Merge pull request #1213 from IntelPython/feature/experimental_target

Add a target inside the experimental module

Author: Diptorup Deb

numba_dpex/experimental/kernel_dispatcher.py

@@ -13,17 +13,19 @@
 from numba.core.compiler import CompileResult
 from numba.core.compiler_lock import global_compiler_lock
 from numba.core.dispatcher import Dispatcher, _FunctionCompiler
+from numba.core.target_extension import dispatcher_registry, target_registry
 from numba.core.typing.typeof import Purpose, typeof
 
 from numba_dpex import config, spirv_generator
-from numba_dpex.core.descriptor import dpex_kernel_target
 from numba_dpex.core.exceptions import (
     ExecutionQueueInferenceError,
     UnsupportedKernelArgumentError,
 )
 from numba_dpex.core.pipelines import kernel_compiler
 from numba_dpex.core.types import DpnpNdArray
 
+from .target import dpex_exp_kernel_target
+
 _KernelModule = namedtuple("_KernelModule", ["kernel_name", "kernel_bitcode"])
 
 _KernelCompileResult = namedtuple(
@@ -177,7 +179,7 @@ class KernelDispatcher(Dispatcher):
 
     """
 
-    targetdescr = dpex_kernel_target
+    targetdescr = dpex_exp_kernel_target
     _fold_args = False
 
     Dispatcher._impl_kinds["kernel"] = _KernelCompiler
@@ -314,3 +316,7 @@ def __call__(self, *args, **kw_args):
         """Functor to launch a kernel."""
 
         raise NotImplementedError
+
+
+_dpex_target = target_registry["dpex_kernel"]
+dispatcher_registry[_dpex_target] = KernelDispatcher

numba_dpex/experimental/target.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""A new target descriptor that includes experimental features that should
+eventually move into the numba_dpex.core.
+"""
+
+from functools import cached_property
+
+from numba.core.descriptors import TargetDescriptor
+
+from numba_dpex.core.descriptor import DpexTargetOptions
+from numba_dpex.core.targets.kernel_target import (
+    DPEX_KERNEL_TARGET_NAME,
+    DpexKernelTargetContext,
+    DpexKernelTypingContext,
+)
+
+
+class DpexExpKernelTypingContext(DpexKernelTypingContext):
+    """Experimental typing context class extending the DpexKernelTypingContext
+    by overriding super class functions for new experimental types.
+
+    A new experimental type may require updating type inference for that type
+    when it is used as an argument, value or attribute in a JIT compiled
+    function. All such experimental functionality should be added here till they
+    are stable enough to be migrated to DpexKernelTypingContext.
+    """
+
+
+#  pylint: disable=W0223
+# FIXME: Remove the pylint disablement once we add an override for
+# get_executable
+class DpexExpKernelTargetContext(DpexKernelTargetContext):
+    """Experimental target context class extending the DpexKernelTargetContext
+    by overriding super class functions for new experimental types.
+
+    A new experimental type may require specific ways for handling the lowering
+    to LLVM IR. All such experimental functionality should be added here till
+    they are stable enough to be migrated to DpexKernelTargetContext.
+    """
+
+
+class DpexExpKernelTarget(TargetDescriptor):
+    """
+    Implements a target descriptor for numba_dpex.kernel decorated functions.
+    """
+
+    options = DpexTargetOptions
+
+    @cached_property
+    def _toplevel_target_context(self):
+        """Lazily-initialized top-level target context, for all threads."""
+        return DpexExpKernelTargetContext(
+            self.typing_context, self._target_name
+        )
+
+    @cached_property
+    def _toplevel_typing_context(self):
+        """Lazily-initialized top-level typing context, for all threads."""
+        return DpexExpKernelTypingContext()
+
+    @property
+    def target_context(self):
+        """
+        The target context used by the Dpex compiler pipeline.
+        """
+        return self._toplevel_target_context
+
+    @property
+    def typing_context(self):
+        """
+        The typing context for used by the Dpex compiler pipeline.
+        """
+        return self._toplevel_typing_context
+
+
+# A global instance of the DpexKernelTarget with the experimental features
+dpex_exp_kernel_target = DpexExpKernelTarget(DPEX_KERNEL_TARGET_NAME)