[FEA] Add method to copy device_buffer back to host memory

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## New Features
 
 - PR #218 Add `_DevicePointer`
+- PR #219 Add method to copy `device_buffer` back to host memory
 - PR #222 Expose free and total memory in Python interface
 
 ## Improvements

include/rmm/device_buffer.hpp

@@ -392,4 +392,29 @@ 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) {
+  if (hb == nullptr) {
+    throw std::runtime_error{"Cannot copy to `nullptr`."};
+  }
+  cudaError_t err = cudaMemcpyAsync(hb,
+                                    db.data(),
+                                    db.size(),
+                                    cudaMemcpyDeviceToHost,
+                                    stream);
+  if (err != cudaSuccess) {
+    throw std::runtime_error{"Failed to copy to host."};
+  }
+}
 }  // namespace rmm

python/rmm/_lib/device_buffer.pxd

@@ -1,4 +1,5 @@
 from libcpp.memory cimport unique_ptr
+from libc.stdint cimport uintptr_t
 
 from rmm._lib.lib cimport cudaStream_t
 
@@ -14,12 +15,19 @@ 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
 
     @staticmethod
     cdef DeviceBuffer c_from_unique_ptr(unique_ptr[device_buffer] ptr)
 
+    cpdef bytes tobytes(self, uintptr_t stream=*)
+
     cdef size_t c_size(self)
     cpdef void resize(self, size_t new_size)
     cpdef size_t capacity(self)

python/rmm/_lib/device_buffer.pyx

@@ -1,6 +1,11 @@
 from libcpp.memory cimport unique_ptr
 from libc.stdint cimport uintptr_t
 
+from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING
+
+from rmm._lib.lib cimport (cudaError_t, cudaSuccess,
+                           cudaStream_t, cudaStreamSynchronize)
+
 
 cdef class DeviceBuffer:
 
@@ -50,6 +55,20 @@ cdef class DeviceBuffer:
         buf.c_obj = move(ptr)
         return buf
 
+    cpdef bytes tobytes(self, uintptr_t stream=0):
+        cdef const device_buffer* dbp = self.c_obj.get()
+        cdef bytes b = PyBytes_FromStringAndSize(NULL, self.c_size())
+        cdef char* p = PyBytes_AS_STRING(b)
+        cdef cudaError_t err
+
+        with nogil:
+            copy_to_host(dbp[0], <void*>p, <cudaStream_t>stream)
+            err = cudaStreamSynchronize(<cudaStream_t>stream)
+        if err != cudaSuccess:
+            raise RuntimeError(f"Stream sync failed with error: {err}")
+
+        return b
+
     cdef size_t c_size(self):
         return self.c_obj.get()[0].size()
 

python/rmm/_lib/lib.pxd

@@ -28,8 +28,13 @@ ctypedef pair[size_t, size_t] memory_pair
 
 cdef extern from * nogil:
 
+    ctypedef enum cudaError_t "cudaError_t":
+        cudaSuccess = 0
+
     ctypedef void* cudaStream_t "cudaStream_t"
 
+    cudaError_t cudaStreamSynchronize(cudaStream_t stream)
+
 
 cdef uintptr_t c_alloc(
     size_t size,

python/rmm/tests/test_rmm.py

@@ -105,6 +105,11 @@ def test_rmm_device_buffer(size):
     assert b.__cuda_array_interface__["typestr"] == "|u1"
     assert b.__cuda_array_interface__["version"] == 0
 
+    # Test conversion to bytes
+    s = b.tobytes()
+    assert isinstance(s, bytes)
+    assert len(s) == len(b)
+
     # Test resizing
     b.resize(2)
     assert b.size == 2

tests/device_buffer_tests.cpp

@@ -111,6 +111,15 @@ 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);