Update to gt4py v1.1.5

model/atmosphere/advection/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # workspace members
   "icon4py-common>=0.0.6",
   # external dependencies
-  "gt4py==1.1.4",
+  "gt4py==1.1.5",
   'packaging>=20.0'
 ]
 description = "ICON advection."

model/atmosphere/diffusion/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # workspace members
   "icon4py-common>=0.0.6",
   # external dependencies
-  "gt4py==1.1.4",
+  "gt4py==1.1.5",
   'packaging>=20.0'
 ]
 description = "ICON diffusion."

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -16,7 +16,6 @@
 
 import gt4py.next as gtx
 import gt4py.next.typing as gtx_typing
-from gt4py.next import allocators as gtx_allocators
 
 import icon4py.model.common.grid.states as grid_states
 import icon4py.model.common.states.prognostic_state as prognostics
@@ -598,7 +597,7 @@ def __init__(
         #   but this requires some changes in gt4py domain inference.
         self.compile_time_connectivities = self._grid.connectivities
 
-    def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocationUtil | None):
+    def _allocate_local_fields(self, allocator: gtx_typing.Allocator | None):
         self.diff_multfac_vn = data_alloc.zero_field(self._grid, dims.KDim, allocator=allocator)
         self.diff_multfac_n2w = data_alloc.zero_field(self._grid, dims.KDim, allocator=allocator)
         self.smag_limit = data_alloc.zero_field(self._grid, dims.KDim, allocator=allocator)

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py

@@ -90,7 +90,7 @@ def geofac_n2s_nbh(self) -> gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CDim], flo
 
 
 def initialize_diffusion_diagnostic_state(
-    grid: icon_grid.IconGrid, allocator: gtx_typing.FieldBufferAllocationUtil
+    grid: icon_grid.IconGrid, allocator: gtx_typing.Allocator
 ) -> DiffusionDiagnosticState:
     hdef_ic = data_alloc.zero_field(
         grid,

model/atmosphere/dycore/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # workspace members
   "icon4py-common>=0.0.6",
   # external dependencies
-  "gt4py==1.1.4",
+  "gt4py==1.1.5",
   'packaging>=20.0'
 ]
 description = "ICON dynamical core."

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_states.py

@@ -329,7 +329,7 @@ class PrepAdvection:
 def initialize_solve_nonhydro_diagnostic_state(
     perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat],
     grid: icon_grid.IconGrid,
-    allocator: gtx_typing.FieldBufferAllocationUtil,
+    allocator: gtx_typing.Allocator,
 ) -> DiagnosticStateNonHydro:
     normal_wind_advective_tendency = common_utils.PredictorCorrectorPair(
         data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, allocator=allocator, dtype=ta.vpfloat),
