Migrate the kernel_dispatcher into _kernel_api_impl.spirv

Author: Diptorup Deb

numba_dpex/_kernel_api_impl/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""The module stores the numba_dpex backends implementing the target-specific
+code generation for the kernel_api Python functions.
+"""

numba_dpex/_kernel_api_impl/spirv/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""A SPIR-V backend to compile the numba_dpex.kernel_api functions to SPIR-V.
+"""

numba_dpex/experimental/kernel_dispatcher.py renamed to numba_dpex/_kernel_api_impl/spirv/dispatcher.py

@@ -40,15 +40,17 @@
 )
 from numba_dpex.core.types import USMNdArray
 from numba_dpex.core.utils import kernel_launcher as kl
+from numba_dpex.experimental.target import (
+    DPEX_KERNEL_EXP_TARGET_NAME,
+    dpex_exp_kernel_target,
+)
 
-from .target import DPEX_KERNEL_EXP_TARGET_NAME, dpex_exp_kernel_target
-
-_KernelCompileResult = namedtuple(
+_SPVKernelCompileResult = namedtuple(
     "_KernelCompileResult", CompileResult._fields + ("kernel_device_ir_module",)
 )
 
 
-class _KernelCompiler(_FunctionCompiler):
+class _SPVKernelCompiler(_FunctionCompiler):
     """A special compiler class used to compile numba_dpex.kernel decorated
     functions.
     """
@@ -188,7 +190,7 @@ def _compile_to_spirv(
             kernel_name=kernel_fn.name, kernel_bitcode=kernel_spirv_module
         )
 
-    def compile(self, args, return_type) -> _KernelCompileResult:
+    def compile(self, args, return_type) -> _SPVKernelCompileResult:
         status, kcres = self._compile_cached(args, return_type)
         if status:
             return kcres
@@ -197,7 +199,7 @@ def compile(self, args, return_type) -> _KernelCompileResult:
 
     def _compile_cached(
         self, args, return_type: types.Type
-    ) -> Tuple[bool, _KernelCompileResult]:
+    ) -> Tuple[bool, _SPVKernelCompileResult]:
         """Compiles the kernel function to bitcode and generates a host-callable
         wrapper to submit the kernel to a SYCL queue.
 
@@ -277,10 +279,10 @@ def _compile_cached(
             self._failed_cache[key] = err
             return False, err
 
-        return True, _KernelCompileResult(*kcres_attrs)
+        return True, _SPVKernelCompileResult(*kcres_attrs)
 
 
-class KernelDispatcher(Dispatcher):
+class SPVKernelDispatcher(Dispatcher):
     """Dispatcher class designed to compile kernel decorated functions. The
     dispatcher inherits the Numba Dispatcher class, but has a different
     compilation strategy. Instead of compiling a kernel decorated function to
@@ -325,7 +327,7 @@ def __init__(
                 targetoptions=targetoptions,
                 pipeline_class=pipeline_class,
             )
-        self._compiler = _KernelCompiler(
+        self._compiler = _SPVKernelCompiler(
             pyfunc,
             self.targetdescr,
             targetoptions,
@@ -426,8 +428,8 @@ def cb_llvm(dur):
                     },
                 ):
                     try:
-                        compiler: _KernelCompiler = self._compiler
-                        kcres: _KernelCompileResult = compiler.compile(
+                        compiler: _SPVKernelCompiler = self._compiler
+                        kcres: _SPVKernelCompileResult = compiler.compile(
                             args, return_type
                         )
                     except errors.ForceLiteralArg as err:
@@ -463,4 +465,4 @@ def __call__(self, *args, **kw_args):
 
 
 _dpex_target = target_registry[DPEX_KERNEL_EXP_TARGET_NAME]
-dispatcher_registry[_dpex_target] = KernelDispatcher
+dispatcher_registry[_dpex_target] = SPVKernelDispatcher

numba_dpex/experimental/__init__.py

@@ -8,6 +8,8 @@
 
 from numba.core.imputils import Registry
 
+from numba_dpex._kernel_api_impl.spirv.dispatcher import SPVKernelDispatcher
+
 # Temporary so that Range and NdRange work in experimental call_kernel
 from numba_dpex.core.boxing import *
 
@@ -17,7 +19,6 @@
     _index_space_id_overloads,
 )
 from .decorators import device_func, kernel
-from .kernel_dispatcher import KernelDispatcher
 from .launcher import call_kernel, call_kernel_async
 from .literal_intenum_type import IntEnumLiteral
 from .models import *
@@ -41,5 +42,5 @@ def dpex_dispatcher_const(context):
     "call_kernel",
     "call_kernel_async",
     "IntEnumLiteral",
-    "KernelDispatcher",
+    "SPVKernelDispatcher",
 ]

numba_dpex/experimental/decorators.py

@@ -15,8 +15,8 @@
     target_registry,
 )
 
+from numba_dpex._kernel_api_impl.spirv.dispatcher import SPVKernelDispatcher
 from numba_dpex.core.targets.kernel_target import CompilationMode
-from numba_dpex.experimental.kernel_dispatcher import KernelDispatcher
 
 from .target import DPEX_KERNEL_EXP_TARGET_NAME
 
@@ -78,7 +78,7 @@ def kernel(func_or_sig=None, **options):
             )
 
     def _kernel_dispatcher(pyfunc):
-        disp: KernelDispatcher = dispatcher(
+        disp: SPVKernelDispatcher = dispatcher(
             pyfunc=pyfunc,
             targetoptions=options,
         )

numba_dpex/experimental/launcher.py

@@ -19,6 +19,10 @@
 from numba.extending import intrinsic
 
 from numba_dpex import dpjit
+from numba_dpex._kernel_api_impl.spirv.dispatcher import (
+    SPVKernelDispatcher,
+    _SPVKernelCompileResult,
+)
 from numba_dpex.core.targets.dpjit_target import DPEX_TARGET_NAME
 from numba_dpex.core.targets.kernel_target import DpexKernelTargetContext
 from numba_dpex.core.types import DpctlSyclEvent, NdRangeType, RangeType
@@ -29,10 +33,6 @@
     ItemType,
     NdItemType,
 )
-from numba_dpex.experimental.kernel_dispatcher import (
-    KernelDispatcher,
-    _KernelCompileResult,
-)
 
 
 class LLRange(NamedTuple):
@@ -156,8 +156,8 @@ def _submit_kernel(  # pylint: disable=too-many-arguments
     # ty_kernel_fn is type specific to exact function, so we can get function
     # directly from type and compile it. Thats why we don't need to get it in
     # codegen
-    kernel_dispatcher: KernelDispatcher = ty_kernel_fn.dispatcher
-    kcres: _KernelCompileResult = kernel_dispatcher.get_compile_result(
+    kernel_dispatcher: SPVKernelDispatcher = ty_kernel_fn.dispatcher
+    kcres: _SPVKernelCompileResult = kernel_dispatcher.get_compile_result(
         types.void(*ty_kernel_args_tuple)  # kernel signature
     )
     kernel_module: kl.SPIRVKernelModule = kcres.kernel_device_ir_module