feat[next]: Compile time domains

src/gt4py/eve/utils.py

@@ -914,7 +914,7 @@ def reset_sequence(self, start: int = 1, *, warn_unsafe: Optional[bool] = None)
         if warn_unsafe is None:
             warn_unsafe = self.warn_unsafe
         if warn_unsafe and start < next(self._counter):
-            warnings.warn("Unsafe reset of UIDGenerator ({self})", stacklevel=2)
+            warnings.warn(f"Unsafe reset of UIDGenerator ({self})", stacklevel=2)
         self._counter = itertools.count(start)
 
         return self

src/gt4py/next/ffront/decorator.py

@@ -86,6 +86,7 @@ class Program:
     static_params: (
         Sequence[str] | None
     )  # if the user requests static params, they will be used later to initialize CompiledPrograms
+    static_domains: bool
 
     @classmethod
     def from_function(
@@ -95,6 +96,7 @@ def from_function(
         grid_type: common.GridType | None = None,
         enable_jit: bool | None = None,
         static_params: Sequence[str] | None = None,
+        static_domains: bool = False,
         connectivities: Optional[
             common.OffsetProvider
         ] = None,  # TODO(ricoh): replace with common.OffsetProviderType once the temporary pass doesn't require the runtime information
@@ -106,6 +108,7 @@ def from_function(
             connectivities=connectivities,
             enable_jit=enable_jit,
             static_params=static_params,
+            static_domains=static_domains,
         )
 
     # TODO(ricoh): linting should become optional, up to the backend.
@@ -170,20 +173,31 @@ def _compiled_programs(self) -> compiled_program.CompiledProgramsPool:
         if self.backend is None or self.backend == eve.NOTHING:
             raise RuntimeError("Cannot compile a program without backend.")
 
-        if self.static_params is None:
-            object.__setattr__(self, "static_params", ())
+        def path_to_expr(path: Sequence[int]) -> str:
+            return "".join(map(lambda idx: f"[{idx}]", path))
+
+        argument_descriptor_mapping: dict[type[arguments.ArgStaticDescriptor], Sequence[str]] = {}
+
+        if self.static_params:
+            argument_descriptor_mapping[arguments.StaticArg] = self.static_params
 
-        argument_descriptor_mapping = {
-            arguments.StaticArg: self.static_params,
-        }
+        if self.static_domains:
+            static_domain_args = []
+            func_type = self.past_stage.past_node.type.definition  # type: ignore[union-attr] # type inference done at this point
+            param_types = func_type.pos_or_kw_args | func_type.kw_only_args
+            for name, type_ in param_types.items():
+                for el_type_, path in type_info.primitive_constituents(type_, with_path_arg=True):
+                    if isinstance(el_type_, ts.FieldType):
+                        static_domain_args.append(f"{name}{path_to_expr(path)}")
+            argument_descriptor_mapping[arguments.FieldDomainDescriptor] = static_domain_args
 
         program_type = self.past_stage.past_node.type
         assert isinstance(program_type, ts_ffront.ProgramType)
         return compiled_program.CompiledProgramsPool(
             backend=self.backend,
             definition_stage=self.definition_stage,
             program_type=program_type,
-            argument_descriptor_mapping=argument_descriptor_mapping,  # type: ignore[arg-type]  # covariant `type[T]` not possible
+            argument_descriptor_mapping=argument_descriptor_mapping,
         )
 
     def with_backend(self, backend: next_backend.Backend) -> Program:
@@ -529,6 +543,7 @@ def program(
     grid_type: common.GridType | None,
     enable_jit: bool | None,
     static_params: Sequence[str] | None,
+    static_domains: bool,
     frozen: bool,
 ) -> Callable[[types.FunctionType], Program]: ...
 
@@ -541,6 +556,7 @@ def program(
     grid_type: common.GridType | None = None,
     enable_jit: bool | None = None,  # only relevant if static_params are set
     static_params: Sequence[str] | None = None,
+    static_domains: bool = False,
     frozen: bool = False,
 ) -> Program | FrozenProgram | Callable[[types.FunctionType], Program | FrozenProgram]:
     """
@@ -569,6 +585,7 @@ def program_inner(definition: types.FunctionType) -> Program:
             ),
             grid_type=grid_type,
             enable_jit=enable_jit,
+            static_domains=static_domains,
             static_params=static_params,
         )
         if frozen:
@@ -703,6 +720,7 @@ def as_program(self, compiletime_args: arguments.CompileTimeArgs) -> Program:
             connectivities=None,
             enable_jit=False,  # TODO(havogt): revisit ProgramFromPast
             static_params=None,  # TODO(havogt): revisit ProgramFromPast
+            static_domains=False,  # TODO(havogt): revisit ProgramFromPast
         )
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:

src/gt4py/next/ffront/past_to_itir.py

@@ -15,7 +15,7 @@
 import devtools
 
 from gt4py.eve import NodeTranslator, traits
-from gt4py.next import common, config, errors, utils as gtx_utils
+from gt4py.next import common, config, errors, utils
 from gt4py.next.ffront import (
     fbuiltins,
     gtcallable,
@@ -28,7 +28,7 @@
 from gt4py.next.ffront.stages import AOT_PRG
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.ir_utils import ir_makers as im
-from gt4py.next.iterator.transforms import remap_symbols
+from gt4py.next.iterator.transforms import remap_symbols, transform_get_domain_range
 from gt4py.next.otf import arguments, stages, workflow
 from gt4py.next.type_system import type_info, type_specifications as ts
 
@@ -102,14 +102,14 @@ def past_to_gtir(inp: AOT_PRG) -> stages.CompilableProgram:
         static_arg_descriptors = inp.args.argument_descriptor_contexts[arguments.StaticArg]
         if not all(
             isinstance(arg_descriptor, arguments.StaticArg)
-            or all(el is None for el in gtx_utils.flatten_nested_tuple(arg_descriptor))  # type: ignore[arg-type]
+            or all(el is None for el in utils.flatten_nested_tuple(arg_descriptor))  # type: ignore[arg-type]
             for arg_descriptor in static_arg_descriptors.values()
         ):
             raise NotImplementedError("Only top-level arguments can be static.")
         static_args = {
             name: im.literal_from_tuple_value(descr.value)  # type: ignore[union-attr]  # type checked above
             for name, descr in static_arg_descriptors.items()
-            if not any(el is None for el in gtx_utils.flatten_nested_tuple(descr))  # type: ignore[arg-type]
+            if not any(el is None for el in utils.flatten_nested_tuple(descr))  # type: ignore[arg-type]
         }
         body = remap_symbols.RemapSymbolRefs().visit(itir_program.body, symbol_map=static_args)
         itir_program = itir.Program(
@@ -120,6 +120,16 @@ def past_to_gtir(inp: AOT_PRG) -> stages.CompilableProgram:
             body=body,
         )
 
+    # TODO(tehrengruber): Put this in a dedicated transformation step.
+    if arguments.FieldDomainDescriptor in inp.args.argument_descriptor_contexts:
+        context = inp.args.argument_descriptor_contexts[arguments.FieldDomainDescriptor]
+        field_domains = {
+            param: utils.tree_map(lambda x: x.domain if x else x)(v) for param, v in context.items()
+        }
+        itir_program = transform_get_domain_range.TransformGetDomainRange.apply(
+            itir_program, sizes=field_domains
+        )
+
     # Translate NamedCollectionTypes to TupleTypes in compile-time args
     args = tuple(ffront_ti.named_collections_to_tuple_types(arg) for arg in inp.args.args)
     kwargs: dict[str, ts.TypeSpec] = {
@@ -443,7 +453,7 @@ def _visit_stencil_call_out_arg(
             "Unexpected 'out' argument. Must be a 'past.Subscript', 'past.Name' or 'past.TupleExpr' node."
         )
 
-        @gtx_utils.tree_map(
+        @utils.tree_map(
             collection_type=ts.COLLECTION_TYPE_SPECS,
             with_path_arg=True,
             unpack=True,

src/gt4py/next/iterator/ir_utils/domain_utils.py

@@ -10,34 +10,27 @@
 
 import dataclasses
 import functools
-from typing import Any, Callable, Iterable, Literal, Mapping, Optional
+import warnings
+from typing import Callable, Iterable, Literal, Optional
+
+import numpy as np
 
 from gt4py.next import common
 from gt4py.next.iterator import builtins, ir as itir
 from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im
-from gt4py.next.iterator.transforms import trace_shifts
+from gt4py.next.iterator.transforms import collapse_tuple, trace_shifts
 from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
 
 
-def _max_domain_sizes_by_location_type(offset_provider: Mapping[str, Any]) -> dict[str, int]:
-    """
-    Extract horizontal domain sizes from an `offset_provider`.
+#: Threshold fraction of domain points which may be added to a domain on translation in order
+#: to have a contiguous domain before a warning is raised.
+NON_CONTIGUOUS_DOMAIN_WARNING_THRESHOLD: float = 1 / 4
 
-    Considers the shape of the neighbor table to get the size of each `source_dim` and the maximum
-    value inside the neighbor table to get the size of each `codomain`.
-    """
-    sizes = dict[str, int]()
-    for provider in offset_provider.values():
-        if common.is_neighbor_connectivity(provider):
-            conn_type = provider.__gt_type__()
-            sizes[conn_type.source_dim.value] = max(
-                sizes.get(conn_type.source_dim.value, 0), provider.ndarray.shape[0]
-            )
-            sizes[conn_type.codomain.value] = max(
-                sizes.get(conn_type.codomain.value, 0),
-                provider.ndarray.max() + 1,  # type: ignore[attr-defined] # TODO(havogt): improve typing for NDArrayObject
-            )
-    return sizes
+#: Skip printing warnings after exceeding `NON_CONTIGUOUS_DOMAIN_WARNING_THRESHOLD` this many times.
+NON_CONTIGUOUS_DOMAIN_MAX_WARNINGS: int = 5
+
+#: Number of warnings raised after exceeding `NON_CONTIGUOUS_DOMAIN_WARNING_THRESHOLD`
+_NON_CONTIGUOUS_DOMAIN_WARNING_COUNTER: int = 0
 
 
 @dataclasses.dataclass(frozen=True)
@@ -114,7 +107,7 @@ def translate(
         offset_provider: common.OffsetProvider | common.OffsetProviderType,
         #: A dictionary mapping axes names to their length. See
         #: func:`gt4py.next.iterator.transforms.infer_domain.infer_expr` for more details.
-        symbolic_domain_sizes: Optional[dict[str, str]] = None,
+        symbolic_domain_sizes: Optional[dict[str, str | itir.Expr]] = None,
     ) -> SymbolicDomain:
         offset_provider_type = common.offset_provider_to_type(offset_provider)
 
@@ -144,28 +137,79 @@ def translate(
                     trace_shifts.Sentinel.ALL_NEIGHBORS,
                     trace_shifts.Sentinel.VALUE,
                 ]
-                horizontal_sizes: dict[str, itir.Expr]
-                if symbolic_domain_sizes is not None:
-                    horizontal_sizes = {
-                        k: im.ensure_expr(v) for k, v in symbolic_domain_sizes.items()
-                    }
-                else:
-                    # note: ugly but cheap re-computation, but should disappear
-                    assert common.is_offset_provider(offset_provider)
-                    horizontal_sizes = {
-                        k: im.literal(str(v), builtins.INTEGER_INDEX_BUILTIN)
-                        for k, v in _max_domain_sizes_by_location_type(offset_provider).items()
-                    }
-
                 old_dim = connectivity_type.source_dim
                 new_dim = connectivity_type.codomain
-
                 assert new_dim not in new_ranges or old_dim == new_dim
+                if symbolic_domain_sizes is not None and new_dim.value in symbolic_domain_sizes:
+                    new_range = SymbolicRange(
+                        im.literal(str(0), builtins.INTEGER_INDEX_BUILTIN),
+                        im.ensure_expr(symbolic_domain_sizes[new_dim.value]),
+                    )
+                else:
+                    assert common.is_offset_provider(offset_provider)
+                    connectivity = offset_provider[off.value]
+                    assert isinstance(connectivity, common.Connectivity)
+                    skip_value = connectivity.skip_value
+
+                    # fold & convert expr into actual integers
+                    range_exprs = new_ranges[old_dim].start, new_ranges[old_dim].stop
+                    range_exprs = tuple(
+                        collapse_tuple.CollapseTuple.apply(
+                            expr,
+                            within_stencil=False,
+                            allow_undeclared_symbols=True,
+                        )
+                        for expr in range_exprs
+                    )  # type: ignore[assignment]  # mypy not smart enough
+                    assert all(isinstance(expr, itir.Literal) for expr in range_exprs)
+                    start, stop = (int(literal.value) for literal in range_exprs)  # type: ignore[attr-defined]  # mypy does not understand assert above
+
+                    nb_index: slice | int
+                    if val in [trace_shifts.Sentinel.ALL_NEIGHBORS, trace_shifts.Sentinel.VALUE]:
+                        nb_index = slice(None)
+                    else:
+                        nb_index = val.value  # type: ignore[assignment]  # assert above
+
+                    accessed = connectivity.ndarray[start:stop, nb_index]
+
+                    if isinstance(val, itir.OffsetLiteral) and np.any(accessed == skip_value):
+                        # TODO(tehrengruber): Turn this into a configurable error. This is currently
+                        #  not possible since some test cases starting from ITIR containing
+                        #  `can_deref` might lead here. The frontend never emits such IR and domain
+                        #  inference runs after we transform reductions into stmts containing
+                        #  `can_deref`.
+                        warnings.warn(
+                            UserWarning(
+                                f"Translating '{self.as_expr()}' using '{shift[0].value}' has "
+                                f"an out-of-bounds access."
+                            ),
+                            stacklevel=2,
+                        )
+
+                    new_start, new_stop = accessed.min(), accessed.max() + 1  # type: ignore[attr-defined]  # TODO(havogt): improve typing for NDArrayObject
+
+                    fraction_accessed = np.unique(accessed).size / (new_stop - new_start)  # type: ignore[call-overload]  # TODO(havogt): improve typing for NDArrayObject
+
+                    if (
+                        fraction_accessed < NON_CONTIGUOUS_DOMAIN_WARNING_THRESHOLD
+                        and (_NON_CONTIGUOUS_DOMAIN_WARNING_COUNTER := +1)
+                        < NON_CONTIGUOUS_DOMAIN_MAX_WARNINGS
+                    ):
+                        warnings.warn(
+                            UserWarning(
+                                f"Translating '{self.as_expr()}' using '{shift[0].value}' requires "
+                                f"computations on many additional points "
+                                f"({round((1 - fraction_accessed) * 100)}%) in order to get a contiguous "
+                                f"domain. Please consider reordering your mesh."
+                            ),
+                            stacklevel=2,
+                        )
+
+                    new_range = SymbolicRange(
+                        im.literal(str(new_start), builtins.INTEGER_INDEX_BUILTIN),
+                        im.literal(str(new_stop), builtins.INTEGER_INDEX_BUILTIN),
+                    )
 
-                new_range = SymbolicRange(
-                    im.literal("0", builtins.INTEGER_INDEX_BUILTIN),
-                    horizontal_sizes[new_dim.value],
-                )
                 new_ranges = dict(
                     (dim, range_) if dim != old_dim else (new_dim, new_range)
                     for dim, range_ in new_ranges.items()

src/gt4py/next/iterator/ir_utils/ir_makers.py

@@ -453,7 +453,7 @@ def _impl(*its: itir.Expr) -> itir.FunCall:
 
 def domain(
     grid_type: Union[common.GridType, str],
-    ranges: dict[common.Dimension, tuple[itir.Expr, itir.Expr]],
+    ranges_or_domain: dict[common.Dimension, tuple[itir.Expr, itir.Expr]] | common.Domain,
 ) -> itir.FunCall:
     """
     >>> IDim = common.Dimension(value="IDim", kind=common.DimensionKind.HORIZONTAL)
@@ -463,6 +463,13 @@ def domain(
     >>> str(domain(common.GridType.UNSTRUCTURED, {IDim: (0, 10), JDim: (0, 20)}))
     'u⟨ IDimₕ: [0, 10[, JDimₕ: [0, 20[ ⟩'
     """
+    if isinstance(ranges_or_domain, common.Domain):
+        domain = ranges_or_domain
+        ranges = {d: (r.start, r.stop) for d, r in zip(domain.dims, domain.ranges)}
+    else:
+        assert isinstance(ranges_or_domain, dict)
+        ranges = ranges_or_domain
+
     if isinstance(grid_type, common.GridType):
         grid_type = f"{grid_type!s}_domain"
     expr = call(grid_type)(

src/gt4py/next/iterator/transforms/constant_folding.py

@@ -77,7 +77,7 @@ class Transformation(enum.Flag):
         # `maximum(maximum(a, 1), 1)` -> `maximum(a, 1)`
         FOLD_MIN_MAX = enum.auto()
 
-        # `maximum(a + 1), a)` -> `a + 1`
+        # `maximum(a + 1, a)` -> `a + 1`
         # `maximum(a + 1, a + (-1))` -> `a + maximum(1, -1)`
         FOLD_MIN_MAX_PLUS = enum.auto()
 

src/gt4py/next/iterator/transforms/global_tmps.py

@@ -312,7 +312,7 @@ def create_global_tmps(
     offset_provider: common.OffsetProvider | common.OffsetProviderType,
     #: A dictionary mapping axes names to their length. See :func:`infer_domain.infer_expr` for
     #: more details.
-    symbolic_domain_sizes: Optional[dict[str, str]] = None,
+    symbolic_domain_sizes: Optional[dict[str, str | itir.Expr]] = None,
     *,
     uids: Optional[eve_utils.UIDGenerator] = None,
 ) -> itir.Program: