Merge branch 'field_storage_interface' of https://github.com/havogt/g…

…t4py into fork_branch

src/gt4py/eve/utils.py

@@ -234,7 +234,7 @@ def itemgetter_(key: Any, default: Any = NOTHING) -> Callable[[Any], Any]:
 class fluid_partial(functools.partial):
     """Create a `functools.partial` with support for multiple applications calling `.partial()`."""
 
-    def partial(self, *args, **kwargs) -> fluid_partial:
+    def partial(self, *args: Any, **kwargs: Any) -> fluid_partial:
         return fluid_partial(self, *args, **kwargs)
 
 
@@ -277,7 +277,7 @@ def with_fluid_partial(  # noqa: F811  # redefinition of unused function
     """
 
     def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        func.partial = fluid_partial(functools.partial, func, *args, **kwargs)
+        func.partial = fluid_partial(functools.partial, func, *args, **kwargs)  # type: ignore[attr-defined]  # add attribute
         return func
 
     return _decorator(func) if func is not None else _decorator

src/gt4py/next/common.py

@@ -80,15 +80,18 @@ def __str__(self):
         return f"{self.value}[{self.kind}]"
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True, init=False)
 class UnitRange(Sequence[int], Set[int]):
     """Range from `start` to `stop` with step size one."""
 
     start: int
     stop: int
 
-    def __post_init__(self):
-        if self.stop <= self.start:
+    def __init__(self, start: core_defs.IntegralScalar, stop: core_defs.IntegralScalar) -> None:
+        if start < stop:
+            object.__setattr__(self, "start", int(start))
+            object.__setattr__(self, "stop", int(stop))
+        else:
             # make UnitRange(0,0) the single empty UnitRange
             object.__setattr__(self, "start", 0)
             object.__setattr__(self, "stop", 0)
@@ -157,14 +160,16 @@ def unit_range(r: RangeLike) -> UnitRange:
         if r.step != 1:
             raise ValueError(f"`UnitRange` requires step size 1, got `{r.step}`.")
         return UnitRange(r.start, r.stop)
+    # TODO(egparedes): use core_defs.IntegralScalar for `isinstance()` checks (see PEP 604)
+    #   once the related mypy bug (#16358) gets fixed
     if (
         isinstance(r, tuple)
         and isinstance(r[0], core_defs.INTEGRAL_TYPES)
         and isinstance(r[1], core_defs.INTEGRAL_TYPES)
     ):
         return UnitRange(r[0], r[1])
     if isinstance(r, core_defs.INTEGRAL_TYPES):
-        return UnitRange(0, r)
+        return UnitRange(0, cast(core_defs.IntegralScalar, r))
     raise ValueError(f"`{r!r}` cannot be interpreted as `UnitRange`.")
 
 

src/gt4py/storage/cartesian/utils.py

@@ -17,7 +17,7 @@
 import collections.abc
 import math
 import numbers
-from typing import Any, Literal, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Final, Literal, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -39,23 +39,33 @@
     cp = None
 
 
+CUPY_DEVICE: Final[Literal[None, core_defs.DeviceType.CUDA, core_defs.DeviceType.ROCM]] = (
+    None
+    if not cp
+    else (core_defs.DeviceType.ROCM if cp.cuda.get_hipcc_path() else core_defs.DeviceType.CUDA)
+)
+
+
 FieldLike = Union["cp.ndarray", np.ndarray, ArrayInterface, CUDAArrayInterface]
 
+assert allocators.is_valid_nplike_allocation_ns(np)
+
 _CPUBufferAllocator = allocators.NDArrayBufferAllocator(
     device_type=core_defs.DeviceType.CPU,
     array_ns=np,
 )
 
+_GPUBufferAllocator: Optional[allocators.NDArrayBufferAllocator] = None
 if cp:
-    cp_device_type = (
-        core_defs.DeviceType.ROCM if cp.cuda.get_hipcc_path() else core_defs.DeviceType.CUDA
-    )
-
     assert allocators.is_valid_nplike_allocation_ns(cp)
-
-    _GPUBufferAllocator = allocators.NDArrayBufferAllocator(device_type=cp_device_type, array_ns=cp)
-else:
-    _GPUBufferAllocator = None
+    if CUPY_DEVICE == core_defs.DeviceType.CUDA:
+        _GPUBufferAllocator = allocators.NDArrayBufferAllocator(
+            device_type=core_defs.DeviceType.CUDA, array_ns=cp
+        )
+    else:
+        _GPUBufferAllocator = allocators.NDArrayBufferAllocator(
+            device_type=core_defs.DeviceType.ROCM, array_ns=cp
+        )
 
 
 def _idx_from_order(order):
@@ -242,7 +252,7 @@ def allocate_gpu(
     buffer = _GPUBufferAllocator.allocate(
         shape,
         core_defs.dtype(dtype),
-        device_id=device.id,
+        device_id=device.device_id,
         layout_map=layout_map,
         byte_alignment=alignment_bytes,
         aligned_index=aligned_index,