feat[next]: Precompilation with static domains

src/gt4py/next/ffront/decorator.py

@@ -163,6 +163,7 @@ def compile(
         | common.OffsetProvider
         | list[common.OffsetProviderType | common.OffsetProvider]
         | None = None,
+        static_domains: dict[str, dict[common.Dimension, tuple[int, int]]] | None = None,
         **static_args: list[xtyping.MaybeNestedInTuple[core_defs.Scalar]],
     ) -> Self:
         """
@@ -204,12 +205,25 @@ def compile(
             for op in offset_provider
         )
 
-        self._compiled_programs.compile(offset_providers=offset_provider, **static_args)
+        if self.compilation_options.static_domains and static_domains is None:
+            raise ValueError(
+                "Static domains option is enabled, but no static domain information was "
+                "provided. Missing required argument 'static_domains'."
+            )
+        if not self.compilation_options.static_domains and static_domains is not None:
+            raise ValueError(
+                "Static domains may not be provided when 'static_domains' option is disabled."
+            )
+
+        self._compiled_programs.compile(
+            offset_providers=offset_provider, static_domains=static_domains, **static_args
+        )
         return self
 
 
 def _field_domain_descriptor_mapping_from_func_type(func_type: ts.FunctionType) -> list[str]:
     static_domain_args = []
+    assert func_type.pos_only_args == []
     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):
@@ -491,6 +505,7 @@ def compile(
         | common.OffsetProvider
         | list[common.OffsetProviderType | common.OffsetProvider]
         | None = None,
+        static_domains: dict[str, dict[common.Dimension, tuple[int, int]]] | None = None,
         **static_args: list[xtyping.MaybeNestedInTuple[core_defs.Scalar]],
     ) -> Self:
         raise NotImplementedError("Compilation of programs with bound arguments is not implemented")

src/gt4py/next/otf/compiled_program.py

@@ -710,11 +710,21 @@ def _compile_variant(
     def compile(
         self,
         offset_providers: list[common.OffsetProvider | common.OffsetProviderType],
+        static_domains: dict[str, dict[common.Dimension, tuple[int, int]]] | None = None,
         **static_args: list[ScalarOrTupleOfScalars],
     ) -> None:
         """
         Compiles the program for all combinations of static arguments and the given 'OffsetProviderType'.
 
+        Args:
+            offset_providers: List of offset providers to compile for.
+            static_domains: Per-field domain specification.
+                Example::
+                    {
+                        "input_field": {IDim: (0, 55), JDim: (0, 30)},
+                        "out": {IHalfDim: (0, 52), JDim: (0, 30)},
+                    }
+
         Note: In case you want to compile for specific combinations of static arguments (instead
         of the combinatoral), you can call compile multiples times.
 
