Add linear indexing overloads for item & nd_item

Author: ZzEeKkAa

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_index_space_id_overloads.py

@@ -16,6 +16,7 @@
     ItemType,
     NdItemType,
 )
+from numba_dpex.kernel_api import Item, NdItem
 from numba_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext
 
 from ..target import DPEX_KERNEL_EXP_TARGET_NAME
@@ -274,3 +275,31 @@ def ol_nd_item_get_group_impl(item):
         return dimensions
 
     return ol_nd_item_get_group_impl
+
+
+def _generate_method_overload(method):
+    """Generates naive method overload with no argument, except self."""
+
+    def ol_method(self):  # pylint: disable=unused-argument
+        return method
+
+    return ol_method
+
+
+def register_jitable_method(type_, method):
+    """
+    Register a regular python method that can be executed by the
+    python interpreter and can be compiled into a nopython
+    function when referenced by other jit'ed functions.
+
+    Same as register_jitable, but for methods with no arguments.
+    """
+    overloaded_method = _generate_method_overload(method)
+    overload_method(type_, method.__name__, target=DPEX_KERNEL_EXP_TARGET_NAME)(
+        overloaded_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)

numba_dpex/kernel_api/index_space_ids.py

@@ -128,14 +128,14 @@ def get_linear_id(self):
         Returns:
             int: The linear id.
         """
-        if len(self._extent) == 1:
-            return self._index[0]
-        if len(self._extent) == 2:
-            return self._index[0] * self._extent[1] + self._index[1]
+        if self.dimensions == 1:
+            return self.get_id(0)
+        if self.dimensions == 2:
+            return self.get_id(0) * self.get_range(1) + self.get_id(1)
         return (
-            (self._index[0] * self._extent[1] * self._extent[2])
-            + (self._index[1] * self._extent[2])
-            + (self._index[2])
+            (self.get_id(0) * self.get_range(1) * self.get_range(2))
+            + (self.get_id(1) * self.get_range(2))
+            + (self.get_id(2))
         )
 
     def get_id(self, idx):
@@ -193,7 +193,24 @@ def get_global_linear_id(self):
         Returns:
             int: The global linear id.
         """
-        return self._global_item.get_linear_id()
+        # Instead of calling self._global_item.get_linear_id(), the linearization
+        # logic is duplicated here so that the method can be JIT compiled by
+        # numba-dpex and works in both Python and Numba nopython modes.
+        if self.dimensions == 1:
+            return self.get_global_id(0)
+        if self.dimensions == 2:
+            return self.get_global_id(0) * self.get_global_range(
+                1
+            ) + self.get_global_id(1)
+        return (
+            (
+                self.get_global_id(0)
+                * self.get_global_range(1)
+                * self.get_global_range(2)
+            )
+            + (self.get_global_id(1) * self.get_global_range(2))
+            + (self.get_global_id(2))
+        )
 
     def get_local_id(self, idx):
         """Get the local id for a specific dimension.
@@ -210,7 +227,24 @@ def get_local_linear_id(self):
         Returns:
             int: The local linear id.
         """
-        return self._local_item.get_linear_id()
+        # Instead of calling self._local_item.get_linear_id(), the linearization
+        # logic is duplicated here so that the method can be JIT compiled by
+        # numba-dpex and works in both Python and Numba nopython modes.
+        if self.dimensions == 1:
+            return self.get_local_id(0)
+        if self.dimensions == 2:
+            return self.get_local_id(0) * self.get_local_range(
+                1
+            ) + self.get_local_id(1)
+        return (
+            (
+                self.get_local_id(0)
+                * self.get_local_range(1)
+                * self.get_local_range(2)
+            )
+            + (self.get_local_id(1) * self.get_local_range(2))
+            + (self.get_local_id(2))
+        )
 
     def get_global_range(self, idx):
         """Get the global range size for a specific dimension.

numba_dpex/tests/experimental/test_index_space_ids.py

@@ -277,18 +277,71 @@ def set_3d_ones_item(item: Item, a):
     j = item.get_id(1)
     k = item.get_id(2)
 
-    # Since we have different sizes for each dimention, wrong order will result
+    # Since we have different sizes for each dimension, wrong order will result
     # that some indexes will be set twice and some won't be set.
     index = i + I_SIZE * (j + J_SIZE * k)
 
     a[index] = 1
 
 
-def test_index_order():
+@dpex_exp.kernel
+def set_3d_ones_item_linear(item: Item, a):
+    # Since we have different sizes for each dimension, wrong order will result
+    # that some indexes will be set twice and some won't be set.
+    index = item.get_linear_id()
+
+    a[index] = 1
+
+
+@dpex_exp.kernel
+def set_3d_ones_nd_item_linear(nd_item: NdItem, a):
+    # Since we have different sizes for each dimension, wrong order will result
+    # that some indexes will be set twice and some won't be set.
+    index = nd_item.get_global_linear_id()
+
+    a[index] = 1
+
+
+@dpex_exp.kernel
+def set_local_3d_ones_nd_item_linear(nd_item: NdItem, a):
+    # Since we have different sizes for each dimension, wrong order will result
+    # that some indexes will be set twice and some won't be set.
+    index = nd_item.get_local_linear_id()
+
+    a[index] = 1
+
+
+@pytest.mark.parametrize("kernel", [set_3d_ones_item, set_3d_ones_item_linear])
+def test_item_index_order(kernel):
+    a = dpnp.zeros(I_SIZE * J_SIZE * K_SIZE, dtype=dpnp.int32)
+
+    dpex_exp.call_kernel(kernel, dpex.Range(I_SIZE, J_SIZE, K_SIZE), a)
+
+    assert np.array_equal(a.asnumpy(), np.ones(a.size, dtype=np.int32))
+
+
+def test_nd_item_index_order():
     a = dpnp.zeros(I_SIZE * J_SIZE * K_SIZE, dtype=dpnp.int32)
 
     dpex_exp.call_kernel(
-        set_3d_ones_item, dpex.Range(I_SIZE, J_SIZE, K_SIZE), a
+        set_3d_ones_nd_item_linear,
+        dpex.NdRange((I_SIZE, J_SIZE, K_SIZE), (1, 1, K_SIZE)),
+        a,
     )
 
     assert np.array_equal(a.asnumpy(), np.ones(a.size, dtype=np.int32))
+
+
+def test_nd_item_local_linear_id():
+    a = dpnp.zeros(I_SIZE * J_SIZE * K_SIZE, dtype=dpnp.int32)
+
+    dpex_exp.call_kernel(
+        set_local_3d_ones_nd_item_linear,
+        dpex.NdRange((I_SIZE, J_SIZE, K_SIZE), (1, 1, K_SIZE)),
+        a,
+    )
+
+    assert np.array_equal(
+        a.asnumpy(),
+        np.array([1] * K_SIZE + [0] * (a.size - K_SIZE), dtype=np.int32),
+    )