Refactoring the indexer classes.

    - Splits up the numba_dpex.core.kernel_interface.indexers
      module and moves the Python classes into a new
      kernel_api.ranges submodule.
    - Moves the numba-dpex type definitions into
      core.types.kernel_api_types
    - Other changes as needed to get everything working

Author: Diptorup Deb

numba_dpex/__init__.py

@@ -99,16 +99,17 @@ def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
 # Re-export types itself
 import numba_dpex.core.types as types  # noqa E402
 from numba_dpex.core import config  # noqa E402
-from numba_dpex.core.kernel_interface.indexers import (  # noqa E402
-    NdRange,
-    Range,
-)
+from numba_dpex.core.kernel_interface import ranges_overloads  # noqa E402
 
 # Re-export all type names
 from numba_dpex.core.types import *  # noqa E402
 from numba_dpex.dpctl_iface import _intrinsic  # noqa E402
 from numba_dpex.dpnp_iface import dpnpimpl  # noqa E402
 
+# Importing NdRange and Range into numba_dpex for
+# backward compatibility
+from numba_dpex.kernel_api import NdRange, Range  # noqa E402
+
 from .core.targets import dpjit_target, kernel_target  # noqa E402
 from .decorators import dpjit, func, kernel  # noqa E402
 from .ocl.stubs import (  # noqa E402

numba_dpex/core/types/range_types.py renamed to numba_dpex/core/boxing.py

@@ -4,55 +4,12 @@
 
 from contextlib import ExitStack
 
-from numba.core import cgutils, errors, types
+from numba.core import cgutils, types
 from numba.core.datamodel import default_manager
 from numba.extending import NativeValue, box, unbox
 
-from ..kernel_interface.indexers import NdRange, Range
-
-
-class RangeType(types.Type):
-    """Numba-dpex type corresponding to
-    :class:`numba_dpex.core.kernel_interface.indexers.Range`
-    """
-
-    def __init__(self, ndim: int):
-        self._ndim = ndim
-        if ndim < 1 or ndim > 3:
-            raise errors.TypingError(
-                "RangeType can only have 1,2, or 3 dimensions"
-            )
-        super(RangeType, self).__init__(name="Range<" + str(ndim) + ">")
-
-    @property
-    def ndim(self):
-        return self._ndim
-
-    @property
-    def key(self):
-        return self._ndim
-
-
-class NdRangeType(types.Type):
-    """Numba-dpex type corresponding to
-    :class:`numba_dpex.core.kernel_interface.indexers.NdRange`
-    """
-
-    def __init__(self, ndim: int):
-        self._ndim = ndim
-        if ndim < 1 or ndim > 3:
-            raise errors.TypingError(
-                "RangeType can only have 1,2, or 3 dimensions"
-            )
-        super(NdRangeType, self).__init__(name="NdRange<" + str(ndim) + ">")
-
-    @property
-    def ndim(self):
-        return self._ndim
-
-    @property
-    def key(self):
-        return self._ndim
+from numba_dpex.core.types import NdRangeType, RangeType
+from numba_dpex.kernel_api import NdRange, Range
 
 
 @unbox(RangeType)

numba_dpex/core/datamodel/models.py

@@ -196,9 +196,7 @@ def __init__(self, dmm, fe_type):
 
 
 class RangeModel(StructModel):
-    """The native data model for a
-    numba_dpex.core.kernel_interface.indexers.Range PyObject.
-    """
+    """The numba_dpex data model for a numba_dpex.kernel_api.Range PyObject."""
 
     def __init__(self, dmm, fe_type):
         members = [
@@ -216,8 +214,8 @@ def flattened_field_count(self):
 
 
 class NdRangeModel(StructModel):
-    """The native data model for a
-    numba_dpex.core.kernel_interface.indexers.NdRange PyObject.
+    """
+    The numba_dpex data model for a numba_dpex.kernel_api.NdRange PyObject.
     """
 
     def __init__(self, dmm, fe_type):

numba_dpex/core/kernel_interface/ranges_overloads.py

@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from llvmlite import ir as llvmir
+from numba.core import cgutils, errors, types
+from numba.core.datamodel import default_manager
+from numba.extending import intrinsic, overload
+
+from numba_dpex.kernel_api import NdRange, Range
+
+# can't import name because of the circular import
+DPEX_TARGET_NAME = "dpex"
+
+
+@intrinsic(target=DPEX_TARGET_NAME)
+def _intrin_range_alloc(typingctx, ty_dim0, ty_dim1, ty_dim2, ty_range):
+    ty_retty = ty_range.instance_type
+    sig = ty_retty(
+        ty_dim0,
+        ty_dim1,
+        ty_dim2,
+        ty_range,
+    )
+
+    def codegen(context, builder, sig, args):
+        typ = sig.return_type
+        dim0, dim1, dim2, _ = args
+        range_struct = cgutils.create_struct_proxy(typ)(context, builder)
+        range_struct.dim0 = dim0
+
+        if not isinstance(sig.args[1], types.NoneType):
+            range_struct.dim1 = dim1
+        else:
+            range_struct.dim1 = llvmir.Constant(
+                llvmir.types.IntType(64), Range.UNDEFINED_DIMENSION
+            )
+
+        if not isinstance(sig.args[2], types.NoneType):
+            range_struct.dim2 = dim2
+        else:
+            range_struct.dim2 = llvmir.Constant(
+                llvmir.types.IntType(64), Range.UNDEFINED_DIMENSION
+            )
+
+        range_struct.ndim = llvmir.Constant(llvmir.types.IntType(64), typ.ndim)
+
+        return range_struct._getvalue()
+
+    return sig, codegen
+
+
+@intrinsic(target=DPEX_TARGET_NAME)
+def _intrin_ndrange_alloc(
+    typingctx, ty_global_range, ty_local_range, ty_ndrange
+):
+    ty_retty = ty_ndrange.instance_type
+    sig = ty_retty(
+        ty_global_range,
+        ty_local_range,
+        ty_ndrange,
+    )
+    range_datamodel = default_manager.lookup(ty_global_range)
+
+    def codegen(context, builder, sig, args):
+        typ = sig.return_type
+
+        global_range, local_range, _ = args
+        ndrange_struct = cgutils.create_struct_proxy(typ)(context, builder)
+        ndrange_struct.ndim = llvmir.Constant(
+            llvmir.types.IntType(64), typ.ndim
+        )
+        ndrange_struct.gdim0 = builder.extract_value(
+            global_range,
+            range_datamodel.get_field_position("dim0"),
+        )
+        ndrange_struct.gdim1 = builder.extract_value(
+            global_range,
+            range_datamodel.get_field_position("dim1"),
+        )
+        ndrange_struct.gdim2 = builder.extract_value(
+            global_range,
+            range_datamodel.get_field_position("dim2"),
+        )
+        ndrange_struct.ldim0 = builder.extract_value(
+            local_range,
+            range_datamodel.get_field_position("dim0"),
+        )
+        ndrange_struct.ldim1 = builder.extract_value(
+            local_range,
+            range_datamodel.get_field_position("dim1"),
+        )
+        ndrange_struct.ldim2 = builder.extract_value(
+            local_range,
+            range_datamodel.get_field_position("dim2"),
+        )
+
+        return ndrange_struct._getvalue()
+
+    return sig, codegen
+
+
+@overload(Range, target=DPEX_TARGET_NAME)
+def _ol_range_init(dim0, dim1=None, dim2=None):
+    """Numba overload of the Range constructor to make it usable inside an
+    njit and dpjit decorated function.
+
+    """
+    from numba_dpex.core.types import RangeType
+
+    ndims = 1
+    ty_optional_dims = (dim1, dim2)
+
+    # A Range should at least have the 0th dimension populated
+    if not isinstance(dim0, types.Integer):
+        raise errors.TypingError(
+            "Expected a Range's dimension should to be an Integer value, but "
+            "encountered " + dim0.name
+        )
+
+    for ty_dim in ty_optional_dims:
+        if isinstance(ty_dim, types.Integer):
+            ndims += 1
+        elif ty_dim is not None:
+            raise errors.TypingError(
+                "Expected a Range's dimension to be an Integer value, "
+                f"but {type(ty_dim)} was provided."
+            )
+
+    ret_ty = RangeType(ndims)
+
+    def impl(dim0, dim1=None, dim2=None):
+        return _intrin_range_alloc(dim0, dim1, dim2, ret_ty)
+
+    return impl
+
+
+@overload(NdRange, target=DPEX_TARGET_NAME)
+def _ol_ndrange_init(global_range, local_range):
+    """Numba overload of the NdRange constructor to make it usable inside an
+    njit and dpjit decorated function.
+
+    """
+    from numba_dpex.core.exceptions import UnmatchedNumberOfRangeDimsError
+    from numba_dpex.core.types import NdRangeType, RangeType
+
+    if not isinstance(global_range, RangeType):
+        raise errors.TypingError(
+            "Only global range values specified as a Range are "
+            "supported inside dpjit"
+        )
+
+    if not isinstance(local_range, RangeType):
+        raise errors.TypingError(
+            "Only local range values specified as a Range are "
+            "supported inside dpjit"
+        )
+
+    if not global_range.ndim == local_range.ndim:
+        raise UnmatchedNumberOfRangeDimsError(
+            kernel_name="",
+            global_ndims=global_range.ndim,
+            local_ndims=local_range.ndim,
+        )
+
+    ret_ty = NdRangeType(global_range.ndim)
+
+    def impl(global_range, local_range):
+        return _intrin_ndrange_alloc(global_range, local_range, ret_ty)
+
+    return impl

numba_dpex/core/types/__init__.py

@@ -5,6 +5,7 @@
 from .array_type import Array
 from .dpctl_types import DpctlSyclEvent, DpctlSyclQueue
 from .dpnp_ndarray_type import DpnpNdArray
+from .kernel_api_types.range_types import NdRangeType, RangeType
 from .numba_types_short_names import (
     b1,
     bool_,
@@ -26,7 +27,6 @@
     uint64,
     void,
 )
-from .range_types import NdRangeType, RangeType
 from .usm_ndarray_type import USMNdArray
 
 usm_ndarray = USMNdArray

numba_dpex/core/types/kernel_api_types/__init__.py

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

numba_dpex/core/types/kernel_api_types/range_types.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import ExitStack
+
+from numba.core import cgutils, errors, types
+
+
+class RangeType(types.Type):
+    """Numba-dpex type corresponding to
+    :class:`numba_dpex.kernel_api.ranges.Range`
+    """
+
+    def __init__(self, ndim: int):
+        self._ndim = ndim
+        if ndim < 1 or ndim > 3:
+            raise errors.TypingError(
+                "RangeType can only have 1,2, or 3 dimensions"
+            )
+        super(RangeType, self).__init__(name="Range<" + str(ndim) + ">")
+
+    @property
+    def ndim(self):
+        return self._ndim
+
+    @property
+    def key(self):
+        return self._ndim
+
+
+class NdRangeType(types.Type):
+    """Numba-dpex type corresponding to
+    :class:`numba_dpex.kernel_api.ranges.NdRange`
+    """
+
+    def __init__(self, ndim: int):
+        self._ndim = ndim
+        if ndim < 1 or ndim > 3:
+            raise errors.TypingError(
+                "RangeType can only have 1,2, or 3 dimensions"
+            )
+        super(NdRangeType, self).__init__(name="NdRange<" + str(ndim) + ">")
+
+    @property
+    def ndim(self):
+        return self._ndim
+
+    @property
+    def key(self):
+        return self._ndim

numba_dpex/core/typing/typeof.py

@@ -8,12 +8,12 @@
 from numba.extending import typeof_impl
 from numba.np import numpy_support
 
+from numba_dpex.kernel_api.ranges import NdRange, Range
 from numba_dpex.utils.constants import address_space
 
-from ..kernel_interface.indexers import NdRange, Range
 from ..types.dpctl_types import DpctlSyclEvent, DpctlSyclQueue
 from ..types.dpnp_ndarray_type import DpnpNdArray
-from ..types.range_types import NdRangeType, RangeType
+from ..types.kernel_api_types.range_types import NdRangeType, RangeType
 from ..types.usm_ndarray_type import USMNdArray
 
 

numba_dpex/core/utils/kernel_launcher.py

@@ -21,7 +21,10 @@
 from numba_dpex.core.exceptions import UnreachableError
 from numba_dpex.core.runtime.context import DpexRTContext
 from numba_dpex.core.types import USMNdArray
-from numba_dpex.core.types.range_types import NdRangeType, RangeType
+from numba_dpex.core.types.kernel_api_types.range_types import (
+    NdRangeType,
+    RangeType,
+)
 from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
 from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum
 from numba_dpex.utils import create_null_ptr

numba_dpex/experimental/__init__.py

@@ -8,6 +8,9 @@
 
 from numba.core.imputils import Registry
 
+# Temporary so that Range and NdRange work in experimental call_kernel
+from numba_dpex.core import boxing
+
 from ._kernel_dpcpp_spirv_overloads import _atomic_ref_overloads
 from .decorators import device_func, kernel
 from .kernel_dispatcher import KernelDispatcher