Add get_group_linear_id to overaloads

Author: ZzEeKkAa

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_index_space_id_overloads.py

@@ -16,7 +16,7 @@
     ItemType,
     NdItemType,
 )
-from numba_dpex.kernel_api import Item, NdItem
+from numba_dpex.kernel_api import Group, Item, NdItem
 from numba_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext
 
 from ..target import DPEX_KERNEL_EXP_TARGET_NAME
@@ -303,3 +303,4 @@ def register_jitable_method(type_, method):
 register_jitable_method(ItemType, Item.get_linear_id)
 register_jitable_method(NdItemType, NdItem.get_global_linear_id)
 register_jitable_method(NdItemType, NdItem.get_local_linear_id)
+register_jitable_method(GroupType, Group.get_group_linear_id)

numba_dpex/kernel_api/index_space_ids.py

@@ -42,14 +42,20 @@ def get_group_id(self, dim):
 
     def get_group_linear_id(self):
         """Returns a linearized version of the work-group index."""
-        if len(self._index) == 1:
-            return self._index[0]
-        if len(self._index) == 2:
-            return self._index[0] * self._group_range[1] + self._index[1]
+        if self.dimensions == 1:
+            return self.get_group_id(0)
+        if self.dimensions == 2:
+            return self.get_group_id(0) * self.get_group_range(
+                1
+            ) + self.get_group_id(1)
         return (
-            (self._index[0] * self._group_range[1] * self._group_range[2])
-            + (self._index[1] * self._group_range[2])
-            + (self._index[2])
+            (
+                self.get_group_id(0)
+                * self.get_group_range(1)
+                * self.get_group_range(2)
+            )
+            + (self.get_group_id(1) * self.get_group_range(2))
+            + (self.get_group_id(2))
         )
 
     def get_group_range(self, dim):

numba_dpex/tests/experimental/test_index_space_ids.py

@@ -87,6 +87,12 @@ def _get_group_id_driver(nditem: NdItem, a):
     a[i] = g.get_group_id(0)
 
 
+def _get_group_linear_id_driver(nditem: NdItem, a):
+    i = nditem.get_global_linear_id()
+    g = nditem.get_group()
+    a[i] = g.get_group_linear_id()
+
+
 def _get_group_range_driver(nditem: NdItem, a):
     i = nditem.get_global_id(0)
     g = nditem.get_group()
@@ -206,21 +212,29 @@ def test_no_item():
     )
 
 
-def test_get_group_id():
-    global_size = 100
-    group_size = 20
-    num_groups = global_size // group_size
+@pytest.mark.parametrize(
+    "driver,rng",
+    [
+        (_get_group_id_driver, dpex.NdRange((_SIZE,), (_GROUP_SIZE,))),
+        (_get_group_linear_id_driver, dpex.NdRange((_SIZE,), (_GROUP_SIZE,))),
+        (
+            _get_group_linear_id_driver,
+            dpex.NdRange((1, 1, _SIZE), (1, 1, _GROUP_SIZE)),
+        ),
+    ],
+)
+def test_get_group_id(driver, rng):
+    num_groups = _SIZE // _GROUP_SIZE
 
-    a = dpnp.empty(global_size, dtype=dpnp.int32)
-    ka = dpnp.empty(global_size, dtype=dpnp.int32)
-    expected = np.empty(global_size, dtype=np.int32)
-    ndrange = NdRange((global_size,), (group_size,))
-    dpex_exp.call_kernel(dpex_exp.kernel(_get_group_id_driver), ndrange, a)
-    kapi_call_kernel(_get_group_id_driver, ndrange, ka)
+    a = dpnp.empty(_SIZE, dtype=dpnp.int32)
+    ka = dpnp.empty(_SIZE, dtype=dpnp.int32)
+    expected = np.empty(_SIZE, dtype=np.int32)
+    dpex_exp.call_kernel(dpex_exp.kernel(driver), rng, a)
+    kapi_call_kernel(driver, rng, ka)
 
     for gid in range(num_groups):
-        for lid in range(group_size):
-            expected[gid * group_size + lid] = gid
+        for lid in range(_GROUP_SIZE):
+            expected[gid * _GROUP_SIZE + lid] = gid
 
     assert np.array_equal(a.asnumpy(), expected)
     assert np.array_equal(ka.asnumpy(), expected)