Merge pull request #1326 from IntelPython/feature/group_index_space_id_class

Extends mock kernel API by more group index functions and adds launchers

Author: Diptorup Deb

numba_dpex/kernel_api/__init__.py

@@ -12,6 +12,7 @@
 from .atomic_ref import AtomicRef
 from .barrier import group_barrier
 from .index_space_ids import Group, Item, NdItem
+from .launcher import call_kernel
 from .memory_enums import AddressSpace, MemoryOrder, MemoryScope
 from .ranges import NdRange, Range
 
@@ -26,4 +27,5 @@
     "NdItem",
     "Item",
     "group_barrier",
+    "call_kernel",
 ]

numba_dpex/kernel_api/index_space_ids.py

@@ -10,7 +10,6 @@
 from .ranges import Range
 
 
-# pylint: disable=too-few-public-methods
 class Group:
     """Analogue to the ``sycl::group`` type."""
 
@@ -25,6 +24,82 @@ def __init__(
         self._local_range = local_range
         self._group_range = group_range
         self._index = index
+        self._leader = False
+
+    def get_group_id(self, dim):
+        """Returns the index of the work-group within the global nd-range for
+        specified dimension.
+
+        Since the work-items in a work-group have a defined position within the
+        global nd-range, the returned group id can be used along with the local
+        id to uniquely identify the work-item in the global nd-range.
+        """
+        if dim > len(self._index) - 1:
+            raise ValueError(
+                "Dimension value is out of bounds for the group index"
+            )
+        return self._index[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]
+        return (
+            (self._index[0] * self._group_range[1] * self._group_range[2])
+            + (self._index[1] * self._group_range[2])
+            + (self._index[2])
+        )
+
+    def get_group_range(self):
+        """Returns a range representing the number of groups in the nd-range."""
+        return self._group_range
+
+    def get_group_linear_range(self):
+        """Return the total number of work-groups in the nd_range."""
+        num_wg = 1
+        for ext in self._group_range:
+            num_wg *= ext
+
+        return num_wg
+
+    def get_local_range(self):
+        """Returns a SYCL range representing all dimensions of the local
+        range. This local range may have been provided by the programmer, or
+        chosen by the SYCL runtime.
+        """
+        return self._local_range
+
+    def get_local_linear_range(self):
+        """Return the total number of work-items in the work-group."""
+        num_wi = 1
+        for ext in self._local_range:
+            num_wi *= ext
+
+        return num_wi
+
+    @property
+    def leader(self):
+        """Return true for exactly one work-item in the work-group, if the
+        calling work-item is the leader of the work-group, and false for all
+        other work-items in the work-group.
+
+        The leader of the work-group is determined during construction of the
+        work-group, and is invariant for the lifetime of the work-group. The
+        leader of the work-group is guaranteed to be the work-item with a
+        local id of 0.
+
+
+        Returns:
+            bool: If the work item is the designated leader of the
+        """
+        return self._leader
+
+    @leader.setter
+    def leader(self, work_item_id):
+        """Sets the leader attribute for the group."""
+        self._leader = work_item_id
 
 
 class Item:
@@ -45,7 +120,7 @@ def get_linear_id(self):
         """
         if len(self._extent) == 1:
             return self._index[0]
-        if len(self._extent) == 1:
+        if len(self._extent) == 2:
             return self._index[0] * self._extent[1] + self._index[1]
         return (
             (self._index[0] * self._extent[1] * self._extent[2])
@@ -90,6 +165,8 @@ def __init__(self, global_item: Item, local_item: Item, group: Group):
         self._global_item = global_item
         self._local_item = local_item
         self._group = group
+        if self.get_local_linear_id() == 0:
+            self._group.leader = True
 
     def get_global_id(self, idx):
         """Get the global id for a specific dimension.

numba_dpex/kernel_api/launcher.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Implementation of mock kernel launcher functions
+"""
+
+from inspect import signature
+from itertools import product
+
+from .index_space_ids import Group, Item, NdItem
+from .ranges import NdRange, Range
+
+
+def _range_kernel_launcher(kernel_fn, index_range, *kernel_args):
+    """Executes a function that mocks a range kernel.
+
+    Converts the range into a set of index tuple that represent an element in
+    the iteration domain over which the kernel will be executed. Then the
+    kernel is called sequentially over that set of indices after each index
+    value is used to construct an Item object.
+
+    Args:
+        kernel_fn : A callable function object
+        index_range (numba_dpex.Range): An instance of a Range object
+
+    Raises:
+        ValueError: If the number of passed in kernel arguments is not the
+        number of function parameters subtracted by one. The first kernel
+        argument is expected to be an Item object.
+    """
+
+    range_sets = [range(ir) for ir in index_range]
+    index_tuples = list(product(*range_sets))
+
+    for idx in index_tuples:
+        it = Item(extent=index_range, index=idx)
+
+        if len(signature(kernel_fn).parameters) - len(kernel_args) != 1:
+            raise ValueError(
+                "Required number of kernel function arguments do not "
+                "match provided number of kernel args"
+            )
+
+        kernel_fn(it, *kernel_args)
+
+
+def _ndrange_kernel_launcher(kernel_fn, index_range, *kernel_args):
+    """Executes a function that mocks a range kernel.
+
+    Args:
+        kernel_fn : A callable function object
+        index_range (numba_dpex.NdRange): An instance of a NdRange object
+
+    Raises:
+        ValueError: If the number of passed in kernel arguments is not the
+        number of function parameters subtracted by one. The first kernel
+        argument is expected to be an Item object.
+    """
+    group_range = tuple(
+        gr // lr
+        for gr, lr in zip(index_range.global_range, index_range.local_range)
+    )
+    local_range_sets = [range(ir) for ir in index_range.local_range]
+    group_range_sets = [range(gr) for gr in group_range]
+    local_index_tuples = list(product(*local_range_sets))
+    group_index_tuples = list(product(*group_range_sets))
+
+    # Loop over the groups (parallel loop)
+    for gidx in group_index_tuples:
+        # loop over work items in the group (parallel loop)
+        for lidx in local_index_tuples:
+            global_id = []
+            # to calculate global indices
+            for dim, gidx_val in enumerate(gidx):
+                global_id.append(
+                    gidx_val * index_range.local_range[dim] + lidx[dim]
+                )
+            # Every NdItem has its own global Item, local Item and Group
+            nditem = NdItem(
+                global_item=Item(
+                    extent=index_range.global_range, index=global_id
+                ),
+                local_item=Item(extent=index_range.local_range, index=lidx),
+                group=Group(
+                    index_range.global_range,
+                    index_range.local_range,
+                    group_range,
+                    gidx,
+                ),
+            )
+
+            if len(signature(kernel_fn).parameters) - len(kernel_args) != 1:
+                raise ValueError(
+                    "Required number of kernel function arguments do not "
+                    "match provided number of kernel args"
+                )
+
+            kernel_fn(nditem, *kernel_args)
+
+
+def call_kernel(kernel_fn, index_range, *kernel_args):
+    """Mocks the launching of a kernel function over either a Range or NdRange.
+
+    Args:
+        kernel_fn : A callable function object
+        index_range (numba_dpex.Range): An instance of a Range object
+
+    Raises:
+        ValueError: If the first positional argument is not callable
+        ValueError: If the second positional argument is not a Range or an
+        Ndrange object
+    """
+    if not callable(kernel_fn):
+        raise ValueError(
+            "Expected the first positional argument to be a function object"
+        )
+    if isinstance(index_range, Range):
+        _range_kernel_launcher(kernel_fn, index_range, *kernel_args)
+    elif isinstance(index_range, NdRange):
+        _ndrange_kernel_launcher(kernel_fn, index_range, *kernel_args)
+    else:
+        raise ValueError(
+            "Expected second positional argument to be Range or NdRange object"
+        )

numba_dpex/kernel_api/ranges.py

@@ -8,6 +8,11 @@
 
 from collections.abc import Iterable
 
+from numba_dpex.core.exceptions import (
+    UnmatchedNumberOfRangeDimsError,
+    UnsupportedGroupWorkItemSizeError,
+)
+
 
 class Range(tuple):
     """A data structure to encapsulate a single kernel launch parameter.
@@ -168,6 +173,22 @@ def __init__(self, global_size, local_size):
                 + "must be of either type Range or Iterable of int's."
             )
 
