Update LocalAccessor host model that contains only shape

Author: ZzEeKkAa

numba_dpex/core/types/kernel_api/local_accessor.py

@@ -2,8 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from numba.core.pythonapi import unbox
-from numba.core.types import Array, Type
+from numba.core import cgutils
+from numba.core.types import Type, UniTuple, intp
+from numba.extending import NativeValue, unbox
 from numba.np import numpy_support
 
 from numba_dpex.core.types import USMNdArray
@@ -69,7 +70,15 @@ def unbox_local_accessor(typ, obj, c):  # pylint: disable=unused-argument
     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)
+    shape = c.pyapi.object_getattr_string(obj, "_shape")
+    local_accessor = cgutils.create_struct_proxy(typ)(c.context, c.builder)
+
+    ty_unituple = UniTuple(intp, typ.ndim)
+    ll_shape = c.unbox(ty_unituple, shape)
+    local_accessor.shape = ll_shape.value
+
+    return NativeValue(
+        c.builder.load(local_accessor._getpointer()),
+        is_error=ll_shape.is_error,
+        cleanup=ll_shape.cleanup,
+    )

numba_dpex/core/utils/kernel_flattened_args_builder.py

@@ -7,6 +7,8 @@
 object.
 """
 
+from functools import reduce
+from math import ceil
 from typing import NamedTuple
 
 import dpctl
@@ -76,7 +78,13 @@ def add_argument(
         arg_packed_llvm_val,
     ):
         """Add flattened representation of a kernel argument."""
-        if isinstance(arg_type, USMNdArray):
+        if isinstance(arg_type, LocalAccessorType):
+            self._kernel_arg_list.extend(
+                self._build_local_accessor_arg(
+                    arg_type, llvm_val=arg_packed_llvm_val
+                )
+            )
+        elif isinstance(arg_type, USMNdArray):
             self._kernel_arg_list.extend(
                 self._build_array_arg(
                     arg_type, llvm_array_val=arg_packed_llvm_val
@@ -262,6 +270,77 @@ def _build_local_accessor_metadata_arg(
             ),
         )
 
+    def _build_local_accessor_arg(self, arg_type: LocalAccessorType, llvm_val):
+        """Creates a list of kernel LLVM Values for an unpacked USMNdArray
+        kernel argument from the local accessor.
+
+        Method generates UsmNdArray fields from local accessor type and value.
+        """
+        # TODO: move extra values build on device side of codegen.
+        ndim = arg_type.ndim
+        la_proxy = cgutils.create_struct_proxy(arg_type)(
+            self._context, self._builder, value=self._builder.load(llvm_val)
+        )
+        shape = cgutils.unpack_tuple(self._builder, la_proxy.shape)
+        ll_size = reduce(self._builder.mul, shape)
+
+        size_ptr = cgutils.alloca_once_value(self._builder, ll_size)
+        itemsize = self._context.get_constant(
+            types.intp, ceil(arg_type.dtype.bitwidth / types.byte.bitwidth)
+        )
+        itemsize_ptr = cgutils.alloca_once_value(self._builder, itemsize)
+
+        kernel_arg_list = []
+
+        kernel_dm = self._kernel_dmm.lookup(arg_type)
+
+        kernel_arg_list.extend(
+            self._build_arg(
+                llvm_val=size_ptr,
+                numba_type=kernel_dm.get_member_fe_type("nitems"),
+            )
+        )
+
+        # Argument itemsize
+        kernel_arg_list.extend(
+            self._build_arg(
+                llvm_val=itemsize_ptr,
+                numba_type=kernel_dm.get_member_fe_type("itemsize"),
+            )
+        )
+
+        # Argument data
+        data_attr_ty = kernel_dm.get_member_fe_type("data")
+
+        kernel_arg_list.extend(
+            self._build_local_accessor_metadata_arg(
+                llvm_val=llvm_val,
+                arg_type=arg_type,
+                data_attr_ty=data_attr_ty,
+            )
+        )
+
+        # Arguments for shape
+        for val in shape:
+            shape_ptr = cgutils.alloca_once_value(self._builder, val)
+            kernel_arg_list.extend(
+                self._build_arg(
+                    llvm_val=shape_ptr,
+                    numba_type=types.int64,
+                )
+            )
+
+        # Arguments for strides
+        for i in range(ndim):
+            kernel_arg_list.extend(
+                self._build_arg(
+                    llvm_val=itemsize_ptr,
+                    numba_type=types.int64,
+                )
+            )
+
+        return kernel_arg_list
+
     def _build_array_arg(self, arg_type, llvm_array_val):
         """Creates a list of LLVM Values for an unpacked USMNdArray kernel
         argument.
