Reuse dpctl.tensor implementation for dpnp.nonzero() (#1330)

Author: antonwolfy

dpnp/backend/include/dpnp_iface.hpp

@@ -57,7 +57,6 @@ typedef ssize_t shape_elem_type;
 
 #include <dpctl_sycl_interface.h>
 
-#include "dpnp_iface_fptr.hpp"
 #include "dpnp_iface_fft.hpp"
 #include "dpnp_iface_random.hpp"
 

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -234,7 +234,6 @@ enum class DPNPFuncName : size_t
     DPNP_FN_NEGATIVE,                     /**< Used in numpy.negative() impl  */
     DPNP_FN_NEGATIVE_EXT,                 /**< Used in numpy.negative() impl, requires extra parameters */
     DPNP_FN_NONZERO,                      /**< Used in numpy.nonzero() impl  */
-    DPNP_FN_NONZERO_EXT,                  /**< Used in numpy.nonzero() impl, requires extra parameters */
     DPNP_FN_NOT_EQUAL_EXT,                /**< Used in numpy.not_equal() impl, requires extra parameters */
     DPNP_FN_ONES,                         /**< Used in numpy.ones() impl */
     DPNP_FN_ONES_LIKE,                    /**< Used in numpy.ones_like() impl */

dpnp/backend/kernels/dpnp_krnl_indexing.cpp

@@ -475,6 +475,7 @@ void dpnp_nonzero_c(const void* in_array1,
                                                             j,
                                                             dep_event_vec_ref);
     DPCTLEvent_WaitAndThrow(event_ref);
+    DPCTLEvent_Delete(event_ref);
 }
 
 template <typename _DataType>
@@ -485,16 +486,6 @@ void (*dpnp_nonzero_default_c)(const void*,
                                const size_t,
                                const size_t) = dpnp_nonzero_c<_DataType>;
 
-template <typename _DataType>
-DPCTLSyclEventRef (*dpnp_nonzero_ext_c)(DPCTLSyclQueueRef,
-                                        const void*,
-                                        void*,
-                                        const size_t,
-                                        const shape_elem_type*,
-                                        const size_t,
-                                        const size_t,
-                                        const DPCTLEventVectorRef) = dpnp_nonzero_c<_DataType>;
-
 template <typename _DataType>
 DPCTLSyclEventRef dpnp_place_c(DPCTLSyclQueueRef q_ref,
                                void* arr_in,
@@ -1021,12 +1012,6 @@ void func_map_init_indexing_func(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_nonzero_default_c<float>};
     fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_nonzero_default_c<double>};
 
-    fmap[DPNPFuncName::DPNP_FN_NONZERO_EXT][eft_BLN][eft_BLN] = {eft_BLN, (void*)dpnp_nonzero_ext_c<bool>};
-    fmap[DPNPFuncName::DPNP_FN_NONZERO_EXT][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_nonzero_ext_c<int32_t>};
-    fmap[DPNPFuncName::DPNP_FN_NONZERO_EXT][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_nonzero_ext_c<int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_NONZERO_EXT][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_nonzero_ext_c<float>};
-    fmap[DPNPFuncName::DPNP_FN_NONZERO_EXT][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_nonzero_ext_c<double>};
-
     fmap[DPNPFuncName::DPNP_FN_PLACE][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_place_default_c<int32_t>};
     fmap[DPNPFuncName::DPNP_FN_PLACE][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_place_default_c<int64_t>};
     fmap[DPNPFuncName::DPNP_FN_PLACE][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_place_default_c<float>};

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -211,7 +211,6 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_NEGATIVE
         DPNP_FN_NEGATIVE_EXT
         DPNP_FN_NONZERO
-        DPNP_FN_NONZERO_EXT
         DPNP_FN_NOT_EQUAL_EXT
         DPNP_FN_ONES
         DPNP_FN_ONES_LIKE

dpnp/dpnp_algo/dpnp_algo_indexing.pyx

@@ -1,7 +1,7 @@
 # cython: language_level=3
 # -*- coding: utf-8 -*-
 # *****************************************************************************
