Add DpctlMDLocalAccessorType

Author: ZzEeKkAa

numba_dpex/core/types/kernel_api/local_accessor.py

@@ -10,6 +10,15 @@
 from numba_dpex.utils import address_space as AddressSpace
 
 
+class DpctlMDLocalAccessorType(Type):
+    """numba-dpex internal type to represent a dpctl SyclInterface type
+    `MDLocalAccessorTy`.
+    """
+
+    def __init__(self):
+        super().__init__(name="DpctlMDLocalAccessor")
+
+
 class LocalAccessorType(USMNdArray):
     """numba-dpex internal type to represent a Python object of
     :class:`numba_dpex.experimental.kernel_iface.LocalAccessor`.

numba_dpex/core/utils/kernel_flattened_args_builder.py

@@ -11,12 +11,15 @@
 
 import dpctl
 from llvmlite import ir as llvmir
-from numba.core import types
+from numba.core import cgutils, types
 from numba.core.cpu import CPUContext
 
 from numba_dpex import utils
 from numba_dpex.core.types import USMNdArray
-from numba_dpex.core.types.kernel_api.local_accessor import LocalAccessorType
+from numba_dpex.core.types.kernel_api.local_accessor import (
+    DpctlMDLocalAccessorType,
+    LocalAccessorType,
+)
 from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum
 
 
@@ -120,40 +123,6 @@ def print_kernel_arg_list(self) -> None:
         for karg in args_list:
             print(f"    {karg.llvm_val} of typeid {karg.typeid}")
 
-    def _allocate_local_accessor_metadata_struct(self):
-        """Allocates a struct into the current function to store the metadata
-        that should be passed to libsyclinterface to allocate a
-        sycl::local_accessor object. The constructor of the sycl::local_accessor
-        class is: local_accessor<Ty, Ndim>(range<Ndims> r).
-
-        For this reason, the struct is allocated as:
-
-        LOCAL_ACCESSOR_MDSTRUCT_TYPE = llvmir.LiteralStructType(
-                [
-                    llvmir.IntType(64), # Ndim (0..3]
-                    llvmir.IntType(32), # typeid
-                    llvmir.IntType(64), # Dim0 extent
-                    llvmir.IntType(64), # Dim1 extent or NULL
-                    llvmir.IntType(64), # Dim2 extent or NULL
-                ]
-            )
-        """
-        local_accessor_mdstruct_type = llvmir.LiteralStructType(
-            [
-                llvmir.IntType(64),
-                llvmir.IntType(32),
-                llvmir.IntType(64),
-                llvmir.IntType(64),
-                llvmir.IntType(64),
-            ]
-        )
-
-        struct_ref = None
-        with self._builder.goto_entry_block():
-            struct_ref = self._builder.alloca(typ=local_accessor_mdstruct_type)
-
-        return struct_ref
-
     def _build_arg(self, llvm_val, numba_type):
         """Returns a KernelArg to be passed to a DPCTLQueue_Submit call.
 
@@ -250,7 +219,7 @@ def _store_val_into_struct(self, struct_ref, index, val):
         )
 
     def _build_local_accessor_metadata_arg(
-        self, llvm_val, arg_type, data_attr_ty
+        self, llvm_val, arg_type: LocalAccessorType, data_attr_ty
     ):
         """Handles the special case of building the kernel argument for the data
         attribute of a kernel_api.LocalAccessor object.
@@ -267,91 +236,27 @@ def _build_local_accessor_metadata_arg(
         handle proper device memory allocation.
         """
 
-        kernel_data_model = self._kernel_dmm.lookup(arg_type)
-        host_data_model = self._context.data_model_manager.lookup(arg_type)
-        shape_member = kernel_data_model.get_member_fe_type("shape")
-        shape_member_pos = host_data_model.get_field_position("shape")
-        ndim = shape_member.count
-
-        mdstruct_ref = self._allocate_local_accessor_metadata_struct()
+        ndim = arg_type.ndim
 
