Add dpnp ufunc specific registry

Author: ZzEeKkAa

numba_dpex/core/targets/dpjit_target.py

@@ -32,6 +32,13 @@ class Dpex(CPU):
 class DpexTypingContext(typing.Context):
     """Custom typing context to support dpjit compilation."""
 
+    def load_additional_registries(self):
+        """Register dpjit specific functions like dpnp ufuncs."""
+        from numba_dpex.core.typing import dpnpdecl
+
+        self.install_registry(dpnpdecl.registry)
+        super().load_additional_registries()
+
 
 class DpexTargetContext(CPUContext):
     def __init__(self, typingctx, target=DPEX_TARGET_NAME):

numba_dpex/core/targets/kernel_target.py

@@ -111,12 +111,16 @@ def load_additional_registries(self):
         """Register the OpenCL API and math and other functions."""
         from numba.core.typing import cmathdecl, enumdecl, npydecl
 
+        from numba_dpex.core.typing import dpnpdecl
+
         from ...ocl import mathdecl, ocldecl
 
         self.install_registry(ocldecl.registry)
         self.install_registry(mathdecl.registry)
         self.install_registry(cmathdecl.registry)
+        # TODO: https://github.com/IntelPython/numba-dpex/issues/1270
         self.install_registry(npydecl.registry)
+        self.install_registry(dpnpdecl.registry)
         self.install_registry(enumdecl.registry)
 
 

numba_dpex/core/typing/dpnpdecl.py

@@ -10,8 +10,24 @@
     NumpyRulesArrayOperator,
     NumpyRulesInplaceArrayOperator,
     NumpyRulesUnaryArrayOperator,
-    infer_global,
 )
+from numba.core.typing.templates import Registry
+
+registry = Registry()
+infer = registry.register
+infer_global = registry.register_global
+infer_getattr = registry.register_attr
+
+
+def _install_operations(cls: NumpyRulesArrayOperator):
+    for op, ufunc_name in cls._op_map.items():
+        infer_global(op)(
+            type(
+                "NumpyRulesArrayOperator_" + ufunc_name,
+                (cls,),
+                dict(key=op),
+            )
+        )
 
 
 class DpnpRulesArrayOperator(NumpyRulesArrayOperator):
@@ -29,13 +45,21 @@ def ufunc(self):
         except:
             pass
 
+    @classmethod
+    def install_operations(cls):
+        _install_operations(cls)
+
 
 class DpnpRulesInplaceArrayOperator(NumpyRulesInplaceArrayOperator):
-    pass
+    @classmethod
+    def install_operations(cls):
+        _install_operations(cls)
 
 
 class DpnpRulesUnaryArrayOperator(NumpyRulesUnaryArrayOperator):
-    pass
+    @classmethod
+    def install_operations(cls):
+        _install_operations(cls)
 
 
 # list of unary ufuncs to register

numba_dpex/dpnp_iface/dpnpimpl.py

@@ -8,7 +8,7 @@
 from numba.core.imputils import Registry
 from numba.np import npyimpl
 
-from numba_dpex.core.typing.dpnpdecl import _unsupported
+from numba_dpex.core.typing import dpnpdecl
 from numba_dpex.dpnp_iface import dpnp_ufunc_db
 
 registry = Registry("dpnpimpl")
@@ -36,11 +36,11 @@ def _register_dpnp_ufuncs():
         )
 
     for _op_map in (
-        npyimpl.npydecl.NumpyRulesUnaryArrayOperator._op_map,
-        npyimpl.npydecl.NumpyRulesArrayOperator._op_map,
+        dpnpdecl.DpnpRulesUnaryArrayOperator._op_map,
+        dpnpdecl.DpnpRulesArrayOperator._op_map,
     ):
         for operator, ufunc_name in _op_map.items():
-            if ufunc_name in _unsupported:
+            if ufunc_name in dpnpdecl._unsupported:
                 continue
             ufunc = getattr(dpnp, ufunc_name)
             kernel = kernels[ufunc]
@@ -57,9 +57,9 @@ def _register_dpnp_ufuncs():
                     "There shouldn't be any non-unary or binary operators"
                 )
 
-    for _op_map in (npyimpl.npydecl.NumpyRulesInplaceArrayOperator._op_map,):
+    for _op_map in (dpnpdecl.DpnpRulesInplaceArrayOperator._op_map,):
         for operator, ufunc_name in _op_map.items():
-            if ufunc_name in _unsupported:
+            if ufunc_name in dpnpdecl._unsupported:
                 continue
             ufunc = getattr(dpnp, ufunc_name)
             kernel = kernels[ufunc]