@@ -453,7 +453,7 @@ def initialize_solve_nonhydro_diagnostic_state(
 
 
 def initialize_prep_advection(
-    grid: icon_grid.IconGrid, allocator: gtx_typing.FieldBufferAllocationUtil
+    grid: icon_grid.IconGrid, allocator: gtx_typing.Allocator
 ) -> PrepAdvection:
     vn_traj = data_alloc.zero_field(
         grid, dims.EdgeDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

@@ -13,7 +13,7 @@
 
 import gt4py.next as gtx
 import gt4py.next.typing as gtx_typing
-from gt4py.next import allocators as gtx_allocators, common as gtx_common
+from gt4py.next import common as gtx_common
 
 import icon4py.model.atmosphere.dycore.solve_nonhydro_stencils as nhsolve_stencils
 import icon4py.model.common.grid.states as grid_states
@@ -110,11 +110,7 @@ class IntermediateFields:
     """
 
     @classmethod
-    def allocate(
-        cls,
-        grid: grid_def.Grid,
-        allocator: gtx_allocators.FieldBufferAllocationUtil | None,
-    ):
+    def allocate(cls, grid: grid_def.Grid, allocator: gtx_typing.Allocator | None):
         return IntermediateFields(
             horizontal_pressure_gradient=data_alloc.zero_field(
                 grid, dims.EdgeDim, dims.KDim, allocator=allocator
@@ -816,7 +812,7 @@ def __init__(
         recomputed or not. The substep length should only change in case of high CFL condition.
         """
 
-    def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocationUtil | None):
+    def _allocate_local_fields(self, allocator: gtx_typing.Allocator | None):
         self.temporal_extrapolation_of_perturbed_exner = data_alloc.zero_field(
             self._grid,
             dims.CellDim,

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py

@@ -11,7 +11,6 @@
 
 import gt4py.next as gtx
 import gt4py.next.typing as gtx_typing
-from gt4py.next import allocators as gtx_allocators  # TODO(havogt): expose in gtx_typing
 
 from icon4py.model.atmosphere.dycore import dycore_states
 from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_horizontal_momentum_equation import (
@@ -181,7 +180,7 @@ def __init__(
             offset_provider=self.grid.connectivities,
         )
 
-    def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocationUtil | None):
+    def _allocate_local_fields(self, allocator: gtx_typing.Allocator | None):
         self._horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field(
             self.grid, dims.EdgeDim, dims.KDim, allocator=allocator, dtype=ta.vpfloat
         )

model/atmosphere/subgrid_scale_physics/microphysics/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # workspace members
   "icon4py-common>=0.0.6",
   # external dependencies
-  "gt4py==1.1.4",
+  "gt4py==1.1.5",
   'packaging>=20.0'
 ]
 description = "ICON microphysics."

model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   "icon4py-common[io]>=0.0.6",
   # external dependencies
   "numpy>=1.23.3",
-  "gt4py==1.1.4",
+  "gt4py==1.1.5",
   "packaging>=20.0"
 ]
 description = "ICON subgrid scale muphys parameterization."

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/common.py

