From fb751823f1a45f0dade26bfdc6f3eb8c4cd99a36 Mon Sep 17 00:00:00 2001 From: jakirkham Date: Tue, 18 Oct 2022 02:22:15 -0700 Subject: [PATCH 1/4] Use `device_buffer` constructors w/memory resource Update the C++ constructors to take `device_memory_resource` arguments. This means that this argument is now required (at least based on the constructors defined here) when calling `device_buffer` constructors. --- python/rmm/_lib/device_buffer.pxd | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pxd b/python/rmm/_lib/device_buffer.pxd index 42c2b4fef..7918b4f92 100644 --- a/python/rmm/_lib/device_buffer.pxd +++ b/python/rmm/_lib/device_buffer.pxd @@ -17,17 +17,23 @@ from libcpp.memory cimport unique_ptr from rmm._cuda.stream cimport Stream from rmm._lib.cuda_stream_view cimport cuda_stream_view -from rmm._lib.memory_resource cimport DeviceMemoryResource +from rmm._lib.memory_resource cimport ( + DeviceMemoryResource, + device_memory_resource, +) cdef extern from "rmm/device_buffer.hpp" namespace "rmm" nogil: cdef cppclass device_buffer: device_buffer() - device_buffer(size_t size, cuda_stream_view stream) except + + device_buffer(size_t size, cuda_stream_view stream, + device_memory_resource* mr) except + device_buffer(const void* source_data, - size_t size, cuda_stream_view stream) except + + size_t size, cuda_stream_view stream, + device_memory_resource* mr) except + device_buffer(const device_buffer buf, - cuda_stream_view stream) except + + cuda_stream_view stream, + device_memory_resource* mr) except + void reserve(size_t new_capacity, cuda_stream_view stream) except + void resize(size_t new_size, cuda_stream_view stream) except + void shrink_to_fit(cuda_stream_view stream) except + From 548b353eced2cd588dc877c020c6efd82371ca73 Mon Sep 17 00:00:00 2001 From: jakirkham Date: Tue, 18 Oct 2022 02:22:16 -0700 Subject: [PATCH 2/4] Use memory resource in `DeviceBuffer` construtor Update the Python constructor to take and handle a `DeviceMemoryResource` argument. Also pass this through to `device_buffer` constructors. --- python/rmm/_lib/device_buffer.pyx | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index ce0009dbf..613e0f2a3 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -47,7 +47,8 @@ cdef class DeviceBuffer: def __cinit__(self, *, uintptr_t ptr=0, size_t size=0, - Stream stream=DEFAULT_STREAM): + Stream stream=DEFAULT_STREAM, + DeviceMemoryResource mr=None): """Construct a ``DeviceBuffer`` with optional size and data pointer Parameters @@ -64,6 +65,9 @@ cdef class DeviceBuffer: scope while the DeviceBuffer is in use. Destroying the underlying stream while the DeviceBuffer is in use will result in undefined behavior. + mr : optional + Memory resource to use to allocate memory for the underlying + ``device_buffer``. Note ---- @@ -77,22 +81,31 @@ cdef class DeviceBuffer: >>> db = rmm.DeviceBuffer(size=5) """ cdef const void* c_ptr + cdef device_memory_resource* c_mr + + # Use default memory resource if none is specified. + # Also get C++ representation to call constructor below. + if mr is None: + mr = get_current_device_resource() + c_mr = mr.get_mr() with nogil: c_ptr = ptr if size == 0: - self.c_obj.reset(new device_buffer()) + self.c_obj.reset(new device_buffer(c_mr)) elif c_ptr == NULL: - self.c_obj.reset(new device_buffer(size, stream.view())) + self.c_obj.reset(new device_buffer(size, stream.view(), c_mr)) else: - self.c_obj.reset(new device_buffer(c_ptr, size, stream.view())) + self.c_obj.reset( + new device_buffer(c_ptr, size, stream.view(), c_mr) + ) if stream.c_is_default(): stream.c_synchronize() # Save a reference to the MR and stream used for allocation - self.mr = get_current_device_resource() + self.mr = mr self.stream = stream def __len__(self): From c777637236bf0215d46bfc94debc28ed298e23b7 Mon Sep 17 00:00:00 2001 From: jakirkham Date: Tue, 18 Oct 2022 02:22:17 -0700 Subject: [PATCH 3/4] Accept memory resource (and stream) in `copy` Changes `copy` to allow a memory resource to be specified (using the default device memory resource if unspecified). Also accepts a `stream` argument (similar to other functions). --- python/rmm/_lib/device_buffer.pyx | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index 613e0f2a3..5dc6d1b55 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -146,7 +146,9 @@ cdef class DeviceBuffer: } return intf - def copy(self): + def copy(self, *, + Stream stream=DEFAULT_STREAM, + DeviceMemoryResource mr=None): """Returns a copy of DeviceBuffer. Returns @@ -165,9 +167,9 @@ cdef class DeviceBuffer: >>> assert db is not db_copy >>> assert db.ptr != db_copy.ptr """ - ret = DeviceBuffer(ptr=self.ptr, size=self.size, stream=self.stream) - ret.mr = self.mr - return ret + return DeviceBuffer( + ptr=self.ptr, size=self.size, stream=stream, mr=mr + ) def __copy__(self): return self.copy() From a66a255697f2a4e735209ce69b8be5b356938461 Mon Sep 17 00:00:00 2001 From: jakirkham Date: Tue, 18 Oct 2022 02:22:18 -0700 Subject: [PATCH 4/4] Add memory resource argument through `to_device` While the `to_device` function already included a memory resource, it didn't use it. Plus other functions calling `to_device` did not use the argument. The change here makes sure `to_device` passes this to the `DeviceBuffer` constructor. Also it makes sure other functions have a default argument, which they set if one is not specified. --- python/rmm/_lib/device_buffer.pxd | 4 ++-- python/rmm/_lib/device_buffer.pyx | 15 +++++++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pxd b/python/rmm/_lib/device_buffer.pxd index 7918b4f92..dc79364f6 100644 --- a/python/rmm/_lib/device_buffer.pxd +++ b/python/rmm/_lib/device_buffer.pxd @@ -59,7 +59,7 @@ cdef class DeviceBuffer: @staticmethod cdef DeviceBuffer c_to_device(const unsigned char[::1] b, - Stream stream=*) + Stream stream=*, DeviceMemoryResource mr=*) cpdef copy_to_host(self, ary=*, Stream stream=*) cpdef copy_from_host(self, ary, Stream stream=*) cpdef copy_from_device(self, cuda_ary, Stream stream=*) @@ -74,7 +74,7 @@ cdef class DeviceBuffer: cdef device_buffer c_release(self) except * cpdef DeviceBuffer to_device(const unsigned char[::1] b, - Stream stream=*) + Stream stream=*, DeviceMemoryResource mr=*) cpdef void copy_ptr_to_host(uintptr_t db, unsigned char[::1] hb, Stream stream=*) except * diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index 5dc6d1b55..59a51066d 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -182,15 +182,17 @@ cdef class DeviceBuffer: @staticmethod cdef DeviceBuffer c_to_device(const unsigned char[::1] b, - Stream stream=DEFAULT_STREAM): + Stream stream=DEFAULT_STREAM, + DeviceMemoryResource mr=None): """Calls ``to_device`` function on arguments provided""" - return to_device(b, stream) + return to_device(b, stream, mr) @staticmethod def to_device(const unsigned char[::1] b, - Stream stream=DEFAULT_STREAM): + Stream stream=DEFAULT_STREAM, + DeviceMemoryResource mr=None): """Calls ``to_device`` function on arguments provided.""" - return to_device(b, stream) + return to_device(b, stream, mr) cpdef copy_to_host(self, ary=None, Stream stream=DEFAULT_STREAM): """Copy from a ``DeviceBuffer`` to a buffer on host. @@ -356,7 +358,8 @@ cdef class DeviceBuffer: @cython.boundscheck(False) cpdef DeviceBuffer to_device(const unsigned char[::1] b, - Stream stream=DEFAULT_STREAM): + Stream stream=DEFAULT_STREAM, + DeviceMemoryResource mr=None): """Return a new ``DeviceBuffer`` with a copy of the data. Parameters @@ -384,7 +387,7 @@ cpdef DeviceBuffer to_device(const unsigned char[::1] b, cdef uintptr_t p = &b[0] cdef size_t s = len(b) - return DeviceBuffer(ptr=p, size=s, stream=stream) + return DeviceBuffer(ptr=p, size=s, stream=stream, mr=mr) @cython.boundscheck(False)