Feature/numpy transcendentals

numba_dppy/dpnp_glue/dpnp_transcendentalsimpl.py

@@ -6,6 +6,33 @@
 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([out.size])
+
+    if print_debug:
+        print("dpnp implementation")
 
 
 @overload(stubs.dpnp.sum)
@@ -25,27 +52,11 @@ def dpnp_sum_impl(a):
     sig = signature(ret_type, types.voidptr, types.voidptr, types.intp)
     dpnp_func = dpnp_ext.dpnp_func("dpnp_" + name, [a.dtype.name, "NONE"], sig)
 
-    def dpnp_impl(a):
-        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)
+    PRINT_DEBUG = dpnp_lowering.DEBUG
 
+    def dpnp_impl(a):
         out = np.empty(1, dtype=a.dtype)
-        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([out.size])
+        common_impl(a, out, dpnp_func, PRINT_DEBUG)
 
         return out[0]
 
@@ -68,28 +79,64 @@ def dpnp_prod_impl(a):
     sig = signature(ret_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):
-        if a.size == 0:
-            raise ValueError("Passed Empty array")
+        out = np.empty(1, dtype=a.dtype)
 
-        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)
+        common_impl(a, out, dpnp_func, PRINT_DEBUG)
+        return out[0]
 
-        out_usm = dpctl_functions.malloc_shared(a.itemsize, sycl_queue)
+    return dpnp_impl
 
-        dpnp_func(a_usm, out_usm, a.size)
 
-        out = np.empty(1, dtype=a.dtype)
-        dpctl_functions.queue_memcpy(
-            sycl_queue, out.ctypes, out_usm, out.size * out.itemsize
-        )
+@overload(stubs.dpnp.nansum)
+def dpnp_nansum_impl(a):
+    name = "nansum"
+    dpnp_lowering.ensure_dpnp(name)
 
-        dpctl_functions.free_with_queue(a_usm, sycl_queue)
-        dpctl_functions.free_with_queue(out_usm, sycl_queue)
+    PRINT_DEBUG = dpnp_lowering.DEBUG
 
-        dpnp_ext._dummy_liveness_func([out.size])
+    def dpnp_impl(a):
+        a_copy = a.copy()
+        a_copy = np.ravel(a_copy)
 
-        return out[0]
+        for i in range(len(a_copy)):
+            if np.isnan(a_copy[i]):
+                a_copy[i] = 0
+
+        result = numba_dppy.dpnp.sum(a_copy)
+        dpnp_ext._dummy_liveness_func([a_copy.size])
+
+        if PRINT_DEBUG:
+            print("dpnp implementation")
+
+        return result
+
+    return dpnp_impl
+
+
+@overload(stubs.dpnp.nanprod)
+def dpnp_nanprod_impl(a):
+    name = "nanprod"
+    dpnp_lowering.ensure_dpnp(name)
+
+    PRINT_DEBUG = dpnp_lowering.DEBUG
+
+    def dpnp_impl(a):
+        a_copy = a.copy()
+        a_copy = np.ravel(a_copy)
+
+        for i in range(len(a_copy)):
+            if np.isnan(a_copy[i]):
+                a_copy[i] = 1
+
+        result = numba_dppy.dpnp.prod(a_copy)
+        dpnp_ext._dummy_liveness_func([a_copy.size])
+
+        if PRINT_DEBUG:
+            print("dpnp implementation")
+
+        return result
 
     return dpnp_impl

numba_dppy/dpnp_glue/stubs.py

@@ -158,3 +158,9 @@ class matrix_rank(Stub):
 
     class eigvals(Stub):
         pass
+
+    class nansum(Stub):
+        pass
+
+    class nanprod(Stub):
+        pass

numba_dppy/ocl/mathimpl.py

@@ -1,5 +1,6 @@
 from __future__ import print_function, absolute_import, division
 import math
+import numpy
 import warnings
 
 from numba.core.imputils import Registry
@@ -20,7 +21,9 @@
 _binary_d_dd = types.float64(types.float64, types.float64)
 
 _binary_f_fi = types.float32(types.float32, types.int32)
+_binary_f_ii = types.float32(types.int32, types.int32)
 _binary_f_fl = types.float32(types.float32, types.int64)
+_binary_f_fd = types.float32(types.float32, types.float64)
 _binary_d_di = types.float64(types.float64, types.int32)
 _binary_d_dl = types.float64(types.float64, types.int64)
 
@@ -31,8 +34,10 @@
     "dd->d": _binary_d_dd,
     "fi->f": _binary_f_fi,
     "fl->f": _binary_f_fl,
+    "ff->f": _binary_f_ff,
     "di->d": _binary_d_di,
     "dl->d": _binary_d_dl,
+    "dd->d": _binary_d_dd,
 }
 
 function_descriptors = {
@@ -69,6 +74,9 @@
     "gamma": (_unary_f_f, _unary_d_d),
     "lgamma": (_unary_f_f, _unary_d_d),
     "ldexp": (_binary_f_fi, _binary_f_fl, _binary_d_di, _binary_d_dl),
+    "hypot": (_binary_f_fi, _binary_f_ff, _binary_d_dl, _binary_d_dd),
+    "exp2": (_unary_f_f, _unary_d_d),
+    "log2": (_unary_f_f, _unary_d_d),
     # unsupported functions listed in the math module documentation:
     # frexp, ldexp, trunc, modf, factorial, fsum
 }
@@ -127,6 +135,9 @@ def core(context, builder, sig, args):
     "lgamma",
     "ldexp",
     "trunc",
+    "hypot",
+    "exp2",
+    "log2",
 ]
 
 
@@ -138,7 +149,10 @@ def function_name_to_supported_decl(name, sig):
         # only symbols present in the math module
         key = getattr(math, name)
     except AttributeError:
