feat[next]: Precompilation with static domains#2483
feat[next]: Precompilation with static domains#2483SF-N wants to merge 8 commits intoGridTools:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Adds support for providing static domain ranges during ahead-of-time compilation so programs decorated with static_domains=True can be precompiled without needing runtime field instances.
Changes:
- Extend
CompiledProgramsPool.compile(...)to optionally constructFieldDomainDescriptorcontexts from user-provided domain ranges. - Plumb a new
static_domainsargument throughProgram.compile(...)to the compiled-program pool. - Add an integration test covering precompilation with static domains for tuple arguments.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 8 comments.
| File | Description |
|---|---|
| tests/next_tests/integration_tests/feature_tests/ffront_tests/test_compiled_program.py | Adjusts compilation parametrization and adds a new test for precompiling with provided static domains. |
| src/gt4py/next/otf/compiled_program.py | Builds FieldDomainDescriptor argument descriptors from provided static domain ranges during precompilation. |
| src/gt4py/next/ffront/decorator.py | Forwards a new static_domains argument through the public compile() API. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_compiled_program.py
Outdated
Show resolved
Hide resolved
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_compiled_program.py
Outdated
Show resolved
Hide resolved
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_compiled_program.py
Outdated
Show resolved
Hide resolved
…into static_domain_precompilation
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| ) | ||
| ), | ||
| } | ||
| if static_domains: |
There was a problem hiding this comment.
static_domains is checked via truthiness (if static_domains:), so passing an empty dict (which is still a non-None value per the public API) will silently skip building FieldDomainDescriptors. This can lead to compiling a variant without any domain descriptors even though the caller explicitly provided static_domains. Consider switching this condition to if static_domains is not None: (and let _build_field_domain_descriptors raise for missing dims as needed).
| if static_domains: | |
| if static_domains is not None: |
| argument_descriptor_dict[arguments.FieldDomainDescriptor] = ( | ||
| _build_field_domain_descriptors(self.program_type, static_domains) # type: ignore[assignment] | ||
| ) |
There was a problem hiding this comment.
There is a # type: ignore[assignment] on the FieldDomainDescriptor insertion because _build_field_domain_descriptors returns dict[str, FieldDomainDescriptor] while argument_descriptor_dict is typed as dict[str, ArgStaticDescriptor]. Consider widening the helper’s return type (e.g., to dict[str, arguments.ArgStaticDescriptor]) or casting at the assignment site to avoid relying on type: ignore here.
Add the option to precompile a program with static domains. Previously this was only supported in JIT mode. Instead of providing the domain of each field argument (which is the format used internally) the user just supplies the range of each dimension which is the most common case. Later we can extend this so that different ranges can be supplied, e.g. for fields that are only defined on a subdomain.
Note that a consequence of this syntax is that it is not possible to have a parameter called
static_domainsand pass a static value to it.