Move choose from cython to backend (#928)

* Move choose from cython to backend

Author: LukichevaPolina

dpnp/backend/include/dpnp_iface.hpp

@@ -600,6 +600,21 @@ INP_DLLEXPORT void dpnp_cov_c(void* array1_in, void* result1, size_t nrows, size
 template <typename _DataType>
 INP_DLLEXPORT void dpnp_det_c(void* array1_in, void* result1, size_t* shape, size_t ndim);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief math library implementation of take function
+ *
+ * @param [out] result        Output array.
+ * @param [in]  array         Input array with data.
+ * @param [in]  choices       Choice arrays.
+ * @param [in]  size          Input array size.
+ * @param [in]  choices_size  Choices size.
+ * @param [in]  choice_size  Choices size.
+ */
+template <typename _DataType1, typename _DataType2>
+INP_DLLEXPORT void
+    dpnp_choose_c(void* result1, void* array1_in, void** choices, size_t size, size_t choices_size, size_t choice_size);
+
 /**
  * @ingroup BACKEND_API
  * @brief math library implementation of det function

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -83,6 +83,7 @@ enum class DPNPFuncName : size_t
     DPNP_FN_CEIL,                     /**< Used in numpy.ceil() implementation  */
     DPNP_FN_CHOLESKY,                 /**< Used in numpy.linalg.cholesky() implementation  */
     DPNP_FN_CONJIGUATE,               /**< Used in numpy.conjugate() implementation  */
+    DPNP_FN_CHOOSE,                   /**< Used in numpy.choose() implementation  */
     DPNP_FN_COPY,                     /**< Used in numpy.copy() implementation  */
     DPNP_FN_COPYSIGN,                 /**< Used in numpy.copysign() implementation  */
     DPNP_FN_COPYTO,                   /**< Used in numpy.copyto() implementation  */

dpnp/backend/kernels/dpnp_krnl_indexing.cpp

@@ -32,6 +32,50 @@
 #include "dpnpc_memory_adapter.hpp"
 #include "queue_sycl.hpp"
 
+template <typename _DataType1, typename _DataType2>
+class dpnp_choose_c_kernel;
+
+template <typename _DataType1, typename _DataType2>
+void dpnp_choose_c(
+    void* result1, void* array1_in, void** choices1, size_t size, size_t choices_size, size_t choice_size)
+{
+    if ((array1_in == nullptr) || (result1 == nullptr) || (choices1 == nullptr))
+    {
+        return;
+    }
+    if (!size || !choices_size || !choice_size)
+    {
+        return;
+    }
+    DPNPC_ptr_adapter<_DataType1> input1_ptr(array1_in, size);
+    _DataType1* array_in = input1_ptr.get_ptr();
+
+    DPNPC_ptr_adapter<_DataType2*> choices_ptr(choices1, choices_size);
+    _DataType2** choices = choices_ptr.get_ptr();
+
+    for (size_t i = 0; i < choices_size; ++i)
+    {
+        DPNPC_ptr_adapter<_DataType2> choice_ptr(choices[i], choice_size);
+        choices[i] = choice_ptr.get_ptr();
+    }
+
+    DPNPC_ptr_adapter<_DataType2> result1_ptr(result1, size, false, true);
+    _DataType2* result = result1_ptr.get_ptr();
+
+    cl::sycl::range<1> gws(size);
+    auto kernel_parallel_for_func = [=](cl::sycl::id<1> global_id) {
+        const size_t idx = global_id[0];
+        result[idx] = choices[array_in[idx]][idx];
+    };
+
+    auto kernel_func = [&](cl::sycl::handler& cgh) {
+        cgh.parallel_for<class dpnp_choose_c_kernel<_DataType1, _DataType2>>(gws, kernel_parallel_for_func);
+    };
+
+    cl::sycl::event event = DPNP_QUEUE.submit(kernel_func);
+    event.wait();
+}
+
 template <typename _DataType>
 class dpnp_diag_indices_c_kernel;
 
@@ -521,6 +565,15 @@ void dpnp_take_c(void* array1_in, const size_t array1_size, void* indices1, void
 
 void func_map_init_indexing_func(func_map_t& fmap)
 {
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_choose_c<int, int>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_INT][eft_LNG] = {eft_LNG, (void*)dpnp_choose_c<int, long>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_INT][eft_FLT] = {eft_FLT, (void*)dpnp_choose_c<int, float>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_INT][eft_DBL] = {eft_DBL, (void*)dpnp_choose_c<int, double>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_LNG][eft_INT] = {eft_INT, (void*)dpnp_choose_c<long, int>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_choose_c<long, long>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_LNG][eft_FLT] = {eft_FLT, (void*)dpnp_choose_c<long, float>};
+    fmap[DPNPFuncName::DPNP_FN_CHOOSE][eft_LNG][eft_DBL] = {eft_DBL, (void*)dpnp_choose_c<long, double>};
+
     fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_diag_indices_c<int>};
     fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_diag_indices_c<long>};
     fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_diag_indices_c<float>};

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -60,6 +60,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_CBRT
         DPNP_FN_CEIL
         DPNP_FN_CHOLESKY
