Adds a compilation mode target option and device_func decorator.

    - A new target option was added for the DpexKernelTarget
      target to compile functions using the experimental
      KernelDispatcher differently based on whether they are
      "kernels" or "device functions".

      kernels have the spir_kernel calling convention, cannot
      return a value, enforce execution queue equivalence, and
      are always compiled down to device IR (SPIR-V).

      device functions have the spir_func calling convention, do
      not have the same restrictions on return value and input
      arguments and are only compiled to LLVM bitcode.

    - A device_func decorator was added to experimental module.
      The new decorator is roughly equivalent to numba_dpex.func
      but uses the new KernelDispatcher and the compilation mode
      of device function.

      The `device_func` decorator is registered to compile
      overloads in DpexExpkernelTarget.

    - In the kernel compilation mode the final LLVM module is now
      "finaliozed" before conversion to SPIR-V. During finalization
      all overload calls are linked into the main (kernel) module
      and optionally inlined.

Author: Diptorup Deb

numba_dpex/core/descriptor.py

@@ -11,6 +11,7 @@
 from .targets.dpjit_target import DPEX_TARGET_NAME, DpexTargetContext
 from .targets.kernel_target import (
     DPEX_KERNEL_TARGET_NAME,
+    CompilationMode,
     DpexKernelTargetContext,
     DpexKernelTypingContext,
 )
@@ -40,13 +41,17 @@ class DpexTargetOptions(CPUTargetOptions):
     release_gil = _option_mapping("release_gil")
     no_compile = _option_mapping("no_compile")
     use_mlir = _option_mapping("use_mlir")
+    _compilation_mode = _option_mapping("_compilation_mode")
 
     def finalize(self, flags, options):
         super().finalize(flags, options)
         _inherit_if_not_set(flags, options, "experimental", False)
         _inherit_if_not_set(flags, options, "release_gil", False)
         _inherit_if_not_set(flags, options, "no_compile", True)
         _inherit_if_not_set(flags, options, "use_mlir", False)
+        _inherit_if_not_set(
+            flags, options, "_compilation_mode", CompilationMode.KERNEL
+        )
 
 
 class DpexKernelTarget(TargetDescriptor):

numba_dpex/core/targets/kernel_target.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+from enum import IntEnum
 from functools import cached_property
 
 import dpnp
@@ -30,6 +31,28 @@
 LLVM_SPIRV_ARGS = 112
 
 
+class CompilationMode(IntEnum):
+    """Flags used to determine how a function should be compiled by the
+    numba_dpex.experimental.dispatcher.KernelDispatcher. Note the functionality
+    will be merged into numba_dpex.core.kernel_interface.dispatcher in the
+    future.
+
+        KERNEL :         Indicates that the function will be compiled into an
+                         LLVM function that has ``spir_kernel`` calling
+                         convention and is compiled down to SPIR-V.
+                         Additionally, the function cannot return any value and
+                         input arguments to the function have to adhere to
+                         "compute follows data" to ensure execution queue
+                         inference.
+        DEVICE_FUNCTION: Indicates that the function will be compiled into an
+                         LLVM function that has ``spir_func`` calling convention
+                         and will be compiled only into LLVM bitcode.
+    """
+
+    KERNEL = 1
+    DEVICE_FUNC = 2
+
+
 class DpexKernelTypingContext(typing.BaseContext):
     """Custom typing context to support kernel compilation.
 

numba_dpex/experimental/decorators.py

@@ -6,6 +6,7 @@
 ready to move to numba_dpex.core.
 """
 import inspect
+from warnings import warn
 
 from numba.core import sigutils
 from numba.core.target_extension import (
@@ -14,6 +15,8 @@
     target_registry,
 )
 
+from numba_dpex.core.targets.kernel_target import CompilationMode
+
 from .target import DPEX_KERNEL_EXP_TARGET_NAME
 
 
