Split the vertically implicit solver calls by vertical dimension

[havogt]

Description

Each field-operator call from the two vertically implicit solver programs now has a single vertical domain.

Before, vertically_implicit_solver_at_{predictor,corrector}_step called one fused field operator with a per-output tuple domain= (KHalfDim for the half-level outputs, KDim for the model-level ones, same index range). That form never states the column the solve_w scans run on: on gtfn/dace it is inferred as the union of the outputs' demands and happens to coincide with the intended one because both ranges are (top, surface-1); on embedded it fails outright (AttributeError: 'tuple' object has no attribute 'dim'). A single shared domain for mixed KDim/KHalfDim outputs is rejected at lowering by the compiled backends, so the calls are split by output dimension instead:

• predictor: _solve_w_at_predictor_step (half levels, writes next_w) + _compute_thermodynamic_variables_at_predictor_step (model levels, writes next_rho, next_exner, next_theta_v, dwdz, exner_dynamical_increment)
• corrector: _solve_w_and_update_vertical_fluxes_at_corrector_step (half levels, writes next_w and the two vertical fluxes) + _compute_thermodynamic_variables_at_corrector_step (model levels)
• the explicit terms and solver coefficients are computed by a shared _compute_explicit_terms_and_solver_coefficients in both calls (no scratch fields; program signatures and callers are unchanged)
• the half-level solver call runs over half levels top+1 .. surface-1, which is the column its scans actually operate on; next_w at the top is written by its own one-level call, like the surface boundary condition already was
• the IAU increments are applied outside the shared operator: leaving the incremented rho term unused behind a runtime is_iau_active gave gtfn 'target_domain' cannot be 'NEVER' in the variant with no static arguments

The embedded xfail (uses_concat_where) stays on the two solver stencil-test classes: the split alone does not make the solver pass on embedded, because the scalar concat_where branches (broadcast(0.0, (CellDim,))) still fail there with Cannot compute length of open 'UnitRange'. #1290 fixes those with value_of_size_* and is expected to rebase on this PR; one more site of that kind is not covered by #1290 yet, the corrector's at_first_substep mass/volume-flux reset.

Verified

• both solver stencil-test classes on embedded: 21/21, on this branch merged with Enable more tests with embedded backend #1290 plus the flux-reset site
• test_solve_nonhydro datatests (test_nonhydro_{predictor,corrector}_step, test_run_solve_nonhydro_{single,multi}_step, test_vertically_implicit_solver_at_{predictor,corrector}_step, MCH_CH_R04B09 + EXCLAIM_APE) on embedded: 12/12, same combination
• predictor stencil-test class on gtfn_cpu: 9/9 on this branch alone
• corrector on gtfn/dace: not run locally (compile does not fit); needs CSCS CI

Open design choice

The second call recomputes the flux divergence (2 neighbor_sum over C2E) and 5 pointwise expressions and re-reads their inputs. The alternative is three calls with three model-level scratch fields (z_beta, z_rho_expl, z_exner_expl) threaded through solve_nonhydro and the bindings. The stencil-test benchmark on the simple grid cannot distinguish the two; a real-grid GPU measurement is needed.

🤖 Generated with Claude Code

https://claude.ai/code/session_015kcU44AVBRvbf5HkLHKpc1

[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.