-# Copyright (c) 2016-2020, Intel Corporation
+# Copyright (c) 2016-2023, Intel Corporation
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -40,7 +40,6 @@ __all__ += [
     "dpnp_diagonal",
     "dpnp_fill_diagonal",
     "dpnp_indices",
-    "dpnp_nonzero",
     "dpnp_place",
     "dpnp_put",
     "dpnp_put_along_axis",
@@ -104,14 +103,6 @@ ctypedef c_dpctl.DPCTLSyclEventRef(*custom_indexing_6in_func_ptr_t)(c_dpctl.DPCT
                                                                     const size_t,
                                                                     const size_t,
                                                                     const c_dpctl.DPCTLEventVectorRef)
-ctypedef c_dpctl.DPCTLSyclEventRef(*fptr_dpnp_nonzero_t)(c_dpctl.DPCTLSyclQueueRef,
-                                                         const void * ,
-                                                         void * ,
-                                                         const size_t,
-                                                         const shape_elem_type * ,
-                                                         const size_t ,
-                                                         const size_t,
-                                                         const c_dpctl.DPCTLEventVectorRef)
 
 
 cpdef utils.dpnp_descriptor dpnp_choose(utils.dpnp_descriptor x1, list choices1):
@@ -316,60 +307,6 @@ cpdef object dpnp_indices(dimensions):
     return dpnp_result
 
 
-cpdef tuple dpnp_nonzero(utils.dpnp_descriptor in_array1):
-    cdef shape_type_c shape_arr = in_array1.shape
-    res_count = in_array1.ndim
-
-    # have to go through array one extra time to count size of result arrays
-    res_size_obj = dpnp_count_nonzero(in_array1)
-    cdef size_t res_size = dpnp.convert_single_elem_array_to_scalar(res_size_obj.get_pyobj())
-
-    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(in_array1.dtype)
-
-    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_NONZERO_EXT, param1_type, param1_type)
-
-    cdef fptr_dpnp_nonzero_t func = <fptr_dpnp_nonzero_t > kernel_data.ptr
-
-    cdef c_dpctl.SyclQueue q
-    cdef c_dpctl.DPCTLSyclQueueRef q_ref
-    cdef c_dpctl.DPCTLSyclEventRef event_ref
-
-    array1_obj = in_array1.get_array()
-
-    res_list = []
-    cdef utils.dpnp_descriptor res_arr
-    cdef shape_type_c result_shape
-    for j in range(res_count):
-        result_shape = utils._object_to_tuple(res_size)
-        res_arr = utils_py.create_output_descriptor_py(result_shape,
-                                                       dpnp.int64,
-                                                       None,
-                                                       device=array1_obj.sycl_device,
-                                                       usm_type=array1_obj.usm_type,
-                                                       sycl_queue=array1_obj.sycl_queue)
-
-        q = <c_dpctl.SyclQueue> res_arr.get_array().sycl_queue
-        q_ref = q.get_queue_ref()
-
-        event_ref = func(q_ref,
-                         in_array1.get_data(),
-                         res_arr.get_data(),
-                         res_arr.size,
-                         shape_arr.data(),
-                         in_array1.ndim,
-                         j,
-                         NULL)  # dep_events_ref
-
-        with nogil: c_dpctl.DPCTLEvent_WaitAndThrow(event_ref)
-        c_dpctl.DPCTLEvent_Delete(event_ref)
-
-        res_list.append(res_arr.get_pyobj())
-
-    result = utils._object_to_tuple(res_list)
-
-    return result
-
-
 cpdef dpnp_place(dpnp_descriptor arr, object mask, dpnp_descriptor vals):
     result_sycl_device, result_usm_type, result_sycl_queue = utils.get_common_usm_allocation(arr, vals)
 

dpnp/dpnp_array.py

@@ -760,7 +760,10 @@ def ndim(self):
         return self._array_obj.ndim
 
  # 'newbyteorder',
- # 'nonzero',
+
+    def nonzero(self):
+        return dpnp.nonzero(self)
+
  # 'partition',
 
     def prod(self, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=True):

dpnp/dpnp_iface_indexing.py

