next: a shifted read of a concat_where result is materialized in dace backend

[havogt]

A concat_where result that is read at a vertical offset tends to be materialized. The usual escape, pushing the shift onto the producer's own inputs so that producer and consumer become pointwise, is blocked by the concat_where.

This is a standard shape in the icon4py dycore: a vertical boundary level supplied by a different field than the interior, consumed by a vertical difference.

@gtx.field_operator
def theta_on_half_levels(interior: CKField, surface: CKField) -> CKField:
    return concat_where(K < 1, surface, interior)

@gtx.field_operator
def vertical_difference(interior: CKField, surface: CKField) -> CKField:
    theta = theta_on_half_levels(interior, surface)
    return theta(Koff[1]) - theta

In the icon4py dycore programs measured here, theta is materialized on both backends, for different reasons. On gtfn fuse_as_fieldop._arg_inline_predicate declines to inline an as_fieldop that is accessed at a non-center shift — not unconditionally, a single argument applied fieldop is inlined regardless of shifts. On dace there is no GTIR fusion stage at all, and on the SDFG the producer and consumer subsets are not pointwise, so MapFusionVertical declines; whether another SDFG level pass could still remove the temporary is open, see below.

The shift can be distributed into the branches, if the condition is translated with it:

as_fieldop(lambda it: deref(shift(K, d)(it)))(concat_where(u<K: [lo, hi[>, a, b))
  ->  concat_where(u<K: [lo - d, hi - d[>, shift_d(a), shift_d(b))

Sunk to the leaves, this removes the barrier and the temporary disappears. Translating the condition is what makes it safe: a branch may be readable on a wider index range than the one it is selected on.

What it buys, measured on GH200 with dace_gpu, icon4py mch_icon-ch1_medium, 1800 calls: 17% on compute_perturbed_quantities_and_interpolation, 0.5189 s to 0.4279 s, and 0.045 to 0.065 s off a 5.70 s model timestep loop. That program is where this shape dominates; nothing else moves beyond the 0.020 s noise floor.

An implementation exists on a fork: havogt#75. Three things are worth settling before proposing it here.

• It needs pruning of a concat_where whose condition selects a branch on an empty region, because the translated condition can produce exactly that. See fix[next]: decide concat_where pruning from the concat_where's domain #2767.
• Profitability is not unconditional. Without a guard declining bindings that transitively feed a scan, the pass was a net model-level regression despite making its target program faster — a scan argument is materialized anyway, so there is no barrier left to remove there.
• Duplication has to be bounded per distinct shift distance rather than per use site, and it matters more on dace, where nothing downstream re-shares copies.

Doing the same thing on the SDFG instead was tried and did not substitute: re-parameterizing the producer map so the subsets line up serves only one offset group, and duplicating producers until they fuse left the transients and the traffic unchanged.

inline_fuser is a much closer fit, and worth a separate look. It inlines the producing dataflow into the consumer's map scope with a symbol_mapping that evaluates the producer at the shifted index, so it does remove the transient rather than just fusing around it — on a synthetic SDFG with one intermediate read at k, k+1 and k+2, the array disappears from sdfg.arrays and is replaced by three scalar transients inside the consumer scope. Two things stop it here:

• it is exported but never called; gt_auto_optimize does not run it, and there is no cost model deciding when inlining is worth the duplicated work;
• an index component is required to have exactly one free symbol (inline_fuser.py:209-211), while gt4py emits i_IDim_gtx_horizontal - __out_IDim_range_0, two symbols. Intersecting with the consumer map's parameters makes it fire.

That would cover shifted single-producer intermediates in general, which is useful independently of concat_where. It does not cover concat_where itself: the intermediate is required to have a single producer (inline_fuser.py:242-243) and a concat node has one per branch. Extending it there means inlining each branch under a condition, and evaluating both branches eagerly is not sound in general — a branch's inputs need not be in bounds outside its own region. That is the information push_shifts still has and an SDFG level pass no longer does, since it runs after domain inference.