Merge pull request #1282 from IntelPython/feature/use_dpnp_specific_registry

Feature/use dpnp specific registry

Author: ZzEeKkAa

numba_dpex/core/descriptor.py

@@ -10,7 +10,11 @@
 
 from numba_dpex.core import config
 
-from .targets.dpjit_target import DPEX_TARGET_NAME, DpexTargetContext
+from .targets.dpjit_target import (
+    DPEX_TARGET_NAME,
+    DpexTargetContext,
+    DpexTypingContext,
+)
 from .targets.kernel_target import (
     DPEX_KERNEL_TARGET_NAME,
     CompilationMode,
@@ -110,7 +114,7 @@ def _toplevel_target_context(self):
     @cached_property
     def _toplevel_typing_context(self):
         # Lazily-initialized top-level typing context, for all threads
-        return typing.Context()
+        return DpexTypingContext()
 
     @property
     def target_context(self):

numba_dpex/core/targets/dpjit_target.py

@@ -7,6 +7,7 @@
 
 from functools import cached_property
 
+from numba.core import typing
 from numba.core.compiler_lock import global_compiler_lock
 from numba.core.cpu import CPUContext
 from numba.core.imputils import Registry
@@ -28,6 +29,17 @@ class Dpex(CPU):
 dpex_function_registry = Registry()
 
 
+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):
         super().__init__(typingctx, target)

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]