Skip to content

Commit 4a347d2

Browse files
RubtsowaRubtsowa
authored
add native version nonzero (#649)
1 parent 7c6e20d commit 4a347d2

File tree

6 files changed

+105
-10
lines changed

6 files changed

+105
-10
lines changed

‎dpnp/backend/include/dpnp_iface.hpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,20 @@ template <typename _DataType>
189189
INP_DLLEXPORT void
190190
dpnp_matmul_c(void* array1, void* array2, void* result1, size_t size_m, size_t size_n, size_t size_k);
191191

192+
/**
193+
* @ingroup BACKEND_API
194+
* @brief Return the indices of the elements that are non-zero.
195+
*
196+
* @param [in] array1 Input array.
197+
* @param [out] result1 Output array.
198+
* @param [in] shape Shape of input array.
199+
* @param [in] ndim Number of elements in shape.
200+
* @param [in] j Number input array.
201+
*/
202+
template <typename _DataType>
203+
INP_DLLEXPORT void
204+
dpnp_nonzero_c(const void* array1, void* result1, const size_t* shape, const size_t ndim, const size_t j);
205+
192206
/**
193207
* @ingroup BACKEND_API
194208
* @brief absolute function.

‎dpnp/backend/include/dpnp_iface_fptr.hpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ enum class DPNPFuncName : size_t
132132
DPNP_FN_MINIMUM, /**< Used in numpy.minimum() implementation */
133133
DPNP_FN_MODF, /**< Used in numpy.modf() implementation */
134134
DPNP_FN_MULTIPLY, /**< Used in numpy.multiply() implementation */
135+
DPNP_FN_NONZERO, /**< Used in numpy.nonzero() implementation */
135136
DPNP_FN_ONES, /**< Used in numpy.ones() implementation */
136137
DPNP_FN_ONES_LIKE, /**< Used in numpy.ones_like() implementation */
137138
DPNP_FN_PLACE, /**< Used in numpy.place() implementation */

‎dpnp/backend/kernels/dpnp_krnl_indexing.cpp

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,50 @@ void dpnp_fill_diagonal_c(void* array1_in, void* val_in, size_t* shape, const si
175175
return;
176176
}
177177

178+
template <typename _DataType>
179+
void dpnp_nonzero_c(const void* in_array1, void* result1, const size_t* shape, const size_t ndim, const size_t j)
180+
{
181+
if ((in_array1 == nullptr) || (result1 == nullptr))
182+
{
183+
return;
184+
}
185+
186+
if (ndim == 0)
187+
{
188+
return;
189+
}
190+
191+
const _DataType* arr = reinterpret_cast<const _DataType*>(in_array1);
192+
long* result = reinterpret_cast<long*>(result1);
193+
194+
size_t size = 1;
195+
for (size_t i = 0; i < ndim; ++i)
196+
{
197+
size *= shape[i];
198+
}
199+
200+
size_t idx = 0;
201+
for (size_t i = 0; i < size; ++i)
202+
{
203+
if (arr[i] != 0)
204+
{
205+
size_t ids[ndim];
206+
size_t ind1 = size;
207+
size_t ind2 = i;
208+
for (size_t k = 0; k < ndim; ++k)
209+
{
210+
ind1 = ind1 / shape[k];
211+
ids[k] = ind2 / ind1;
212+
ind2 = ind2 % ind1;
213+
}
214+
215+
result[idx] = ids[j];
216+
idx += 1;
217+
}
218+
}
219+
return;
220+
}
221+
178222
template <typename _DataType>
179223
void dpnp_place_c(void* arr_in, long* mask_in, void* vals_in, const size_t arr_size, const size_t vals_size)
180224
{
@@ -441,6 +485,11 @@ void func_map_init_indexing_func(func_map_t& fmap)
441485
fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_fill_diagonal_c<float>};
442486
fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_fill_diagonal_c<double>};
443487

