Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use memory resource in DeviceBuffer #1132

Draft
wants to merge 4 commits into
base: branch-23.06
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions python/rmm/_lib/device_buffer.pxd
Original file line number Diff line number Diff line change
@@ -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 +
@@ -53,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=*)
@@ -68,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 *
48 changes: 33 additions & 15 deletions python/rmm/_lib/device_buffer.pyx
Original file line number Diff line number Diff line change
@@ -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 = <const void*>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):
@@ -133,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
@@ -152,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()
@@ -167,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.
@@ -341,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
@@ -369,7 +387,7 @@ cpdef DeviceBuffer to_device(const unsigned char[::1] b,

cdef uintptr_t p = <uintptr_t>&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)