Add Cython/Python copy_to_host and to_device

CHANGELOG.md

@@ -8,6 +8,7 @@
 - PR #261 Add `__bytes__` method to `DeviceBuffer`
 - PR #262 Moved device memory resource files to `mr/device` directory
 - PR #266 Drop `rmm.auto_device`
+- PR #268 Add Cython/Python `copy_to_host` and `to_device`
 
 ## Improvements
 

CMakeLists.txt

@@ -109,7 +109,6 @@ link_directories("${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES}" # CMAKE_CUDA_IMPLICIT
 
 add_library(rmm SHARED
             src/rmm.cpp
-            src/device_buffer.cpp
             src/memory_manager.cpp
             thirdparty/cnmem/src/cnmem.cpp
             src/mr/device/default_memory_resource.cpp)

include/rmm/device_buffer.hpp

@@ -392,17 +392,4 @@ class device_buffer {
       mr::get_default_resource()};  ///< The memory resource used to
                                     ///< allocate/deallocate device memory
 };
-
-/**--------------------------------------------------------------------------*
- * @brief Copies rmm::device_buffer to a preallocated host buffer.
- *
- * Copies device memory asynchronously on the specified stream
- *
- * @throws std::runtime_error if `hb` is `nullptr` or copy fails
- *
- * @param db `rmm::device_buffer` to copy to host
- * @param hb host allocated buffer to copy data to
- * @param stream CUDA stream on which the device to host copy will be performed
- *-------------------------------------------------------------------------**/
-void copy_to_host(const device_buffer& db, void* hb, cudaStream_t stream = 0);
 }  // namespace rmm

python/rmm/_lib/device_buffer.pxd

@@ -20,7 +20,7 @@
 from libcpp.memory cimport unique_ptr
 from libc.stdint cimport uintptr_t
 
-from rmm._lib.lib cimport cudaStream_t
+from rmm._lib.lib cimport cudaStream_t, cudaMemcpyAsync, cudaMemcpyDeviceToHost
 
 cdef extern from "rmm/device_buffer.hpp" namespace "rmm" nogil:
     cdef cppclass device_buffer:
@@ -37,10 +37,6 @@ cdef extern from "rmm/device_buffer.hpp" namespace "rmm" nogil:
         size_t size()
         size_t capacity()
 
-    void copy_to_host(const device_buffer& db, void* hb) except +
-    void copy_to_host(const device_buffer& db,
-                      void* hb,
-                      cudaStream_t stream) except +
 
 cdef class DeviceBuffer:
     cdef unique_ptr[device_buffer] c_obj
@@ -49,8 +45,9 @@ cdef class DeviceBuffer:
     cdef DeviceBuffer c_from_unique_ptr(unique_ptr[device_buffer] ptr)
 
     @staticmethod
