Implement np.cumsum and np.cumprod in kernel by dpnp (#258)

Author: vlad-perevezentsev

numba_dppy/dpnp_glue/dpnp_array_ops_impl.py

@@ -17,12 +17,90 @@
 from numba.core.typing import signature
 from . import stubs
 import numba_dppy.dpnp_glue as dpnp_lowering
-from numba.core.extending import overload
+from numba.core.extending import overload, register_jitable
 import numpy as np
 from numba_dppy import dpctl_functions
 import numba_dppy
 
 
+@register_jitable
+def common_impl(a, out, dpnp_func, print_debug):
+    if a.size == 0:
+        raise ValueError("Passed Empty array")
+
+    sycl_queue = dpctl_functions.get_current_queue()
+    a_usm = dpctl_functions.malloc_shared(a.size * a.itemsize, sycl_queue)
+    dpctl_functions.queue_memcpy(sycl_queue, a_usm, a.ctypes, a.size * a.itemsize)
+
+    out_usm = dpctl_functions.malloc_shared(a.itemsize, sycl_queue)
+
+    dpnp_func(a_usm, out_usm, a.size)
+
+    dpctl_functions.queue_memcpy(
+        sycl_queue, out.ctypes, out_usm, out.size * out.itemsize
+    )
+
+    dpctl_functions.free_with_queue(a_usm, sycl_queue)
+    dpctl_functions.free_with_queue(out_usm, sycl_queue)
+
+    dpnp_ext._dummy_liveness_func([a.size, out.size])
+
+    if print_debug:
+        print("dpnp implementation")
+
+
+@overload(stubs.dpnp.cumsum)
+def dpnp_cumsum_impl(a):
+    name = "cumsum"
+    dpnp_lowering.ensure_dpnp(name)
+
+    res_type = types.void
+    """
+    dpnp source:
+    https://github.com/IntelPython/dpnp/blob/0.5.1/dpnp/backend/kernels/dpnp_krnl_mathematical.cpp#L135
+    Function declaration:
+    void dpnp_cumsum_c(void* array1_in, void* result1, size_t size)
+    """
+    sig = signature(res_type, types.voidptr, types.voidptr, types.intp)
+    dpnp_func = dpnp_ext.dpnp_func("dpnp_" + name, [a.dtype.name, "NONE"], sig)
+
+    PRINT_DEBUG = dpnp_lowering.DEBUG
+
+    def dpnp_impl(a):
+        out = np.arange(a.size, dtype=a.dtype)
+        common_impl(a, out, dpnp_func, PRINT_DEBUG)
+
+        return out
+
+    return dpnp_impl
+
+
+@overload(stubs.dpnp.cumprod)
+def dpnp_cumprod_impl(a):
+    name = "cumprod"
+    dpnp_lowering.ensure_dpnp(name)
+
+    res_type = types.void
+    """
+    dpnp source:
+    https://github.com/IntelPython/dpnp/blob/0.5.1/dpnp/backend/kernels/dpnp_krnl_mathematical.cpp#L110
+    Function declaration:
+    void dpnp_cumprod_c(void* array1_in, void* result1, size_t size)
+    """
+    sig = signature(res_type, types.voidptr, types.voidptr, types.intp)
+    dpnp_func = dpnp_ext.dpnp_func("dpnp_" + name, [a.dtype.name, "NONE"], sig)
+
+    PRINT_DEBUG = dpnp_lowering.DEBUG
+
+    def dpnp_impl(a):
+        out = np.arange(a.size, dtype=a.dtype)
+        common_impl(a, out, dpnp_func, PRINT_DEBUG)
+
+        return out
+
+    return dpnp_impl
+
+
 @overload(stubs.dpnp.sort)
 def dpnp_sort_impl(a):
     name = "sort"

numba_dppy/dpnp_glue/dpnp_fptr_interface.pyx

@@ -30,6 +30,8 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_COS
         DPNP_FN_COSH
         DPNP_FN_COV
+        DPNP_FN_CUMPROD
+        DPNP_FN_CUMSUM
         DPNP_FN_DEGREES
         DPNP_FN_DET
         DPNP_FN_DIVIDE
@@ -201,6 +203,10 @@ cdef DPNPFuncName get_DPNPFuncName_from_str(name):
         return DPNPFuncName.DPNP_FN_DET
     elif name == "dpnp_matrix_rank":
         return DPNPFuncName.DPNP_FN_MATRIX_RANK
+    elif name == "dpnp_cumsum":
+        return DPNPFuncName.DPNP_FN_CUMSUM
+    elif name == "dpnp_cumprod":
+        return DPNPFuncName.DPNP_FN_CUMPROD
     elif name == "dpnp_sort":
         return DPNPFuncName.DPNP_FN_SORT
     else:

numba_dppy/dpnp_glue/stubs.py

@@ -179,5 +179,11 @@ class nansum(Stub):
     class nanprod(Stub):
         pass
 
+    class cumsum(Stub):
+        pass
+
+    class cumprod(Stub):
+        pass
+
     class sort(Stub):
         pass

numba_dppy/rename_numpy_functions_pass.py

@@ -83,6 +83,8 @@
     "prod": (["numpy"], "prod"),
     "sum": (["numpy"], "sum"),
     # array ops
+    "cumsum": (["numpy"], "cumsum"),
+    "cumprod": (["numpy"], "cumprod"),
     "sort": (["numpy"], "sort"),
 }
 

numba_dppy/tests/test_dpnp_functions.py

@@ -1408,6 +1408,38 @@ def f(a, b):
 class Testdpnp_array_ops_functions(unittest.TestCase):
     tys = [np.int32, np.uint32, np.int64, np.uint64, np.float, np.double]
 
+    def test_cumsum(self):
+        @njit
+        def f(a):
+            c = np.cumsum(a)
+            return c
+
+        with assert_dpnp_implementaion():
+            self.assertTrue(
+                check_for_different_datatypes(f, np.cumsum, [10], 1, self.tys, True)
+            )
+            self.assertTrue(check_for_dimensions(f, np.cumsum, [10, 2], self.tys, True))
+            self.assertTrue(
+                check_for_dimensions(f, np.cumsum, [10, 2, 3], self.tys, True)
+            )
+
+    def test_cumprod(self):
+        @njit
+        def f(a):
+            c = np.cumprod(a)
+            return c
+
+        with assert_dpnp_implementaion():
+            self.assertTrue(
+                check_for_different_datatypes(f, np.cumprod, [10], 1, self.tys, True)
+            )
+            self.assertTrue(
+                check_for_dimensions(f, np.cumprod, [10, 2], self.tys, True)
+            )
+            self.assertTrue(
+                check_for_dimensions(f, np.cumprod, [10, 2, 3], self.tys, True)
+            )
+
     def test_sort(self):
         @njit
         def f(a):