feat[next]: Support for Array Api namespace as allocator

src/gt4py/next/allocators.py

@@ -60,7 +60,7 @@ def __gt_allocate__(
         dtype: core_defs.DType[core_defs.ScalarT],
         device_id: int = 0,
         aligned_index: Optional[Sequence[common.NamedIndex]] = None,  # absolute position
-    ) -> core_allocators.TensorBuffer[core_defs.DeviceTypeT, core_defs.ScalarT]: ...
+    ) -> core_defs.NDArrayObject: ...
 
 
 def is_field_allocator(obj: Any) -> TypeGuard[FieldBufferAllocatorProtocol]:
@@ -160,15 +160,15 @@ def __gt_allocate__(
         dtype: core_defs.DType[core_defs.ScalarT],
         device_id: int = 0,
         aligned_index: Optional[Sequence[common.NamedIndex]] = None,  # absolute position
-    ) -> core_allocators.TensorBuffer[core_defs.DeviceTypeT, core_defs.ScalarT]:
+    ) -> core_defs.NDArrayObject:
         shape = domain.shape
         layout_map = self.layout_mapper(domain.dims)
         # TODO(egparedes): add support for non-empty aligned index values
         assert aligned_index is None
 
         return self.buffer_allocator.allocate(
             shape, dtype, device_id, layout_map, self.byte_alignment, aligned_index
-        )
+        ).ndarray
 
 
 if TYPE_CHECKING:
@@ -242,7 +242,7 @@ def __gt_allocate__(
         dtype: core_defs.DType[core_defs.ScalarT],
         device_id: int = 0,
         aligned_index: Optional[Sequence[common.NamedIndex]] = None,  # absolute position
-    ) -> core_allocators.TensorBuffer[core_defs.DeviceTypeT, core_defs.ScalarT]:
+    ) -> core_defs.NDArrayObject:
         raise self.exception
 
 
@@ -292,16 +292,28 @@ def __init__(self) -> None:
 )
 
 
