Apply builder pattern to kernel launcher

Author: ZzEeKkAa

.pre-commit-config.yaml

@@ -52,7 +52,7 @@ repos:
     - id: pylint
       name: pylint
       entry: pylint
-      files: ^numba_dpex/experimental
+      files: ^numba_dpex/experimental|^numba_dpex/core/utils/kernel_launcher.py
       language: system
       types: [python]
       require_serial: true

numba_dpex/core/parfors/kernel_builder.py

@@ -40,7 +40,7 @@ def __init__(
         signature,
         kernel_args,
         kernel_arg_types,
-        queue,
+        queue: dpctl.SyclQueue,
     ):
         self.name = name
         self.kernel = kernel
@@ -244,7 +244,7 @@ def create_kernel_for_parfor(
     has_aliases,
     races,
     parfor_outputs,
-):
+) -> ParforKernel:
     """
     Creates a numba_dpex.kernel function for a parfor node.
 
@@ -422,7 +422,7 @@ def create_kernel_for_parfor(
     # arrays are on same device. We can take the queue from the first input
     # array and use that to compile the kernel.
 
-    exec_queue = None
+    exec_queue: dpctl.SyclQueue = None
 
     for arg in parfor_args:
         obj = typemap[arg]

numba_dpex/core/parfors/parfor_lowerer.py

@@ -13,12 +13,15 @@
 )
 
 from numba_dpex import config
-from numba_dpex.core.datamodel.models import dpex_data_model_manager as dpex_dmm
 from numba_dpex.core.parfors.reduction_helper import (
     ReductionHelper,
     ReductionKernelVariables,
 )
 from numba_dpex.core.utils.kernel_launcher import KernelLaunchIRBuilder
+from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
+from numba_dpex.core.datamodel.models import (
+    dpex_data_model_manager as kernel_dmm,
+)
 
 from ..exceptions import UnsupportedParforError
 from ..types.dpnp_ndarray_type import DpnpNdArray
@@ -28,11 +31,6 @@
     create_reduction_remainder_kernel_for_parfor,
 )
 
-_KernelArgs = namedtuple(
-    "_KernelArgs",
-    ["num_flattened_args", "arg_vals", "arg_types"],
-)
-
 
 # A global list of kernels to keep the objects alive indefinitely.
 keep_alive_kernels = []
@@ -68,11 +66,8 @@ def _getvar(lowerer, x):
         var_val = lowerer.varmap[x]
 
     if var_val:
-        if not isinstance(var_val.type, llvmir.PointerType):
-            with lowerer.builder.goto_entry_block():
-                var_val_ptr = lowerer.builder.alloca(var_val.type)
-            lowerer.builder.store(var_val, var_val_ptr)
-            return var_val_ptr
+        if isinstance(var_val.type, llvmir.PointerType):
+            return lowerer.builder.load(var_val)
         else:
             return var_val
     else:
@@ -91,56 +86,6 @@ class ParforLowerImpl:
     for a parfor and submits it to a queue.
     """
 
