argmin, argmax, refactoring of sorting funcs

dpnp/backend.pxd

@@ -32,6 +32,8 @@ from dpnp.dparray cimport dparray, dparray_shape_type
 cdef extern from "backend/backend_iface_fptr.hpp" namespace "DPNPFuncName":  # need this namespace for Enum import
     cdef enum DPNPFuncName "DPNPFuncName":
         DPNP_FN_ADD
+        DPNP_FN_ARGMAX
+        DPNP_FN_ARGMIN
         DPNP_FN_DOT
 
 cdef extern from "backend/backend_iface_fptr.hpp" namespace "DPNPFuncType":  # need this namespace for Enum import
@@ -122,6 +124,9 @@ cdef extern from "backend/backend_iface.hpp":
     void custom_argsort_c[_DataType, _idx_DataType](void * array, void * result, size_t size)
     void custom_sort_c[_DataType](void * array, void * result, size_t size)
 
+    # Sorting routines
+    void custom_argmax_c[_DataType, _idx_DataType](void * array, void * result, size_t size)
+    void custom_argmin_c[_DataType, _idx_DataType](void * array, void * result, size_t size)
 
 cpdef dparray dpnp_remainder(dparray array1, int scalar)
 cpdef dparray dpnp_astype(dparray array1, dtype_target)