-def allocate(
-    domain: common.DomainLike,
+class ConcreteAllocator(Protocol):
+    def __call__(
+        domain: common.DomainLike,
+        dtype: core_defs.DType[core_defs.ScalarT],
+        *,
+        aligned_index: Optional[Sequence[common.NamedIndex]],
+        allocator: FieldBufferAllocationUtil,
+        device: core_defs.Device,
+    ) -> core_defs.NDArrayObject: ...
+
+
+def make_concrete_allocator(
+    domain: common.DomainLike,  # TODO: there is an inconsistency between DomainLike and concrete DType, probably accept either (Domain, DType) or (DomainLike, DTypeLike). anyway this is not meant to be user-facing
     dtype: core_defs.DType[core_defs.ScalarT],
     *,
     aligned_index: Optional[Sequence[common.NamedIndex]] = None,
     allocator: Optional[FieldBufferAllocationUtil] = None,
     device: Optional[core_defs.Device] = None,
-) -> core_allocators.TensorBuffer:
+) -> ConcreteAllocator:
     """
-    Allocate a TensorBuffer for the given domain and device or allocator.
+    TODO: docstring
+    Allocate an NDArrayObject for the given domain and device or allocator.
 
     The arguments `device` and `allocator` are mutually exclusive.
     If `device` is specified, the corresponding default allocator
@@ -334,9 +346,20 @@ def allocate(
     elif device.device_type != actual_allocator.__gt_device_type__:
         raise ValueError(f"Device '{device}' and allocator '{actual_allocator}' are incompatible.")
 
-    return actual_allocator.__gt_allocate__(
-        domain=common.domain(domain),
-        dtype=dtype,
-        device_id=device.device_id,
-        aligned_index=aligned_index,
-    )
+    def allocate(
+        domain: common.DomainLike = domain,
+        dtype: core_defs.DType[core_defs.ScalarT] = dtype,
+        *,
+        aligned_index: Optional[Sequence[common.NamedIndex]] = aligned_index,
+        allocator: FieldBufferAllocationUtil = actual_allocator,
+        device: core_defs.Device = device,
+    ) -> core_defs.NDArrayObject:
+        # TODO check how to get from FieldBufferAllocationUtil to FieldBufferAllocatorProtocol
+        return allocator.__gt_allocate__(
+            domain=common.domain(domain),
+            dtype=dtype,
+            device_id=device.device_id,
+            aligned_index=aligned_index,
+        )
+
+    return allocate

src/gt4py/next/common.py

@@ -950,6 +950,7 @@ def _field(
     /,
     *,
     domain: Optional[DomainLike] = None,
+    allocator: Optional[Any] = None,  # TODO: resolve the type annotation
     dtype: Optional[core_defs.DType] = None,
 ) -> Field:
     raise NotImplementedError
@@ -963,6 +964,7 @@ def _connectivity(
     codomain: Dimension,
     *,
     domain: Optional[DomainLike] = None,
+    allocator: Optional[Any] = None,  # TODO: resolve the type annotation
     dtype: Optional[core_defs.DType] = None,
     skip_value: Optional[core_defs.IntegralScalar] = None,
 ) -> Connectivity:

src/gt4py/next/constructors.py

@@ -11,6 +11,8 @@
 from collections.abc import Mapping, Sequence
 from typing import Optional, cast
 
+from typing_extensions import NotRequired, TypedDict
+
 import gt4py._core.definitions as core_defs
 import gt4py.eve as eve
 import gt4py.eve.extended_typing as xtyping
@@ -77,10 +79,11 @@ def empty(
     dtype = core_defs.dtype(dtype)
     if allocator is None and device is None:
         device = core_defs.Device(core_defs.DeviceType.CPU, device_id=0)
-    buffer = next_allocators.allocate(
+    allocate = next_allocators.make_concrete_allocator(
         domain, dtype, aligned_index=aligned_index, allocator=allocator, device=device
     )
-    res = common._field(buffer.ndarray, domain=domain)
+    buffer = allocate()
+    res = common._field(buffer, domain=domain, allocator=allocate)
     assert isinstance(res, common.MutableField)
     assert isinstance(res, nd_array_field.NdArrayField)
     return res
@@ -349,12 +352,51 @@ def as_connectivity(
 
     if (allocator is None) and (device is None) and xtyping.supports_dlpack(data):
         device = core_defs.Device(*data.__dlpack_device__())
-    buffer = next_allocators.allocate(actual_domain, dtype, allocator=allocator, device=device)
+    allocate = next_allocators.make_concrete_allocator(
+        actual_domain, dtype, allocator=allocator, device=device
+    )
+    buffer = allocate()
     # TODO(havogt): consider adding MutableNDArrayObject
-    buffer.ndarray[...] = storage_utils.asarray(data)  # type: ignore[index]
+    buffer[...] = storage_utils.asarray(data)  # type: ignore[index]
     connectivity_field = common._connectivity(
-        buffer.ndarray, codomain=codomain, domain=actual_domain, skip_value=skip_value
+        buffer, codomain=codomain, domain=actual_domain, skip_value=skip_value, allocator=allocate
     )
     assert isinstance(connectivity_field, nd_array_field.NdArrayConnectivityField)
 
     return connectivity_field
+
+
+_like_field = None  # for more descriptive function signature in editors
+
+
+class AllocatorParams(TypedDict):
+    domain: NotRequired[common.DomainLike]
+    dtype: NotRequired[core_defs.DType[core_defs.ScalarT],]
+    aligned_index: NotRequired[Sequence[common.NamedIndex]]
+    allocator: NotRequired[next_allocators.FieldBufferAllocatorProtocol]
+    device: NotRequired[core_defs.Device]
+
+
+def empty_like(
+    field: nd_array_field.NdArrayField,
+    *,
+    domain: Optional[common.DomainLike] = _like_field,
+    dtype: Optional[core_defs.DTypeLike] = _like_field,
+    aligned_index: Optional[Sequence[common.NamedIndex]] = _like_field,
+    allocator: Optional[next_allocators.FieldBufferAllocationUtil] = _like_field,
+    device: Optional[core_defs.Device] = _like_field,
+) -> nd_array_field.NdArrayField:
+    kwargs: AllocatorParams = {}
+    if domain is not None:
+        kwargs["domain"] = domain
+    if dtype is not None:
+        kwargs["dtype"] = core_defs.dtype(dtype)
+    if aligned_index is not None:
+        kwargs["aligned_index"] = aligned_index
+    if allocator is not eve.NOTHING:
+        kwargs["allocator"] = allocator
+    if device is not eve.NOTHING:
+        kwargs["device"] = device
+    if field._allocator is None:
+        raise ValueError("'Field' does not have an allocator.")  # TODO discuss if this is possible
+    return field._allocator(**kwargs)

src/gt4py/next/embedded/nd_array_field.py

@@ -29,7 +29,7 @@
     TypeVar,
     cast,
 )
-from gt4py.next import common
+from gt4py.next import allocators, common
 from gt4py.next.embedded import (
     common as embedded_common,
     context as embedded_context,
@@ -116,6 +116,7 @@ class NdArrayField(
 
     _domain: common.Domain
     _ndarray: core_defs.NDArrayObject
+    _allocator: Optional[allocators.ConcreteAllocator]
 
     array_ns: ClassVar[ModuleType]  # TODO(havogt) introduce a NDArrayNamespace protocol
 
@@ -167,6 +168,9 @@ def from_array(
         /,
         *,
         domain: common.DomainLike,
+        allocator: Optional[
+            allocators.ConcreteAllocator
+        ] = None,  # TODO: maybe an NDArrayField always has an allocator?
         dtype: Optional[core_defs.DTypeLike] = None,
     ) -> NdArrayField:
         domain = common.domain(domain)
@@ -184,7 +188,7 @@ def from_array(
         assert len(domain) == array.ndim
         assert all(s == 1 or len(r) == s for r, s in zip(domain.ranges, array.shape))
 
-        return cls(domain, array)
+        return cls(domain, array, allocator)
 
     def premap(
         self: NdArrayField,
@@ -513,6 +517,7 @@ def from_array(  # type: ignore[override]
         codomain: common.DimT,
         *,
         domain: common.DomainLike,
+        allocator: Optional[allocators.ConcreteAllocator] = None,
         dtype: Optional[core_defs.DTypeLike] = None,
         skip_value: Optional[core_defs.IntegralScalar] = None,
     ) -> NdArrayConnectivityField:
@@ -533,7 +538,7 @@ def from_array(  # type: ignore[override]
 
         assert isinstance(codomain, common.Dimension)
 
-        return cls(domain, array, codomain, _skip_value=skip_value)
+        return cls(domain, array, allocator, codomain, _skip_value=skip_value)
 
     def inverse_image(self, image_range: common.UnitRange | common.NamedRange) -> common.Domain:
         cache_key = hash((id(self.ndarray), self.domain, image_range))