Add more tests to cover different use cases

Author: antonwolfy

tests/test_dlpack.py

@@ -0,0 +1,74 @@
+import numpy
+import pytest
+from numpy.testing import assert_array_equal
+
+import dpnp
+
+from .helper import (
+    get_all_dtypes,
+)
+
+device_oneAPI = 14  # DLDeviceType.kDLOneAPI
+
+
+class TestDLPack:
+    @pytest.mark.parametrize("stream", [None, 1])
+    def test_stream(self, stream):
+        x = dpnp.arange(5)
+        x.__dlpack__(stream=stream)
+
+    @pytest.mark.parametrize("copy", [True, None, False])
+    def test_copy(self, copy):
+        x = dpnp.arange(5)
+        x.__dlpack__(copy=copy)
+
+    def test_wrong_copy(self):
+        x = dpnp.arange(5)
+        x.__dlpack__(copy=dpnp.array([1, 2, 3]))
+
+    @pytest.mark.parametrize("xp", [dpnp, numpy])
+    @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True))
+    def test_dtype_passthrough(self, xp, dt):
+        x = xp.arange(5).astype(dt)
+        y = xp.from_dlpack(x)
+
+        assert y.dtype == x.dtype
+        assert_array_equal(x, y)
+
+    @pytest.mark.parametrize("xp", [dpnp, numpy])
+    def test_non_contiguous(self, xp):
+        x = xp.arange(25).reshape((5, 5))
+
+        y1 = x[0]
+        assert_array_equal(y1, xp.from_dlpack(y1))
+
+        y2 = x[:, 0]
+        assert_array_equal(y2, xp.from_dlpack(y2))
+
+        y3 = x[1, :]
+        assert_array_equal(y3, xp.from_dlpack(y3))
+
+        y4 = x[1]
+        assert_array_equal(y4, xp.from_dlpack(y4))
+
+        y5 = xp.diagonal(x).copy()
+        assert_array_equal(y5, xp.from_dlpack(y5))
+
+    def test_device(self):
+        x = dpnp.arange(5)
+        assert x.__dlpack_device__()[0] == device_oneAPI
+        y = dpnp.from_dlpack(x)
+        assert y.__dlpack_device__()[0] == device_oneAPI
+        z = y[::2]
+        assert z.__dlpack_device__()[0] == device_oneAPI
+
+    def test_ndim0(self):
+        x = dpnp.array(1.0)
+        y = dpnp.from_dlpack(x)
+        assert_array_equal(x, y)
+
+    def test_device(self):
+        x = dpnp.arange(5)
+        y = dpnp.from_dlpack(x, device=x.__dlpack_device__())
+        assert x.device == y.device
+        assert x.get_array()._pointer == y.get_array()._pointer

tests/third_party/cupy/core_tests/test_dlpack.py

