next: how to do domain inference for dynamic shift

[havogt]

@field_operator
def foo(..., k_offset):
    return compute_k_field(...)(as_offset(Koff, k_offset))

lowers to

return as_fieldop(lambda k_offset, k_field: deref(shift(Koff, deref(k_offset))), out_domain)(k_offset, compute_k_field(...))

How can we compute the domain that is needed for compute_k_field?

a) take the out_domain: probably wrong in some cases (e.g. between k and khalf) b) take the domain of k_offset (this is the same as a), as the domain of k_offset is taken from out_domain c) user has to annotate the domain explicitly (either explicit (requires to pass in the domain bounds -> ugly) or relative to a domain or field; both introduce slicing-like syntax and requires clear concept definition) d) inline everything below for now and find a better way to express the pattern

Note: in icon4py all as_offsets are in the vertical

[havogt]

Does a) work in icon4py?

Probably broken in https://github.com/C2SM/icon4py/blob/1acfb78134f09d254457ea97ae81423da54c28bb/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py#L38 where we compute a single output level, but access multiple levels with an as_offset shift.

[tehrengruber]

Sketch for domain inference:

• Forward propagate the domain from the inputs. This should be unconditionally possible in a single pass. The domain of field(as_offset(Ioff, offset_field)) is intersection(domain(field), domain(offset_field)). If the input IR does not have any Out-Of-Bounds memory accesses so does direct execution of the resulting IR here (since we only compute values were the input is defined).
• Backward propagate the domain. In case we encounter a stencil with a dynamic shift as_fieldop(dynamic_shift_stencil, forward_inferred_domain) we can either
  • Stop propagation to inputs of the as_fieldop. Since the domains of the inputs are inferred already by the forward propagation all domains are inferred and we can execute as usual, but might compute (way) too much. This is trivial to implement.
  • Introduce a static bound, i.e. known at compile time, on the (relative) offsets in the offset_field. Given the domain where the as_fieldop is read, let's call it accessed_domain, the resulting domain is intersection(forward_inferred_domain, extend(accessed_domain, max_relative_offset))
  • Instead of a static offset use a similar approach as in the inverse image the resulting domain is intersection(forward_inferred_domain, domain(min(offset_field_with_absolute_indices), max(offset_field_with_absolute_indices)).