Various improvements to the DpexTarget context.

    - Implement refresh and a dpex target registry.
    - Initialize the dpjit dispacther and runtime sub-modules
      when dpex loads.
    - doxstrings etc.

Author: Diptorup Deb

numba_dpex/__init__.py

@@ -5,10 +5,14 @@
 """
 The numba-dpex extension module adds data-parallel offload support to Numba.
 """
-import numba.testing
 
+import numba_dpex.core.dpjit_dispatcher
 import numba_dpex.core.offload_dispatcher
 
+# Initialize the _dpexrt_python extension
+import numba_dpex.core.runtime
+import numba_dpex.core.targets.dpjit_target
+
 # Re-export types itself
 import numba_dpex.core.types as types
 from numba_dpex.core.kernel_interface.utils import *

numba_dpex/core/dpjit_dispatcher.py

@@ -12,6 +12,14 @@
 
 
 class DpjitDispatcher(dispatcher.Dispatcher):
+    """A dpex.djit-specific dispatcher.
+
+    The DpjitDispatcher sets the targetdescr string to "dpex" so that Numba's
+    Dispatcher can lookup the global target_registry with that string and
+    correctly use the DpexTarget context.
+
+    """
+
     targetdescr = dpex_target
 
     def __init__(

numba_dpex/core/targets/dpjit_target.py

@@ -5,7 +5,11 @@
 """Defines the target and typing contexts for numba_dpex's dpjit decorator.
 """
 
+from numba.core import utils
+from numba.core.codegen import JITCPUCodegen
+from numba.core.compiler_lock import global_compiler_lock
 from numba.core.cpu import CPUContext
+from numba.core.imputils import Registry, RegistryLoader
 from numba.core.target_extension import CPU, target_registry
 
 
@@ -19,7 +23,38 @@ class Dpex(CPU):
 # permits lookup and reference in user space by the string "dpex"
 target_registry[DPEX_TARGET_NAME] = Dpex
 
+# This is the function registry for the dpu, it just has one registry, this one!
+dpex_function_registry = Registry()
+
 
 class DpexTargetContext(CPUContext):
     def __init__(self, typingctx, target=DPEX_TARGET_NAME):
         super().__init__(typingctx, target)
+
+    @global_compiler_lock
+    def init(self):
+        self.is32bit = utils.MACHINE_BITS == 32
+        self._internal_codegen = JITCPUCodegen("numba.exec")
+        self.lower_extensions = {}
+        # Initialize NRT runtime
+        # rtsys.initialize(self)
+        self.refresh()
+
+    @utils.cached_property
+    def dpexrt(self):
+        from numba_dpex.core.runtime.context import DpexRTContext
+
+        return DpexRTContext(self)
+
+    def refresh(self):
+        registry = dpex_function_registry
+        try:
+            loader = self._registries[registry]
+        except KeyError:
+            loader = RegistryLoader(registry)
+            self._registries[registry] = loader
+        self.install_registry(registry)
+        # Also refresh typing context, since @overload declarations can
+        # affect it.
+        self.typing_context.refresh()
+        super().refresh()

numba_dpex/decorators.py

@@ -166,6 +166,12 @@ def dpjit(*args, **kws):
         del kws["forceobj"]
     kws.update({"nopython": True})
     kws.update({"pipeline_class": OffloadCompiler})
+
+    # FIXME: When trying to use dpex's target context, overloads do not work
+    # properly. We will turn on dpex target once the issue is fixed.
+
+    # kws.update({"_target": "dpex"})
+
     return decorators.jit(*args, **kws)