From fdbd8fce02971197ceca107b7ebdc8ae8cabc889 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 20 Nov 2020 03:49:16 -0800 Subject: [PATCH 1/6] Support enabling PTDS via CUDA_PTDS environment variable --- python/rmm/_lib/cuda_stream_view.pyx | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/rmm/_lib/cuda_stream_view.pyx b/python/rmm/_lib/cuda_stream_view.pyx index 817d091f2..2f5b4d66e 100644 --- a/python/rmm/_lib/cuda_stream_view.pyx +++ b/python/rmm/_lib/cuda_stream_view.pyx @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os + from libc.stdint cimport uintptr_t @@ -26,7 +28,10 @@ cdef class CudaStreamView: CUDA stream to wrap, default 0 """ if (stream == 0): - self.c_obj.reset(new cuda_stream_view()) + if int(os.environ.get("CUDA_PTDS", "0")) != 0: + self.c_obj.reset(new cuda_stream_view(2)) + else: + self.c_obj.reset(new cuda_stream_view()) else: self.c_obj.reset(new cuda_stream_view(stream)) From d7891178b1c6d3bf52b28862e76efab2e71a19ea Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 23 Nov 2020 04:31:37 -0800 Subject: [PATCH 2/6] Expose cudaStreamLegacy/cudaStreamPerThread in Python --- python/rmm/_lib/cuda_stream_view.pyx | 3 ++- python/rmm/_lib/lib.pxd | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/rmm/_lib/cuda_stream_view.pyx b/python/rmm/_lib/cuda_stream_view.pyx index 2f5b4d66e..fde7e12ae 100644 --- 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(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 index ba023a802..4e4e0a604 100644 --- 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 From 1ee52afcc5ae9ef3c02d4ddc117f3a8a688042b4 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 23 Nov 2020 04:33:09 -0800 Subject: [PATCH 3/6] Rename PTDS environment variable to RMM_PER_THREAD_DEFAULT_STREAM --- python/rmm/_lib/cuda_stream_view.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/_lib/cuda_stream_view.pyx b/python/rmm/_lib/cuda_stream_view.pyx index fde7e12ae..2145e689b 100644 --- a/python/rmm/_lib/cuda_stream_view.pyx +++ b/python/rmm/_lib/cuda_stream_view.pyx @@ -29,7 +29,7 @@ cdef class CudaStreamView: CUDA stream to wrap, default 0 """ if (stream == 0): - if int(os.environ.get("CUDA_PTDS", "0")) != 0: + if int(os.environ.get("RMM_PER_THREAD_DEFAULT_STREAM", "0")) != 0: self.c_obj.reset(new cuda_stream_view(cudaStreamPerThread)) else: self.c_obj.reset(new cuda_stream_view()) From d6c635315c4906cf2a0362bc24e2d37dbf75017b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 23 Nov 2020 11:32:23 -0800 Subject: [PATCH 4/6] Switch Python stream types to cudaStream enum --- python/rmm/_lib/cuda_stream_view.pyx | 13 +++++++++---- python/rmm/_lib/lib.pxd | 3 --- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/python/rmm/_lib/cuda_stream_view.pyx b/python/rmm/_lib/cuda_stream_view.pyx index 2145e689b..2faaf7f30 100644 --- a/python/rmm/_lib/cuda_stream_view.pyx +++ b/python/rmm/_lib/cuda_stream_view.pyx @@ -15,12 +15,17 @@ import os from libc.stdint cimport uintptr_t -from rmm._lib.lib cimport cudaStreamPerThread + + +cpdef enum cudaStream: + cudaStreamDefault = 0 + cudaStreamLegacy = 1 + cudaStreamPerThread = 2 cdef class CudaStreamView: - def __cinit__(self, uintptr_t stream=0): + def __cinit__(self, cudaStream stream=cudaStreamDefault): """Construct a view of the specified CUDA stream Parameters @@ -28,9 +33,9 @@ cdef class CudaStreamView: stream : uintptr_t, optional CUDA stream to wrap, default 0 """ - if (stream == 0): + if (stream == cudaStreamDefault): if int(os.environ.get("RMM_PER_THREAD_DEFAULT_STREAM", "0")) != 0: - self.c_obj.reset(new cuda_stream_view(cudaStreamPerThread)) + 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 index 4e4e0a604..ba023a802 100644 --- a/python/rmm/_lib/lib.pxd +++ b/python/rmm/_lib/lib.pxd @@ -25,9 +25,6 @@ cdef extern from * nogil: ctypedef void* cudaStream_t "cudaStream_t" - cudaStream_t cudaStreamLegacy - cudaStream_t cudaStreamPerThread - ctypedef enum cudaMemcpyKind "cudaMemcpyKind": cudaMemcpyHostToHost = 0 cudaMemcpyHostToDevice = 1 From b1d6aad2539384b1d29a0a44df83553839a04c0e Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 6 Jan 2021 09:31:41 -0800 Subject: [PATCH 5/6] Update RMM_PER_THREAD_DEFAULT_STREAM usage --- python/rmm/_cuda/stream.pyx | 7 ++++++- python/setup.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/rmm/_cuda/stream.pyx b/python/rmm/_cuda/stream.pyx index 0f6c5ab19..541524124 100644 --- a/python/rmm/_cuda/stream.pyx +++ b/python/rmm/_cuda/stream.pyx @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os + from libc.stdint cimport uintptr_t from libcpp cimport bool @@ -117,7 +119,10 @@ cdef class Stream: self._cuda_stream, self._owner = stream._cuda_stream, stream._owner -DEFAULT_STREAM = Stream._from_cudaStream_t(cuda_stream_default.value()) +if int(os.environ.get("RMM_PER_THREAD_DEFAULT_STREAM", "0")) != 0: + DEFAULT_STREAM = Stream._from_cudaStream_t(cuda_stream_per_thread.value()) +else: + DEFAULT_STREAM = Stream._from_cudaStream_t(cuda_stream_default.value()) LEGACY_DEFAULT_STREAM = Stream._from_cudaStream_t(cuda_stream_legacy.value()) PER_THREAD_DEFAULT_STREAM = Stream._from_cudaStream_t( cuda_stream_per_thread.value() diff --git a/python/setup.py b/python/setup.py index 49444f13f..d413a5ac3 100644 --- a/python/setup.py +++ b/python/setup.py @@ -64,7 +64,7 @@ def get_cuda_version_from_header(cuda_include_dir): cwd = os.getcwd() preprocess_files = ["gpu.pxd"] -supported_cuda_versions = {"10.1", "10.2", "11.0"} +supported_cuda_versions = {"10.1", "10.2", "11.0", "11.1"} for file_p in preprocess_files: pxi_file = ".".join(file_p.split(".")[:-1]) From 7ede244fb14cde6f1f46b2f7a825a5387d3befc3 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 6 Jan 2021 09:33:40 -0800 Subject: [PATCH 6/6] Revert unnecessary change to setup.py --- python/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/setup.py b/python/setup.py index d413a5ac3..49444f13f 100644 --- a/python/setup.py +++ b/python/setup.py @@ -64,7 +64,7 @@ def get_cuda_version_from_header(cuda_include_dir): cwd = os.getcwd() preprocess_files = ["gpu.pxd"] -supported_cuda_versions = {"10.1", "10.2", "11.0", "11.1"} +supported_cuda_versions = {"10.1", "10.2", "11.0"} for file_p in preprocess_files: pxi_file = ".".join(file_p.split(".")[:-1])