@@ -34,7 +34,7 @@ def _calc_dz(z: np.ndarray) -> np.ndarray:
 
 def _as_field_from_nc(
     dataset: netCDF4.Dataset,
-    allocator: gtx_typing.FieldBufferAllocationUtil,
+    allocator: gtx_typing.Allocator,
     varname: str,
     optional: bool = False,
     dtype: np.dtype | None = None,
@@ -96,7 +96,7 @@ def q(self) -> Q:
     def load(
         cls,
         filename: pathlib.Path | str,
-        allocator: gtx_typing.FieldBufferAllocationUtil,
+        allocator: gtx_typing.Allocator,
         dtype=np.float64,
     ) -> None:
         with netCDF4.Dataset(filename, mode="r") as ncfile:
@@ -146,7 +146,7 @@ class GraupelOutput:
     @classmethod
     def allocate(
         cls,
-        allocator: gtx_typing.FieldBufferAllocationUtil,
+        allocator: gtx_typing.Allocator,
         domain: gtx.Domain,
         references: dict[str, gtx.Field] | None = None,
     ):
@@ -169,7 +169,7 @@ def allocate(
         )
 
     @classmethod
-    def load(cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil):
+    def load(cls, filename: pathlib.Path | str, allocator: gtx_typing.Allocator):
         with netCDF4.Dataset(filename, mode="r") as ncfile:
             field_from_nc = functools.partial(_as_field_from_nc, ncfile, allocator)
             return cls(

model/common/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
   'Topic :: Scientific/Engineering :: Physics'
 ]
 dependencies = [
-  'gt4py==1.1.4',
+  'gt4py==1.1.5',
   'packaging>=20.0',
   'typing-extensions>=4.11.0',
   'numpy>=1.23.3',

model/common/src/icon4py/model/common/grid/base.py

@@ -13,7 +13,8 @@
 from types import ModuleType
 
 import gt4py.next as gtx
-from gt4py.next import allocators as gtx_allocators, common as gtx_common
+import gt4py.next.typing as gtx_typing
+from gt4py.next import common as gtx_common
 
 from icon4py.model.common import dimension as dims
 from icon4py.model.common.grid import horizontal as h_grid
@@ -169,7 +170,7 @@ def construct_connectivity(
     table: data_alloc.NDArray,
     skip_value: int | None = None,
     *,
-    allocator: gtx_allocators.FieldBufferAllocationUtil | None = None,
+    allocator: gtx_typing.Allocator | None = None,
     replace_skip_values: bool = False,
 ):
     from_dim, dim = offset.target

model/common/src/icon4py/model/common/grid/geometry.py

@@ -866,7 +866,7 @@ def create_auxiliary_coordinate_arrays_for_orientation(
     cell_lon: fa.CellField[ta.wpfloat],
     edge_lat: fa.EdgeField[ta.wpfloat],
     edge_lon: fa.EdgeField[ta.wpfloat],
-    allocator: gtx_typing.FieldBufferAllocationUtil | None,
+    allocator: gtx_typing.Allocator | None,
 ) -> tuple[
     fa.EdgeField[ta.wpfloat],
     fa.EdgeField[ta.wpfloat],

model/common/src/icon4py/model/common/grid/grid_manager.py

@@ -121,7 +121,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is FileNotFoundError:
             raise FileNotFoundError(f"gridfile {self._file_name} not found, aborting")
 
-    def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_values: bool):
+    def __call__(self, allocator: gtx_typing.Allocator, keep_skip_values: bool):
         if not self._reader:
             self.open()
 
@@ -141,7 +141,7 @@ def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_va
 
     def _read_coordinates(
         self,
-        allocator: gtx_typing.FieldBufferAllocationUtil,
+        allocator: gtx_typing.Allocator,
         geometry_type: base.GeometryType,
     ) -> CoordinateDict:
         coordinates = {
@@ -249,7 +249,7 @@ def _read_coordinates(
 
     def _read_geometry_fields(
         self,
-        allocator: gtx_typing.FieldBufferAllocationUtil,
+        allocator: gtx_typing.Allocator,
     ) -> GeometryDict:
         return {
             # TODO(halungge): still needs to ported, values from "our" grid files contains (wrong) values:
@@ -311,7 +311,7 @@ def _read_grid_refinement_fields(
         self,
         *,
         decomposition_info: decomposition.DecompositionInfo | None = None,
-        allocator: gtx_typing.FieldBufferAllocationUtil,
+        allocator: gtx_typing.Allocator,
     ) -> dict[gtx.Dimension, gtx.Field]:
         """
         Reads the refinement control fields from the grid file.
@@ -354,7 +354,7 @@ def coordinates(self) -> CoordinateDict:
 
     def _construct_grid(
         self,
-        allocator: gtx_typing.FieldBufferAllocationUtil | None,
+        allocator: gtx_typing.Allocator | None,
         with_skip_values: bool,
         geometry_type: base.GeometryType,
     ) -> icon.IconGrid:

model/common/src/icon4py/model/common/grid/icon.py

@@ -12,7 +12,7 @@
 from typing import Final, TypeVar
 
 import gt4py.next as gtx
-from gt4py.next import allocators as gtx_allocators
+import gt4py.next.typing as gtx_typing
 
 from icon4py.model.common import constants, dimension as dims
 from icon4py.model.common.grid import base, horizontal as h_grid
@@ -182,7 +182,7 @@ def _should_replace_skip_values(
 
 def icon_grid(
     id_: str,
-    allocator: gtx_allocators.FieldBufferAllocationUtil | None,
+    allocator: gtx_typing.Allocator | None,
     config: base.GridConfig,
     neighbor_tables: dict[gtx.FieldOffset, data_alloc.NDArray],
     start_index: Callable[[h_grid.Domain], gtx.int32],

model/common/src/icon4py/model/common/grid/simple.py

@@ -413,7 +413,7 @@ def c2e2c2e2c_table(self):
 
 def simple_grid(
     *,
-    allocator: gtx_typing.FieldBufferAllocationUtil | None = None,
+    allocator: gtx_typing.Allocator | None = None,
     num_levels: int = DEFAULT_NUM_LEVELS,
 ) -> base.Grid:
     """

model/common/src/icon4py/model/common/grid/vertical.py

@@ -294,7 +294,7 @@ def _determine_end_index_of_flat_layers(
 
 
 def _read_vct_a_and_vct_b_from_file(
-    file_path: pathlib.Path, num_levels: int, allocator: gtx_typing.FieldBufferAllocationUtil
+    file_path: pathlib.Path, num_levels: int, allocator: gtx_typing.Allocator
 ) -> tuple[fa.KField, fa.KField]:
     """
     Read vct_a and vct_b from a file.
@@ -340,7 +340,7 @@ def _read_vct_a_and_vct_b_from_file(
 
 
 def _compute_vct_a_and_vct_b(  # noqa: PLR0912 [too-many-branches]
-    vertical_config: VerticalGridConfig, allocator: gtx_typing.FieldBufferAllocationUtil
+    vertical_config: VerticalGridConfig, allocator: gtx_typing.Allocator
 ) -> tuple[fa.KField, fa.KField]:
     """
     Compute vct_a and vct_b.
@@ -529,7 +529,7 @@ def _compute_vct_a_and_vct_b(  # noqa: PLR0912 [too-many-branches]
 
 
 def get_vct_a_and_vct_b(
-    vertical_config: VerticalGridConfig, allocator: gtx_typing.FieldBufferAllocationUtil
+    vertical_config: VerticalGridConfig, allocator: gtx_typing.Allocator
 ) -> tuple[fa.KField, fa.KField]:
     """
     get vct_a and vct_b.

model/common/src/icon4py/model/common/model_backends.py

@@ -9,7 +9,7 @@
 
 import gt4py.next as gtx
 import gt4py.next.typing as gtx_typing
-from gt4py.next import allocators as gtx_allocators, backend as gtx_backend
+from gt4py.next import backend as gtx_backend, custom_layout_allocators as gtx_allocators
 from gt4py.next.program_processors.runners import dace as gtx_dace, gtfn
 
 
@@ -68,6 +68,7 @@ def make_custom_dace_backend(
     optimization_args: dict[str, Any] | None = None,
     use_metrics: bool = True,
     use_zero_origin: bool = False,
+    use_max_domain_range_on_unstructured_shift: bool | None = True,
     **_,
 ) -> gtx_typing.Backend:
     """Customize the dace backend with the given configuration parameters.
@@ -82,6 +83,12 @@ def make_custom_dace_backend(
             the SDFG auto-optimize pipeline.
         use_metrics: Add SDFG instrumentation to collect the metric for stencil
             compute time.
+        use_max_domain_range_on_unstructured_shift: When True, compute `as_fieldop`
+            expressions everywhere. Otherwise, when all connectivities are given
+            at compile time, infer the minimal domain of all `as_fieldop` statically.
+
+    TODO(edopao): We should use the default `use_max_domain_range_on_unstructured_shift=None`
+        once the minimal static domain produces the correct result.
 
     Returns:
         A dace backend with custom configuration for the target device.
@@ -95,6 +102,7 @@ def make_custom_dace_backend(
         optimization_args=optimization_args,
         use_metrics=use_metrics,
         use_zero_origin=use_zero_origin,
+        use_max_domain_range_on_unstructured_shift=use_max_domain_range_on_unstructured_shift,
     )
 
 

model/common/src/icon4py/model/common/states/diagnostic_state.py

@@ -59,7 +59,7 @@ class DiagnosticMetricState:
 
 def initialize_diagnostic_state(
     grid: icon_grid.IconGrid,
-    allocator: gtx_typing.FieldBufferAllocationUtil,
+    allocator: gtx_typing.Allocator,
 ) -> DiagnosticState:
     """Initialize the diagnostic state with zero fields."""
     pressure = data_alloc.zero_field(

model/common/src/icon4py/model/common/states/prognostic_state.py

@@ -43,7 +43,7 @@ def w_1(self) -> fa.CellField[ta.wpfloat]:
 
 def initialize_prognostic_state(
     grid: icon_grid.IconGrid,
-    allocator: gtx_typing.FieldBufferAllocationUtil,
+    allocator: gtx_typing.Allocator,
 ) -> PrognosticState:
     """Initialize the prognostic state with zero fields."""
     rho = data_alloc.zero_field(