← all records

Audit ICON4Py programs with new GT4Py features

pitch-c58f20 · in Runtime improvements

  • No No-gos section. The pitch template asks for one — it is what keeps the appetite honest.

Problem

Most ICON4Py programs were written against limitations that recent GT4Py versions have lifted: K-staggering, multiple output domains, custom collections. Some of these are mostly syntactic, but multiple output domains changes the GT4Py representation itself. The goal there is to expose more optimization potential to the backends -- but because the structure changes, performance may temporarily get worse where optimization passes no longer trigger.

Underneath it sits the unresolved question of why AMD and NVIDIA GPUs differ so much, which these changes give a fresh angle on.

Appetite

Full cycle. Eight person-weeks is what the five named people can actually give between them rather than five full-timers: Hannes full, Nikki half, Edoardo two weeks, and Ioannis and Christoph with no figure in the cycle-38 roster at all.

Solution

K-staggering is the critical part, because it is all or nothing: once KHalfDim is introduced for a field, every producer and consumer must move at the same time. Performance matches main from the next GT4Py version. Full support needs a fix for domain inference of the as_offset case, though that is not blocking -- the single case where the bug triggers has a workaround that translates KHalfDim back to KDim locally.

Multiple output domains: one field_operator call per program, which is the granularity GT4Py optimizes at, and in some cases that removes program-level temporaries.

Additional performance work: a thread block size sweep on A100 alongside the ongoing MI300 and GH200 work; workgroup-to-XCD assignment on AMD MI300A in DaCe; and exploring whether limiting Map fusion helps, since simpler kernels put less pressure on memory and cache.

Rabbit holes

Improving the multiple-output-domain syntax is not part of this. Readability regressions are accepted for now, and ideas for improvement can be discussed during the cycle.

No-gos

Progress

  • K-staggering: #1429 merged 02.09.2026
  • Velocity advection merged into two programs with per-output domains: #1460, merged 21.09.2026
  • Vertically implicit solver: #1468, open. Its state is not well defined in GT4Py after the staggering -- the scan range is deduced from the vertical dimension and there are now two of those. It happens to work in compiled backends and is broken in embedded. Ioannis also found a regression in vertically_implicit_solver_at_corrector_step after the KHalfDim merge.
  • Tracer advection (Nikki)
  • Diffusion (Edoardo)