Expose cudaStreamLegacy/cudaStreamPerThread in Python

rapidsai · Nov 23, 2020 · d789117 · d789117
1 parent fdbd8fc
commit d789117
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/python/rmm/_lib/cuda_stream_view.pyx b/python/rmm/_lib/cuda_stream_view.pyx
@@ -15,6 +15,7 @@
 import os
 
 from libc.stdint cimport uintptr_t
+from rmm._lib.lib cimport cudaStreamPerThread
 
 
 cdef class CudaStreamView:
@@ -29,7 +30,7 @@ cdef class CudaStreamView:
         """
         if (stream == 0):
             if int(os.environ.get("CUDA_PTDS", "0")) != 0:
-                self.c_obj.reset(new cuda_stream_view(<cudaStream_t>2))
+                self.c_obj.reset(new cuda_stream_view(cudaStreamPerThread))
             else:
                 self.c_obj.reset(new cuda_stream_view())
         else:

diff --git a/python/rmm/_lib/lib.pxd b/python/rmm/_lib/lib.pxd
@@ -25,6 +25,9 @@ cdef extern from * nogil:
 
     ctypedef void* cudaStream_t "cudaStream_t"
 
+    cudaStream_t cudaStreamLegacy
+    cudaStream_t cudaStreamPerThread
+
     ctypedef enum cudaMemcpyKind "cudaMemcpyKind":
         cudaMemcpyHostToHost = 0
         cudaMemcpyHostToDevice = 1