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.

Sign up for GitHub

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

Jump to bottom

Use memory resource in DeviceBuffer #1132

Draft
wants to merge 4 commits into
base: branch-23.06
Choose a base branch
from
Draft

Use memory resource in DeviceBuffer #1132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fb75182
Use `device_buffer` constructors w/memory resource
jakirkham Oct 18, 2022
548b353
Use memory resource in `DeviceBuffer` construtor
jakirkham Oct 18, 2022
c777637
Accept memory resource (and stream) in `copy`
jakirkham Oct 18, 2022
a66a255
Add memory resource argument through `to_device`
jakirkham Oct 18, 2022
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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 +
Expand All @@ -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=*)
Expand All @@ -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 *
Expand Down
48 changes: 33 additions & 15 deletions python/rmm/_lib/device_buffer.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
----
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down