Move stream functions from CUDAContext to CUDAStream (#14110)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14110

I'm planning to move CUDAStream to c10/cuda, without also moving
CUDAContext, and so it's most convenient if these definitions
are in the actual header file in question.

Reviewed By: smessmer

Differential Revision: D13104693

fbshipit-source-id: 23ce492003091adadaa5ca6a17124213005046c2

diff --git a/aten/src/ATen/cuda/CUDAContext.cpp b/aten/src/ATen/cuda/CUDAContext.cpp
index a250619..e3df328 100644 (file)
--- a/aten/src/ATen/cuda/CUDAContext.cpp
+++ b/aten/src/ATen/cuda/CUDAContext.cpp
@@ -16,27 +16,6 @@ cudaDeviceProp* getDeviceProperties(int64_t device) {
   return THCState_getDeviceProperties(at::globalContext().getTHCState(), (int)device);
 }
 
-/* Streams */
-CUDAStream getStreamFromPool(
-  const bool isHighPriority
-, int64_t device) {
-  return CUDAStream(impl::CUDAStream_getStreamFromPool(isHighPriority, device));
-}
-
-CUDAStream getDefaultCUDAStream(int64_t device) {
-  return CUDAStream(impl::CUDAStream_getDefaultStream(device));
-}
-CUDAStream getCurrentCUDAStream(int64_t device) {
-  return CUDAStream(impl::CUDAStream_getCurrentStream(device));
-}
-
-void setCurrentCUDAStream(CUDAStream stream) {
-  impl::CUDAStream_setStream(stream.internals());
-}
-void uncheckedSetCurrentCUDAStream(CUDAStream stream) {
-  impl::CUDAStream_uncheckedSetStream(stream.internals());
-}
-
 Allocator* getCUDADeviceAllocator() {
   return at::globalContext().getTHCState()->cudaDeviceAllocator;
 }

diff --git a/aten/src/ATen/cuda/CUDAContext.h b/aten/src/ATen/cuda/CUDAContext.h
index a0792bb..413e3a3 100644 (file)
--- a/aten/src/ATen/cuda/CUDAContext.h
+++ b/aten/src/ATen/cuda/CUDAContext.h
@@ -47,27 +47,6 @@ CAFFE2_API int warp_size();
 
 CAFFE2_API cudaDeviceProp* getDeviceProperties(int64_t device);
 
-/* Streams */
-
-/**
- * Get a new stream from the CUDA stream pool.  You can think of this
- * as "creating" a new stream, but no such creation actually happens;
- * instead, streams are preallocated from the pool and returned in a
- * round-robin fashion.
- *
- * You can request a stream from the high priority pool by setting
- * isHighPriority to true, or a stream for a specific device by setting device
- * (defaulting to the current CUDA stream.)
- */
-CAFFE2_API CUDAStream
-getStreamFromPool(const bool isHighPriority = false, int64_t device = -1);
-
-CAFFE2_API CUDAStream getDefaultCUDAStream(int64_t device = -1);
-CAFFE2_API CUDAStream getCurrentCUDAStream(int64_t device = -1);
-
-CAFFE2_API void setCurrentCUDAStream(CUDAStream stream);
-CAFFE2_API void uncheckedSetCurrentCUDAStream(CUDAStream stream);
-
 CAFFE2_API Allocator* getCUDADeviceAllocator();
 
 /* Handles */

diff --git a/aten/src/ATen/cuda/CUDAStream.cpp b/aten/src/ATen/cuda/CUDAStream.cpp
index 317ebe7..0e4a4e3 100644 (file)
--- a/aten/src/ATen/cuda/CUDAStream.cpp
+++ b/aten/src/ATen/cuda/CUDAStream.cpp
@@ -323,5 +323,27 @@ impl::CUDAStreamInternals* CUDAStream::internals() const {
   }
 }
 
+/* Streams */
+CUDAStream getStreamFromPool(
+  const bool isHighPriority
+, int64_t device) {
+  return CUDAStream(impl::CUDAStream_getStreamFromPool(isHighPriority, device));
+}
+
+CUDAStream getDefaultCUDAStream(int64_t device) {
+  return CUDAStream(impl::CUDAStream_getDefaultStream(device));
+}
+CUDAStream getCurrentCUDAStream(int64_t device) {
+  return CUDAStream(impl::CUDAStream_getCurrentStream(device));
+}
+
+void setCurrentCUDAStream(CUDAStream stream) {
+  impl::CUDAStream_setStream(stream.internals());
+}
+void uncheckedSetCurrentCUDAStream(CUDAStream stream) {
+  impl::CUDAStream_uncheckedSetStream(stream.internals());
+}
+
+
 } // namespace cuda
 } // namespace at

diff --git a/aten/src/ATen/cuda/CUDAStream.h b/aten/src/ATen/cuda/CUDAStream.h
index 1e66ea0..4f0b257 100644 (file)
--- a/aten/src/ATen/cuda/CUDAStream.h
+++ b/aten/src/ATen/cuda/CUDAStream.h
@@ -13,8 +13,6 @@
 * A CUDAStream interface. See CUDAStream.cpp for implementation details.
 *
 * Includes the CUDAStream convenience class and a pointer-based stream API.
-*
-* The ATen/cuda/CUDAContext interface should be preferred when working with streams.
 */
 
 /*
@@ -108,5 +106,25 @@ private:
   Stream stream_;
 };
 
+/**
+ * Get a new stream from the CUDA stream pool.  You can think of this
+ * as "creating" a new stream, but no such creation actually happens;
+ * instead, streams are preallocated from the pool and returned in a
+ * round-robin fashion.
+ *
+ * You can request a stream from the high priority pool by setting
+ * isHighPriority to true, or a stream for a specific device by setting device
+ * (defaulting to the current CUDA stream.)
+ */
+CAFFE2_API CUDAStream
+getStreamFromPool(const bool isHighPriority = false, int64_t device = -1);
+
+CAFFE2_API CUDAStream getDefaultCUDAStream(int64_t device = -1);
+CAFFE2_API CUDAStream getCurrentCUDAStream(int64_t device = -1);
+
+CAFFE2_API void setCurrentCUDAStream(CUDAStream stream);
+CAFFE2_API void uncheckedSetCurrentCUDAStream(CUDAStream stream);
+
+
 } // namespace cuda
 } // namespace at