@@ -198,3 +203,9 @@ Sorting functions
 """
 cpdef dparray dpnp_argsort(dparray array1)
 cpdef dparray dpnp_sort(dparray array1)
+
+"""
+Searching functions
+"""
+cpdef dparray dpnp_argmax(dparray array1)
+cpdef dparray dpnp_argmin(dparray array1)

dpnp/backend.pyx

@@ -56,6 +56,7 @@ __all__ = [
 include "backend_logic.pyx"
 include "backend_manipulation.pyx"
 include "backend_mathematical.pyx"
+include "backend_searching.pyx"
 include "backend_sorting.pyx"
 include "backend_statistics.pyx"
 include "backend_trigonometric.pyx"

dpnp/backend/backend_iface.hpp

@@ -198,6 +198,34 @@ INP_DLLEXPORT void custom_sort_c(void* array, void* result, size_t size);
 template <typename _DataType>
 INP_DLLEXPORT void custom_cov_c(void* array1_in, void* result1, const std::vector<long>& input_shape);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief MKL implementation of argmax function
+ *
+ * @param [in]  array   Input array with data.
+ *
+ * @param [out] result  Output array with indeces.
+ *
+ * @param [in]  size    Number of elements in input array.
+ *
+ */
+template <typename _DataType, typename _idx_DataType>
+INP_DLLEXPORT void custom_argmax_c(void* array, void* result, size_t size);
+
+/**
+ * @ingroup BACKEND_API
+ * @brief MKL implementation of argmin function
+ *
+ * @param [in]  array   Input array with data.
+ *
+ * @param [out] result  Output array with indeces.
+ *
+ * @param [in]  size    Number of elements in input array.
+ *
+ */
+template <typename _DataType, typename _idx_DataType>
+INP_DLLEXPORT void custom_argmin_c(void* array, void* result, size_t size);
+
 #if 0 // Example for OpenCL kernel
 template <typename _DataType>
 void custom_dgemm_c_opencl(void* array_1, void* array_2, void* result_1, size_t size);

dpnp/backend/backend_iface_fptr.cpp

@@ -81,6 +81,62 @@ func_map_t func_map =
       }
     }
   },
+  {DPNPFuncName::DPNP_FN_ARGMAX,
+    { // T1. First template parameter
+      {DPNPFuncType::DPNP_FT_INT,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmax_c<int, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmax_c<int, long>}}
+        }
+      },
+      {DPNPFuncType::DPNP_FT_LONG,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmax_c<long, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmax_c<long, long>}}
+        }
+      },
+      {DPNPFuncType::DPNP_FT_FLOAT,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmax_c<float, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmax_c<float, long>}}
+        }
+      },
+      {DPNPFuncType::DPNP_FT_DOUBLE,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmax_c<double, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmax_c<double, long>}}
+        }
+      }
+    }
+  },
+  {DPNPFuncName::DPNP_FN_ARGMIN,
+    { // T1. First template parameter
+      {DPNPFuncType::DPNP_FT_INT,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmin_c<int, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmin_c<int, long>}}
+        }
+      },
+      {DPNPFuncType::DPNP_FT_LONG,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmin_c<long, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmin_c<long, long>}}
+        }
+      },
+      {DPNPFuncType::DPNP_FT_FLOAT,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmin_c<float, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmin_c<float, long>}}
+        }
+      },
+      {DPNPFuncType::DPNP_FT_DOUBLE,
+        { // T2. Second template parameter
+          {DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_argmin_c<double, int>}},
+          {DPNPFuncType::DPNP_FT_LONG, {DPNPFuncType::DPNP_FT_LONG, (void*)custom_argmin_c<double, long>}}
+        }
+      }
+    }
+  },
   {DPNPFuncName::DPNP_FN_DOT,
     {
       {DPNPFuncType::DPNP_FT_INT, {{DPNPFuncType::DPNP_FT_INT, {DPNPFuncType::DPNP_FT_INT, (void*)custom_blas_dot_c<int>}}}},

dpnp/backend/backend_iface_fptr.hpp

@@ -59,10 +59,12 @@
  */
 enum class DPNPFuncName : size_t
 {
-    DPNP_FN_NONE, /**< Very first element of the enumeration */
-    DPNP_FN_ADD,  /**< Used in numpy.add() implementation  */
-    DPNP_FN_DOT,  /**< Used in numpy.dot() implementation  */
-    DPNP_FN_LAST  /**< The latest element of the enumeration */
+    DPNP_FN_NONE,   /**< Very first element of the enumeration */
+    DPNP_FN_ADD,    /**< Used in numpy.add() implementation  */
+    DPNP_FN_ARGMAX, /**< Used in numpy.argmax() implementation  */
+    DPNP_FN_ARGMIN, /**< Used in numpy.argmin() implementation  */
+    DPNP_FN_DOT,    /**< Used in numpy.dot() implementation  */
+    DPNP_FN_LAST    /**< The latest element of the enumeration */
 };
 
 /**

dpnp/backend/custom_kernels.cpp

@@ -120,95 +120,6 @@ void custom_blas_dot_c(void* array1_in, void* array2_in, void* result1, size_t s
 template void custom_blas_dot_c<long>(void* array1_in, void* array2_in, void* result1, size_t size);
 template void custom_blas_dot_c<int>(void* array1_in, void* array2_in, void* result1, size_t size);
 
-template <typename _DataType, typename _idx_DataType>
-struct _argsort_less
-{
-    _argsort_less(_DataType* data_ptr)
-    {
-        _data_ptr = data_ptr;
-    }
-
-    inline bool operator()(const _idx_DataType& idx1, const _idx_DataType& idx2)
-    {
-        return (_data_ptr[idx1] < _data_ptr[idx2]);
-    }
-
-private:
-    _DataType* _data_ptr = nullptr;
-};
-
-template <typename _DataType, typename _idx_DataType>
-class custom_argsort_c_kernel;
-
-template <typename _DataType, typename _idx_DataType>
-void custom_argsort_c(void* array1_in, void* result1, size_t size)
-{
-    _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
-    _idx_DataType* result = reinterpret_cast<_idx_DataType*>(result1);
-
-    for (size_t i = 0; i < size; ++i)
-    {
-        result[i] = i;
-    }
-
-    auto queue = DPNP_QUEUE;
-
-    auto policy =
-        oneapi::dpl::execution::make_device_policy<class custom_argsort_c_kernel<_DataType, _idx_DataType>>(queue);
-
-    std::sort(policy, result, result + size, _argsort_less<_DataType, _idx_DataType>(array_1));
-
-    queue.wait_and_throw(); // looks like it is necessary to sync after call of pstl
-}
-
-template void custom_argsort_c<double, long>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<float, long>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<long, long>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<int, long>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<double, int>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<float, int>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<long, int>(void* array1_in, void* result1, size_t size);
-template void custom_argsort_c<int, int>(void* array1_in, void* result1, size_t size);
-
-template <typename _DataType>
-struct _sort_less
-{
-    inline bool operator()(const _DataType& val1, const _DataType& val2)
-    {
-        return (val1 < val2);
-    }
-};
-
-template <typename _DataType>
-class custom_sort_c_kernel;
-
-template <typename _DataType>
-void custom_sort_c(void* array1_in, void* result1, size_t size)
-{
-    _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
-    _DataType* result = reinterpret_cast<_DataType*>(result1);
-
-    for (size_t i = 0; i < size; ++i)
-    {
-        result[i] = array_1[i];
-    }
-
-    auto queue = DPNP_QUEUE;
-
-    auto policy = oneapi::dpl::execution::make_device_policy<class custom_sort_c_kernel<_DataType>>(queue);
-
-    // fails without explicitly specifying of comparator or with std::less during kernels compilation
-    // affects other kernels
-    std::sort(policy, result, result + size, _sort_less<_DataType>());
-
-    queue.wait_and_throw(); // looks like it is necessary to sync after call of pstl
-}
-
-template void custom_sort_c<double>(void* array1_in, void* result1, size_t size);
-template void custom_sort_c<float>(void* array1_in, void* result1, size_t size);
-template void custom_sort_c<long>(void* array1_in, void* result1, size_t size);
-template void custom_sort_c<int>(void* array1_in, void* result1, size_t size);
-
 #if 0 // Example for OpenCL kernel
 #include <map>
 #include <typeindex>

dpnp/backend/custom_kernels_searching.cpp

@@ -0,0 +1,96 @@
+//*****************************************************************************
+// Copyright (c) 2016-2020, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#include <iostream>
+
+#include <backend/backend_iface.hpp>
+#include "backend_pstl.hpp"
+#include "queue_sycl.hpp"
+
+template <typename _DataType, typename _idx_DataType>
+class custom_argmax_c_kernel;
+
+template <typename _DataType, typename _idx_DataType>
+void custom_argmax_c(void* array1_in, void* result1, size_t size)
+{
+    _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
+    _idx_DataType* result = reinterpret_cast<_idx_DataType*>(result1);
+
+    auto queue = DPNP_QUEUE;
+
+    auto policy =
+        oneapi::dpl::execution::make_device_policy<class custom_argmax_c_kernel<_DataType, _idx_DataType>>(DPNP_QUEUE);
+
+    _DataType* res = std::max_element(policy, array_1, array_1 + size);
+    policy.queue().wait();
+
+    result[0] = std::distance(array_1, res);
+
+#if 0
+    std::cout << "result " << result[0] << "\n";
+#endif
+}
+
+template void custom_argmax_c<double, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<float, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<long, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<int, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<double, int>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<float, int>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<long, int>(void* array1_in, void* result1, size_t size);
+template void custom_argmax_c<int, int>(void* array1_in, void* result1, size_t size);
+
+template <typename _DataType, typename _idx_DataType>
+class custom_argmin_c_kernel;
+
+template <typename _DataType, typename _idx_DataType>
+void custom_argmin_c(void* array1_in, void* result1, size_t size)
+{
+    _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
+    _idx_DataType* result = reinterpret_cast<_idx_DataType*>(result1);
+
+    auto queue = DPNP_QUEUE;
+
+    auto policy =
+        oneapi::dpl::execution::make_device_policy<class custom_argmin_c_kernel<_DataType, _idx_DataType>>(DPNP_QUEUE);
+
+    _DataType* res = std::min_element(policy, array_1, array_1 + size);
+    policy.queue().wait();
+
+    result[0] = std::distance(array_1, res);
+
+#if 0
+    std::cout << "result " << result[0] << "\n";
+#endif
+}
+
+template void custom_argmin_c<double, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<float, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<long, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<int, long>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<double, int>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<float, int>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<long, int>(void* array1_in, void* result1, size_t size);
+template void custom_argmin_c<int, int>(void* array1_in, void* result1, size_t size);