Add a decorator-level option to set inlining threshold.

    - A new target option "inlining threshold" was added to
      DpexKernelTarget to define how the LLVM inlining passes
      should optimize the final codegen library.

      The decorator-level option will supersede any global
      configuration setting.

Author: Diptorup Deb

numba_dpex/core/codegen.py

@@ -29,6 +29,25 @@ class SPIRVCodeLibrary(CPUCodeLibrary):
     def _optimize_functions(self, ll_module):
         pass
 
+    @property
+    def inline_threshold(self):
+        """The inlining threshold value to be used to optimize the final library"""
+        if hasattr(self, "_inline_threshold"):
+            return self._inline_threshold
+        else:
+            return 0
+
+    @inline_threshold.setter
+    def inline_threshold(self, value: int):
+        """Returns the current inlining threshold level for the library."""
+        if value < 0 or value > 3:
+            logging.warning(
+                "Unsupported inline threshold. Set a value between 0 and 3"
+            )
+            self._inline_threshold = 0
+        else:
+            self._inline_threshold = value
+
     def _optimize_final_module(self):
         # Run some lightweight optimization to simplify the module.
         pmb = ll.PassManagerBuilder()
@@ -43,12 +62,38 @@ def _optimize_final_module(self):
             )
 
         pmb.disable_unit_at_a_time = False
+
         if config.INLINE_THRESHOLD is not None:
-            logging.warning(
-                "Setting INLINE_THRESHOLD leads to very aggressive "
-                + "optimizations that may produce incorrect binary."
-            )
-            pmb.inlining_threshold = config.INLINE_THRESHOLD
+            # Check if a decorator-level inline threshold was set and use that
+            # instead of the global configuration.
+            if (
+                hasattr(self, "_inline_threshold")
+                and self._inline_threshold > 0
+                and self._inline_threshold <= 3
+            ):
+                logging.warning(
+                    "Setting INLINE_THRESHOLD leads to very aggressive "
+                    + "optimizations that may produce incorrect binary."
+                )
+                pmb.inlining_threshold = self._inline_threshold
+            elif not hasattr(self, "_inline_threshold"):
+                logging.warning(
+                    "Setting INLINE_THRESHOLD leads to very aggressive "
+                    + "optimizations that may produce incorrect binary."
+                )
+                pmb.inlining_threshold = config.INLINE_THRESHOLD
+        else:
+            if (
+                hasattr(self, "_inline_threshold")
+                and self._inline_threshold > 0
+                and self._inline_threshold <= 3
+            ):
+                logging.warning(
+                    "Setting INLINE_THRESHOLD leads to very aggressive "
+                    + "optimizations that may produce incorrect binary."
+                )
+                pmb.inlining_threshold = self._inline_threshold
+
         pmb.disable_unroll_loops = True
         pmb.loop_vectorize = False
         pmb.slp_vectorize = False

numba_dpex/core/descriptor.py

@@ -41,6 +41,7 @@ class DpexTargetOptions(CPUTargetOptions):
     release_gil = _option_mapping("release_gil")
     no_compile = _option_mapping("no_compile")
     use_mlir = _option_mapping("use_mlir")
+    inline_threshold = _option_mapping("inline_threshold")
     _compilation_mode = _option_mapping("_compilation_mode")
 
     def finalize(self, flags, options):
@@ -49,6 +50,7 @@ def finalize(self, flags, options):
         _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, "inline_threshold", 0)
         _inherit_if_not_set(
             flags, options, "_compilation_mode", CompilationMode.KERNEL
         )

numba_dpex/experimental/kernel_dispatcher.py

@@ -84,7 +84,14 @@ def _compile_to_spirv(
         kernel_fn = kernel_targetctx.prepare_spir_kernel(
             kernel_func, kernel_fndesc.argtypes
         )
-
+        # If the inline_threshold option was set then set the property in the
+        # kernel_library to force inlining ``overload`` calls into a kernel.
+        try:
+            kernel_library.inline_threshold = self.targetoptions[
+                "inline_threshold"
+            ]
+        except KeyError:
+            pass
         # 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.