Enable constructing DeviceBuffers using a non-default memory resource and correctly manage memory resource lifetime in DeviceBuffer

python/rmm/_lib/device_buffer.pxd

@@ -17,15 +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 +
+        device_buffer(const void* source_data,
+                      size_t size, cuda_stream_view stream,
+                      device_memory_resource* mr) except +
         void resize(size_t new_size, cuda_stream_view stream) except +
         void shrink_to_fit(cuda_stream_view stream) except +
         void* data()

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=None,
+                  DeviceMemoryResource mr=None):
         """Construct a ``DeviceBuffer`` with optional size and data pointer
 
         Parameters
@@ -64,9 +65,13 @@ 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 : DeviceMemoryResource, optional
+            The memory resource to use for the allocation. If not
+            provided, the memory resource returned from
+            ``get_current_device_resource()`` is used.
 
-        Note
-        ----
+        Notes
+        -----
 
         If the pointer passed is non-null and ``stream`` is the default stream,
         it is synchronized after the copy. However if a non-default ``stream``
@@ -77,25 +82,36 @@ cdef class DeviceBuffer:
         >>> import rmm
         >>> db = rmm.DeviceBuffer(size=5)
         """
-        cdef const void* c_ptr
+        cdef:
+            const void* c_ptr
+            cuda_stream_view c_stream_view
+            device_memory_resource* c_mr
+
+        if stream is None:
+            stream = DEFAULT_STREAM
+        if mr is None:
+            mr = get_current_device_resource()
+
+        # Save a reference to the MR and stream used for allocation
+        self.mr = mr
+        self.stream = stream
 
         with nogil:
             c_ptr = <const void*>ptr
+            c_stream_view = stream.view()
+            c_mr = mr.c_obj.get()
 
             if size == 0:
                 self.c_obj.reset(new device_buffer())
             elif c_ptr == NULL:
-                self.c_obj.reset(new device_buffer(size, stream.view()))
+                self.c_obj.reset(new device_buffer(size, c_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, c_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.stream = stream
-
     def __len__(self):
         return self.size
 

python/rmm/tests/test_rmm.py

@@ -716,3 +716,10 @@ def test_dev_buf_circle_ref_dealloc():
     # deallocate `dbuf1` (which needs the MR alive), a segfault occurs.
 
     gc.collect()
+
+
+def test_device_buffer_nondefault_mr():
+    # test creating a DeviceBuffer with a non-default memory resource
+    dbuf = rmm.DeviceBuffer(size=100_000, mr=rmm.mr.CudaMemoryResource())
+    del dbuf
+    gc.collect()