-        # Store the number of dimensions in the local accessor
-        self._store_val_into_struct(
-            mdstruct_ref,
-            index=0,
-            val=self._context.get_constant(types.int64, ndim),
-        )
-        # Get the underlying dtype of the data (a CPointer) attribute of a
-        # local_accessor object
-        self._store_val_into_struct(
-            mdstruct_ref,
-            index=1,
-            val=numba_type_to_dpctl_typenum(self._context, data_attr_ty.dtype),
-        )
-        # Extract and store the shape values from array into mdstruct
-        shape_attr = self._builder.gep(
-            llvm_val,
-            [
-                self._context.get_constant(types.int32, 0),
-                self._context.get_constant(types.int32, shape_member_pos),
-            ],
-        )
-        # Store the extent of the 1st dimension of the local accessor
-        dim0_shape_ext = self._builder.gep(
-            shape_attr,
-            [
-                self._context.get_constant(types.int32, 0),
-                self._context.get_constant(types.int32, 0),
-            ],
+        md_proxy = cgutils.create_struct_proxy(DpctlMDLocalAccessorType())(
+            self._context,
+            self._builder,
         )
-        self._store_val_into_struct(
-            mdstruct_ref,
-            index=2,
-            val=self._builder.load(dim0_shape_ext),
+        la_proxy = cgutils.create_struct_proxy(arg_type)(
+            self._context, self._builder, value=self._builder.load(llvm_val)
         )
 
-        if ndim == 2:
-            dim1_shape_ext = self._builder.gep(
-                shape_attr,
-                [
-                    self._context.get_constant(types.int32, 0),
-                    self._context.get_constant(types.int32, 1),
-                ],
-            )
-            self._store_val_into_struct(
-                mdstruct_ref,
-                index=3,
-                val=self._builder.load(dim1_shape_ext),
-            )
-        else:
-            self._store_val_into_struct(
-                mdstruct_ref,
-                index=3,
-                val=self._context.get_constant(types.int64, 1),
-            )
-
-        if ndim == 3:
-            dim2_shape_ext = self._builder.gep(
-                shape_attr,
-                [
-                    self._context.get_constant(types.int32, 0),
-                    self._context.get_constant(types.int32, 2),
-                ],
-            )
-            self._store_val_into_struct(
-                mdstruct_ref,
-                index=4,
-                val=self._builder.load(dim2_shape_ext),
-            )
-        else:
-            self._store_val_into_struct(
-                mdstruct_ref,
-                index=4,
-                val=self._context.get_constant(types.int64, 1),
-            )
+        md_proxy.ndim = self._context.get_constant(types.int64, ndim)
+        md_proxy.dpctl_type_id = numba_type_to_dpctl_typenum(
+            self._context, data_attr_ty.dtype
+        )
+        for i, val in enumerate(
+            cgutils.unpack_tuple(self._builder, la_proxy.shape)
+        ):
+            setattr(md_proxy, f"dim{i}", val)
 
         return self._build_arg(
-            llvm_val=mdstruct_ref,
+            llvm_val=md_proxy._getpointer(),
             numba_type=LocalAccessorType(
                 ndim, dpctl.tensor.dtype(data_attr_ty.dtype.name)
             ),

numba_dpex/experimental/models.py

@@ -19,7 +19,10 @@
 )
 
 from ..core.types.kernel_api.atomic_ref import AtomicRefType
-from ..core.types.kernel_api.local_accessor import LocalAccessorType
+from ..core.types.kernel_api.local_accessor import (
+    DpctlMDLocalAccessorType,
+    LocalAccessorType,
+)
 from .types import KernelDispatcherType
 
 
@@ -45,6 +48,26 @@ def __init__(self, dmm, fe_type):
         super().__init__(dmm, fe_type, members)
 
 
+class DpctlMDLocalAccessorModel(StructModel):
+    """Data model to represent DpctlMDLocalAccessorType.
+
+    Must be the same structure as
+    dpctl/syclinterface/dpctl_sycl_queue_interface.h::MDLocalAccessor.
+
+    Structure intended to be used only on host side of the kernel call.
+    """
+
+    def __init__(self, dmm, fe_type):
+        members = [
+            ("ndim", types.size_t),
+            ("dpctl_type_id", types.int32),
+            ("dim0", types.size_t),
+            ("dim1", types.size_t),
+            ("dim2", types.size_t),
+        ]
+        super().__init__(dmm, fe_type, members)
+
+
 def _init_exp_data_model_manager() -> DataModelManager:
     """Initializes a DpexExpKernelTarget-specific data model manager.
 
@@ -89,6 +112,9 @@ def _init_exp_data_model_manager() -> DataModelManager:
 # Register the NdItemType type
 register_model(NdItemType)(EmptyStructModel)
 
+# 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)