Skip to content

Commit

Permalink
Use device_buffer constructors w/memory resource
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jakirkham committed Oct 18, 2022
1 parent c690822 commit fc8e3ed
Showing 1 changed file with 11 additions and 5 deletions.
16 changes: 11 additions & 5 deletions python/rmm/_lib/device_buffer.pxd
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 (
device_memory_resource,
DeviceMemoryResource,
)


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 Down

0 comments on commit fc8e3ed

Please sign in to comment.