-        return None
+        try:
+            key = getattr(numpy, name)
+        except:
+            return None
 
     fn = _mk_fn_decl(name, sig)
     # lower(key, *sig.args)(fn)

numba_dppy/rename_numpy_functions_pass.py

@@ -11,28 +11,19 @@
 
 
 rewrite_function_name_map = {
-    "sum": (["numpy"], "sum"),
-    "eig": (["linalg"], "eig"),
-    "prod": (["numpy"], "prod"),
-    "max": (["numpy"], "max"),
+    # numpy
     "amax": (["numpy"], "amax"),
-    "min": (["numpy"], "min"),
     "amin": (["numpy"], "amin"),
-    "mean": (["numpy"], "mean"),
-    "median": (["numpy"], "median"),
     "argmax": (["numpy"], "argmax"),
     "argmin": (["numpy"], "argmin"),
     "argsort": (["numpy"], "argsort"),
     "cov": (["numpy"], "cov"),
-    "dot": (["numpy"], "dot"),
-    "matmul": (["numpy"], "matmul"),
-    "random_sample": (["random"], "random_sample"),
-    "ranf": (["random"], "ranf"),
-    "sample": (["random"], "sample"),
-    "random": (["random"], "random"),
-    "rand": (["random"], "rand"),
-    "randint": (["random"], "randint"),
-    "random_integers": (["random"], "random_integers"),
+    "max": (["numpy"], "max"),
+    "mean": (["numpy"], "mean"),
+    "median": (["numpy"], "median"),
+    "min": (["numpy"], "min"),
+    "vdot": (["numpy"], "vdot"),
+    # random
     "beta": (["random"], "beta"),
     "binomial": (["random"], "binomial"),
     "chisquare": (["random"], "chisquare"),
@@ -48,20 +39,35 @@
     "negative_binomial": (["random"], "negative_binomial"),
     "normal": (["random"], "normal"),
     "poisson": (["random"], "poisson"),
+    "rand": (["random"], "rand"),
+    "randint": (["random"], "randint"),
+    "random_integers": (["random"], "random_integers"),
+    "random_sample": (["random"], "random_sample"),
+    "random": (["random"], "random"),
+    "ranf": (["random"], "ranf"),
     "rayleigh": (["random"], "rayleigh"),
+    "sample": (["random"], "sample"),
     "standard_cauchy": (["random"], "standard_cauchy"),
     "standard_exponential": (["random"], "standard_exponential"),
     "standard_gamma": (["random"], "standard_gamma"),
     "standard_normal": (["random"], "standard_normal"),
     "uniform": (["random"], "uniform"),
     "weibull": (["random"], "weibull"),
-    "vdot": (["numpy"], "vdot"),
+    # linalg
     "cholesky": (["linalg"], "cholesky"),
     "det": (["linalg"], "det"),
-    "multi_dot": (["linalg"], "multi_dot"),
+    "dot": (["numpy"], "dot"),
+    "eig": (["linalg"], "eig"),
+    "eigvals": (["linalg"], "eigvals"),
+    "matmul": (["numpy"], "matmul"),
     "matrix_power": (["linalg"], "matrix_power"),
     "matrix_rank": (["linalg"], "matrix_rank"),
-    "eigvals": (["linalg"], "eigvals"),
+    "multi_dot": (["linalg"], "multi_dot"),
+    # transcendentals
+    "nanprod": (["numpy"], "nanprod"),
+    "nansum": (["numpy"], "nansum"),
+    "prod": (["numpy"], "prod"),
+    "sum": (["numpy"], "sum"),
 }
 
 

numba_dppy/target.py

@@ -77,43 +77,6 @@ def _init_data_model_manager():
 spirv_data_model_manager = _init_data_model_manager()
 
 
-def _replace_numpy_ufunc_with_opencl_supported_functions():
-    from numba.np.ufunc_db import _ufunc_db as ufunc_db
-    from numba_dppy.ocl.mathimpl import lower_ocl_impl, sig_mapper
-
-    ufuncs = [
-        ("fabs", np.fabs),
-        ("exp", np.exp),
-        ("log", np.log),
-        ("log10", np.log10),
-        ("expm1", np.expm1),
-        ("log1p", np.log1p),
-        ("sqrt", np.sqrt),
-        ("sin", np.sin),
-        ("cos", np.cos),
-        ("tan", np.tan),
-        ("asin", np.arcsin),
-        ("acos", np.arccos),
-        ("atan", np.arctan),
-        ("atan2", np.arctan2),
-        ("sinh", np.sinh),
-        ("cosh", np.cosh),
-        ("tanh", np.tanh),
-        ("asinh", np.arcsinh),
-        ("acosh", np.arccosh),
-        ("atanh", np.arctanh),
-        ("ldexp", np.ldexp),
-        ("floor", np.floor),
-        ("ceil", np.ceil),
-        ("trunc", np.trunc),
-    ]
-
-    for name, ufunc in ufuncs:
-        for sig in ufunc_db[ufunc].keys():
-            if sig in sig_mapper and (name, sig_mapper[sig]) in lower_ocl_impl:
-                ufunc_db[ufunc][sig] = lower_ocl_impl[(name, sig_mapper[sig])]
-
-
 class DPPYTargetContext(BaseContext):
     implement_powi_as_math_call = True
     generic_addrspace = SPIR_GENERIC_ADDRSPACE
@@ -168,6 +131,9 @@ def replace_numpy_ufunc_with_opencl_supported_functions(self):
             ("floor", np.floor),
             ("ceil", np.ceil),
             ("trunc", np.trunc),
+            ("hypot", np.hypot),
+            ("exp2", np.exp2),
+            ("log2", np.log2),
         ]
 
         for name, ufunc in ufuncs: