feat[next]: Compile time domains

[SF-N]

This PR adds support for domain sizes and domain inference at compile time. As a result temporary fields are only allocated on the domains that they are read on instead of everywhere.

@gtx.program(static_domains=True)
def testee(...):
        ...

Internally compile time domains are realized by the addition of a new argument descriptor FieldDomainDescriptor which holds the domain of the corresponding field argument in the form of a gtx.Domain. The output domains of SetAt statements are replaced with these static domains using the TransformGetDomainRange pass introduced in #2151. The domain inference is extended to use the actual values stored in the connectivity to deduce the required domain on which temporaries need to be computed. This domain is potentially non-contiguous and, since we only support contiguous domains, extended to the minimal contiguous domain. In case this extension leads to many unnecessary computations (default is 25% of the total size) a warning is emitted. The current default can be configured by setting NON_CONTIGUOUS_DOMAIN_WARNING_THRESHOLD in gt4py.next.iterator.ir_utils.domain_utils. The old behavior to compute everywhere is still used in case either a) the static_domain option is False or b) the backend is configured to compute everywhere by setting use_max_domain_range_on_unstructured_shift to True in GTFNTranslationStep or DaCeTranslator.

Note that compile time domains are only supported in jit mode right now, as we need to change the compile interface, which will happen in a later PR.

[tehrengruber]

Suggested change

	symbolic_domain_sizes: Optional[dict[str, str | itir.Expr]],
	symbolic_domain_sizes: Optional[dict[str, itir.Expr]],

Try if this is possible everywhere with low effort.

[havogt]

Still open?

[tehrengruber]

@havogt The case we discussed where a program does not create any temporaries, for which using static values for the domain has no benefits, but leads to unnecessary recompilation when changing the number of nodes, can not be easily detected in a way that would allow us to keep dynamic domain sizes. Such programs can only be identified after field operator fusion, which is after the replacement of the get_domain_range calls and domain inference. As soon as we introduce symbolic domain operations (e.g. union, intersection) we can keep the domain inference symbolic and only after operator fusion make everything static. As such I would postpone working on this case and stay with the workaround of the user selectively using static_domains to avoid unnecessary recompilation. Most likely this is anyway only important for PMAP-GO where we have a good amount of small operators without temporaries.