Add dlpack support with tests and docstrings

Author: vlad-perevezentsev

dpnp/dpnp_array.py

@@ -140,7 +140,7 @@ def __bool__(self):
         return self._array_obj.__bool__()
 
  # '__class__',
- 
+
     def __complex__(self):
         return self._array_obj.__complex__()
 
@@ -153,6 +153,12 @@ def __complex__(self):
  # '__divmod__',
  # '__doc__',
 
+    def __dlpack__(self, stream=None):
+        return self._array_obj.__dlpack__(stream=stream)
+
+    def __dlpack_device__(self):
+        return self._array_obj.__dlpack_device__()
+
     def __eq__(self, other):
         return dpnp.equal(self, other)
 
@@ -190,7 +196,7 @@ def __gt__(self, other):
  # '__imatmul__',
  # '__imod__',
  # '__imul__',
- 
+
     def __index__(self):
         return self._array_obj.__index__()
 

dpnp/dpnp_iface.py

@@ -64,6 +64,7 @@
     "default_float_type",
     "dpnp_queue_initialize",
     "dpnp_queue_is_cpu",
+    "from_dlpack",
     "get_dpnp_descriptor",
     "get_include",
     "get_normalized_queue_device"
@@ -222,6 +223,30 @@ def default_float_type(device=None, sycl_queue=None):
     return map_dtype_to_device(float64, _sycl_queue.sycl_device)
 
 
+def from_dlpack(obj):
+    """
+    Create a dpnp array from a Python object implementing the ``__dlpack__``
+    protocol.
+
+    See https://dmlc.github.io/dlpack/latest/ for more details.
+
+    Parameters
+    ----------
+    obj : A Python object representing an array that implements the ``__dlpack__``
+        and ``__dlpack_device__`` methods.
+
+    Returns
+    -------
+    array : dpnp_array
+
+    """
+
+    usm_ary = dpt.from_dlpack(obj)
+    dpnp_ary = dpnp_array.__new__(dpnp_array)
+    dpnp_ary._array_obj = usm_ary
+    return dpnp_ary
+
+
 def get_dpnp_descriptor(ext_obj,
                         copy_when_strides=True,
                         copy_when_nondefault_queue=True,

tests/test_dparray.py

@@ -23,6 +23,25 @@ def test_astype(arr, arr_dtype, res_dtype):
     assert_array_equal(expected, result)
 
 
+@pytest.mark.parametrize("arr_dtype", get_all_dtypes())
+@pytest.mark.parametrize("shape", [tuple(), (2,), (3, 0, 1), (2, 2, 2)])
+def test_from_dlpack(arr_dtype,shape):
+    X = dpnp.empty(shape=shape,dtype=arr_dtype)
+    Y = dpnp.from_dlpack(X)
+    assert_array_equal(X, Y)
+    assert X.__dlpack_device__() == Y.__dlpack_device__()
+    assert X.shape == Y.shape
+    assert X.dtype == Y.dtype or (
+        str(X.dtype) == "bool" and str(Y.dtype) == "uint8"
+    )
+    assert X.sycl_device == Y.sycl_device
+    assert X.usm_type == Y.usm_type
+    if Y.ndim:
+        V = Y[::-1]
+        W = dpnp.from_dlpack(V)
+        assert V.strides == W.strides
+
+
 @pytest.mark.parametrize("arr_dtype", get_all_dtypes())
 @pytest.mark.parametrize("arr",
                          [[-2, -1, 0, 1, 2], [[-2, -1], [1, 2]], []],