Add item and nd_item style of kernel programing

Author: ZzEeKkAa

numba_dpex/experimental/__init__.py

@@ -11,7 +11,10 @@
 # Temporary so that Range and NdRange work in experimental call_kernel
 from numba_dpex.core.boxing import *
 
-from ._kernel_dpcpp_spirv_overloads import _atomic_ref_overloads
+from ._kernel_dpcpp_spirv_overloads import (
+    _atomic_ref_overloads,
+    _index_space_id_overloads,
+)
 from .decorators import device_func, kernel
 from .kernel_dispatcher import KernelDispatcher
 from .launcher import call_kernel, call_kernel_async

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_index_space_id_overloads.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Implements the SPIR-V overloads for the kernel_api.items class methods.
+"""
+
+from numba.core import types
+from numba.core.errors import TypingError
+from numba.extending import intrinsic, overload_method
+
+from numba_dpex.experimental.core.types.kernel_api.items import (
+    ItemType,
+    NdItemType,
+)
+from numba_dpex.ocl._declare_function import _declare_function
+
+from ..target import DPEX_KERNEL_EXP_TARGET_NAME
+
+
+@intrinsic(target=DPEX_KERNEL_EXP_TARGET_NAME)
+def _intrinsic_get_global_id(
+    ty_context, ty_dim  # pylint: disable=unused-argument
+):
+    """Generates instruction for spirv i64 get_global_id(i32) call."""
+    sig = types.int64(types.int32)
+
+    def _intrinsic_exchange_gen(context, builder, sig, args):
+        [dim] = args
+        get_global_id = _declare_function(
+            # unsigned int - is what demangler returns from IR instruction
+            # generated by dpcpp. However int32 is passed as an argument.
+            # Most likely it does not matter, since argument can be from 0 to 3.
+            # TODO: https://github.com/IntelPython/numba-dpex/issues/936
+            # Numba generates unnecessary checks because of the type mismatch.
+            context,
+            builder,
+            "get_global_id",
+            sig,
+            ["unsigned int"],
+        )
+        res = builder.call(get_global_id, [dim])
+        return context.cast(builder, res, types.uintp, types.intp)
+
+    return sig, _intrinsic_exchange_gen
+
+
+@overload_method(ItemType, "get_id", target=DPEX_KERNEL_EXP_TARGET_NAME)
+def ol_item_get_id(item, dim):
+    """SPIR-V overload for
+    :meth:`numba_dpex.kernel_api.Item.get_id`.
+
+    Generates the same LLVM IR instruction as dpcpp for the
+    `sycl::item::get_id` function.
+
+    Raises:
+        TypingError: When argument is not an integer.
+    """
+    if not isinstance(item, ItemType):
+        raise TypingError("Only Item is supported")
+
+    if not isinstance(dim, types.Integer):
+        raise TypingError("Only integers supported")
+
+    # pylint: disable=unused-argument
+    def ol_get_id(item, dim):
+        # pylint: disable=no-value-for-parameter
+        return _intrinsic_get_global_id(dim)
+
+    return ol_get_id
+
+
+@overload_method(
+    NdItemType, "get_global_id", target=DPEX_KERNEL_EXP_TARGET_NAME
+)
+def ol_nd_item_get_global_id(nd_item, dim):
+    """SPIR-V overload for
+    :meth:`numba_dpex.kernel_api.NdItem.get_global_id`.
+
+    Generates the same LLVM IR instruction as dpcpp for the
+    `sycl::nd_item::get_global_id` function.
+
+    Raises:
+        TypingError: When argument is not an integer.
+    """
+    if not isinstance(nd_item, NdItemType):
+        raise TypingError("Only NdItem is supported")
+
+    if not isinstance(dim, types.Integer):
+        raise TypingError("Only integers supported")
+
+    # pylint: disable=unused-argument
+    def ol_get_global_id(nd_item, dim):
+        # pylint: disable=no-value-for-parameter
+        return _intrinsic_get_global_id(dim)
+
+    return ol_get_global_id

numba_dpex/experimental/core/__init__.py

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

numba_dpex/experimental/core/types/__init__.py

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

numba_dpex/experimental/core/types/kernel_api/__init__.py

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

numba_dpex/experimental/core/types/kernel_api/items.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Defines numba types for Item and NdItem classes"""
+
+from numba.core import errors, types
+
+
+class ItemType(types.Type):
+    """Numba-dpex type corresponding to :class:`numba_dpex.kernel_api.Item`"""
+
+    def __init__(self, ndim: int):
+        self._ndim = ndim
+        if ndim < 1 or ndim > 3:
+            raise errors.TypingError(
+                "ItemType can only have 1, 2 or 3 dimensions"
+            )
+        super().__init__(name="Item<" + str(ndim) + ">")
+
+    @property
+    def ndim(self):
+        """Returns number of dimensions"""
+        return self._ndim
+
+    @property
+    def key(self):
+        """Numba type specific overload"""
+        return self._ndim
+
+    def cast_python_value(self, args):
+        raise NotImplementedError
+
+
+class NdItemType(types.Type):
+    """Numba-dpex type corresponding to :class:`numba_dpex.kernel_api.NdItem`"""
+
+    def __init__(self, ndim: int):
+        self._ndim = ndim
+        if ndim < 1 or ndim > 3:
+            raise errors.TypingError(
+                "ItemType can only have 1, 2 or 3 dimensions"
+            )
+        super().__init__(name="NdItem<" + str(ndim) + ">")
+
+    @property
+    def ndim(self):
+        """Returns number of dimensions"""
+        return self._ndim
+
+    @property
+    def key(self):
+        """Numba type specific overload"""
+        return self._ndim
+
+    def cast_python_value(self, args):
+        raise NotImplementedError

