Add nd_item to reduction template

Author: ZzEeKkAa

numba_dpex/core/parfors/kernel_templates/reduction_template.py

@@ -55,7 +55,7 @@ def _generate_kernel_stub_as_string(self):
 
         gufunc_txt = ""
         gufunc_txt += "def " + self._kernel_name
-        gufunc_txt += "(" + (", ".join(self._kernel_params)) + "):\n"
+        gufunc_txt += "(nd_item, " + (", ".join(self._kernel_params)) + "):\n"
         global_id_dim = 0
         for_loop_dim = self._parfor_dim
 
@@ -64,14 +64,17 @@ def _generate_kernel_stub_as_string(self):
         else:
             global_id_dim = self._parfor_dim
 
+        gufunc_txt += "    group = nd_item.get_group()\n"
         for dim in range(global_id_dim):
             dstr = str(dim)
             gufunc_txt += (
-                f"    {self._ivar_names[dim]} = dpex.get_global_id({dstr})\n"
+                f"    {self._ivar_names[dim]} = nd_item.get_global_id({dstr})\n"
             )
-            gufunc_txt += f"    local_id{dim} = dpex.get_local_id({dstr})\n"
-            gufunc_txt += f"    local_size{dim} = dpex.get_local_size({dstr})\n"
-            gufunc_txt += f"    group_id{dim} = dpex.get_group_id({dstr})\n"
+            gufunc_txt += f"    local_id{dim} = nd_item.get_local_id({dstr})\n"
+            gufunc_txt += (
+                f"    local_size{dim} = group.get_local_range({dstr})\n"
+            )
+            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:

numba_dpex/core/parfors/reduction_kernel_builder.py

@@ -20,6 +20,7 @@
 from numba.core.typing import signature
 
 from numba_dpex.core.types import DpctlSyclQueue
+from numba_dpex.core.types.kernel_api.index_space_ids import NdItemType
 
 from .kernel_builder import _print_body  # saved for debug
 from .kernel_builder import (
@@ -48,6 +49,7 @@ def create_reduction_main_kernel_for_parfor(
     """
 
     loc = parfor_node.init_block.loc
+    parfor_dim = len(parfor_node.loop_nests)
 
     for race in parfor_node.races:
         msg = (
@@ -84,7 +86,7 @@ def create_reduction_main_kernel_for_parfor(
         sentinel_name=sentinel_name,
         loop_ranges=loop_ranges,
         param_dict=reductionKernelVar.param_dict,
-        parfor_dim=len(parfor_node.loop_nests),
+        parfor_dim=parfor_dim,
         redvars=reductionKernelVar.parfor_redvars,
         parfor_args=reductionKernelVar.parfor_params,
         parfor_reddict=parfor_reddict,
@@ -136,6 +138,13 @@ def create_reduction_main_kernel_for_parfor(
     if not has_aliases:
         flags.noalias = True
 
+    # The first argument to a range kernel is a kernel_api.NdItem object. The
+    # ``NdItem`` object is used by the kernel_api.spirv backend to generate the
+    # correct SPIR-V indexing instructions. Since, the argument is not something
+    # available originally in the kernel_param_types, we add it at this point to
+    # make sure the kernel signature matches the actual generated code.
+    ty_item = NdItemType(parfor_dim)
+    kernel_param_types = (ty_item, *kernel_param_types)
     kernel_sig = signature(types.none, *kernel_param_types)
 
     # FIXME: A better design is required so that we do not have to create a