@@ -290,22 +369,13 @@ def _build_array_arg(self, arg_type, llvm_array_val):
         data_attr_pos = host_data_model.get_field_position("data")
         data_attr_ty = kernel_data_model.get_member_fe_type("data")
 
-        if isinstance(arg_type, LocalAccessorType):
-            kernel_arg_list.extend(
-                self._build_local_accessor_metadata_arg(
-                    llvm_val=llvm_array_val,
-                    arg_type=arg_type,
-                    data_attr_ty=data_attr_ty,
-                )
-            )
-        else:
-            kernel_arg_list.extend(
-                self._build_collections_attr_arg(
-                    llvm_val=llvm_array_val,
-                    attr_index=data_attr_pos,
-                    attr_type=data_attr_ty,
-                )
+        kernel_arg_list.extend(
+            self._build_collections_attr_arg(
+                llvm_val=llvm_array_val,
+                attr_index=data_attr_pos,
+                attr_type=data_attr_ty,
             )
+        )
         # Arguments for shape
         kernel_arg_list.extend(
             self._build_unituple_member_arg(

numba_dpex/experimental/models.py

@@ -8,10 +8,11 @@
 
 from numba.core import types
 from numba.core.datamodel import DataModelManager, models
-from numba.core.datamodel.models import ArrayModel, StructModel
+from numba.core.datamodel.models import StructModel
 from numba.core.extending import register_model
 
 import numba_dpex.core.datamodel.models as dpex_core_models
+from numba_dpex.core.datamodel.models import USMArrayDeviceModel
 from numba_dpex.core.types.kernel_api.index_space_ids import (
     GroupType,
     ItemType,
@@ -68,6 +69,17 @@ def __init__(self, dmm, fe_type):
         super().__init__(dmm, fe_type, members)
 
 
+class LocalAccessorModel(StructModel):
+    """Data model for the LocalAccessor type when used in a host-only function."""
+
+    def __init__(self, dmm, fe_type):
+        ndim = fe_type.ndim
+        members = [
+            ("shape", types.UniTuple(types.intp, ndim)),
+        ]
+        super().__init__(dmm, fe_type, members)
+
+
 def _init_exp_data_model_manager() -> DataModelManager:
     """Initializes a DpexExpKernelTarget-specific data model manager.
 
@@ -84,7 +96,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 LocalAccessorType type
+    dmm.register(LocalAccessorType, USMArrayDeviceModel)
 
     # Register the GroupType type
     dmm.register(GroupType, EmptyStructModel)
@@ -115,6 +128,5 @@ def _init_exp_data_model_manager() -> DataModelManager:
 # Register the MDLocalAccessorType type
 register_model(DpctlMDLocalAccessorType)(DpctlMDLocalAccessorModel)
 
-# 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)
+# Register the LocalAccessorType type
+register_model(LocalAccessorType)(LocalAccessorModel)

numba_dpex/experimental/typeof.py

@@ -97,4 +97,4 @@ def typeof_local_accessor(val: LocalAccessor, c) -> LocalAccessorType:
     Returns: LocalAccessorType object corresponding to the LocalAccessor object.
     """
     # pylint: disable=protected-access
-    return LocalAccessorType(ndim=val._data.ndim, dtype=val._data.dtype)
+    return LocalAccessorType(ndim=len(val._shape), dtype=val._dtype)

numba_dpex/kernel_api/local_accessor.py

@@ -35,19 +35,17 @@ def __init__(self, shape, dtype) -> None:
             if hasattr(shape, "tolist"):
                 fn = getattr(shape, "tolist")
                 if callable(fn):
-                    self._shape = shape.tolist()
+                    self._shape = tuple(shape.tolist())
             else:
                 try:
-                    self._shape = [
-                        shape,
-                    ]
+                    self._shape = (shape,)
                 except Exception as e:
                     raise TypeError(
                         "Argument shape must a non-negative integer, "
                         "or a list/tuple of such integers."
                     ) from e
         else:
-            self._shape = list(shape)
+            self._shape = tuple(shape)
 
         # Make sure shape is made up a supported types
         if not self._verify_positive_integral_list(self._shape):
@@ -118,7 +116,9 @@ class is designed in a way to not have any data container backing up the
     """
 
     def __init__(self, local_accessor: LocalAccessor):
-        self._data = local_accessor._data
+        self._data = numpy.empty(
+            local_accessor._shape, dtype=local_accessor._dtype
+        )
 
     def __getitem__(self, idx_obj):
         """Returns the value stored at the position represented by idx_obj in