make range functions respect current stream (#21619)

Natalia Gimelshein · facebook-github-bot · commit cf5c3bb3fe4c · 2019-06-11T15:46:48.000-07:00
Summary: Stream is not respected on range/linspace/logspace functions, which contributes to #21589 (this is not a complete solution for that issue). Pull Request resolved: #21619 Differential Revision: D15769666 Pulled By: ezyang fbshipit-source-id: 7c036f7aecb3119430c4d432775cad98a5028fa8
diff --git a/aten/src/ATen/native/cuda/RangeFactories.cu b/aten/src/ATen/native/cuda/RangeFactories.cu
@@ -2,11 +2,13 @@
 #include <ATen/NativeFunctions.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/cuda/Exceptions.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <cmath>
 #include <limits>
 
 #include <thrust/device_ptr.h>
 #include <thrust/sequence.h>
+#include <thrust/execution_policy.h>
 
 namespace at {
 namespace native {
@@ -56,7 +58,9 @@ Tensor& linspace_cuda_out(Tensor& result, Scalar start, Scalar end, int64_t step
       scalar_t step = (scalar_end - scalar_start) / static_cast<scalar_t>(steps - 1);
       LinspaceOp<scalar_t> linspace_method(scalar_start, step);
       thrust::device_ptr<scalar_t> data_(r.data<scalar_t>());
-      thrust::tabulate(data_, data_ + steps, linspace_method);
+      cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+      auto policy = thrust::cuda::par.on(stream);
+      thrust::tabulate(policy, data_, data_ + steps, linspace_method);
     });
   }
 
@@ -87,7 +91,9 @@ Tensor& logspace_cuda_out(Tensor& result, Scalar start, Scalar end, int64_t step
       scalar_t step = (scalar_end - scalar_start) / static_cast<scalar_t>(steps - 1);
       LogspaceOp<scalar_t> logspace_method(scalar_start, step, scalar_base);
       thrust::device_ptr<scalar_t> data_(r.data<scalar_t>());
-      thrust::tabulate(data_, data_ + steps, logspace_method);
+      cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+      auto policy = thrust::cuda::par.on(stream);
+      thrust::tabulate(policy, data_, data_ + steps, logspace_method);
     });
   }
 
@@ -117,8 +123,10 @@ Tensor& range_cuda_out(Tensor& result, Scalar start, Scalar end, Scalar step) {
     }
     Tensor r = result.is_contiguous() ? result : result.contiguous();
     LinspaceOp<scalar_t, accscalar_t> linspace_method(xstart, xstep);
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    auto policy = thrust::cuda::par.on(stream);
     thrust::device_ptr<scalar_t> data_ptr(r.data<scalar_t>());
-    thrust::tabulate(data_ptr, data_ptr + size, linspace_method);
+    thrust::tabulate(policy, data_ptr, data_ptr + size, linspace_method);
 
     if (!result.is_contiguous()) {
       result.copy_(r);
@@ -168,8 +176,10 @@ Tensor& arange_cuda_out(Tensor& result, Scalar start, Scalar end, Scalar step) {
     }
     Tensor r = result.is_contiguous() ? result : result.contiguous();
     LinspaceOp<scalar_t, accscalar_t> linspace_method(xstart, xstep);
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    auto policy = thrust::cuda::par.on(stream);
     thrust::device_ptr<scalar_t> data_ptr(r.data<scalar_t>());
-    thrust::tabulate(data_ptr, data_ptr + size, linspace_method);
+    thrust::tabulate(policy, data_ptr, data_ptr + size, linspace_method);
 
     if (!result.is_contiguous()) {
       result.copy_(r);