numba_dpex/experimental/launcher.py

@@ -6,6 +6,8 @@
 from either CPython or a numba_dpex.dpjit decorated function.
 """
 
+import warnings
+from inspect import signature
 from typing import NamedTuple, Union
 
 import dpctl
@@ -23,6 +25,10 @@
 from numba_dpex.core.utils import kernel_launcher as kl
 from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
 from numba_dpex.dpctl_iface.wrappers import wrap_event_reference
+from numba_dpex.experimental.core.types.kernel_api.items import (
+    ItemType,
+    NdItemType,
+)
 from numba_dpex.experimental.kernel_dispatcher import (
     KernelDispatcher,
     _KernelCompileResult,
@@ -129,13 +135,30 @@ def _submit_kernel(  # pylint: disable=too-many-arguments
     else:
         sig = ty_return(ty_kernel_fn, ty_index_space, ty_kernel_args_tuple)
 
-    kernel_sig = types.void(*ty_kernel_args_tuple)
+    # Add Item/NdItem as a first argument to kernel arguments list. It is
+    # an empty struct so any other modifications at kernel submission are not
+    # needed.
+    if len(signature(ty_kernel_fn.dispatcher.py_func).parameters) > len(
+        ty_kernel_args_tuple
+    ):
+        if isinstance(ty_index_space, RangeType):
+            ty_item = ItemType(ty_index_space.ndim)
+        else:
+            ty_item = NdItemType(ty_index_space.ndim)
+
+        ty_kernel_args_tuple = (ty_item, *ty_kernel_args_tuple)
+    else:
+        warnings.warn(
+            "Kernels without item/nd_item will be not supported in the future",
+            DeprecationWarning,
+        )
+
     # ty_kernel_fn is type specific to exact function, so we can get function
     # directly from type and compile it. Thats why we don't need to get it in
     # codegen
     kernel_dispatcher: KernelDispatcher = ty_kernel_fn.dispatcher
     kcres: _KernelCompileResult = kernel_dispatcher.get_compile_result(
-        kernel_sig
+        types.void(*ty_kernel_args_tuple)  # kernel signature
     )
     kernel_module: kl.SPIRVKernelModule = kcres.kernel_device_ir_module
     kernel_targetctx: DpexKernelTargetContext = kernel_dispatcher.targetctx

numba_dpex/experimental/models.py

@@ -13,6 +13,10 @@
 from numba.core.extending import register_model
 
 import numba_dpex.core.datamodel.models as dpex_core_models
+from numba_dpex.experimental.core.types.kernel_api.items import (
+    ItemType,
+    NdItemType,
+)
 
 from .dpcpp_types import AtomicRefType
 from .literal_intenum_type import IntEnumLiteral
@@ -43,6 +47,15 @@ def __init__(self, dmm, fe_type):
         super().__init__(dmm, fe_type, be_type)
 
 
+class EmptyStructModel(StructModel):
+    """Data model that does not take space. Intended to be used with types that
+    are presented only at typing stage and not represented physically."""
+
+    def __init__(self, dmm, fe_type):
+        members = []
+        super().__init__(dmm, fe_type, members)
+
+
 def _init_exp_data_model_manager() -> DataModelManager:
     """Initializes a DpexExpKernelTarget-specific data model manager.
 
