Full test coverage for Buffer.__dlpack__, __dlpack_device__

Author: rwgk

cuda_core/cuda/core/experimental/_memory.py

@@ -179,14 +179,14 @@ def __dlpack__(
         # Note: we ignore the stream argument entirely (as if it is -1).
         # It is the user's responsibility to maintain stream order.
         if dl_device is not None:
-            raise BufferError("Sorry not supported: dl_device other than None")
+            raise BufferError("Sorry, not supported: dl_device other than None")
         if copy is True:
-            raise BufferError("Sorry not supported: copy=True")
+            raise BufferError("Sorry, not supported: copy=True")
         if max_version is None:
             versioned = False
         else:
             if not isinstance(max_version, tuple) or len(max_version) != 2:
-                raise RuntimeError(f"Expected max_version Tuple[int, int], got {max_version}")
+                raise BufferError(f"Expected max_version Tuple[int, int], got {max_version}")
             versioned = max_version >= (1, 0)
         capsule = make_py_capsule(self, versioned)
         return capsule

cuda_core/tests/test_memory.py

@@ -13,8 +13,10 @@
 
 import ctypes
 
+import pytest
+
 from cuda.core.experimental import Device
-from cuda.core.experimental._memory import Buffer, MemoryResource
+from cuda.core.experimental._memory import Buffer, DLDeviceType, MemoryResource
 from cuda.core.experimental._utils import handle_return
 
 
@@ -116,6 +118,12 @@ def device_id(self) -> int:
         raise RuntimeError("the pinned memory resource is not bound to any GPU")
 
 
+class NullMemoryResource(DummyHostMemoryResource):
+    @property
+    def is_host_accessible(self) -> bool:
+        return False
+
+
 def buffer_initialization(dummy_mr: MemoryResource):
     buffer = dummy_mr.allocate(size=1024)
     assert buffer.handle != 0
@@ -211,3 +219,46 @@ def test_buffer_close():
     buffer_close(DummyHostMemoryResource())
     buffer_close(DummyUnifiedMemoryResource(device))
     buffer_close(DummyPinnedMemoryResource(device))
+
+
+def test_buffer_dunder_dlpack():
+    device = Device()
+    device.set_current()
+    dummy_mr = DummyDeviceMemoryResource(device)
+    buffer = dummy_mr.allocate(size=1024)
+    capsule = buffer.__dlpack__()
+    assert "dltensor" in repr(capsule)
+    capsule = buffer.__dlpack__(max_version=(1, 0))
+    assert "dltensor" in repr(capsule)
+    with pytest.raises(BufferError, match=r"^Sorry, not supported: dl_device other than None$"):
+        buffer.__dlpack__(dl_device=[])
+    with pytest.raises(BufferError, match=r"^Sorry, not supported: copy=True$"):
+        buffer.__dlpack__(copy=True)
+    with pytest.raises(BufferError, match=r"^Expected max_version Tuple\[int, int\], got \[\]$"):
+        buffer.__dlpack__(max_version=[])
+    with pytest.raises(BufferError, match=r"^Expected max_version Tuple\[int, int\], got \(9, 8, 7\)$"):
+        buffer.__dlpack__(max_version=(9, 8, 7))
+
+
+@pytest.mark.parametrize(
+    ("DummyMR", "expected"),
+    [
+        (DummyDeviceMemoryResource, (DLDeviceType.kDLCUDA, 0)),
+        (DummyHostMemoryResource, (DLDeviceType.kDLCPU, 0)),
+        (DummyUnifiedMemoryResource, (DLDeviceType.kDLCUDAHost, 0)),
+        (DummyPinnedMemoryResource, (DLDeviceType.kDLCUDAHost, 0)),
+    ],
+)
+def test_buffer_dunder_dlpack_device_success(DummyMR, expected):
+    device = Device()
+    device.set_current()
+    dummy_mr = DummyMR() if DummyMR is DummyHostMemoryResource else DummyMR(device)
+    buffer = dummy_mr.allocate(size=1024)
+    assert buffer.__dlpack_device__() == expected
+
+
+def test_buffer_dunder_dlpack_device_failure():
+    dummy_mr = NullMemoryResource()
+    buffer = dummy_mr.allocate(size=1024)
+    with pytest.raises(BufferError, match=r"^buffer is neither device-accessible nor host-accessible$"):
+        buffer.__dlpack_device__()