Merge pull request #1419 from IntelPython/remove/core.kernel_interface

Refactors and removes core.kernel interface

Author: Diptorup Deb

numba_dpex/__init__.py

@@ -84,7 +84,7 @@ def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
 import numba_dpex.core.types as types  # noqa E402
 from numba_dpex.core import boxing  # noqa E402
 from numba_dpex.core import config  # noqa E402
-from numba_dpex.core.kernel_interface import ranges_overloads  # noqa E402
+from numba_dpex.core.overloads import ranges_overloads  # noqa E402
 
 # Re-export all type names
 from numba_dpex.core.types import *  # noqa E402

numba_dpex/core/kernel_interface/arrayobj.py

@@ -1,6 +1,3 @@
 # SPDX-FileCopyrightText: 2020 - 2024 Intel Corporation
 #
 # SPDX-License-Identifier: Apache-2.0
-
-"""Defines the interface for kernel compilation using numba-dpex.
-"""

numba_dpex/core/kernel_interface/__init__.py renamed to numba_dpex/core/overloads/__init__.py

@@ -16,10 +16,10 @@
 from numba.np.arrayobj import make_array
 from numba.np.numpy_support import is_nonelike
 
-from numba_dpex.core.kernel_interface.arrayobj import (
+from numba_dpex.core.types import DpnpNdArray
+from numba_dpex.kernel_api_impl.spirv.arrayobj import (
     _getitem_array_generic as kernel_getitem_array_generic,
 )
-from numba_dpex.core.types import DpnpNdArray
 from numba_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext
 
 from ._intrinsic import (

numba_dpex/core/kernel_interface/ranges_overloads.py renamed to numba_dpex/core/overloads/ranges_overloads.py

@@ -1,30 +1,149 @@
+# SPDX-FileCopyrightText: 2012 - 2024 Anaconda Inc.
 # SPDX-FileCopyrightText: 2024 Intel Corporation
 #
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-License-Identifier: BSD-2-Clause
 
 """Contains SPIR-V specific array functions."""
 
-import operator
-from functools import reduce
 from typing import Union
 
 import llvmlite.ir as llvmir
 from llvmlite.ir.builder import IRBuilder
 from numba.core import cgutils, errors, types
 from numba.core.base import BaseContext
+from numba.np.arrayobj import (
+    basic_indexing,
+    get_itemsize,
+    load_item,
+    make_array,
+)
 
-from numba_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext
-from numba_dpex.ocl.oclimpl import _get_target_data
+from numba_dpex.core.types import USMNdArray
 
 
-def get_itemsize(context: SPIRVTargetContext, array_type: types.Array):
+def populate_array(
+    arraystruct, data, shape, strides, itemsize
+):  # pylint: disable=too-many-arguments,too-many-locals
     """
-    Return the item size for the given array or buffer type.
-    Same as numba.np.arrayobj.get_itemsize, but using spirv data.
+    Helper function for populating array structures.
+
+    The function is copied from upstream Numba and modified to support the
+    USMNdArray data type that uses a different data model on SYCL devices
+    than the upstream types.Array data type. USMNdArray data model does not
+    have the ``parent`` and ``meminfo`` fields. This function intended to be
+    used only in the SPIRVKernelTarget.
+
+    *shape* and *strides* can be Python tuples or LLVM arrays.
+    """
+    context = arraystruct._context  # pylint: disable=protected-access
+    builder = arraystruct._builder  # pylint: disable=protected-access
+    datamodel = arraystruct._datamodel  # pylint: disable=protected-access
+    # doesn't matter what this array type instance is, it's just to get the
+    # fields for the data model of the standard array type in this context
+    standard_array = USMNdArray(ndim=1, layout="C", dtype=types.float64)
+    standard_array_type_datamodel = context.data_model_manager[standard_array]
+    required_fields = set(standard_array_type_datamodel._fields)
+    datamodel_fields = set(datamodel._fields)
+    # Make sure that the presented array object has a data model that is
+    # close enough to an array for this function to proceed.
+    if (required_fields & datamodel_fields) != required_fields:
+        missing = required_fields - datamodel_fields
+        msg = (
+            f"The datamodel for type {arraystruct} is missing "
+            f"field{'s' if len(missing) > 1 else ''} {missing}."
+        )
+        raise ValueError(msg)
+
+    intp_t = context.get_value_type(types.intp)
+    if isinstance(shape, (tuple, list)):
+        shape = cgutils.pack_array(builder, shape, intp_t)
+    if isinstance(strides, (tuple, list)):
+        strides = cgutils.pack_array(builder, strides, intp_t)
+    if isinstance(itemsize, int):
+        itemsize = intp_t(itemsize)
+
+    attrs = {
+        "shape": shape,
+        "strides": strides,
+        "data": data,
+        "itemsize": itemsize,
+    }
+
+    # Calc num of items from shape
+    nitems = context.get_constant(types.intp, 1)
+    unpacked_shape = cgutils.unpack_tuple(builder, shape, shape.type.count)
+    # (note empty shape => 0d array therefore nitems = 1)
+    for axlen in unpacked_shape:
+        nitems = builder.mul(nitems, axlen, flags=["nsw"])
+    attrs["nitems"] = nitems
+
+    # Make sure that we have all the fields
+    got_fields = set(attrs.keys())
+    if got_fields != required_fields:
+        raise ValueError(f"missing {required_fields - got_fields}")
+
+    # Set field value
+    for k, v in attrs.items():
+        setattr(arraystruct, k, v)
+
+    return arraystruct
+
+
+def make_view(
+    context, builder, ary, return_type, data, shapes, strides
+):  # pylint: disable=too-many-arguments
+    """
+    Build a view over the given array with the given parameters.
+
+    This is analog of numpy.np.arrayobj.make_view without parent and
+    meminfo fields, because they don't make sense on device. This function
+    intended to be used only in kernel targets.
+    """
+    retary = make_array(return_type)(context, builder)
+    context.populate_array(
+        retary, data=data, shape=shapes, strides=strides, itemsize=ary.itemsize
+    )
+    return retary
+
+
+def _getitem_array_generic(
+    context, builder, return_type, aryty, ary, index_types, indices
+):  # pylint: disable=too-many-arguments
     """
-    targetdata = _get_target_data(context)
-    lldtype = context.get_data_type(array_type.dtype)
-    return lldtype.get_abi_size(targetdata)
+    Return the result of indexing *ary* with the given *indices*,
+    returning either a scalar or a view.
+
+    This is analog of numpy.np.arrayobj._getitem_array_generic without parent
+    and meminfo fields, because they don't make sense on device. This function
+    intended to be used only in kernel targets.
+    """
+    dataptr, view_shapes, view_strides = basic_indexing(
+        context,
+        builder,
+        aryty,
+        ary,
+        index_types,
+        indices,
+        boundscheck=context.enable_boundscheck,
+    )
+
+    if isinstance(return_type, types.Buffer):
+        # Build array view
+        retary = make_view(
+            context,
+            builder,
+            ary,
+            return_type,
+            dataptr,
+            view_shapes,
+            view_strides,
+        )
+        return retary._getvalue()  # pylint: disable=protected-access
+
+    # Load scalar from 0-d result
+    assert not view_shapes
+    return load_item(context, builder, aryty, dataptr)
 
 
 def require_literal(literal_type: types.Type):
@@ -46,15 +165,22 @@ def require_literal(literal_type: types.Type):
             )
 
 
-def make_spirv_array(  # pylint: disable=too-many-arguments
-    context: SPIRVTargetContext,
+def np_cfarray(  # pylint: disable=too-many-arguments
+    context: BaseContext,
     builder: IRBuilder,
     ty_array: types.Array,
     ty_shape: Union[types.IntegerLiteral, types.BaseTuple],
     shape: llvmir.Value,
     data: llvmir.Value,
 ):
-    """Makes SPIR-V array and fills it data."""
+    """Makes numpy-like array and fills it's data depending on the context's
+    datamodel.
+
+    Generic version of numba.np.arrayobj.np_cfarray so that it can be used
+    not only as intrinsic, but inside instruction generation.
+
+    TODO: upstream changes to numba.
+    """
     # Create array object
     ary = context.make_array(ty_array)(context, builder)
 
@@ -92,32 +218,3 @@ def make_spirv_array(  # pylint: disable=too-many-arguments
     )
 
     return ary
-
-
-def allocate_array_data_on_stack(
-    context: BaseContext,
-    builder: IRBuilder,
-    ty_array: types.Array,
-    ty_shape: Union[types.IntegerLiteral, types.BaseTuple],
-):
-    """Allocates flat array of given shape on the stack."""
-    if not isinstance(ty_shape, types.BaseTuple):
-        ty_shape = (ty_shape,)
-
-    return cgutils.alloca_once(
-        builder,
-        context.get_data_type(ty_array.dtype),
-        size=reduce(operator.mul, [s.literal_value for s in ty_shape]),
-    )
-
-
-def make_spirv_generic_array_on_stack(
-    context: SPIRVTargetContext,
-    builder: IRBuilder,
-    ty_array: types.Array,
-    ty_shape: Union[types.IntegerLiteral, types.BaseTuple],
-    shape: llvmir.Value,
-):
-    """Makes SPIR-V array of given shape with empty data."""
-    data = allocate_array_data_on_stack(context, builder, ty_array, ty_shape)
-    return make_spirv_array(context, builder, ty_array, ty_shape, shape, data)