@@ -60,10 +73,22 @@ def _init_exp_data_model_manager() -> DataModelManager:
     dmm.register(IntEnumLiteral, IntEnumLiteralModel)
     dmm.register(AtomicRefType, AtomicRefModel)
 
+    # Register the ItemType type
+    dmm.register(ItemType, EmptyStructModel)
+
+    # Register the NdItemType type
+    dmm.register(NdItemType, EmptyStructModel)
+
     return dmm
 
 
 exp_dmm = _init_exp_data_model_manager()
 
 # Register any new type that should go into numba.core.datamodel.default_manager
 register_model(KernelDispatcherType)(models.OpaqueModel)
+
+# Register the ItemType type
+register_model(ItemType)(EmptyStructModel)
+
+# Register the NdItemType type
+register_model(NdItemType)(EmptyStructModel)

numba_dpex/experimental/typeof.py

@@ -9,7 +9,11 @@
 
 from numba.extending import typeof_impl
 
-from numba_dpex.kernel_api import AtomicRef
+from numba_dpex.experimental.core.types.kernel_api.items import (
+    ItemType,
+    NdItemType,
+)
+from numba_dpex.kernel_api import AtomicRef, Item, NdItem
 
 from .dpcpp_types import AtomicRefType
 
@@ -34,3 +38,33 @@ def typeof_atomic_ref(val: AtomicRef, ctx) -> AtomicRefType:
         memory_scope=val.memory_scope.value,
         address_space=val.address_space.value,
     )
+
+
+@typeof_impl.register(Item)
+def typeof_item(val: Item, c):
+    """Registers the type inference implementation function for a
+    numba_dpex.kernel_api.Item PyObject.
+
+    Args:
+        val : An instance of numba_dpex.kernel_api.Item.
+        c : Unused argument used to be consistent with Numba API.
+
+    Returns: A numba_dpex.experimental.core.types.kernel_api.items.ItemType
+        instance.
+    """
+    return ItemType(val.ndim)
+
+
+@typeof_impl.register(NdItem)
+def typeof_nditem(val, c):
+    """Registers the type inference implementation function for a
+    numba_dpex.kernel_api.NdItem PyObject.
+
+    Args:
+        val : An instance of numba_dpex.kernel_api.NdItem.
+        c : Unused argument used to be consistent with Numba API.
+
+    Returns: A numba_dpex.experimental.core.types.kernel_api.items.NdItemType
+        instance.
+    """
+    return NdItemType(val.ndim)

numba_dpex/kernel_api/__init__.py

@@ -10,6 +10,7 @@
 """
 
 from .atomic_ref import AtomicRef
+from .index_space_ids import Item, NdItem
 from .memory_enums import AddressSpace, MemoryOrder, MemoryScope
 from .ranges import NdRange, Range
 
@@ -20,4 +21,6 @@
     "MemoryScope",
     "NdRange",
     "Range",
+    "NdItem",
+    "Item",
 ]