488+
fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_nonzero_c<int>};
489+
fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_nonzero_c<long>};
490+
fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_nonzero_c<float>};
491+
fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_nonzero_c<double>};
492+
444493
fmap[DPNPFuncName::DPNP_FN_PLACE][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_place_c<int>};
445494
fmap[DPNPFuncName::DPNP_FN_PLACE][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_place_c<long>};
446495
fmap[DPNPFuncName::DPNP_FN_PLACE][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_place_c<float>};

‎dpnp/dpnp_algo/dpnp_algo.pxd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName": # need this na
105105
DPNP_FN_MINIMUM
106106
DPNP_FN_MODF
107107
DPNP_FN_MULTIPLY
108+
DPNP_FN_NONZERO
108109
DPNP_FN_ONES
109110
DPNP_FN_ONES_LIKE
110111
DPNP_FN_PLACE

‎dpnp/dpnp_algo/dpnp_algo_indexing.pyx

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ ctypedef void(*custom_indexing_3in_func_ptr_t)(void * , void * , void * , const
6666
ctypedef void(*custom_indexing_3in_with_axis_func_ptr_t)(void * , void * , void * , const size_t, size_t * , const size_t,
6767
const size_t, const size_t,)
6868
ctypedef void(*custom_indexing_6in_func_ptr_t)(void *, void * , void * , const size_t, const size_t, const size_t)
69+
ctypedef void(*fptr_dpnp_nonzero_t)(const void * , void * , const size_t * , const size_t , const size_t)
6970

7071

7172
cpdef dparray dpnp_choose(input, choices):
@@ -170,18 +171,21 @@ cpdef tuple dpnp_nonzero(dparray in_array1):
170171
# have to go through array one extra time to count size of result arrays
171172
res_size = count_nonzero(in_array1)
172173

174+
cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(in_array1.dtype)
175+
176+
cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_NONZERO, param1_type, param1_type)
177+
178+
cdef fptr_dpnp_nonzero_t func = <fptr_dpnp_nonzero_t > kernel_data.ptr
179+
173180
res_list = []
174-
for i in range(res_count):
175-
res_list.append(dparray((res_size, ), dtype=dpnp.int64))
176-
result = _object_to_tuple(res_list)
181+
for j in range(res_count):
182+
res_arr = dparray((res_size, ), dtype=dpnp.int64)
177183

178-
idx = 0
179-
for i in range(in_array1.size):
180-
if in_array1[i] != 0:
181-
ids = get_axis_indeces(i, in_array1.shape)
182-
for j in range(res_count):
183-
result[j][idx] = ids[j]
184-
idx = idx + 1
184+
func(in_array1.get_data(), res_arr.get_data(), < size_t * > in_array1._dparray_shape.data(), in_array1.ndim, j)
185+
186+
res_list.append(res_arr)
187+
188+
result = _object_to_tuple(res_list)
185189

186190
return result
187191

‎tests/test_indexing.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,32 @@ def test_indices(dimension):
8383
numpy.testing.assert_array_equal(expected, result)
8484

8585

86+
@pytest.mark.parametrize("array",
87+
[[],
88+
[[0, 0], [0, 0]],
89+
[[1, 0], [1, 0]],
90+
[[1, 2], [3, 4]],
91+
[[0, 1, 2], [3, 0, 5], [6, 7, 0]],
92+
[[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]],
93+
[[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]],
94+
[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [
95+
[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]],
96+
ids=['[]',
97+
'[[0, 0], [0, 0]]',
98+
'[[1, 0], [1, 0]]',
99+
'[[1, 2], [3, 4]]',
100+
'[[0, 1, 2], [3, 0, 5], [6, 7, 0]]',
101+
'[[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]]',
102+
'[[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]]',
103+
'[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]'])
104+
def test_nonzero(array):
105+
a = numpy.array(array)
106+
ia = dpnp.array(array)
107+
expected = numpy.nonzero(a)
108+
result = dpnp.nonzero(ia)
109+
numpy.testing.assert_array_equal(expected, result)
110+
111+
86112
@pytest.mark.parametrize("vals",
87113
[[100, 200],
88114
(100, 200)],

0 commit comments

Comments
 (0)