Merge VelocityAdvection into two programs with per-output domains

[havogt]

Part of the Audit ICON4Py programs with new GT4Py features project. First unit of work: VelocityAdvection in its final shape.

What this is for

A field_operator call is the granularity at which the GT4Py optimization pipeline works. When two calls are separate, the values passing between them are program-level fields the pipeline must treat as materialized and visible. Merging them moves that work inside one unit of optimization, so the pipeline becomes free to fuse it, keep it in registers, or still materialize it.

This is not a claim about kernel count or memory — both remain the pipeline's decision and neither is measured here. Performance will be measured separately, on santis; the pitch expects the new structure may temporarily hurt performance until optimization passes adapt.

What changed

VelocityAdvection issued 5 program calls across its two run methods; it now issues 2, each a single field_operator call with a per-output domain.

• run_predictor_step: 3 programs → compute_velocity_advection_in_predictor_step (9 outputs)
• run_corrector_step: 2 programs → compute_velocity_advection_in_corrector_step (3 outputs)

_allocate_local_fields goes from 3 fields to 1. horizontal_advection_of_w_at_edges_on_half_levels (ICON z_v_grad_w) and contravariant_corrected_w_at_cells_on_model_levels (ICON z_w_con_c_full) are no longer program fields; GT4Py derives their extents from the neighbour reads of their consumers. The predictor's vertical-momentum operator now computes the horizontal advection of w inline, which is how the corrector already did it.

Both host max_vertical_cfl reductions moved after the fused call — safe because apply_extra_diffusion_on_vn = True is a hardcoded literal at both sites and nothing between the calls consumed the reduction.

Three outputs got their own narrowed domain, each derived from mo_velocity_advection.f90 rather than from the previous icon4py bounds:

Output	Domain	Fortran
contravariant_correction_at_edges_on_model_levels	KDim: (nflatlev, nlev)	DO jk = nflatlev_jg, nlev
vertical_wind_advective_tendency (both programs)	KHalfDim: (1, nlev)	DO jk = 2, nlev
vn_on_half_levels	KHalfDim: (0, nlev+1)	absorbs the separate top-level extrapolation

Every other per-output domain is the one its predecessor program used, checked against the Fortran and against the neighbour reach of each consumer. The wider-than-necessary cell domains were deliberately not narrowed to match their consumers: that slack is what keeps the C2E/E2C reads in bounds.