-    cdef DeviceBuffer c_frombytes(const unsigned char[::1] b,
+    cdef DeviceBuffer c_to_device(const unsigned char[::1] b,
                                   uintptr_t stream=*)
+    cpdef copy_to_host(self, unsigned char[::1] hb=*, uintptr_t stream=*)
     cpdef bytes tobytes(self, uintptr_t stream=*)
 
     cdef size_t c_size(self)
@@ -59,5 +56,11 @@ cdef class DeviceBuffer:
     cdef void* c_data(self)
 
 
+cpdef DeviceBuffer to_device(const unsigned char[::1] b, uintptr_t stream=*)
+cpdef void copy_ptr_to_host(uintptr_t db,
+                            unsigned char[::1] hb,
+                            uintptr_t stream=*) nogil except *
+
+
 cdef extern from "<utility>" namespace "std" nogil:
     cdef unique_ptr[device_buffer] move(unique_ptr[device_buffer])

python/rmm/_lib/device_buffer.pyx

@@ -17,6 +17,9 @@
 # cython: embedsignature = True
 # cython: language_level = 3
 
+
+import numpy as np
+
 from libcpp.memory cimport unique_ptr
 from libc.stdint cimport uintptr_t
 
@@ -71,7 +74,7 @@ cdef class DeviceBuffer:
         return self.tobytes()
 
     def __setstate__(self, state):
-        cdef DeviceBuffer other = DeviceBuffer.c_frombytes(state)
+        cdef DeviceBuffer other = DeviceBuffer.c_to_device(state)
         self.c_obj = move(other.c_obj)
 
     @property
@@ -92,39 +95,61 @@ cdef class DeviceBuffer:
         return buf
 
     @staticmethod
-    @cython.boundscheck(False)
-    cdef DeviceBuffer c_frombytes(const unsigned char[::1] b,
+    cdef DeviceBuffer c_to_device(const unsigned char[::1] b,
                                   uintptr_t stream=0):
-        if b is None:
-            raise TypeError(
-                "Argument 'b' has incorrect type"
-                " (expected bytes, got NoneType)"
+        """Calls ``to_device`` function on arguments provided"""
+        return to_device(b, stream)
+
+    @staticmethod
+    def to_device(const unsigned char[::1] b, uintptr_t stream=0):
+        """Calls ``to_device`` function on arguments provided"""
+        return to_device(b, stream)
+
+    cpdef copy_to_host(self, unsigned char[::1] hb=None, uintptr_t stream=0):
+        """Copy from a ``DeviceBuffer`` to a buffer on host
+
+        Parameters
+        ----------
+        hb : ``bytes``-like buffer to write into
+        stream : CUDA stream to use for copying, default 0
+
+        Examples
+        --------
+        >>> import rmm
+        >>> db = rmm.DeviceBuffer.to_device(b"abc")
+        >>> hb = bytearray(db.nbytes)
+        >>> db.copy_to_host(hb)
+        >>> print(hb)
+        bytearray(b'abc')
+        >>> hb = db.copy_to_host()
+        >>> print(hb)
+        bytearray(b'abc')
+        """
+        cdef const device_buffer* dbp = self.c_obj.get()
+        cdef size_t s = dbp.size()
+
+        if hb is None:
+            # NumPy leverages huge pages under-the-hood,
+            # which speeds up the copy from device to host.
+            hb = np.empty((s,), dtype="u1")
+        elif len(hb) < s:
+            raise ValueError(
+                "Argument `hb` is too small. Need space for %i bytes." % s
             )
 
-        cdef uintptr_t p = <uintptr_t>&b[0]
-        cdef size_t s = len(b)
-        return DeviceBuffer(ptr=p, size=s, stream=stream)
+        with nogil:
+            copy_ptr_to_host(<uintptr_t>dbp.data(), hb[:s], stream)
 
-    @staticmethod
-    def frombytes(const unsigned char[::1] b, uintptr_t stream=0):
-        return DeviceBuffer.c_frombytes(b, stream)
+        return hb
 
     cpdef bytes tobytes(self, uintptr_t stream=0):
         cdef const device_buffer* dbp = self.c_obj.get()
         cdef size_t s = dbp.size()
-        if s == 0:
-            return b""
 
         cdef bytes b = PyBytes_FromStringAndSize(NULL, s)
-        cdef void* p = <void*>PyBytes_AS_STRING(b)
-        cdef cudaStream_t c_stream
-        cdef cudaError_t err
-        with nogil:
-            c_stream = <cudaStream_t>stream
-            copy_to_host(dbp[0], p, c_stream)
-            err = cudaStreamSynchronize(c_stream)
-        if err != cudaSuccess:
-            raise RuntimeError(f"Stream sync failed with error: {err}")
+        cdef unsigned char* p = <unsigned char*>PyBytes_AS_STRING(b)
+        cdef unsigned char[::1] mv = (<unsigned char[:(s + 1):1]>p)[:s]
+        self.copy_to_host(mv, stream)
 
         return b
 
@@ -139,3 +164,79 @@ cdef class DeviceBuffer:
 
     cdef void* c_data(self):
         return self.c_obj.get()[0].data()
+
+
+@cython.boundscheck(False)
+cpdef DeviceBuffer to_device(const unsigned char[::1] b, uintptr_t stream=0):
+    """Return a new ``DeviceBuffer`` with a copy of the data
+
+    Parameters
+    ----------
+    b : ``bytes``-like data on host to copy to device
+    stream : CUDA stream to use for copying, default 0
+
+    Returns
+    -------
+    ``DeviceBuffer`` with copy of data from host
+
+    Examples
+    --------
+    >>> import rmm
+    >>> db = rmm._lib.device_buffer.to_device(b"abc")
+    >>> print(bytes(db))
+    b'abc'
+    """
+
+    if b is None:
+        raise TypeError(
+            "Argument 'b' has incorrect type"
+            " (expected bytes-like, got NoneType)"
+        )
+
+    cdef uintptr_t p = <uintptr_t>&b[0]
+    cdef size_t s = len(b)
+    return DeviceBuffer(ptr=p, size=s, stream=stream)
+
+
+@cython.boundscheck(False)
+cpdef void copy_ptr_to_host(uintptr_t db,
+                            unsigned char[::1] hb,
+                            uintptr_t stream=0) nogil except *:
+    """Copy from a device pointer to a buffer on host
+
+    Parameters
+    ----------
+    db : pointer to data on device to copy
+    hb : ``bytes``-like buffer to write into
+    stream : CUDA stream to use for copying, default 0
+
+    Examples
+    --------
+    >>> import rmm
+    >>> db = rmm.DeviceBuffer.to_device(b"abc")
+    >>> hb = bytearray(db.nbytes)
+    >>> rmm._lib.device_buffer.copy_ptr_to_host(db.ptr, hb)
+    >>> print(hb)
+    bytearray(b'abc')
+    """
+
+    if hb is None:
+        with gil:
+            raise TypeError(
+                "Argument `hb` has incorrect type"
+                " (expected bytes-like, got NoneType)"
+            )
+
+    cdef cudaError_t err
+
+    err = cudaMemcpyAsync(<void*>&hb[0], <const void*>db, len(hb),
+                          cudaMemcpyDeviceToHost, <cudaStream_t>stream)
+    if err != cudaSuccess:
+        with gil:
+            raise RuntimeError(f"Memcpy failed with error: {err}")
+
+    if stream == 0:
+        err = cudaStreamSynchronize(<cudaStream_t>stream)
+        if err != cudaSuccess:
+            with gil:
+                raise RuntimeError(f"Stream sync failed with error: {err}")

python/rmm/_lib/lib.pxd

@@ -33,6 +33,17 @@ cdef extern from * nogil:
 
     ctypedef void* cudaStream_t "cudaStream_t"
 
+    ctypedef enum cudaMemcpyKind "cudaMemcpyKind":
+        cudaMemcpyHostToHost = 0
+        cudaMemcpyHostToDevice = 1
+        cudaMemcpyDeviceToHost = 2
+        cudaMemcpyDeviceToDevice = 3
+        cudaMemcpyDefault = 4
+
+    cudaError_t cudaMemcpyAsync(void* dst, const void* src, size_t count,
+                                cudaMemcpyKind kind)
+    cudaError_t cudaMemcpyAsync(void* dst, const void* src, size_t count,
+                                cudaMemcpyKind kind, cudaStream_t stream)
     cudaError_t cudaStreamSynchronize(cudaStream_t stream)
 
 

python/rmm/tests/test_rmm.py

@@ -114,7 +114,7 @@ def test_rmm_device_buffer(size):
     assert len(s) == len(b)
 
     # Test conversion from bytes
-    b2 = rmm.DeviceBuffer.frombytes(s)
+    b2 = rmm.DeviceBuffer.to_device(s)
     assert isinstance(b2, rmm.DeviceBuffer)
     assert len(b2) == len(s)
 
@@ -124,6 +124,32 @@ def test_rmm_device_buffer(size):
     assert b.capacity() >= b.size
 
 
+@pytest.mark.parametrize(
+    "hb",
+    [
+        b"abc",
+        bytearray(b"abc"),
+        memoryview(b"abc"),
+        np.asarray(memoryview(b"abc")),
+        np.arange(3, dtype="u1"),
+    ],
+)
+def test_rmm_device_buffer_memoryview_roundtrip(hb):
+    mv = memoryview(hb)
+    db = rmm.DeviceBuffer.to_device(hb)
+    hb2 = db.copy_to_host()
+    mv2 = memoryview(hb2)
+    assert mv == mv2
+    hb3 = bytearray(mv.nbytes)
+    hb3 = db.copy_to_host(hb3)
+    mv3 = memoryview(hb3)
+    assert mv == mv3
+    hb4 = np.empty_like(mv)
+    hb4 = db.copy_to_host(hb4)
+    mv4 = memoryview(hb4)
+    assert mv == mv4
+
+
 @pytest.mark.parametrize(
     "hb",
     [
@@ -146,19 +172,19 @@ def test_rmm_device_buffer_bytes_roundtrip(hb):
         mv = memoryview(hb)
     except TypeError:
         with pytest.raises(TypeError):
-            rmm.DeviceBuffer.frombytes(hb)
+            rmm.DeviceBuffer.to_device(hb)
     else:
         if mv.format != "B":
             with pytest.raises(ValueError):
-                rmm.DeviceBuffer.frombytes(hb)
+                rmm.DeviceBuffer.to_device(hb)
         elif len(mv.strides) != 1:
             with pytest.raises(ValueError):
-                rmm.DeviceBuffer.frombytes(hb)
+                rmm.DeviceBuffer.to_device(hb)
         elif mv.strides[0] != 1:
             with pytest.raises(ValueError):
-                rmm.DeviceBuffer.frombytes(hb)
+                rmm.DeviceBuffer.to_device(hb)
         else:
-            db = rmm.DeviceBuffer.frombytes(hb)
+            db = rmm.DeviceBuffer.to_device(hb)
             hb2 = db.tobytes()
             mv2 = memoryview(hb2)
             assert mv == mv2
@@ -169,7 +195,7 @@ def test_rmm_device_buffer_bytes_roundtrip(hb):
 
 @pytest.mark.parametrize("hb", [b"", b"123", b"abc"])
 def test_rmm_device_buffer_pickle_roundtrip(hb):
-    db = rmm.DeviceBuffer.frombytes(hb)
+    db = rmm.DeviceBuffer.to_device(hb)
     pb = pickle.dumps(db)
     del db
     db2 = pickle.loads(pb)