@@ -724,6 +734,59 @@ def compile(
             pool.compile(static_arg0=[0], static_arg1=[2]).compile(static_arg=[1], static_arg1=[3])
                 will compile for (0,2), (1,3)
         """
+
+        def _build_field_domain_descriptors(
+            program_type: ts_ffront.ProgramType,
+            static_domains: dict[str, dict[common.Dimension, tuple[int, int]]],
+        ) -> dict[str, arguments.FieldDomainDescriptor]:
+            field_domain_descriptors: dict[str, arguments.FieldDomainDescriptor] = {}
+            assert program_type.definition.pos_only_args == []
+
+            matched_keys = set()
+            param_types = (
+                program_type.definition.pos_or_kw_args | program_type.definition.kw_only_args
+            )
+            for arg_name, arg_type_ in param_types.items():
+                for el_type_, path in type_info.primitive_constituents(
+                    arg_type_, with_path_arg=True
+                ):
+                    if isinstance(el_type_, ts.FieldType):
+                        path_as_expr = "".join(map(lambda idx: f"[{idx}]", path))
+                        field_expr = f"{arg_name}{path_as_expr}"
+
+                        if field_expr not in static_domains:
+                            raise ValueError(
+                                f"Missing static domain for field '{field_expr}'. "
+                                f"Expected domains for all field parameters. "
+                                f"Available keys in static_domains: {list(static_domains.keys())}"
+                            )
+
+                        field_domain = static_domains[field_expr]
+                        matched_keys.add(field_expr)
+
+                        expected_dims = set(el_type_.dims)
+                        provided_dims = set(field_domain.keys())
+                        if expected_dims != provided_dims:
+                            raise ValueError(
+                                f"Domain dimension mismatch for field '{field_expr}': "
+                                f"field has dimensions {set(d.value for d in expected_dims)}, "
+                                f"but static_domains provides {set(d.value for d in provided_dims)}."
+                            )
+
+                        domain_ranges = {dim: field_domain[dim] for dim in el_type_.dims}
+                        field_domain_descriptors[field_expr] = arguments.FieldDomainDescriptor(
+                            common.domain(domain_ranges)
+                        )
+
+            extra_keys = set(static_domains.keys()) - matched_keys
+            if extra_keys:
+                raise ValueError(
+                    f"static_domains contains keys that do not correspond to any field parameter: "
+                    f"{list(extra_keys)}. Valid field expressions are: {list(matched_keys)}"
+                )
+
+            return field_domain_descriptors
+
         for offset_provider in offset_providers:  # not included in product for better type checking
             for static_values in itertools.product(*static_args.values()):
                 # The argument descriptors built here must be kept in sync with the
@@ -738,6 +801,10 @@ def compile(
                             strict=True,
                         )
                     )
+                if static_domains:
+                    argument_descriptors[arguments.FieldDomainDescriptor] = (
+                        _build_field_domain_descriptors(self.program_type, static_domains)  # type: ignore[assignment]
+                    )
                 self._compile_variant(
                     argument_descriptors=argument_descriptors,
                     offset_provider=offset_provider,

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_compiled_program.py

@@ -39,7 +39,6 @@
     skip_value_mesh,
 )
 
-from gt4py.next.otf import arguments
 
 _raise_on_compile = mock.Mock()
 _raise_on_compile.compile.side_effect = AssertionError("This function should never be called.")
@@ -66,6 +65,7 @@ def testee(a: cases.IField, b: cases.IField, out: cases.IField):
         testee_op(a, b, out=out)
 
     wrap_in_program = request.param
+
     if wrap_in_program:
         return testee
     else:
@@ -1005,7 +1005,8 @@ def compile(self, program, compile_time_args):
 
 @pytest.mark.uses_tuple_args
 @pytest.mark.uses_tuple_returns
-def test_compile_variants_decorator_static_domains(cartesian_case):
+@pytest.mark.parametrize("precompile", [True, False], ids=["precompile", "run"])
+def test_compile_variants_decorator_static_domains(cartesian_case, precompile):
     if cartesian_case.backend is None:
         pytest.skip("Embedded compiled program doesn't make sense.")
 
@@ -1033,6 +1034,22 @@ def testee(inp: tuple[cases.IField, cases.IField, float], out: NamedTupleNamedCo
 
     inp = cases.allocate(cartesian_case, testee, "inp")()
     out = cases.allocate(cartesian_case, testee, "out")()
+    if precompile:
+        static_domain = {cases.IDim: (0, cartesian_case.default_sizes[cases.IDim])}
+        testee.compile(
+            offset_provider=cartesian_case.offset_provider,
+            static_domains={
+                "inp[0]": static_domain,
+                "inp[1]": static_domain,
+                "out[0]": static_domain,
+                "out[1]": static_domain,
+            },
+        )
+        gtx.wait_for_compilation()
+    else:
+        testee(inp, out, offset_provider={})
+        assert np.allclose(inp[0].ndarray, out[0].ndarray)
+        assert np.allclose(inp[1].ndarray, out[1].ndarray)
 
     testee(inp, out, offset_provider={})
     assert np.allclose(inp[0].ndarray, out[0].ndarray)