-    def _build_kernel_arglist(
-        self, kernel_fn, lowerer, kernel_builder: KernelLaunchIRBuilder
-    ):
-        """Creates local variables for all the arguments and the argument types
-        that are passes to the kernel function.
-
-        Args:
-            kernel_fn: Kernel function to be launched.
-            lowerer: The Numba lowerer used to generate the LLVM IR
-
-        Raises:
-            AssertionError: If the LLVM IR Value for an argument defined in
-            Numba IR is not found.
-        """
-        num_flattened_args = 0
-
-        # Compute number of args to be passed to the kernel. Note that the
-        # actual number of kernel arguments is greater than the count of
-        # kernel_fn.kernel_args as arrays get flattened.
-        for arg_type in kernel_fn.kernel_arg_types:
-            if isinstance(arg_type, DpnpNdArray):
-                datamodel = dpex_dmm.lookup(arg_type)
-                num_flattened_args += datamodel.flattened_field_count
-            elif arg_type == types.complex64 or arg_type == types.complex128:
-                num_flattened_args += 2
-            else:
-                num_flattened_args += 1
-
-        # Create LLVM values for the kernel args list and kernel arg types list
-        args_list = kernel_builder.allocate_kernel_arg_array(num_flattened_args)
-        args_ty_list = kernel_builder.allocate_kernel_arg_ty_array(
-            num_flattened_args
-        )
-        callargs_ptrs = []
-        for arg in kernel_fn.kernel_args:
-            callargs_ptrs.append(_getvar(lowerer, arg))
-
-        kernel_builder.populate_kernel_args_and_args_ty_arrays(
-            kernel_argtys=kernel_fn.kernel_arg_types,
-            callargs_ptrs=callargs_ptrs,
-            args_list=args_list,
-            args_ty_list=args_ty_list,
-        )
-
-        return _KernelArgs(
-            num_flattened_args=num_flattened_args,
-            arg_vals=args_list,
-            arg_types=args_ty_list,
-        )
-
     def _loop_ranges(
         self,
         lowerer,
@@ -163,7 +108,10 @@ def _loop_ranges(
                     "non-unit strides are not yet supported."
                 )
             global_range.append(stop)
-
+        # For now the local_range is always an empty list as numba_dpex always
+        # submits kernels generated for parfor nodes as range kernels.
+        # The provision is kept here if in future there is newer functionality
+        # to submit these kernels as ndrange.
         local_range = []
 
         return global_range, local_range
@@ -215,31 +163,34 @@ def _submit_parfor_kernel(
         # Ensure that the Python arguments are kept alive for the duration of
         # the kernel execution
         keep_alive_kernels.append(kernel_fn.kernel)
-        kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)
+        kl_builder = KernelLaunchIRBuilder(
+            lowerer.context, lowerer.builder, kernel_dmm
+        )
+
+        queue_ref = kl_builder.get_queue(exec_queue=kernel_fn.queue)
 
-        ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)
-        args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)
+        kernel_args = []
+        for arg in kernel_fn.kernel_args:
+            kernel_args.append(_getvar(lowerer, arg))
 
         kernel_ref_addr = kernel_fn.kernel.addressof_ref()
         kernel_ref = lowerer.builder.inttoptr(
             lowerer.context.get_constant(types.uintp, kernel_ref_addr),
             cgutils.voidptr_t,
         )
-        curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)
-
-        # Submit a synchronous kernel
-        kernel_builder.submit_sycl_kernel(
-            sycl_kernel_ref=kernel_ref,
-            sycl_queue_ref=curr_queue_ref,
-            total_kernel_args=args.num_flattened_args,
-            arg_list=args.arg_vals,
-            arg_ty_list=args.arg_types,
-            global_range=global_range,
-            local_range=local_range,
+
+        kl_builder.set_kernel(kernel_ref)
+        kl_builder.set_queue(queue_ref)
+        kl_builder.set_range(global_range, local_range)
+        kl_builder.set_arguments(
+            kernel_fn.kernel_arg_types, kernel_args=kernel_args
         )
+        kl_builder.set_dependant_event_list(dep_events=[])
+        event_ref = kl_builder.submit()
 
-        # At this point we can free the DPCTLSyclQueueRef (curr_queue)
-        kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)
+        sycl.dpctl_event_wait(lowerer.builder, event_ref)
+        sycl.dpctl_event_delete(lowerer.builder, event_ref)
+        sycl.dpctl_queue_delete(lowerer.builder, queue_ref)
 
     def _reduction_codegen(
         self,

numba_dpex/core/parfors/reduction_helper.py

@@ -17,6 +17,9 @@
 from numba.parfors.parfor_lowering_utils import ParforLoweringBuilder
 
 from numba_dpex import utils
+from numba_dpex.core.datamodel.models import (
+    dpex_data_model_manager as kernel_dmm,
+)
 from numba_dpex.core.utils.kernel_launcher import KernelLaunchIRBuilder
 from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
 
@@ -395,11 +398,13 @@ def work_group_size(self):
 
     def copy_final_sum_to_host(self, parfor_kernel):
         lowerer = self.lowerer
-        ir_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)
+        kl_builder = KernelLaunchIRBuilder(
+            lowerer.context, lowerer.builder, kernel_dmm
+        )
 
         # Create a local variable storing a pointer to a DPCTLSyclQueueRef
         # pointer.
-        curr_queue = ir_builder.get_queue(exec_queue=parfor_kernel.queue)
+        queue_ref = kl_builder.get_queue(exec_queue=parfor_kernel.queue)
 
         builder = lowerer.builder
         context = lowerer.context
@@ -433,7 +438,7 @@ def copy_final_sum_to_host(self, parfor_kernel):
             )
 
             args = [
-                builder.load(curr_queue),
+                queue_ref,
                 dest,
                 src,
                 builder.load(item_size),
@@ -443,4 +448,4 @@ def copy_final_sum_to_host(self, parfor_kernel):
             sycl.dpctl_event_wait(builder, event_ref)
             sycl.dpctl_event_delete(builder, event_ref)
 
-        ir_builder.free_queue(ptr_to_sycl_queue_ref=curr_queue)
+        sycl.dpctl_queue_delete(builder, queue_ref)