+        DPNP_FN_CHOOSE
         DPNP_FN_CONJIGUATE
         DPNP_FN_COPY
         DPNP_FN_COPYSIGN

dpnp/dpnp_algo/dpnp_algo.pyx

@@ -33,6 +33,7 @@ and the rest of the library
 """
 
 from libc.time cimport time, time_t
+from libcpp.vector cimport vector
 import dpnp
 import dpnp.config as config
 import dpnp.dpnp_utils as utils_py

dpnp/dpnp_algo/dpnp_algo_indexing.pyx

@@ -54,6 +54,7 @@ __all__ += [
     "dpnp_triu_indices_from"
 ]
 
+ctypedef void(*fptr_dpnp_choose_t)(void * , void * , void **, size_t, size_t, size_t)
 ctypedef void(*fptr_dpnp_diag_indices)(void*, size_t)
 ctypedef void(*custom_indexing_2in_1out_func_ptr_t)(void *, const size_t, void * , void * , size_t)
 ctypedef void(*custom_indexing_2in_1out_func_ptr_t_)(void * , const size_t, void * , const size_t, size_t * , size_t * , const size_t)
@@ -65,12 +66,33 @@ ctypedef void(*custom_indexing_6in_func_ptr_t)(void *, void * , void * , const s
 ctypedef void(*fptr_dpnp_nonzero_t)(const void * , void * , const size_t, const size_t * , const size_t , const size_t)
 
 
-cpdef utils.dpnp_descriptor dpnp_choose(object input, list choices):
-    cdef shape_type_c obj_shape = utils._object_to_tuple(len(input))
-    cdef utils.dpnp_descriptor res_array = utils_py.create_output_descriptor_py(obj_shape, choices[0].dtype, None)
+cpdef utils.dpnp_descriptor dpnp_choose(utils.dpnp_descriptor input, list choices1):
+    cdef vector[void * ] choices
+    cdef utils.dpnp_descriptor choice
+    for desc in choices1:
+        choice = desc
+        choices.push_back(choice.get_data())
+
+    cdef shape_type_c input_shape = input.shape
+    cdef size_t choice_size = choices1[0].size
+        
+    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input.dtype)
+
+    cdef DPNPFuncType param2_type = dpnp_dtype_to_DPNPFuncType(choices1[0].dtype)
+
+    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_CHOOSE, param1_type, param2_type)
+
+    cdef utils.dpnp_descriptor res_array = utils.create_output_descriptor(input_shape, kernel_data.return_type, None)
+
+    cdef fptr_dpnp_choose_t func = <fptr_dpnp_choose_t> kernel_data.ptr
+
+    func(res_array.get_data(),
+         input.get_data(),
+         choices.data(),
+         input_shape[0],
+         choices.size(),
+         choice_size)
 
-    for i in range(len(input)):
-        res_array.get_pyobj()[i] = (choices[input[i]])[i]
     return res_array
 
 
@@ -299,7 +321,7 @@ cpdef utils.dpnp_descriptor dpnp_select(list condlist, list choicelist, default)
 
 cpdef utils.dpnp_descriptor dpnp_take(utils.dpnp_descriptor input, utils.dpnp_descriptor indices):
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input.dtype)
-
+    
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_TAKE, param1_type, param1_type)
 
     cdef utils.dpnp_descriptor result = utils.create_output_descriptor(indices.shape, kernel_data.return_type, None)

dpnp/dpnp_iface_indexing.py

@@ -81,35 +81,32 @@ def choose(x1, choices, out=None, mode='raise'):
     --------
     :obj:`take_along_axis` : Preferable if choices is an array.
     """
-    if not use_origin_backend(x1):
-        if not isinstance(x1, list):
-            pass
-        elif not isinstance(choices, list):
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    
+    choices_list =[]
+    for choice in choices:
+        choices_list.append(dpnp.get_dpnp_descriptor(choice))
+
+    if x1_desc:
+        if any(not desc for desc in choices_list):
             pass
         elif out is not None:
             pass
         elif mode != 'raise':
             pass
+        elif any(not choices[0].dtype == choice.dtype for choice in choices):
+            pass
+        elif not len(choices_list):
+            pass
         else:
-            val = True
-            len_ = len(x1)
-            size_ = choices[0].size
-            for i in range(len(choices)):
-                if choices[i].size != size_ or choices[i].size != len_:
-                    val = False
-                    break
-            if not val:
+            size = x1_desc.size
+            choices_size = choices_list[0].size
+            if any(choice.size != choices_size or choice.size != size for choice in choices):
+                pass
+            elif any(x >= choices_size for x in dpnp.asnumpy(x1)):
                 pass
             else:
-                val = True
-                for i in range(len_):
-                    if x1[i] >= size_:
-                        val = False
-                        break
-                if not val:
-                    pass
-                else:
-                    return dpnp_choose(x1, choices).get_pyobj()
+                return dpnp_choose(x1_desc, choices_list).get_pyobj()
 
     return call_origin(numpy.choose, x1, choices, out, mode)