@@ -2,7 +2,7 @@
 # distutils: language = c++
 # -*- coding: utf-8 -*-
 # *****************************************************************************
-# Copyright (c) 2016-2020, Intel Corporation
+# Copyright (c) 2016-2023, Intel Corporation
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -39,14 +39,14 @@
 
 """
 
-
-import collections
-
 from dpnp.dpnp_algo import *
 from dpnp.dpnp_utils import *
 
 import dpnp
+from dpnp.dpnp_array import dpnp_array
+
 import numpy
+import dpctl.tensor as dpt
 
 
 __all__ = [
@@ -286,15 +286,21 @@ def indices(dimensions, dtype=int, sparse=False):
     return call_origin(numpy.indices, dimensions, dtype, sparse)
 
 
-def nonzero(x1):
+def nonzero(x, /):
     """
     Return the indices of the elements that are non-zero.
 
     For full documentation refer to :obj:`numpy.nonzero`.
 
+    Returns
+    -------
+    y : tuple[dpnp.ndarray]
+        Indices of elements that are non-zero.
+    
     Limitations
     -----------
-    Input array is supported as :obj:`dpnp.ndarray`.
+    Parameters `x` is supported as either :class:`dpnp.ndarray`
+    or :class:`dpctl.tensor.usm_ndarray`.
     Otherwise the function will be executed sequentially on CPU.
     Input array data types are limited by supported DPNP :ref:`Data types`.
 
@@ -329,11 +335,11 @@ def nonzero(x1):
 
     """
 
-    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
-    if x1_desc:
-        return dpnp_nonzero(x1_desc)
+    if isinstance(x, dpnp_array) or isinstance(x, dpt.usm_ndarray):
+        dpt_array = x.get_array() if isinstance(x, dpnp_array) else x
+        return tuple(dpnp_array._create_from_usm_ndarray(y) for y in dpt.nonzero(dpt_array))
 
-    return call_origin(numpy.nonzero, x1)
+    return call_origin(numpy.nonzero, x)
 
 
 def place(x1, mask, vals):

dpnp/dpnp_iface_searching.py

@@ -44,10 +44,7 @@
 from dpnp.dpnp_utils import *
 
 import dpnp
-from dpnp.dpnp_array import dpnp_array
-
 import numpy
-import dpctl.tensor as dpt
 
 
 __all__ = [
@@ -221,13 +218,7 @@ def where(condition, x=None, y=None, /):
     if missing == 1:
         raise ValueError("Must provide both 'x' and 'y' or neither.")
     elif missing == 2:
-        # TODO: rework through dpnp.nonzero() once ready
-        # return dpnp.nonzero(condition)
-        if isinstance(condition, dpnp_array):
-            return dpt.nonzero(condition.get_array())
-
-        if isinstance(condition, dpt.usm_ndarray):
-            return dpt.nonzero(condition)
+        return dpnp.nonzero(condition)
     elif missing == 0:
         # get USM type and queue to copy scalar from the host memory into a USM allocation
         usm_type, queue = get_usm_allocations([condition, x, y])

tests/skipped_tests_gpu.tbl

@@ -52,13 +52,6 @@ tests/test_sycl_queue.py::test_modf[level_zero:gpu:0]
 tests/test_sycl_queue.py::test_1in_1out[opencl:gpu:0-trapz-data19]
 tests/test_sycl_queue.py::test_1in_1out[opencl:cpu:0-trapz-data19]
 
-tests/test_indexing.py::test_nonzero[[[1, 0], [1, 0]]]
-tests/test_indexing.py::test_nonzero[[[1, 2], [3, 4]]]
-tests/test_indexing.py::test_nonzero[[[0, 1, 2], [3, 0, 5], [6, 7, 0]]]
-tests/test_indexing.py::test_nonzero[[[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]]]
-tests/test_indexing.py::test_nonzero[[[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]]]
-tests/test_indexing.py::test_nonzero[[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]]
-
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_no_axis
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPlace_param_3_{n_vals=1, shape=(7,)}::test_place
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPlace_param_4_{n_vals=1, shape=(2, 3)}::test_place