@@ -30,6 +33,14 @@ def kernel(func_or_sig=None, **options):
     """
 
     dispatcher = resolve_dispatcher_from_str(DPEX_KERNEL_EXP_TARGET_NAME)
+    if "_compilation_mode" in options:
+        user_compilation_mode = options["_compilation_mode"]
+        warn(
+            "_compilation_mode is an internal flag that should not be set "
+            "in the decorator. The decorator defined option "
+            f"{user_compilation_mode} is going to be ignored."
+        )
+    options["_compilation_mode"] = CompilationMode.KERNEL
 
     # FIXME: The options need to be evaluated and checked here like it is
     # done in numba.core.decorators.jit
@@ -80,4 +91,44 @@ def _specialized_kernel_dispatcher(pyfunc):
     return _kernel_dispatcher(func)
 
 
-jit_registry[target_registry[DPEX_KERNEL_EXP_TARGET_NAME]] = kernel
+def device_func(func_or_sig=None, **options):
+    """Generates a function with a device-only calling convention, e.g.,
+    spir_func for SPIR-V based devices.
+
+    The decorator is used to compile overloads in the DpexKernelTarget and
+    users should use the decorator to define functions that are only callable
+    from inside another device_func or a kernel.
+
+    A device_func is not compiled down to device binary IR and instead left as
+    LLVM IR. It is done so that the function can be inlined fully into the
+    kernel module from where it is used at the LLVM level, leading to more
+    optimization opportunities.
+
+    Returns:
+        KernelDispatcher: A KernelDispatcher instance with the
+        _compilation_mode option set to DEVICE_FUNC.
+    """
+    dispatcher = resolve_dispatcher_from_str(DPEX_KERNEL_EXP_TARGET_NAME)
+
+    if "_compilation_mode" in options:
+        user_compilation_mode = options["_compilation_mode"]
+        warn(
+            "_compilation_mode is an internal flag that should not be set "
+            "in the decorator. The decorator defined option "
+            f"{user_compilation_mode} is going to be ignored."
+        )
+    options["_compilation_mode"] = CompilationMode.DEVICE_FUNC
+
+    def _kernel_dispatcher(pyfunc):
+        return dispatcher(
+            pyfunc=pyfunc,
+            targetoptions=options,
+        )
+
+    if func_or_sig is None:
+        return _kernel_dispatcher
+
+    return _kernel_dispatcher(func_or_sig)
+
+
+jit_registry[target_registry[DPEX_KERNEL_EXP_TARGET_NAME]] = device_func

numba_dpex/experimental/kernel_dispatcher.py

@@ -15,14 +15,17 @@
 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.types import void
 from numba.core.typing.typeof import Purpose, typeof
 
 from numba_dpex import config, spirv_generator
 from numba_dpex.core.exceptions import (
     ExecutionQueueInferenceError,
+    KernelHasReturnValueError,
     UnsupportedKernelArgumentError,
 )
 from numba_dpex.core.pipelines import kernel_compiler
+from numba_dpex.core.targets.kernel_target import CompilationMode
 from numba_dpex.core.types import DpnpNdArray
 
 from .target import DPEX_KERNEL_EXP_TARGET_NAME, dpex_exp_kernel_target
@@ -82,9 +85,10 @@ def _compile_to_spirv(
             kernel_func, kernel_fndesc.argtypes
         )
 
-        # makes sure that the spir_func is completely inlined into the
-        # spir_kernel wrapper
-        kernel_library.optimize_final_module()
+        # Call finalize on the LLVM module. Finalization will result in
+        # all linking libraries getting linked together and final optimization
+        # including inlining of functions if an inlining level is specified.
+        kernel_library.finalize()
         # Compiled the LLVM IR to SPIR-V
         kernel_spirv_module = spirv_generator.llvm_to_spirv(
             kernel_targetctx,
@@ -144,9 +148,15 @@ def _compile_cached(
         try:
             cres: CompileResult = self._compile_core(args, return_type)
 
-            kernel_device_ir_module = self._compile_to_spirv(
-                cres.library, cres.fndesc, cres.target_context
-            )
+            if (
+                self.targetoptions["_compilation_mode"]
+                == CompilationMode.KERNEL
+            ):
+                kernel_device_ir_module: _KernelModule = self._compile_to_spirv(
+                    cres.library, cres.fndesc, cres.target_context
+                )
+            else:
+                kernel_device_ir_module = None
 
             kcres_attrs = []
 
@@ -282,12 +292,28 @@ def cb_llvm(dur):
             with self._compiling_counter:
                 args, return_type = sigutils.normalize_signature(sig)
 
-                try:
-                    self._compiler.check_queue_equivalence_of_args(
-                        self._kernel_name, args
-                    )
-                except ExecutionQueueInferenceError as eqie:
-                    raise eqie
+                if (
+                    self.targetoptions["_compilation_mode"]
+                    == CompilationMode.KERNEL
+                ):
+                    # Compute follows data based queue equivalence is only
+                    # evaluated for kernel functions whose arguments are
+                    # supposed to be arrays. For device_func decorated
+                    # functions, the arguments can be scalar and we skip queue
+                    # equivalence check.
+                    try:
+                        self._compiler.check_queue_equivalence_of_args(
+                            self._kernel_name, args
+                        )
+                    except ExecutionQueueInferenceError as eqie:
+                        raise eqie
+
+                    # A function being compiled in the KERNEL compilation mode
+                    # cannot have a non-void return value
+                    if return_type and return_type != void:
+                        raise KernelHasReturnValueError(
+                            kernel_name=None, return_type=return_type, sig=sig
+                        )
 
                 # Don't recompile if signature already exists
                 existing = self.overloads.get(tuple(args))