Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactors and removes core.kernel interface #1419

Merged
merged 4 commits into from
Apr 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion numba_dpex/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
137 changes: 0 additions & 137 deletions numba_dpex/core/kernel_interface/arrayobj.py

This file was deleted.

Original file line number Diff line number Diff line change
@@ -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.
"""
4 changes: 2 additions & 2 deletions numba_dpex/dpnp_iface/arrayobj.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
181 changes: 139 additions & 42 deletions numba_dpex/kernel_api_impl/spirv/arrayobj.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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)

Expand Down Expand Up @@ -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)
Loading
Loading