Migrate parfor local accessor to new api

Author: ZzEeKkAa

numba_dpex/core/parfors/kernel_builder.py

@@ -44,13 +44,17 @@ def __init__(
         kernel_args,
         kernel_arg_types,
         queue: dpctl.SyclQueue,
+        local_accessors=None,
+        work_group_size=None,
     ):
         self.name = name
         self.kernel = kernel
         self.signature = signature
         self.kernel_args = kernel_args
         self.kernel_arg_types = kernel_arg_types
         self.queue = queue
+        self.local_accessors = local_accessors
+        self.work_group_size = work_group_size
 
 
 def _print_block(block):

numba_dpex/core/parfors/kernel_templates/reduction_template.py

@@ -30,8 +30,8 @@ def __init__(
         parfor_args,
         parfor_reddict,
         redvars_dict,
+        local_accessors_dict,
         typemap,
-        work_group_size,
     ) -> None:
         self._kernel_name = kernel_name
         self._kernel_params = kernel_params
@@ -44,8 +44,8 @@ def __init__(
         self._parfor_args = parfor_args
         self._parfor_reddict = parfor_reddict
         self._redvars_dict = redvars_dict
+        self._local_accessors_dict = local_accessors_dict
         self._typemap = typemap
-        self._work_group_size = work_group_size
 
         self._kernel_txt = self._generate_kernel_stub_as_string()
         self._kernel_ir = self._generate_kernel_ir()
@@ -76,13 +76,6 @@ def _generate_kernel_stub_as_string(self):
             )
             gufunc_txt += f"    group_id{dim} = group.get_group_id({dstr})\n"
 
-        # Allocate local_sums arrays for each reduction variable.
-        for redvar in self._redvars:
-            rtyp = str(self._typemap[redvar])
-            redvar = self._redvars_dict[redvar]
-            gufunc_txt += f"    local_sums_{redvar} = \
-                dpex.local.array({self._work_group_size}, dpnp.{rtyp})\n"
-
         for dim in range(global_id_dim, for_loop_dim):
             for indent in range(1 + (dim - global_id_dim)):
                 gufunc_txt += "    "

numba_dpex/core/parfors/parfor_lowerer.py

@@ -169,8 +169,24 @@ def _submit_parfor_kernel(
         queue_ref = kl_builder.get_queue(exec_queue=kernel_fn.queue)
 
         kernel_args = []
-        for arg in kernel_fn.kernel_args:
-            kernel_args.append(_getvar(lowerer, arg))
+        for i, arg in enumerate(kernel_fn.kernel_args):
+            if (
+                kernel_fn.local_accessors is not None
+                and arg in kernel_fn.local_accessors
+            ):
+                wg_size = lowerer.context.get_constant(
+                    types.intp, kernel_fn.work_group_size
+                )
+                la_shape = cgutils.pack_array(lowerer.builder, [wg_size])
+                arg_ty = kernel_fn.kernel_arg_types[i]
+                la = cgutils.create_struct_proxy(arg_ty)(
+                    lowerer.context,
+                    lowerer.builder,
+                )
+                la.shape = la_shape
+                kernel_args.append(la._getvalue())
+            else:
+                kernel_args.append(_getvar(lowerer, arg))
 
         kernel_ref_addr = kernel_fn.kernel.addressof_ref()
         kernel_ref = lowerer.builder.inttoptr(

numba_dpex/core/parfors/reduction_kernel_builder.py

@@ -19,8 +19,10 @@
 )
 from numba.core.typing import signature
 
+from numba_dpex.core.parfors.reduction_helper import ReductionKernelVariables
 from numba_dpex.core.types import DpctlSyclQueue
 from numba_dpex.core.types.kernel_api.index_space_ids import NdItemType
+from numba_dpex.core.types.kernel_api.local_accessor import LocalAccessorType
 
 from .kernel_builder import _print_body  # saved for debug
 from .kernel_builder import (
@@ -41,7 +43,7 @@ def create_reduction_main_kernel_for_parfor(
     typemap,
     flags,
     has_aliases,
-    reductionKernelVar,
+    reductionKernelVar: ReductionKernelVariables,
     parfor_reddict=None,
 ):
     """
@@ -79,20 +81,35 @@ def create_reduction_main_kernel_for_parfor(
         except KeyError:
             pass
 
+    parfor_params = reductionKernelVar.parfor_params.copy()
+    parfor_legalized_params = reductionKernelVar.parfor_legalized_params.copy()
+    parfor_param_types = reductionKernelVar.param_types.copy()
+    local_accessors_dict = {}
+    for k, v in reductionKernelVar.redvars_legal_dict.items():
+        la_var = "local_sums_" + v
+        local_accessors_dict[k] = la_var
+        idx = reductionKernelVar.parfor_params.index(k)
+        arr_ty = reductionKernelVar.param_types[idx]
+        la_ty = LocalAccessorType(parfor_dim, arr_ty.dtype)
+
+        parfor_params.append(la_var)
+        parfor_legalized_params.append(la_var)
+        parfor_param_types.append(la_ty)
+
     kernel_template = TreeReduceIntermediateKernelTemplate(
         kernel_name=kernel_name,
-        kernel_params=reductionKernelVar.parfor_legalized_params,
+        kernel_params=parfor_legalized_params,
         ivar_names=reductionKernelVar.legal_loop_indices,
         sentinel_name=sentinel_name,
         loop_ranges=loop_ranges,
         param_dict=reductionKernelVar.param_dict,
         parfor_dim=parfor_dim,
         redvars=reductionKernelVar.parfor_redvars,
-        parfor_args=reductionKernelVar.parfor_params,
+        parfor_args=parfor_params,
         parfor_reddict=parfor_reddict,
         redvars_dict=reductionKernelVar.redvars_legal_dict,
+        local_accessors_dict=local_accessors_dict,
         typemap=typemap,
-        work_group_size=reductionKernelVar.work_group_size,
     )
     kernel_ir = kernel_template.kernel_ir
 
@@ -116,7 +133,7 @@ def create_reduction_main_kernel_for_parfor(
             new_var_dict[name] = mk_unique_var(name)
 
     replace_var_names(kernel_ir.blocks, new_var_dict)
-    kernel_param_types = reductionKernelVar.param_types
+    kernel_param_types = parfor_param_types
     kernel_stub_last_label = max(kernel_ir.blocks.keys()) + 1
     # Add kernel stub last label to each parfor.loop_body label to prevent
     # label conflicts.
@@ -164,13 +181,20 @@ def create_reduction_main_kernel_for_parfor(
 
     flags.noalias = old_alias
 
+    parfor_params = (
+        reductionKernelVar.parfor_params.copy()
+        + parfor_params[len(reductionKernelVar.parfor_params) :]  # noqa: $203
+    )
+
     return ParforKernel(
         name=kernel_name,
         kernel=sycl_kernel,
         signature=kernel_sig,
-        kernel_args=reductionKernelVar.parfor_params,
-        kernel_arg_types=reductionKernelVar.func_arg_types,
+        kernel_args=parfor_params,
+        kernel_arg_types=parfor_param_types,
         queue=exec_queue,
+        local_accessors=set(local_accessors_dict.values()),
+        work_group_size=reductionKernelVar.work_group_size,
     )