Add numba typing infrastructure for LocalAccessor

Author: Diptorup Deb

numba_dpex/core/types/kernel_api/local_accessor.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from numba.core.pythonapi import unbox
+from numba.core.types import Array, Type
+from numba.np import numpy_support
+
+from numba_dpex.core.types import USMNdArray
+from numba_dpex.utils import address_space as AddressSpace
+
+
+class LocalAccessorType(USMNdArray):
+    """numba-dpex internal type to represent a Python object of
+    :class:`numba_dpex.experimental.kernel_iface.LocalAccessor`.
+    """
+
+    def __init__(self, ndim, dtype):
+        try:
+            if isinstance(dtype, Type):
+                parsed_dtype = dtype
+            else:
+                parsed_dtype = numpy_support.from_dtype(dtype)
+        except NotImplementedError as exc:
+            raise ValueError(f"Unsupported array dtype: {dtype}") from exc
+
+        type_name = (
+            f"LocalAccessor(dtype={parsed_dtype}, ndim={ndim}, "
+            f"address_space={AddressSpace.LOCAL})"
+        )
+
+        super().__init__(
+            ndim=ndim,
+            layout="C",
+            dtype=parsed_dtype,
+            addrspace=AddressSpace.LOCAL,
+            name=type_name,
+        )
+
+    def cast_python_value(self, args):
+        """The helper function is not overloaded and using it on the
+        LocalAccessorType throws a NotImplementedError.
+        """
+        raise NotImplementedError
+
+
+@unbox(LocalAccessorType)
+def unbox_local_accessor(typ, obj, c):  # pylint: disable=unused-argument
+    """Unboxes a Python LocalAccessor PyObject* into a numba-dpex internal
+    representation.
+
+    A LocalAccessor object is represented internally in numba-dpex with the
+    same data model as a numpy.ndarray. It is done as a LocalAccessor object
+    serves only as a placeholder type when passed to ``call_kernel`` and the
+    data buffer should never be accessed inside a host-side compiled function
+    such as ``call_kernel``.
+
+    When a LocalAccessor object is passed as an argument to a kernel function
+    it uses the USMArrayDeviceModel. Doing so allows numba-dpex to correctly
+    generate the kernel signature passing in a pointer in the local address
+    space.
+    """
+
+    nparrobj = c.pyapi.object_getattr_string(obj, "data")
+    nparrtype = Array(typ.dtype, typ.ndim, typ.layout, readonly=False)
+    return c.unbox(nparrtype, nparrobj)

numba_dpex/experimental/__init__.py

@@ -12,6 +12,7 @@
 from numba_dpex.core.boxing import *
 from numba_dpex.kernel_api_impl.spirv.dispatcher import SPIRVKernelDispatcher
 
+from . import typeof
 from ._kernel_dpcpp_spirv_overloads import (
     _atomic_fence_overloads,
     _atomic_ref_overloads,

numba_dpex/experimental/models.py

@@ -8,7 +8,7 @@
 
 from numba.core import types
 from numba.core.datamodel import DataModelManager, models
-from numba.core.datamodel.models import StructModel
+from numba.core.datamodel.models import ArrayModel, StructModel
 from numba.core.extending import register_model
 
 import numba_dpex.core.datamodel.models as dpex_core_models
@@ -19,6 +19,7 @@
 )
 
 from ..core.types.kernel_api.atomic_ref import AtomicRefType
+from ..core.types.kernel_api.local_accessor import LocalAccessorType
 from .types import KernelDispatcherType
 
 
@@ -60,6 +61,8 @@ def _init_exp_data_model_manager() -> DataModelManager:
     # Register the types and data model in the DpexExpTargetContext
     dmm.register(AtomicRefType, AtomicRefModel)
 
+    dmm.register(LocalAccessorType, dpex_core_models.USMArrayDeviceModel)
+
     # Register the GroupType type
     dmm.register(GroupType, EmptyStructModel)
 
@@ -85,3 +88,7 @@ def _init_exp_data_model_manager() -> DataModelManager:
 
 # Register the NdItemType type
 register_model(NdItemType)(EmptyStructModel)
+
+# The LocalAccessorType is registered with the EmptyStructModel in the default
+# data manager so that its attributes are not accessible inside dpjit.
+register_model(LocalAccessorType)(ArrayModel)

numba_dpex/experimental/typeof.py

@@ -14,7 +14,8 @@
     ItemType,
     NdItemType,
 )
-from numba_dpex.kernel_api import AtomicRef, Group, Item, NdItem
+from numba_dpex.core.types.kernel_api.local_accessor import LocalAccessorType
+from numba_dpex.kernel_api import AtomicRef, Group, Item, LocalAccessor, NdItem
 
 from ..core.types.kernel_api.atomic_ref import AtomicRefType
 
@@ -84,3 +85,16 @@ def typeof_nditem(val: NdItem, c):
         instance.
     """
     return NdItemType(val.dimensions)
+
+
+@typeof_impl.register(LocalAccessor)
+def typeof_local_accessor(val: LocalAccessor, c) -> LocalAccessorType:
+    """Returns a ``numba_dpex.experimental.dpctpp_types.LocalAccessorType``
+    instance for a Python LocalAccessor object.
+    Args:
+        val (LocalAccessor): Instance of the LocalAccessor type.
+        c : Numba typing context used for type inference.
+    Returns: LocalAccessorType object corresponding to the LocalAccessor object.
+    """
+
+    return LocalAccessorType(ndim=val.data.ndim, dtype=val.data.dtype)