Fixed issues in slicing and array construction

These were discovered while preparing for customer presentation

Author: oleksandr-pavlyk

dpctl/tensor/_slicing.pxi

@@ -17,6 +17,23 @@
 import numbers
 
 
+cdef Py_ssize_t _slice_len(
+    Py_ssize_t sl_start,
+    Py_ssize_t sl_stop,
+    Py_ssize_t sl_step
+):
+    """
+    Compute len(range(sl_start, sl_stop, sl_step))
+    """
+    if sl_start == sl_stop:
+        return 0
+    if sl_step > 0:
+        # 1 + argmax k such htat sl_start + sl_step*k < sl_stop
+        return 1 + ((sl_stop - sl_start - 1) // sl_step)
+    else:
+        return 1 + ((sl_stop - sl_start + 1) // sl_step)
+
+
 cdef object _basic_slice_meta(object ind, tuple shape,
                               tuple strides, Py_ssize_t offset):
     """
@@ -33,9 +50,9 @@ cdef object _basic_slice_meta(object ind, tuple shape,
         return ((1,) + shape, (0,) + strides, offset)
     elif isinstance(ind, slice):
         sl_start, sl_stop, sl_step = ind.indices(shape[0])
-        sh0 = (sl_stop - sl_start) // sl_step
+        sh0 = _slice_len(sl_start, sl_stop, sl_step)
         str0 = sl_step * strides[0]
-        new_strides = strides if (sl_step == 1) else (str0,) + strides[1:]
+        new_strides = strides if (sl_step == 1 or sh0 == 0) else (str0,) + strides[1:]
         return (
             (sh0, ) + shape[1:],
             new_strides,
@@ -101,8 +118,8 @@ cdef object _basic_slice_meta(object ind, tuple shape,
             elif isinstance(ind_i, slice):
                 k_new = k + 1
                 sl_start, sl_stop, sl_step = ind_i.indices(shape[k])
-                sh_i = (sl_stop - sl_start) // sl_step
-                str_i = sl_step * strides[k]
+                sh_i = _slice_len(sl_start, sl_stop, sl_step)
+                str_i = (1 if sh_i == 0 else sl_step) * strides[k]
                 new_shape.append(sh_i)
                 new_strides.append(str_i)
                 new_offset = new_offset + sl_start * strides[k]

dpctl/tensor/_stride_utils.pxi

@@ -72,7 +72,7 @@ cdef int _from_input_shape_strides(
 
     # 0-d array
     if (nd == 0):
-        contig[0] = USM_ARRAY_C_CONTIGUOUS
+        contig[0] = (USM_ARRAY_C_CONTIGUOUS | USM_ARRAY_F_CONTIGUOUS)
         nelems[0] = 1
         min_disp[0] = 0
         max_disp[0] = 0
@@ -88,17 +88,28 @@ cdef int _from_input_shape_strides(
         shape_arr[i] = <Py_ssize_t> shape[i]
         elem_count *= shape_arr[i]
     if elem_count == 0:
-        contig[0] = USM_ARRAY_C_CONTIGUOUS
+        contig[0] = (USM_ARRAY_C_CONTIGUOUS | USM_ARRAY_F_CONTIGUOUS)
         nelems[0] = 1
         min_disp[0] = 0
         max_disp[0] = 0
-        strides_ptr[0] = <Py_ssize_t *>(<size_t>0)
+        if strides is None:
+            strides_ptr[0] = <Py_ssize_t *>(<size_t>0)
+        else:
+            strides_arr = <Py_ssize_t*>PyMem_Malloc(nd * sizeof(Py_ssize_t))
+            if (not strides_arr):
+                PyMem_Free(shape_ptr[0]);
+                shape_ptr[0] = <Py_ssize_t *>(<size_t>0)
+                return ERROR_MALLOC
+            strides_ptr[0] = strides_arr
+            for i in range(0, nd):
+                strides_arr[i] = <Py_ssize_t> strides[i]
         return 0
     nelems[0] = elem_count
-
     if (strides is None):
         # no need to allocate and populate strides
         if (int(order) not in [ord('C'), ord('F'), ord('c'), ord('f')]):
+            PyMem_Free(shape_ptr[0]);
+            shape_ptr[0] = <Py_ssize_t *>(<size_t>0)
             return ERROR_INCORRECT_ORDER
         if order == <char> ord('C') or order == <char> ord('c'):
             contig[0] = USM_ARRAY_C_CONTIGUOUS
@@ -112,6 +123,8 @@ cdef int _from_input_shape_strides(
           and len(strides) == nd):
         strides_arr = <Py_ssize_t*>PyMem_Malloc(nd * sizeof(Py_ssize_t))
         if (not strides_arr):
+            PyMem_Free(shape_ptr[0]);
+            shape_ptr[0] = <Py_ssize_t *>(<size_t>0)
             return ERROR_MALLOC
         strides_ptr[0] = strides_arr
         for i in range(0, nd):
@@ -143,6 +156,8 @@ cdef int _from_input_shape_strides(
             contig[0] = 0  # non-contiguous
         return 0
     else:
+        PyMem_Free(shape_ptr[0]);
+        shape_ptr[0] = <Py_ssize_t *>(<size_t>0)
         return ERROR_UNEXPECTED_STRIDES
     # return ERROR_INTERNAL
 

dpctl/tensor/_usmarray.pyx

@@ -475,7 +475,8 @@ cdef class usm_ndarray:
         cdef usm_ndarray res
 
         res = usm_ndarray.__new__(
-            usm_ndarray, _meta[0],
+            usm_ndarray,
+            _meta[0],
             dtype=_make_typestr(self.typenum_),
             strides=_meta[1],
             buffer=self.base_,

dpctl/tests/test_usm_ndarray_ctor.py

@@ -17,12 +17,14 @@
 import numbers
 
 import numpy as np
+import numpy.lib.stride_tricks as np_st
 import pytest
 
 import dpctl
 
 # import dpctl.memory as dpmem
 import dpctl.tensor as dpt
+from dpctl.tensor._usmarray import Device
 
 
 @pytest.mark.parametrize(
@@ -112,6 +114,8 @@ def test_properties():
         (2, 2, None, 3, 4),
         (Ellipsis,),
         (None, slice(0, None, 2), Ellipsis, slice(0, None, 3)),
+        (None, slice(1, None, 2), Ellipsis, slice(1, None, 3)),
+        (None, slice(None, -1, -2), Ellipsis, slice(2, None, 3)),
         (
             slice(None, None, -1),
             slice(None, None, -1),
@@ -121,10 +125,86 @@ def test_properties():
     ],
 )
 def test_basic_slice(ind):
-    X = dpt.usm_ndarray((2 * 3, 2 * 4, 3 * 5, 3 * 7), dtype="u1")
+    X = dpt.usm_ndarray((2 * 3, 2 * 4, 3 * 5, 2 * 7), dtype="u1")
     Xnp = np.empty(X.shape, dtype=X.dtype)
     S = X[ind]
     Snp = Xnp[ind]
     assert S.shape == Snp.shape
     assert S.strides == Snp.strides
     assert S.dtype == X.dtype
+
+
+def _from_numpy(np_ary, device=None, usm_type="shared"):
+    if type(np_ary) is np.ndarray:
+        if np_ary.flags["FORC"]:
+            x = np_ary
+        else:
+            x = np.ascontiguous(np_ary)
+        R = dpt.usm_ndarray(
+            np_ary.shape,
+            dtype=np_ary.dtype,
+            buffer=usm_type,
+            buffer_ctor_kwargs={
+                "queue": Device.create_device(device).sycl_queue
+            },
+        )
+        R.usm_data.copy_from_host(x.reshape((-1)).view("|u1"))
+        return R
+    else:
+        raise ValueError("Expected numpy.ndarray, got {}".format(type(np_ary)))
+
+
+def _to_numpy(usm_ary):
+    if type(usm_ary) is dpt.usm_ndarray:
+        usm_buf = usm_ary.usm_data
+        s = usm_buf.nbytes
+        host_buf = usm_buf.copy_to_host().view(usm_ary.dtype)
+        usm_ary_itemsize = usm_ary.itemsize
+        R_offset = (
+            usm_ary.__sycl_usm_array_interface__["offset"] * usm_ary_itemsize
+        )
+        R = np.ndarray((s,), dtype="u1", buffer=host_buf)
+        R = R[R_offset:].view(usm_ary.dtype)
+        R_strides = (usm_ary_itemsize * si for si in usm_ary.strides)
+        return np_st.as_strided(R, shape=usm_ary.shape, strides=R_strides)
+    else:
+        raise ValueError(
+            "Expected dpctl.tensor.usm_ndarray, got {}".format(type(usm_ary))
+        )
+
+
+def test_slice_constructor_1d():
+    Xh = np.arange(37, dtype="i4")
+    Xusm = _from_numpy(Xh, device="gpu", usm_type="device")
+    for ind in [
+        slice(1, None, 2),
+        slice(0, None, 3),
+        slice(1, None, 3),
+        slice(2, None, 3),
+        slice(None, None, -1),
+        slice(-2, 2, -2),
+        slice(-1, 1, -2),
+        slice(None, None, -13),
+    ]:
+        assert np.array_equal(
+            _to_numpy(Xusm[ind]), Xh[ind]
+        ), "Failed for {}".format(ind)
+
+
+def test_slice_constructor_3d():
+    Xh = np.empty((37, 24, 35), dtype="i4")
+    Xusm = _from_numpy(Xh, device="gpu", usm_type="device")
+    for ind in [
+        slice(1, None, 2),
+        slice(0, None, 3),
+        slice(1, None, 3),
+        slice(2, None, 3),
+        slice(None, None, -1),
+        slice(-2, 2, -2),
+        slice(-1, 1, -2),
+        slice(None, None, -13),
+        (slice(None, None, -2), Ellipsis, None, 15),
+    ]:
+        assert np.array_equal(
+            _to_numpy(Xusm[ind]), Xh[ind]
+        ), "Failed for {}".format(ind)