diff --git a/python/rmm/docs/guide.md b/python/rmm/docs/guide.md index c7e940497..4b9d305ab 100644 --- a/python/rmm/docs/guide.md +++ b/python/rmm/docs/guide.md @@ -93,7 +93,7 @@ destination device ID and stream are optional parameters. ```python >>> import rmm >>> rmm.reinitialize(managed_memory=True) ->>> from rmm._cuda.stream import Stream +>>> from rmm.pylibrmm.stream import Stream >>> stream = Stream() >>> buf = rmm.DeviceBuffer(size=100, stream=stream) >>> buf.prefetch(device=3, stream=stream) # prefetch to device on stream. diff --git a/python/rmm/rmm/_cuda/CMakeLists.txt b/python/rmm/rmm/_cuda/CMakeLists.txt index 7759432d3..1617ead7f 100644 --- a/python/rmm/rmm/_cuda/CMakeLists.txt +++ b/python/rmm/rmm/_cuda/CMakeLists.txt @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2022-2024, NVIDIA CORPORATION. +# Copyright (c) 2022-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -12,7 +12,6 @@ # the License. # ============================================================================= -set(cython_sources stream.pyx) set(linked_libraries rmm::rmm cpp_logger) rapids_cython_create_modules(SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}" diff --git a/python/rmm/rmm/_cuda/stream.pxd b/python/rmm/rmm/_cuda/stream.pxd index 219b75864..026830162 100644 --- a/python/rmm/rmm/_cuda/stream.pxd +++ b/python/rmm/rmm/_cuda/stream.pxd @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2024, NVIDIA CORPORATION. +# Copyright (c) 2020-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,22 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from cuda.bindings.cyruntime cimport cudaStream_t -from libc.stdint cimport uintptr_t -from libcpp cimport bool - -from rmm.librmm.cuda_stream_view cimport cuda_stream_view - - -cdef class Stream: - cdef cudaStream_t _cuda_stream - cdef object _owner - - @staticmethod - cdef Stream _from_cudaStream_t(cudaStream_t s, object owner=*) - - cdef cuda_stream_view view(self) except * nogil - cdef void c_synchronize(self) except * nogil - cdef bool c_is_default(self) except * nogil - cdef void _init_with_new_cuda_stream(self) except * - cdef void _init_from_stream(self, Stream stream) except * +from rmm.pylibrmm.stream cimport Stream diff --git a/python/rmm/rmm/_cuda/stream.py b/python/rmm/rmm/_cuda/stream.py new file mode 100644 index 000000000..f912b9ae6 --- /dev/null +++ b/python/rmm/rmm/_cuda/stream.py @@ -0,0 +1,35 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings + +from rmm.pylibrmm.stream import ( + DEFAULT_STREAM, + LEGACY_DEFAULT_STREAM, + PER_THREAD_DEFAULT_STREAM, + Stream, +) + +__all__ = [ + "DEFAULT_STREAM", + "LEGACY_DEFAULT_STREAM", + "PER_THREAD_DEFAULT_STREAM", + "Stream", +] + +warnings.warn( + "The `rmm._cuda.stream` module is deprecated in 25.02 and will be removed in a future release. Use `rmm.pylibrmm.stream` instead.", + FutureWarning, + stacklevel=2, +) diff --git a/python/rmm/rmm/allocators/cupy.py b/python/rmm/rmm/allocators/cupy.py index 780ff2abf..47210e8d6 100644 --- a/python/rmm/rmm/allocators/cupy.py +++ b/python/rmm/rmm/allocators/cupy.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2024, NVIDIA CORPORATION. +# Copyright (c) 2023-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from rmm import pylibrmm -from rmm._cuda.stream import Stream +from rmm.pylibrmm.stream import Stream try: import cupy diff --git a/python/rmm/rmm/pylibrmm/CMakeLists.txt b/python/rmm/rmm/pylibrmm/CMakeLists.txt index 0012cb93d..1be58c32e 100644 --- a/python/rmm/rmm/pylibrmm/CMakeLists.txt +++ b/python/rmm/rmm/pylibrmm/CMakeLists.txt @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2022-2024, NVIDIA CORPORATION. +# Copyright (c) 2022-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -12,7 +12,8 @@ # the License. # ============================================================================= -set(cython_sources device_buffer.pyx logger.pyx memory_resource.pyx cuda_stream.pyx helper.pyx) +set(cython_sources device_buffer.pyx logger.pyx memory_resource.pyx cuda_stream.pyx helper.pyx + stream.pyx) set(linked_libraries rmm::rmm cpp_logger) # Build all of the Cython targets diff --git a/python/rmm/rmm/pylibrmm/device_buffer.pxd b/python/rmm/rmm/pylibrmm/device_buffer.pxd index a0d287423..295c2494e 100644 --- a/python/rmm/rmm/pylibrmm/device_buffer.pxd +++ b/python/rmm/rmm/pylibrmm/device_buffer.pxd @@ -1,4 +1,4 @@ -# Copyright (c) 2019-2024, NVIDIA CORPORATION. +# Copyright (c) 2019-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,9 +15,9 @@ from libc.stdint cimport uintptr_t from libcpp.memory cimport unique_ptr -from rmm._cuda.stream cimport Stream from rmm.librmm.device_buffer cimport device_buffer from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource +from rmm.pylibrmm.stream cimport Stream cdef class DeviceBuffer: diff --git a/python/rmm/rmm/pylibrmm/device_buffer.pyx b/python/rmm/rmm/pylibrmm/device_buffer.pyx index 9343d77c1..0d6825ed8 100644 --- a/python/rmm/rmm/pylibrmm/device_buffer.pyx +++ b/python/rmm/rmm/pylibrmm/device_buffer.pyx @@ -1,4 +1,4 @@ -# Copyright (c) 2019-2024, NVIDIA CORPORATION. +# Copyright (c) 2019-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,9 +19,9 @@ from libc.stdint cimport uintptr_t from libcpp.memory cimport unique_ptr from libcpp.utility cimport move -from rmm._cuda.stream cimport Stream +from rmm.pylibrmm.stream cimport Stream -from rmm._cuda.stream import DEFAULT_STREAM +from rmm.pylibrmm.stream import DEFAULT_STREAM from cuda.bindings.cyruntime cimport ( cudaError, diff --git a/python/rmm/rmm/pylibrmm/memory_resource.pyx b/python/rmm/rmm/pylibrmm/memory_resource.pyx index ca9ee01ac..0189a58b5 100644 --- a/python/rmm/rmm/pylibrmm/memory_resource.pyx +++ b/python/rmm/rmm/pylibrmm/memory_resource.pyx @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2024, NVIDIA CORPORATION. +# Copyright (c) 2020-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,9 +32,9 @@ from cuda.bindings.runtime import cudaError_t from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice -from rmm._cuda.stream cimport Stream +from rmm.pylibrmm.stream cimport Stream -from rmm._cuda.stream import DEFAULT_STREAM +from rmm.pylibrmm.stream import DEFAULT_STREAM from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rmm.librmm.per_device_resource cimport ( diff --git a/python/rmm/rmm/pylibrmm/stream.pxd b/python/rmm/rmm/pylibrmm/stream.pxd new file mode 100644 index 000000000..219b75864 --- /dev/null +++ b/python/rmm/rmm/pylibrmm/stream.pxd @@ -0,0 +1,33 @@ +# Copyright (c) 2020-2024, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from cuda.bindings.cyruntime cimport cudaStream_t +from libc.stdint cimport uintptr_t +from libcpp cimport bool + +from rmm.librmm.cuda_stream_view cimport cuda_stream_view + + +cdef class Stream: + cdef cudaStream_t _cuda_stream + cdef object _owner + + @staticmethod + cdef Stream _from_cudaStream_t(cudaStream_t s, object owner=*) + + cdef cuda_stream_view view(self) except * nogil + cdef void c_synchronize(self) except * nogil + cdef bool c_is_default(self) except * nogil + cdef void _init_with_new_cuda_stream(self) except * + cdef void _init_from_stream(self, Stream stream) except * diff --git a/python/rmm/rmm/_cuda/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx similarity index 100% rename from python/rmm/rmm/_cuda/stream.pyx rename to python/rmm/rmm/pylibrmm/stream.pyx diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py index de363e4b4..ee02d5d0e 100644 --- a/python/rmm/rmm/tests/test_rmm.py +++ b/python/rmm/rmm/tests/test_rmm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2024, NVIDIA CORPORATION. +# Copyright (c) 2020-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,10 +26,10 @@ from numba import cuda import rmm -import rmm._cuda.stream from rmm.allocators.cupy import rmm_cupy_allocator from rmm.allocators.numba import RMMNumbaManager from rmm.pylibrmm.logger import level_enum +from rmm.pylibrmm.stream import Stream cuda.set_memory_manager(RMMNumbaManager) @@ -348,8 +348,8 @@ def test_rmm_device_buffer_prefetch(pool, managed): def test_rmm_pool_numba_stream(stream): rmm.reinitialize(pool_allocator=True) - stream = rmm._cuda.stream.Stream(stream) - a = rmm.pylibrmm.device_buffer.DeviceBuffer(size=3, stream=stream) + stream = Stream(stream) + a = rmm.DeviceBuffer(size=3, stream=stream) assert a.size == 3 assert a.ptr != 0 @@ -600,7 +600,7 @@ def test_mr_devicebuffer_lifetime(): ) # Creates a new non-default stream - stream = rmm._cuda.stream.Stream() + stream = rmm.pylibrmm.stream.Stream() # Allocate DeviceBuffer with Pool and Stream a = rmm.DeviceBuffer(size=10, stream=stream) @@ -695,7 +695,7 @@ def test_cuda_async_memory_resource_stream(nelems): # with a non-default stream works mr = rmm.mr.CudaAsyncMemoryResource() rmm.mr.set_current_device_resource(mr) - stream = rmm._cuda.stream.Stream() + stream = Stream() expected = np.full(nelems, 5, dtype="u1") dbuf = rmm.DeviceBuffer.to_device(expected, stream=stream) result = np.asarray(dbuf.copy_to_host())