add kernel for cumsum and cumprod functions (#557)

Author: KsanaKozlova

dpnp/backend/include/dpnp_iface.hpp

@@ -169,6 +169,30 @@ INP_DLLEXPORT void dpnp_dot_c(void* array1, void* array2, void* result1, size_t
 template <typename _DataType_input1, typename _DataType_input2, typename _DataType_output>
 INP_DLLEXPORT void dpnp_cross_c(void* array1_in, void* array2_in, void* result1, size_t size);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief Custom implementation of cumprod function
+ *
+ * @param [in]  array1_in  Input array.
+ * @param [out] result1    Output array.
+ * @param [in]  size       Number of elements in input arrays.
+ *
+ */
+template <typename _DataType_input, typename _DataType_output>
+INP_DLLEXPORT void dpnp_cumprod_c(void* array1_in, void* result1, size_t size);
+
+/**
+ * @ingroup BACKEND_API
+ * @brief Custom implementation of cumsum function
+ *
+ * @param [in]  array1_in  Input array.
+ * @param [out] result1    Output array.
+ * @param [in]  size       Number of elements in input arrays.
+ *
+ */
+template <typename _DataType_input, typename _DataType_output>
+INP_DLLEXPORT void dpnp_cumsum_c(void* array1_in, void* result1, size_t size);
+
 /**
  * @ingroup BACKEND_API
  * @brief Sum of array elements

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -84,6 +84,8 @@ enum class DPNPFuncName : size_t
     DPNP_FN_COSH,                     /**< Used in numpy.cosh() implementation  */
     DPNP_FN_COV,                      /**< Used in numpy.cov() implementation  */
     DPNP_FN_CROSS,                    /**< Used in numpy.cross() implementation  */
+    DPNP_FN_CUMPROD,                  /**< Used in numpy.cumprod() implementation  */
+    DPNP_FN_CUMSUM,                   /**< Used in numpy.cumsum() implementation  */
     DPNP_FN_DEGREES,                  /**< Used in numpy.degrees() implementation  */
     DPNP_FN_DET,                      /**< Used in numpy.linalg.det() implementation  */
     DPNP_FN_DIVIDE,                   /**< Used in numpy.divide() implementation  */

dpnp/backend/kernels/dpnp_krnl_mathematical.cpp

@@ -103,6 +103,56 @@ void dpnp_cross_c(void* array1_in, void* array2_in, void* result1, size_t size)
     return;
 }
 
+template <typename _KernelNameSpecialization1, typename _KernelNameSpecialization2>
+class dpnp_cumprod_c_kernel;
+
+template <typename _DataType_input, typename _DataType_output>
+void dpnp_cumprod_c(void* array1_in, void* result1, size_t size)
+{
+    if (!size)
+    {
+        return;
+    }
+
+    _DataType_input* array1 = reinterpret_cast<_DataType_input*>(array1_in);
+    _DataType_output* result = reinterpret_cast<_DataType_output*>(result1);
+
+    _DataType_output cur_res = 1;
+
+    for (size_t i = 0; i < size; ++i)
+    {
+        cur_res *= array1[i];
+        result[i] = cur_res;
+    }
+
+    return;
+}
+
+template <typename _KernelNameSpecialization1, typename _KernelNameSpecialization2>
+class dpnp_cumsum_c_kernel;
+
+template <typename _DataType_input, typename _DataType_output>
+void dpnp_cumsum_c(void* array1_in, void* result1, size_t size)
+{
+    if (!size)
+    {
+        return;
+    }
+
+    _DataType_input* array1 = reinterpret_cast<_DataType_input*>(array1_in);
+    _DataType_output* result = reinterpret_cast<_DataType_output*>(result1);
+
+    _DataType_output cur_res = 0;
+
+    for (size_t i = 0; i < size; ++i)
+    {
+        cur_res += array1[i];
+        result[i] = cur_res;
+    }
+
+    return;
+}
+
 template <typename _KernelNameSpecialization1, typename _KernelNameSpecialization2, typename _KernelNameSpecialization3>
 class dpnp_floor_divide_c_kernel;
 
@@ -251,6 +301,16 @@ void func_map_init_mathematical(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_CROSS][eft_DBL][eft_FLT] = {eft_DBL, (void*)dpnp_cross_c<double, float, double>};
     fmap[DPNPFuncName::DPNP_FN_CROSS][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_cross_c<double, double, double>};
 
+    fmap[DPNPFuncName::DPNP_FN_CUMPROD][eft_INT][eft_INT] = {eft_LNG, (void*)dpnp_cumprod_c<int, long>};
+    fmap[DPNPFuncName::DPNP_FN_CUMPROD][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_cumprod_c<long, long>};
+    fmap[DPNPFuncName::DPNP_FN_CUMPROD][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_cumprod_c<float, float>};
+    fmap[DPNPFuncName::DPNP_FN_CUMPROD][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_cumprod_c<double, double>};
+
+    fmap[DPNPFuncName::DPNP_FN_CUMSUM][eft_INT][eft_INT] = {eft_LNG, (void*)dpnp_cumsum_c<int, long>};
+    fmap[DPNPFuncName::DPNP_FN_CUMSUM][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_cumsum_c<long, long>};
+    fmap[DPNPFuncName::DPNP_FN_CUMSUM][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_cumsum_c<float, float>};
+    fmap[DPNPFuncName::DPNP_FN_CUMSUM][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_cumsum_c<double, double>};
+
     fmap[DPNPFuncName::DPNP_FN_FLOOR_DIVIDE][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_floor_divide_c<int, int, int>};
     fmap[DPNPFuncName::DPNP_FN_FLOOR_DIVIDE][eft_INT][eft_LNG] = {eft_LNG, (void*)dpnp_floor_divide_c<int, long, long>};
     fmap[DPNPFuncName::DPNP_FN_FLOOR_DIVIDE][eft_INT][eft_FLT] = {eft_DBL,

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -57,6 +57,8 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_COSH
         DPNP_FN_COV
         DPNP_FN_CROSS
+        DPNP_FN_CUMPROD
+        DPNP_FN_CUMSUM
         DPNP_FN_DEGREES
         DPNP_FN_DET
         DPNP_FN_DIVIDE

dpnp/dpnp_algo/dpnp_algo_mathematical.pyx

@@ -63,6 +63,8 @@ __all__ += [
     "dpnp_minimum",
     "dpnp_modf",
     "dpnp_multiply",
+    "dpnp_nancumprod",
+    "dpnp_nancumsum",
     "dpnp_nanprod",
     "dpnp_nansum",
     "dpnp_negative",
@@ -139,54 +141,28 @@ cpdef dparray dpnp_cross(dparray x1, dparray x2):
     return call_fptr_2in_1out(DPNP_FN_CROSS, x1, x2, x1.shape)
 
 
-cpdef dparray dpnp_cumprod(dparray x1, bint usenan=False):
+cpdef dparray dpnp_cumprod(dparray x1):
+    #instead of x1.shape, (x1.size, ) is passed to the function
+    # due to the following:
+    # >>> import numpy
+    # >>> a = numpy.array([[1, 2], [2, 3]])
+    # >>> res = numpy.cumprod(a)
+    # >>> res.shape
+    # (4,)
 
-    types_map = {
-        dpnp.int32: dpnp.int64,
-        dpnp.int64: dpnp.int64,
-        dpnp.float32: dpnp.float32,
-        dpnp.float64: dpnp.float64
-    }
-
-    res_type = types_map[x1.dtype.type]
-
-    cdef dparray result = dparray(x1.size, dtype=res_type)
-
-    cur_res = 1
-
-    for i in range(result.size):
-
-        if not usenan or not dpnp.isnan(x1[i]):
-            cur_res *= x1[i]
-
-        result._setitem_scalar(i, cur_res)
-
-    return result
-
-
-cpdef dparray dpnp_cumsum(dparray x1, bint usenan=False):
-
-    types_map = {
-        dpnp.int32: dpnp.int64,
-        dpnp.int64: dpnp.int64,
-        dpnp.float32: dpnp.float32,
-        dpnp.float64: dpnp.float64
-    }
-
-    res_type = types_map[x1.dtype.type]
+    return call_fptr_1in_1out(DPNP_FN_CUMPROD, x1, (x1.size,))
 
-    cdef dparray result = dparray(x1.size, dtype=res_type)
 
-    cur_res = 0
+cpdef dparray dpnp_cumsum(dparray x1):
+    #instead of x1.shape, (x1.size, ) is passed to the function
+    # due to the following:
+    # >>> import numpy
+    # >>> a = numpy.array([[1, 2], [2, 3]])
+    # >>> res = numpy.cumsum(a)
+    # >>> res.shape
+    # (4,)
 
-    for i in range(result.size):
-
-        if not usenan or not dpnp.isnan(x1[i]):
-            cur_res += x1[i]
-
-        result._setitem_scalar(i, cur_res)
-
-    return result
+    return call_fptr_1in_1out(DPNP_FN_CUMSUM, x1, (x1.size,))
 
 
 cpdef dparray dpnp_diff(dparray input, int n):
@@ -344,6 +320,28 @@ cpdef dparray dpnp_multiply(dparray x1, x2):
         return call_fptr_2in_1out(DPNP_FN_MULTIPLY, x1, x2, x1.shape)
 
 
+cpdef dparray dpnp_nancumprod(dparray x1):
+
+    cur_x1 = dpnp.copy(x1)
+
+    for i in range(cur_x1.size):
+        if dpnp.isnan(cur_x1[i]):
+            cur_x1._setitem_scalar(i, 1)
+
+    return dpnp_cumprod(cur_x1)
+
+
+cpdef dparray dpnp_nancumsum(dparray x1):
+
+    cur_x1 = dpnp.copy(x1)
+
+    for i in range(cur_x1.size):
+        if dpnp.isnan(cur_x1[i]):
+            cur_x1._setitem_scalar(i, 0)
+
+    return dpnp_cumsum(cur_x1)
+
+
 cpdef dpnp_nanprod(dparray x1):
     cdef dparray result = dparray(x1.shape, dtype=x1.dtype)
 

dpnp/dpnp_iface_mathematical.py

@@ -1065,7 +1065,7 @@ def nancumprod(x1, **kwargs):
         if not isinstance(x1, dparray):
             pass
         else:
-            return dpnp_cumprod(x1, usenan=True)
+            return dpnp_nancumprod(x1)
 
     return call_origin(numpy.nancumprod, x1, **kwargs)
 
@@ -1103,7 +1103,7 @@ def nancumsum(x1, **kwargs):
         if not isinstance(x1, dparray):
             pass
         else:
-            return dpnp_cumsum(x1, usenan=True)
+            return dpnp_nancumsum(x1)
 
     return call_origin(numpy.nancumsum, x1, **kwargs)
 

tests/test_mathematical.py

@@ -62,8 +62,8 @@ def test_nancumprod(array):
     a = numpy.array(array)
     ia = inp.array(a)
 
-    result = inp.nancumsum(ia)
-    expected = numpy.nancumsum(a)
+    result = inp.nancumprod(ia)
+    expected = numpy.nancumprod(a)
     numpy.testing.assert_array_equal(expected, result)