Aliasing (fields both read and written in the same statement, the gt4py#2817 hazard) drops from 3 statements / 5 fields to 1 statement / 2 fields. Removing the two concat_where(1 <= KHalfDim, …) guards meant the operator stops reading the incoming tendency altogether, which cleared the corrector's aliasing entirely.

The four superseded programs are deleted along with their stencil tests; every field_operator they used survives and is exercised through the merged programs.

Verification

Suite	Backend	Result
test_compute_velocity_advection.py	gtfn_cpu	12 passed
test_compute_velocity_advection.py	dace_cpu	11 passed, 1 xfailed
test_velocity_advection.py --datatest-only	gtfn_cpu / dace_cpu	10 passed each
test_solve_nonhydro.py (4 dycore cases)	gtfn_cpu / dace_cpu	8 passed each

No --skip-stenciltest-verification anywhere. The embedded backend is not evidence here — these tests xfail on it, and an xfailing stencil test never evaluates its numpy reference, so no claim above rests on it.

Savepoint comparisons lost, both inherent to the demotions: z_v_grad_w and z_w_con_c_full. To compensate, z_vt_ie, z_kin_hor_e, z_w_concorr_me and a cfl_clipping mask comparison were added to the component tests, where none of them existed before.

Known issues

One xfail, strict=True. One dace_cpu variant hits assert read_right or read_left in splitting_tools.py:515, reached from SplitMemlet.can_be_applied: both flags come from (a < b) == True on SymPy relationals, which is undecidable — hence False both ways — when horizontal bounds are symbolic while vertical ones are static. The merge exposes this rather than causing it; the three predecessor programs pass the same variants on 4a064f23c. The model is unaffected: setup_program binds horizontal and vertical bounds together, so the fully-static shape is what ever gets compiled, and it passes on dace, as do both dace datatest suites. A gt4py-only reproducer is being prepared separately.

A pre-existing icon4py/ICON divergence, now documented in a test. When skip_compute_predictor_vertical_advection=True, icon4py leaves z_vt_ie untouched while ICON sets half level 0 unconditionally (mo_velocity_advection.f90:300, outside the IF (.NOT. lvn_only)). 29020 edges differ, level 0 only, on MCH_CH_R04B09. ICON itself notes that level is unused. Unchanged by this PR — the new assertion compares from level 1 and names the divergence.

scripts/python/compare_icon4py_openacc.py loses four stale entries: its model is one icon4py program per OpenACC kernel, and each merged program now aggregates several. Restoring that comparison needs a decision from whoever owns the perf tooling.

Module consolidation and dead-code removal

The 15 field_operators with no live consumer outside these two programs moved into compute_velocity_advection.py (now 1300 lines, 17 operators + 2 programs). compute_advection_in_{vertical,horizontal}_momentum_equation.py are deleted; compute_diagnostics_from_normal_wind.py is down to 36 lines, holding the two operators compute_horizontal_velocity_quantities.py still imports. This is pure code motion — every moved operator's AST is byte-identical before and after, and the generated C++ was checked to be byte-identical by diff against a separate clean build cache.

Separately, five superseded modules are deleted with their stencil tests: the old _mo_velocity_advection_stencil_05/14/15 ports (compute_horizontal_kinetic_energy, interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges, interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges, compute_maximum_cfl_and_clip_contravariant_vertical_velocity, interpolate_contravariant_vertical_velocity_to_full_levels). Nothing in any run path referenced them; they were kept alive only by their own tests and by each other.

Two of them defined operators whose names collide with live ones but whose implementations had diverged, so those tests were asserting behaviour the model does not have: the dead _interpolate_contravariant_vertical_velocity_to_full_levels has no concat_where and takes no nlev, so it does not handle the top level; the dead _compute_horizontal_kinetic_energy returns a 3-tuple where the live one returns a single field. Four numpy reference helpers from the deleted tests did have live consumers and were relocated rather than removed.

Not in scope

Narrowing vertical_cfl to KHalfDim: (max(2, nrdmax-2), nlev-3) is left open — it needs the host reduction slice narrowed in step, which is a decision that has not been made. No performance measurement or tuning.

Review follow-up

Three further commits address the review comments.

The vertical wind tendency is computed as three named contributions rather than accumulated through operators that each took the running total as an argument. _add_interpolated_horizontal_advection_of_w and _add_extra_diffusion_for_w_con_approaching_cfl now compute their contribution and are renamed _compute_*, with the cfl_clipping & owner_mask gate moved to the call site; _add_vertical_advection_of_w_to_advective_vertical_wind_tendency never added anything and is _compute_vertical_advection_of_w.

This is numerically unchanged including under mixed precision, which is the reason the additions stay sequenced instead of collapsing into one sum: ICON's ddt_w_adv_pc is a REAL(vp) array that is assigned once and then accumulated into, so each contribution is rounded to vpfloat before the next is added. Hoisting those conversions would drop roundings ICON performs. A TODO in the source asks whether that was deliberate on the Fortran side.

Computing the CFL is separated from clipping w. _compute_maximum_cfl_and_clip_contravariant_vertical_velocity did both and returned a three-tuple; it is split into _compute_cfl and _clip_contravariant_corrected_w, so the enclosing operator reads as correct w, compute the clipping, clip. The clip now sits outside the concat_where that restricts the CFL to the damping range — behaviour-preserving because cfl_clipping is False outside it, so both wheres degenerate to passing the velocity through, which is what the else-branch returned. _compute_horizontal_advection_of_w also moves next to the interpolation that is its only consumer.

One tangential wind stencil instead of two. _compute_tangential_wind_vp was _compute_tangential_wind_wp wrapped in a single astype, which made a usage search return a second definition rather than a consumer. There is now one operator and one program, both working precision, and the three call sites that want variable precision do the conversion where a reader can see it. The compute_tangential_wind_vp program goes with the operator it wrapped — a program body is a field operator call with out= and domain=, so the astype cannot live there; nothing outside its own stencil test used it.

Two review points remain open: finding a pattern that avoids passing tangential_wind_on_half_levels in purely to pass it through under skip_compute_predictor_vertical_advection, which looks like it needs GT4Py support rather than a local change; and whether the three CFL steps should be inlined at both call sites rather than kept in a shared wrapper.

---

🤖 Generated with Claude Code

https://claude.ai/code/session_017s66WxborSWhPes8eHxiSF

[havogt]

cscs-ci run default;MODEL_SUBPACKAGES=dycore:common:tracer_advection;SESSIONS=model

[havogt]

Overview of the current state https://claude.ai/code/artifact/1441c859-e2ed-49b0-a517-6613716af1e3

[nfarabullini]

just some comments

[nfarabullini]

these are incredibly long names, can you shorten?

[nfarabullini]

The comment doesn't match the action

[havogt]

I think it does, but maybe the comment is not clear enough. See the corresponding comment in the predictor field_operator.

[nfarabullini]

can we make this cleaner and define these above?

[havogt]

cscs-ci run default;MODEL_SUBPACKAGES=dycore:common:tracer_advection;SESSIONS=model

[havogt]

cscs-ci run default;MODEL_SUBPACKAGES=dycore:common:tracer_advection;SESSIONS=model

[github-actions]

When developing, you can test your changes on CSCS CI before merge with the default pipeline: cscs-ci run default. This will run a default subset of tests.

You can pass options to override pipeline variables, for example:

• cscs-ci run default;BACKENDS=gtfn_cpu;LEVELS=unit
• cscs-ci run default;MODEL_SUBPACKAGES=common:driver;SESSIONS=model

Avoid running the pipeline for all tests when you are developing.

Available options are:

• SESSIONS: model, model_mpi, or tools (correspond to nox sessions)
• MODEL_SUBSETS: datatest, basic, or stencils (correspond to nox session selections)
• MODEL_SUBPACKAGES: subpackages for non-MPI tests (last component, e.g. diffusion, driver)
• MODEL_MPI_SUBPACKAGES: subpackages for MPI tests (as above)
• BACKENDS: backends
• GRIDS: grids for stencil tests (simple, icon_regional, or icon_global)
• LEVELS: testing level for non-stencil tests (unit or integration)

For each option, all can be used as a shorthand for all possible values of that variable, e.g. LEVELS=all.

Multiple values can be given to each option with : used as the separator (; separates options and , separates pipelines).

See scripts/python/generate_ci_pipeline.py and noxfile.py for available values for each option.

The all pipeline can be run with cscs-ci run all. This will run all icon4py tests in CSCS CI which can be expensive. This pipeline runs on a schedule on main, and can be run when extensive validation is needed (e.g. before releases).

Merging

Once your PR is approved and ready for merging, add it to the merge queue. The merge CSCS CI pipeline will run automatically on the merge-queue branch and must pass before the PR is merged. A dummy merge check will be triggered on the PR itself since it's required to add a PR to the merge queue.

Optional Tests

To run benchmarks you can use:

• cscs-ci run benchmark-bencher

For more detailed information please look at CI in the EXCLAIM universe.