Add PrivateArray kernel_api

Author: ZzEeKkAa

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

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_private_array_overloads.py

@@ -0,0 +1,152 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Implements the SPIR-V overloads for the kernel_api.AtomicRef class methods.
+"""
+
+import operator
+from functools import reduce
+
+from numba.core import cgutils, errors, types
+from numba.core.cpu import CPUContext
+from numba.extending import intrinsic, overload
+from numba.np.arrayobj import get_itemsize as get_itemsize_np
+
+from numba_dpex.core.types import USMNdArray
+from numba_dpex.kernel_api import PrivateArray
+from numba_dpex.ocl.oclimpl import _get_target_data
+from numba_dpex.utils import address_space as AddressSpace
+
+from ..target import DPEX_KERNEL_EXP_TARGET_NAME
+
+
+def get_itemsize_spirv(context, array_type):
+    """
+    Return the item size for the given array or buffer type.
+    """
+    targetdata = _get_target_data(context)
+    lldtype = context.get_data_type(array_type.dtype)
+    return lldtype.get_abi_size(targetdata)
+
+
+def require_literal(literal_type):
+    if not hasattr(literal_type, "__len__"):
+        if not isinstance(literal_type, types.Literal):
+            raise errors.TypingError("requires literal type")
+        return
+
+    for i in range(len(literal_type)):
+        if not isinstance(literal_type[i], types.Literal):
+            raise errors.TypingError("requires literal type")
+
+
+@intrinsic(target=DPEX_KERNEL_EXP_TARGET_NAME)
+def _intrinsic_private_array_ctor(ty_context, ty_shape, ty_dtype):
+    require_literal(ty_shape)
+
+    if not isinstance(ty_dtype, types.DType):
+        raise errors.TypingError("Second argument must be instance of DType")
+
+    ndim = 1
+    if hasattr(ty_shape, "__len__"):
+        ndim = len(ty_shape)
+
+    ty_array = USMNdArray(
+        dtype=ty_dtype.dtype,
+        ndim=ndim,
+        layout="C",
+        addrspace=AddressSpace.PRIVATE,
+    )
+
+    sig = ty_array(ty_shape, ty_dtype)
+
+    def codegen(context: CPUContext, builder, sig, args):
+        shape = args[0]
+        ty_shape = sig.args[0]
+        ty_array = sig.return_type
+        # Create array object
+        ary = context.make_array(ty_array)(context, builder)
+
+        itemsize = get_itemsize_spirv(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()
+
+        dataptr = cgutils.alloca_once(
+            builder,
+            context.get_data_type(ty_array.dtype),
+            size=reduce(operator.mul, [s.literal_value for s in ty_shape]),
+        )
+
+        context.populate_array(
+            ary,
+            data=dataptr,
+            shape=shapes,
+            strides=strides,
+            itemsize=ll_itemsize,
+        )
+
+        return ary._getvalue()
+
+    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.AtomicRef`.
+
+    Raises:
+        errors.TypingError: If the `ref` argument is not a UsmNdArray type.
+        errors.TypingError: If the dtype of the `ref` is not supported in an
+        AtomicRef.
+        errors.TypingError: If the device does not support atomic operations on
+        the dtype of the `ref`.
+        errors.TypingError: If the `memory_order`, `address_type`, or
+        `memory_scope` arguments could not be parsed as integer literals.
+        errors.TypingError: If the `address_space` argument is different from
+        the address space attribute of the `ref` argument.
+        errors.TypingError: If the address space is PRIVATE.
+
+    """
+
+    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

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",
 ]

numba_dpex/kernel_api/private_array.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Implements a Python analogue to SYCL's local_accessor class. The class is
+intended to be used in pure Python code when prototyping a kernel function
+and to be passed to an actual kernel function for local memory allocation.
+"""
+
+from dpctl.tensor import usm_ndarray
+
+KernelUseOnlyError = NotImplementedError("Only for use inside kernel")
+
+
+class PrivateArray(usm_ndarray):
+    """
+    The ``LocalAccessor`` class is analogous to SYCL's ``local_accessor``
+    class. The class acts a s proxy to allocating device local memory and
+    accessing that memory from within a :func:`numba_dpex.kernel` decorated
+    function.
+    """
+
+    def __init__(self, shape, dtype) -> None:
+        """Creates a new LocalAccessor instance of the given shape and dtype."""
+
+        raise KernelUseOnlyError
+
+    def __getitem__(self, idx_obj):
+        """Returns the value stored at the position represented by idx_obj in
+        the self._data ndarray.
+        """
+
+        raise KernelUseOnlyError
+
+    def __setitem__(self, idx_obj, val):
+        """Assigns a new value to the position represented by idx_obj in
+        the self._data ndarray.
+        """
+
+        raise KernelUseOnlyError

numba_dpex/tests/experimental/test_private_array.py

@@ -0,0 +1,44 @@
+import dpnp
+import numpy as np
+import pytest
+
+import numba_dpex.experimental as dpex_exp
+from numba_dpex.kernel_api import Item, PrivateArray, Range
+
+
+@dpex_exp.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
+
+    a[i] = 0
+    for j in range(10):
+        a[i] += p[j]
+
+
+@dpex_exp.kernel
+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
+
+    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]
+)
+def test_private_array(kernel):
+    a = dpnp.empty(10, dtype=dpnp.float32)
+    dpex_exp.call_kernel(kernel, Range(a.size), a)
+
+    want = np.full(a.size, (0 + 9) * 10 / 2, dtype=np.float32)
+
+    assert np.array_equal(want, a.asnumpy())