@@ -0,0 +1,120 @@
+import unittest
+
+import dpctl
+import dpctl.tensor._dlpack as dlp
+import numpy
+import pytest
+
+import dpnp as cupy
+from tests.third_party.cupy import testing
+
+
+def _gen_array(dtype, alloc_q=None):
+    if cupy.issubdtype(dtype, numpy.unsignedinteger):
+        array = cupy.random.randint(
+            0, 10, size=(2, 3), sycl_queue=alloc_q
+        ).astype(dtype)
+    elif cupy.issubdtype(dtype, cupy.integer):
+        array = cupy.random.randint(
+            -10, 10, size=(2, 3), sycl_queue=alloc_q
+        ).astype(dtype)
+    elif cupy.issubdtype(dtype, cupy.floating):
+        array = cupy.random.rand(2, 3, sycl_queue=alloc_q).astype(dtype)
+    elif cupy.issubdtype(dtype, cupy.complexfloating):
+        array = cupy.random.random((2, 3), sycl_queue=alloc_q).astype(dtype)
+    elif dtype == cupy.bool_:
+        array = cupy.random.randint(
+            0, 2, size=(2, 3), sycl_queue=alloc_q
+        ).astype(cupy.bool_)
+    else:
+        assert False, f"unrecognized dtype: {dtype}"
+    return array
+
+
+class TestDLPackConversion(unittest.TestCase):
+    @testing.for_all_dtypes(no_bool=False)
+    def test_conversion(self, dtype):
+        orig_array = _gen_array(dtype)
+        tensor = orig_array.__dlpack__()
+        out_array = dlp.from_dlpack_capsule(tensor)
+        testing.assert_array_equal(orig_array, out_array)
+        assert orig_array.get_array()._pointer == out_array._pointer
+
+
+@testing.parameterize(*testing.product({"memory": ("device", "managed")}))
+class TestNewDLPackConversion(unittest.TestCase):
+    def _get_stream(self, stream_name):
+        if stream_name == "null":
+            return dpctl.SyclQueue()
+        return dpctl.SyclQueue()
+
+    @testing.for_all_dtypes(no_bool=False)
+    def test_conversion(self, dtype):
+        orig_array = _gen_array(dtype)
+        out_array = cupy.from_dlpack(orig_array)
+        testing.assert_array_equal(orig_array, out_array)
+        assert orig_array.get_array()._pointer == out_array.get_array()._pointer
+
+    def test_stream(self):
+        allowed_streams = ["null", True]
+
+        # stream order is automatically established via DLPack protocol
+        for src_s in [self._get_stream(s) for s in allowed_streams]:
+            for dst_s in [self._get_stream(s) for s in allowed_streams]:
+                orig_array = _gen_array(cupy.float32, alloc_q=src_s)
+                dltensor = orig_array.__dlpack__(stream=orig_array)
+
+                out_array = dlp.from_dlpack_capsule(dltensor)
+                out_array = cupy.from_dlpack(out_array, device=dst_s)
+                testing.assert_array_equal(orig_array, out_array)
+                assert (
+                    orig_array.get_array()._pointer
+                    == out_array.get_array()._pointer
+                )
+
+
+class TestDLTensorMemory(unittest.TestCase):
+    # def setUp(self):
+    #     self.old_pool = cupy.get_default_memory_pool()
+    #     self.pool = cupy.cuda.MemoryPool()
+    #     cupy.cuda.set_allocator(self.pool.malloc)
+
+    # def tearDown(self):
+    #     self.pool.free_all_blocks()
+    #     cupy.cuda.set_allocator(self.old_pool.malloc)
+
+    def test_deleter(self):
+        # memory is freed when tensor is deleted, as it's not consumed
+        array = cupy.empty(10)
+        tensor = array.__dlpack__()
+        # str(tensor): <capsule object "dltensor" at 0x7f7c4c835330>
+        assert '"dltensor"' in str(tensor)
+        # assert self.pool.n_free_blocks() == 0
+        # del array
+        # assert self.pool.n_free_blocks() == 0
+        # del tensor
+        # assert self.pool.n_free_blocks() == 1
+
+    def test_deleter2(self):
+        # memory is freed when array2 is deleted, as tensor is consumed
+        array = cupy.empty(10)
+        tensor = array.__dlpack__()
+        assert '"dltensor"' in str(tensor)
+        array2 = dlp.from_dlpack_capsule(tensor)
+        assert '"used_dltensor"' in str(tensor)
+        # assert self.pool.n_free_blocks() == 0
+        # del array
+        # assert self.pool.n_free_blocks() == 0
+        # del array2
+        # assert self.pool.n_free_blocks() == 1
+        # del tensor
+        # assert self.pool.n_free_blocks() == 1
+
+    def test_multiple_consumption_error(self):
+        # Prevent segfault, see #3611
+        array = cupy.empty(10)
+        tensor = array.__dlpack__()
+        array2 = dlp.from_dlpack_capsule(tensor)
+        with pytest.raises(ValueError) as e:
+            array3 = dlp.from_dlpack_capsule(tensor)
+        assert "consumed multiple times" in str(e.value)