DPPy target will retain a copy of ufunc_db for replacing the implementation of numpy ufuncs with that of OpenCL's when such function exists

Author: reazulhoque

numba/dppy/target.py

@@ -97,7 +97,30 @@ def init(self):
                                 .SPIR_DATA_LAYOUT[utils.MACHINE_BITS]))
         # Override data model manager to SPIR model
         self.data_model_manager = spirv_data_model_manager
-        self.done_once = False
+
+        from numba.np.ufunc_db import _ufunc_db as ufunc_db, _lazy_init_db
+        import copy
+        _lazy_init_db()
+        self.ufunc_db = copy.deepcopy(ufunc_db)
+
+
+    def replace_numpy_ufunc_with_opencl_supported_functions(self):
+        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 self.ufunc_db[ufunc].keys():
+                if sig in sig_mapper and (name, sig_mapper[sig]) in lower_ocl_impl:
+                    self.ufunc_db[ufunc][sig] = lower_ocl_impl[(name, sig_mapper[sig])]
+
 
     def load_additional_registries(self):
         from .ocl import oclimpl, mathimpl
@@ -111,9 +134,7 @@ def load_additional_registries(self):
             functions we will redirect some of NUMBA's NumPy
             ufunc with OpenCL's.
         """
-        if not self.done_once:
-            _replace_numpy_ufunc_with_opencl_supported_functions()
-            self.done_once = True
+        self.replace_numpy_ufunc_with_opencl_supported_functions()
 
 
     @cached_property

numba/np/npyimpl.py

@@ -412,7 +412,12 @@ def __init__(self, context, builder, outer_sig):
             super(_KernelImpl, self).__init__(context, builder, outer_sig)
             loop = ufunc_find_matching_loop(
                 ufunc, outer_sig.args + (outer_sig.return_type,))
-            self.fn = ufunc_db.get_ufunc_info(ufunc).get(loop.ufunc_sig)
+
+            if hasattr(context, 'ufunc_db'):
+                self.fn = context.ufunc_db[ufunc].get(loop.ufunc_sig)
+            else:
+                self.fn = ufunc_db.get_ufunc_info(ufunc).get(loop.ufunc_sig)
+
             self.inner_sig = typing.signature(
                 *(loop.outputs + loop.inputs))