Add frombytes to convert bytes-like to DeviceBuffer

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## New Features
 
+- PR #253 Add `frombytes` to convert `bytes`-like to `DeviceBuffer`
+
 ## Improvements
 
 - PR #246 Type `DeviceBuffer` arguments to `__cinit__`

python/rmm/_lib/device_buffer.pxd

@@ -29,6 +29,9 @@ cdef class DeviceBuffer:
     @staticmethod
     cdef DeviceBuffer c_from_unique_ptr(unique_ptr[device_buffer] ptr)
 
+    @staticmethod
+    cdef DeviceBuffer c_frombytes(const unsigned char[::1] b,
+                                  uintptr_t stream=*)
     cpdef bytes tobytes(self, uintptr_t stream=*)
 
     cdef size_t c_size(self)

python/rmm/_lib/device_buffer.pyx

@@ -6,6 +6,8 @@ from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING
 from rmm._lib.lib cimport (cudaError_t, cudaSuccess,
                            cudaStream_t, cudaStreamSynchronize)
 
+cimport cython
+
 
 cdef class DeviceBuffer:
 
@@ -57,17 +59,36 @@ cdef class DeviceBuffer:
         buf.c_obj = move(ptr)
         return buf
 
+    @staticmethod
+    @cython.boundscheck(False)
+    cdef DeviceBuffer c_frombytes(const unsigned char[::1] b,
+                                  uintptr_t stream=0):
+        if b is None:
+            raise TypeError(
+                "Argument 'b' has incorrect type"
+                " (expected bytes, got NoneType)"
+            )
+
+        cdef uintptr_t p = <uintptr_t>&b[0]
+        cdef size_t s = len(b)
+        return DeviceBuffer(ptr=p, size=s, stream=stream)
+
+    @staticmethod
+    def frombytes(const unsigned char[::1] b, uintptr_t stream=0):
+        return DeviceBuffer.c_frombytes(b, stream)
+
     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 cudaStream_t c_stream = <cudaStream_t>stream
         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:

python/rmm/tests/test_rmm.py

@@ -110,12 +110,56 @@ def test_rmm_device_buffer(size):
     assert isinstance(s, bytes)
     assert len(s) == len(b)
 
+    # Test conversion from bytes
+    b2 = rmm.DeviceBuffer.frombytes(s)
+    assert isinstance(b2, rmm.DeviceBuffer)
+    assert len(b2) == len(s)
+
     # Test resizing
     b.resize(2)
     assert b.size == 2
     assert b.capacity() >= b.size
 
 
+@pytest.mark.parametrize(
+    "hb",
+    [
+        None,
+        "abc",
+        123,
+        b"",
+        np.ones((2,), "u2"),
+        np.ones((2, 2), "u1"),
+        np.ones(4, "u1")[::2],
+        b"abc",
+        bytearray(b"abc"),
+        memoryview(b"abc"),
+        np.arange(3, dtype="u1"),
+    ],
+)
+def test_rmm_device_buffer_bytes_roundtrip(hb):
+    try:
+        mv = memoryview(hb)
+    except TypeError:
+        with pytest.raises(TypeError):
+            rmm.DeviceBuffer.frombytes(hb)
+    else:
+        if mv.format != "B":
+            with pytest.raises(ValueError):
+                rmm.DeviceBuffer.frombytes(hb)
+        elif len(mv.strides) != 1:
+            with pytest.raises(ValueError):
+                rmm.DeviceBuffer.frombytes(hb)
+        elif mv.strides[0] != 1:
+            with pytest.raises(ValueError):
+                rmm.DeviceBuffer.frombytes(hb)
+        else:
+            db = rmm.DeviceBuffer.frombytes(hb)
+            hb2 = db.tobytes()
+            mv2 = memoryview(hb2)
+            assert mv == mv2
+
+
 def test_rmm_cupy_allocator():
     cupy = pytest.importorskip("cupy")