Add Cython/Python copy_to_host and to_device

CHANGELOG.md

@@ -6,6 +6,7 @@
 - PR #252 Add `__sizeof__` method to `DeviceBuffer`
 - PR #258 Define pickling behavior for `DeviceBuffer`
 - PR #261 Add `__bytes__` method to `DeviceBuffer`
+- PR #268 Add Cython/Python `copy_to_host` function
 
 ## 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/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
@@ -51,6 +47,7 @@ cdef class DeviceBuffer:
     @staticmethod
     cdef DeviceBuffer c_frombytes(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,10 @@ cdef class DeviceBuffer:
     cdef void* c_data(self)
 
 
+cpdef void copy_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,11 @@
 # cython: embedsignature = True
 # cython: language_level = 3
 
+
+import warnings
+
+import numpy as np
+
 from libcpp.memory cimport unique_ptr
 from libc.stdint cimport uintptr_t
 
@@ -109,22 +114,37 @@ cdef class DeviceBuffer:
     def frombytes(const unsigned char[::1] b, uintptr_t stream=0):
         return DeviceBuffer.c_frombytes(b, stream)
 
+    cpdef copy_to_host(self, unsigned char[::1] hb=None, uintptr_t stream=0):
+        cdef const device_buffer* dbp = self.c_obj.get()
+        cdef size_t s = dbp.size()
+
+        if hb is None:
+            hb = np.empty((s,), dtype="u1")
+        elif len(hb) < s:
+            raise ValueError(
+                "Argument `hb` is too small. Need space for %i bytes." % s
+            )
+        elif len(hb) > s:
+            hb = hb[:s]
+            warnings.warn(
+                "Argument `hb` larger than needed."
+                " Will fill only first %i bytes." % s,
+                RuntimeWarning
+            )
+
+        with nogil:
+            copy_to_host(<uintptr_t>dbp.data(), hb, 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 +159,28 @@ cdef class DeviceBuffer:
 
     cdef void* c_data(self):
         return self.c_obj.get()[0].data()
+
+
+@cython.boundscheck(False)
+cpdef void copy_to_host(uintptr_t db,
+                        unsigned char[::1] hb,
+                        uintptr_t stream=0) nogil except *:
+    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}")
+
+    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

@@ -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.frombytes(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",
     [

tests/device_buffer_tests.cpp

@@ -111,15 +111,6 @@ TYPED_TEST(DeviceBufferTest, CopyFromRawDevicePointer) {
   EXPECT_EQ(cudaSuccess, cudaFree(device_memory));
 }
 
-TYPED_TEST(DeviceBufferTest, CopyToRawHostPointer) {
-  rmm::device_buffer buff(this->size);
-  std::vector<uint8_t> host_data(this->size);
-  uint8_t* host_data_ptr = host_data.data();
-  rmm::copy_to_host(buff, host_data_ptr);
-  EXPECT_EQ(0, buff.stream());
-  // TODO check for equality between the contents of the two allocations
-}
-
 TYPED_TEST(DeviceBufferTest, CopyFromRawHostPointer) {
   std::vector<uint8_t> host_data(this->size);
   rmm::device_buffer buff(static_cast<void*>(host_data.data()), this->size);