Merge pull request #578 from IntelPython/scalar_copy

Author: oleksandr-pavlyk

dpctl/memory/_sycl_usm_array_interface_utils.pxi

@@ -2,7 +2,7 @@
 
 cdef bint _valid_usm_ptr_and_context(DPCTLSyclUSMRef ptr, SyclContext ctx):
     usm_type = _Memory.get_pointer_type(ptr, ctx)
-    return usm_type in (b'shared', b'device', b'host')
+    return usm_type in (b"shared", b"device", b"host")
 
 
 cdef DPCTLSyclQueueRef _queue_ref_copy_from_SyclQueue(
@@ -49,7 +49,7 @@ cdef DPCTLSyclQueueRef get_queue_ref_from_ptr_and_syclobj(
     elif pycapsule.PyCapsule_IsValid(syclobj, "SyclContextRef"):
         ctx = <SyclContext>SyclContext(syclobj)
         return _queue_ref_copy_from_USMRef_and_SyclContext(ptr, ctx)
-    elif hasattr(syclobj, '_get_capsule'):
+    elif hasattr(syclobj, "_get_capsule"):
         cap = syclobj._get_capsule()
         if pycapsule.PyCapsule_IsValid(cap, "SyclQueueRef"):
             q = SyclQueue(cap)
@@ -166,8 +166,8 @@ cdef class _USMBufferData:
         nd = len(ary_shape)
         try:
             dt = np.dtype(ary_typestr)
-            if (dt.hasobject or not (np.issubdtype(dt.type, np.integer) or
-                                     np.issubdtype(dt.type, np.inexact))):
+            if (dt.hasobject or not (np.issubdtype(dt.type, np.number) or
+                                     dt.type is np.bool_)):
                 DPCTLQueue_Delete(QRef)
                 raise TypeError("Only integer types, floating and complex "
                                 "floating types are supported.")

dpctl/tensor/_slicing.pxi

@@ -44,6 +44,9 @@ cdef object _basic_slice_meta(object ind, tuple shape,
     Raises IndexError for invalid index `ind`, and NotImplementedError
     if `ind` is an array.
     """
+    is_integral = lambda x: (
+        isinstance(x, numbers.Integral) or callable(getattr(x, "__index__", None))
+    )
     if ind is Ellipsis:
         return (shape, strides, offset)
     elif ind is None:
@@ -58,7 +61,8 @@ cdef object _basic_slice_meta(object ind, tuple shape,
             new_strides,
             offset + sl_start * strides[0]
         )
-    elif isinstance(ind, numbers.Integral):
+    elif is_integral(ind):
+        ind = ind.__index__()
         if 0 <= ind < shape[0]:
             return (shape[1:], strides[1:], offset + ind * strides[0])
         elif -shape[0] <= ind < 0:
@@ -82,7 +86,7 @@ cdef object _basic_slice_meta(object ind, tuple shape,
                 ellipses_count = ellipses_count + 1
             elif isinstance(i, slice):
                 axes_referenced = axes_referenced + 1
-            elif isinstance(i, numbers.Integral):
+            elif is_integral(i):
                 explicit_index = explicit_index + 1
                 axes_referenced = axes_referenced + 1
             elif isinstance(i, list):
@@ -124,7 +128,8 @@ cdef object _basic_slice_meta(object ind, tuple shape,
                 new_strides.append(str_i)
                 new_offset = new_offset + sl_start * strides[k]
                 k = k_new
-            elif isinstance(ind_i, numbers.Integral):
+            elif is_integral(ind_i):
+                ind_i = ind_i.__index__()
                 if 0 <= ind_i < shape[k]:
                     k_new = k + 1
                     new_offset = new_offset + ind_i * strides[k]

dpctl/tensor/_usmarray.pyx

@@ -491,6 +491,54 @@ cdef class usm_ndarray:
         res.flags_ |= (self.flags_ & USM_ARRAY_WRITEABLE)
         return res
 
+    def __bool__(self):
+        if self.size == 1:
+            mem_view = dpmem.as_usm_memory(self)
+            return mem_view.copy_to_host().view(self.dtype).__bool__()
+
+        if self.size == 0:
+            raise ValueError(
+                "The truth value of an empty array is ambiguous"
+            )
+
+        raise ValueError(
+            "The truth value of an array with more than one element is "
+            "ambiguous. Use a.any() or a.all()"
+        )
+
+    def __float__(self):
+        if self.size == 1:
+            mem_view = dpmem.as_usm_memory(self)
+            return mem_view.copy_to_host().view(self.dtype).__float__()
+
+        raise ValueError(
+            "only size-1 arrays can be converted to Python scalars"
+        )
+
+    def __complex__(self):
+        if self.size == 1:
+            mem_view = dpmem.as_usm_memory(self)
+            return mem_view.copy_to_host().view(self.dtype).__complex__()
+
+        raise ValueError(
+            "only size-1 arrays can be converted to Python scalars"
+        )
+
+    def __int__(self):
+        if self.size == 1:
+            mem_view = dpmem.as_usm_memory(self)
+            return mem_view.copy_to_host().view(self.dtype).__int__()
+
+        raise ValueError(
+            "only size-1 arrays can be converted to Python scalars"
+        )
+
+    def __index__(self):
+        if np.issubdtype(self.dtype, np.integer):
+            return int(self)
+
+        raise IndexError("only integer arrays are valid indices")
+
     def to_device(self, target_device):
         """
         Transfer array to target device

dpctl/tests/test_usm_ndarray_ctor.py

@@ -114,6 +114,64 @@ def test_properties():
     assert isinstance(X.ndim, numbers.Integral)
 
 
+@pytest.mark.parametrize("func", [bool, float, int, complex])
+@pytest.mark.parametrize("shape", [tuple(), (1,), (1, 1), (1, 1, 1)])
+@pytest.mark.parametrize("dtype", ["|b1", "|u2", "|f4", "|i8"])
+def test_copy_scalar_with_func(func, shape, dtype):
+    X = dpt.usm_ndarray(shape, dtype=dtype)
+    Y = np.arange(1, X.size + 1, dtype=dtype).reshape(shape)
+    X.usm_data.copy_from_host(Y.reshape(-1).view("|u1"))
+    assert func(X) == func(Y)
+
+
+@pytest.mark.parametrize(
+    "method", ["__bool__", "__float__", "__int__", "__complex__"]
+)
+@pytest.mark.parametrize("shape", [tuple(), (1,), (1, 1), (1, 1, 1)])
+@pytest.mark.parametrize("dtype", ["|b1", "|u2", "|f4", "|i8"])
+def test_copy_scalar_with_method(method, shape, dtype):
+    X = dpt.usm_ndarray(shape, dtype=dtype)
+    Y = np.arange(1, X.size + 1, dtype=dtype).reshape(shape)
+    X.usm_data.copy_from_host(Y.reshape(-1).view("|u1"))
+    assert getattr(X, method)() == getattr(Y, method)()
+
+
+@pytest.mark.parametrize("func", [bool, float, int, complex])
+@pytest.mark.parametrize("shape", [(2,), (1, 2), (3, 4, 5), (0,)])
+def test_copy_scalar_invalid_shape(func, shape):
+    X = dpt.usm_ndarray(shape)
+    with pytest.raises(ValueError):
+        func(X)
+
+
+@pytest.mark.parametrize("shape", [(1,), (1, 1), (1, 1, 1)])
+@pytest.mark.parametrize("index_dtype", ["|i8"])
+def test_usm_ndarray_as_index(shape, index_dtype):
+    X = dpt.usm_ndarray(shape, dtype=index_dtype)
+    Xnp = np.arange(1, X.size + 1, dtype=index_dtype).reshape(shape)
+    X.usm_data.copy_from_host(Xnp.reshape(-1).view("|u1"))
+    Y = np.arange(X.size + 1)
+    assert Y[X] == Y[1]
+
+
+@pytest.mark.parametrize("shape", [(2,), (1, 2), (3, 4, 5), (0,)])
+@pytest.mark.parametrize("index_dtype", ["|i8"])
+def test_usm_ndarray_as_index_invalid_shape(shape, index_dtype):
+    X = dpt.usm_ndarray(shape, dtype=index_dtype)
+    Y = np.arange(X.size + 1)
+    with pytest.raises(IndexError):
+        Y[X]
+
+
+@pytest.mark.parametrize("shape", [(1,), (1, 1), (1, 1, 1)])
+@pytest.mark.parametrize("index_dtype", ["|f8"])
+def test_usm_ndarray_as_index_invalid_dtype(shape, index_dtype):
+    X = dpt.usm_ndarray(shape, dtype=index_dtype)
+    Y = np.arange(X.size + 1)
+    with pytest.raises(IndexError):
+        Y[X]
+
+
 @pytest.mark.parametrize(
     "ind",
     [
@@ -251,6 +309,14 @@ def test_slicing_basic():
         Xusm[:, -128]
     with pytest.raises(TypeError):
         Xusm[{1, 2, 3, 4, 5, 6, 7}]
+    X = dpt.usm_ndarray(10, "u1")
+    X.usm_data.copy_from_host(b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09")
+    int(
+        X[X[2]]
+    )  # check that objects with __index__ method can be used as indices
+    Xh = dpm.as_usm_memory(X[X[2] : X[5]]).copy_to_host()
+    Xnp = np.arange(0, 10, dtype="u1")
+    assert np.array_equal(Xh, Xnp[Xnp[2] : Xnp[5]])
 
 
 def test_ctor_invalid_shape():
@@ -291,3 +357,19 @@ def test_usm_ndarray_props():
     except dpctl.SyclQueueCreationError:
         pytest.skip("Sycl device CPU was not detected")
     Xusm.to_device("cpu")
+
+
+def test_datapi_device():
+    X = dpt.usm_ndarray(1)
+    dev_t = type(X.device)
+    with pytest.raises(TypeError):
+        dev_t()
+    dev_t.create_device(X.device)
+    dev_t.create_device(X.sycl_queue)
+    dev_t.create_device(X.sycl_device)
+    dev_t.create_device(X.sycl_device.filter_string)
+    dev_t.create_device(None)
+    X.device.sycl_context
+    X.device.sycl_queue
+    X.device.sycl_device
+    repr(X.device)