Move nonzero() to dpctl_ext/tensor and reuse it in dpnp

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -41,6 +41,7 @@
 )
 from dpctl_ext.tensor._indexing_functions import (
     extract,
+    nonzero,
     place,
     put,
     take,
@@ -57,6 +58,7 @@
     "extract",
     "from_numpy",
     "full",
+    "nonzero",
     "place",
     "put",
     "reshape",

dpctl_ext/tensor/_copy_utils.py

@@ -200,6 +200,37 @@ def _extract_impl(ary, ary_mask, axis=0):
     return dst
 
 
+def _nonzero_impl(ary):
+    if not isinstance(ary, dpt.usm_ndarray):
+        raise TypeError(
+            f"Expecting type dpctl.tensor.usm_ndarray, got {type(ary)}"
+        )
+    exec_q = ary.sycl_queue
+    usm_type = ary.usm_type
+    mask_nelems = ary.size
+    cumsum_dt = dpt.int32 if mask_nelems < int32_t_max else dpt.int64
+    cumsum = dpt.empty(
+        mask_nelems, dtype=cumsum_dt, sycl_queue=exec_q, order="C"
+    )
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
+    dep_evs = _manager.submitted_events
+    mask_count = ti.mask_positions(
+        ary, cumsum, sycl_queue=exec_q, depends=dep_evs
+    )
+    indexes_dt = ti.default_device_index_type(exec_q.sycl_device)
+    indexes = dpt.empty(
+        (ary.ndim, mask_count),
+        dtype=indexes_dt,
+        usm_type=usm_type,
+        sycl_queue=exec_q,
+        order="C",
+    )
+    hev, nz_ev = ti._nonzero(cumsum, indexes, ary.shape, exec_q)
+    res = tuple(indexes[i, :] for i in range(ary.ndim))
+    _manager.add_event_pair(hev, nz_ev)
+    return res
+
+
 def from_numpy(np_ary, /, *, device=None, usm_type="device", sycl_queue=None):
     """
     from_numpy(arg, device=None, usm_type="device", sycl_queue=None)

dpctl_ext/tensor/_indexing_functions.py

@@ -39,6 +39,7 @@
 
 from ._copy_utils import (
     _extract_impl,
+    _nonzero_impl,
 )
 from ._numpy_helper import normalize_axis_index
 
@@ -98,6 +99,33 @@ def extract(condition, arr):
     return _extract_impl(arr, condition)
 
 
+def nonzero(arr):
+    """nonzero(arr)
+
+    Return the indices of non-zero elements.
+
+    Returns a tuple of usm_ndarrays, one for each dimension
+    of ``arr``, containing the indices of the non-zero elements
+    in that dimension. The values of ``arr`` are always tested in
+    row-major, C-style order.
+
+    Args:
+        arr (usm_ndarray):
+            Input array, which has non-zero array rank.
+
+    Returns:
+        Tuple[usm_ndarray, ...]:
+            Indices of non-zero array elements.
+    """
+    if not isinstance(arr, dpt.usm_ndarray):
+        raise TypeError(
+            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(arr)}"
+        )
+    if arr.ndim == 0:
+        raise ValueError("Array of positive rank is expected")
+    return _nonzero_impl(arr)
+
+
 def place(arr, mask, vals):
     """place(arr, mask, vals)
 

dpnp/dpnp_iface_indexing.py

@@ -817,7 +817,7 @@ def extract(condition, a):
         usm_a = dpt_ext.reshape(usm_a, -1)
         usm_cond = dpt_ext.reshape(usm_cond, -1)
 
-        usm_res = dpt_ext.take(usm_a, dpt.nonzero(usm_cond)[0])
+        usm_res = dpt_ext.take(usm_a, dpt_ext.nonzero(usm_cond)[0])
     else:
         if usm_cond.shape != usm_a.shape:
             usm_a = dpt_ext.reshape(usm_a, -1)
@@ -1546,7 +1546,7 @@ def nonzero(a):
 
     usm_a = dpnp.get_usm_ndarray(a)
     return tuple(
-        dpnp_array._create_from_usm_ndarray(y) for y in dpt.nonzero(usm_a)
+        dpnp_array._create_from_usm_ndarray(y) for y in dpt_ext.nonzero(usm_a)
     )