+        if len(self._local_range) != len(self._global_range):
+            raise UnmatchedNumberOfRangeDimsError(
+                kernel_name="",
+                global_ndims=len(self._global_range),
+                local_ndims=len(self._local_range),
+            )
+
+        for i, _ in enumerate(self._global_range):
+            if self._global_range[i] % self._local_range[i] != 0:
+                raise UnsupportedGroupWorkItemSizeError(
+                    kernel_name="",
+                    dim=i,
+                    work_groups=self._global_range[i],
+                    work_items=self._local_range[i],
+                )
+
     @property
     def global_range(self):
         """Accessor for global_range.

numba_dpex/tests/kernel_api/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0

numba_dpex/tests/kernel_api/test_ndrange_kernel_launch.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy
+
+from numba_dpex import kernel_api as kapi
+
+
+def test_range_kernel_call1D():
+    def vecadd(item: kapi.NdItem, a, b, c):
+        idx = item.get_global_id(0)
+        c[idx] = a[idx] + b[idx]
+
+    a = numpy.ones(100)
+    b = numpy.ones(100)
+    c = numpy.empty(100)
+
+    kapi.call_kernel(vecadd, kapi.NdRange((100,), (20,)), a, b, c)
+
+    assert numpy.allclose(c, a + b)
+
+
+def test_range_kernel_call2D():
+    def vecadd(item: kapi.NdItem, a, b, c):
+        idx = item.get_global_id(0)
+        jdx = item.get_global_id(1)
+        c[idx, jdx] = a[idx, jdx] + b[idx, jdx]
+
+    a = numpy.ones((10, 10))
+    b = numpy.ones((10, 10))
+    c = numpy.empty((10, 10))
+
+    kapi.call_kernel(vecadd, kapi.NdRange((10, 10), (2, 2)), a, b, c)
+
+    assert numpy.allclose(c, a + b)
+
+
+def test_range_kernel_call3D():
+    def vecadd(item: kapi.Item, a, b, c):
+        idx = item.get_global_id(0)
+        jdx = item.get_global_id(1)
+        kdx = item.get_global_id(2)
+        c[idx, jdx, kdx] = a[idx, jdx, kdx] + b[idx, jdx, kdx]
+
+    a = numpy.ones((8, 8, 8))
+    b = numpy.ones((8, 8, 8))
+    c = numpy.empty((8, 8, 8))
+
+    kapi.call_kernel(vecadd, kapi.NdRange((8, 8, 8), (2, 2, 2)), a, b, c)
+
+    assert numpy.allclose(c, a + b)

numba_dpex/tests/kernel_api/test_range_kernel_launch.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy
+
+from numba_dpex import kernel_api as kapi
+
+
+def test_range_kernel_call1D():
+    def vecadd(item: kapi.Item, a, b, c):
+        idx = item.get_id(0)
+        c[idx] = a[idx] + b[idx]
+
+    a = numpy.ones(100)
+    b = numpy.ones(100)
+    c = numpy.empty(100)
+
+    kapi.call_kernel(vecadd, kapi.Range(100), a, b, c)
+
+    assert numpy.allclose(c, a + b)
+
+
+def test_range_kernel_call2D():
+    def vecadd(item: kapi.Item, a, b, c):
+        idx = item.get_id(0)
+        jdx = item.get_id(1)
+        c[idx, jdx] = a[idx, jdx] + b[idx, jdx]
+
+    a = numpy.ones((10, 10))
+    b = numpy.ones((10, 10))
+    c = numpy.empty((10, 10))
+
+    kapi.call_kernel(vecadd, kapi.Range(10, 10), a, b, c)
+
+    assert numpy.allclose(c, a + b)
+
+
+def test_range_kernel_call3D():
+    def vecadd(item: kapi.Item, a, b, c):
+        idx = item.get_id(0)
+        jdx = item.get_id(1)
+        kdx = item.get_id(2)
+        c[idx, jdx, kdx] = a[idx, jdx, kdx] + b[idx, jdx, kdx]
+
+    a = numpy.ones((5, 5, 5))
+    b = numpy.ones((5, 5, 5))
+    c = numpy.empty((5, 5, 5))
+
+    kapi.call_kernel(vecadd, kapi.Range(5, 5, 5), a, b, c)
+
+    assert numpy.allclose(c, a + b)