diff --git a/numba_dpex/experimental/__init__.py b/numba_dpex/experimental/__init__.py index aa5b4a9646..f764e40843 100644 --- a/numba_dpex/experimental/__init__.py +++ b/numba_dpex/experimental/__init__.py @@ -17,6 +17,7 @@ _atomic_ref_overloads, _group_barrier_overloads, _index_space_id_overloads, + _private_array_overloads, ) from .decorators import device_func, kernel from .launcher import call_kernel, call_kernel_async diff --git a/numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_private_array_overloads.py b/numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_private_array_overloads.py new file mode 100644 index 0000000000..66f7b7e835 --- /dev/null +++ b/numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_private_array_overloads.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: 2024 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Implements the SPIR-V overloads for the kernel_api.PrivateArray class. +""" + + +import llvmlite.ir as llvmir +from llvmlite.ir.builder import IRBuilder +from numba.core.typing.npydecl import parse_dtype as _ty_parse_dtype +from numba.core.typing.npydecl import parse_shape as _ty_parse_shape +from numba.core.typing.templates import Signature +from numba.extending import intrinsic, overload + +from numba_dpex.core.types import USMNdArray +from numba_dpex.experimental.target import DpexExpKernelTypingContext +from numba_dpex.kernel_api import PrivateArray +from numba_dpex.kernel_api_impl.spirv.arrayobj import ( + make_spirv_generic_array_on_stack, + require_literal, +) +from numba_dpex.utils import address_space as AddressSpace + +from ..target import DPEX_KERNEL_EXP_TARGET_NAME + + +@intrinsic(target=DPEX_KERNEL_EXP_TARGET_NAME) +def _intrinsic_private_array_ctor( + ty_context, ty_shape, ty_dtype # pylint: disable=unused-argument +): + require_literal(ty_shape) + + ty_array = USMNdArray( + dtype=_ty_parse_dtype(ty_dtype), + ndim=_ty_parse_shape(ty_shape), + layout="C", + addrspace=AddressSpace.PRIVATE, + ) + + sig = ty_array(ty_shape, ty_dtype) + + def codegen( + context: DpexExpKernelTypingContext, + builder: IRBuilder, + sig: Signature, + args: list[llvmir.Value], + ): + shape = args[0] + ty_shape = sig.args[0] + ty_array = sig.return_type + + ary = make_spirv_generic_array_on_stack( + context, builder, ty_array, ty_shape, shape + ) + return ary._getvalue() # pylint: disable=protected-access + + return ( + sig, + codegen, + ) + + +@overload( + PrivateArray, + prefer_literal=True, + target=DPEX_KERNEL_EXP_TARGET_NAME, +) +def ol_private_array_ctor( + shape, + dtype, +): + """Overload of the constructor for the class + class:`numba_dpex.kernel_api.PrivateArray`. + + Raises: + errors.TypingError: If the shape argument is not a shape compatible + type. + errors.TypingError: If the dtype argument is not a dtype compatible + type. + """ + + def ol_private_array_ctor_impl( + shape, + dtype, + ): + # pylint: disable=no-value-for-parameter + return _intrinsic_private_array_ctor(shape, dtype) + + return ol_private_array_ctor_impl diff --git a/numba_dpex/kernel_api/__init__.py b/numba_dpex/kernel_api/__init__.py index 1172dd324c..4ff9ec742a 100644 --- a/numba_dpex/kernel_api/__init__.py +++ b/numba_dpex/kernel_api/__init__.py @@ -15,6 +15,7 @@ from .index_space_ids import Group, Item, NdItem from .launcher import call_kernel from .memory_enums import AddressSpace, MemoryOrder, MemoryScope +from .private_array import PrivateArray from .ranges import NdRange, Range __all__ = [ @@ -28,6 +29,7 @@ "Group", "NdItem", "Item", + "PrivateArray", "group_barrier", "call_kernel", ] diff --git a/numba_dpex/kernel_api/private_array.py b/numba_dpex/kernel_api/private_array.py new file mode 100644 index 0000000000..7393cc71b8 --- /dev/null +++ b/numba_dpex/kernel_api/private_array.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: 2024 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +"""Implements a simple array intended to be used inside kernel work item. +Implementation is intended to be used in pure Python code when prototyping a +kernel function. +""" + +from numpy import ndarray + + +class PrivateArray: + """ + The ``PrivateArray`` class is an simple version of array intended to be used + inside kernel work item. + """ + + def __init__(self, shape, dtype) -> None: + """Creates a new PrivateArray instance of the given shape and dtype.""" + + self._data = ndarray(shape=shape, dtype=dtype) + + def __getitem__(self, idx_obj): + """Returns the value stored at the position represented by idx_obj in + the self._data ndarray. + """ + + return self._data[idx_obj] + + def __setitem__(self, idx_obj, val): + """Assigns a new value to the position represented by idx_obj in + the self._data ndarray. + """ + + self._data[idx_obj] = val diff --git a/numba_dpex/kernel_api_impl/spirv/arrayobj.py b/numba_dpex/kernel_api_impl/spirv/arrayobj.py new file mode 100644 index 0000000000..0b0b5fbed9 --- /dev/null +++ b/numba_dpex/kernel_api_impl/spirv/arrayobj.py @@ -0,0 +1,117 @@ +"""Contains spriv 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_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext +from numba_dpex.ocl.oclimpl import _get_target_data + + +def get_itemsize(context: SPIRVTargetContext, array_type: types.Array): + """ + Return the item size for the given array or buffer type. + Same as numba.np.arrayobj.get_itemsize, but using spirv data. + """ + targetdata = _get_target_data(context) + lldtype = context.get_data_type(array_type.dtype) + return lldtype.get_abi_size(targetdata) + + +def require_literal(literal_type: types.Type): + """Checks if the numba type is Literal. If iterable object is passed, + checks that every element is Literal. + + Raises: + TypingError: When argument is not Iterable. + """ + if not hasattr(literal_type, "__len__"): + if not isinstance(literal_type, types.Literal): + raise errors.TypingError("requires literal type") + return + + for i, _ in enumerate(literal_type): + if not isinstance(literal_type[i], types.Literal): + raise errors.TypingError("requires literal type") + + +def make_spirv_array( # pylint: disable=too-many-arguments + context: SPIRVTargetContext, + 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.""" + # Create array object + ary = context.make_array(ty_array)(context, builder) + + itemsize = get_itemsize(context, ty_array) + ll_itemsize = cgutils.intp_t(itemsize) + + if isinstance(ty_shape, types.BaseTuple): + shapes = cgutils.unpack_tuple(builder, shape) + else: + ty_shape = (ty_shape,) + shapes = (shape,) + shapes = [ + context.cast(builder, value, fromty, types.intp) + for fromty, value in zip(ty_shape, shapes) + ] + + off = ll_itemsize + strides = [] + if ty_array.layout == "F": + for s in shapes: + strides.append(off) + off = builder.mul(off, s) + else: + for s in reversed(shapes): + strides.append(off) + off = builder.mul(off, s) + strides.reverse() + + context.populate_array( + ary, + data=data, + shape=shapes, + strides=strides, + itemsize=ll_itemsize, + ) + + 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) diff --git a/numba_dpex/tests/experimental/test_private_array.py b/numba_dpex/tests/experimental/test_private_array.py new file mode 100644 index 0000000000..fcbf69b825 --- /dev/null +++ b/numba_dpex/tests/experimental/test_private_array.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: 2024 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +import dpnp +import numpy as np +import pytest + +import numba_dpex.experimental as dpex_exp +from numba_dpex.kernel_api import Item, PrivateArray, Range +from numba_dpex.kernel_api import call_kernel as kapi_call_kernel + + +def private_array_kernel(item: Item, a): + i = item.get_linear_id() + p = PrivateArray(10, a.dtype) + + for j in range(10): + p[j] = j * j + + a[i] = 0 + for j in range(10): + a[i] += p[j] + + +def private_2d_array_kernel(item: Item, a): + i = item.get_linear_id() + p = PrivateArray(shape=(5, 2), dtype=a.dtype) + + for j in range(10): + p[j % 5, j // 5] = j * j + + a[i] = 0 + for j in range(10): + a[i] += p[j % 5, j // 5] + + +@pytest.mark.parametrize( + "kernel", [private_array_kernel, private_2d_array_kernel] +) +@pytest.mark.parametrize( + "call_kernel, decorator", + [(dpex_exp.call_kernel, dpex_exp.kernel), (kapi_call_kernel, lambda a: a)], +) +def test_private_array(call_kernel, decorator, kernel): + kernel = decorator(kernel) + + a = dpnp.empty(10, dtype=dpnp.float32) + call_kernel(kernel, Range(a.size), a) + + # sum of squares from 1 to n: n*(n+1)*(2*n+1)/6 + want = np.full(a.size, (9) * (9 + 1) * (2 * 9 + 1) / 6, dtype=np.float32) + + assert np.array_equal(want, a.asnumpy())