Combine compute_perturbed_quantities_and_interpolation further #896

muellch · 2025-10-08T11:59:45Z

further in-lining of stencils 1 to 13 to less field_operators calls within the program:
compute_perturbed_quantities_and_interpolation

Co-authored-by: Hannes Vogt <[email protected]>

Fix bug in the input arguments of `IconSerialDataProvider.from_grid_savepoint` in `initialization_utils.py` of the driver package after this PR #829 is merged.

On the mch-ch1_medium experiment - this is 4% faster on the full timestep compared to the version without `ravel`; - within fluctuations there is no difference between no reduction and this version.

Co-authored-by: Edoardo Paone <[email protected]>

delete top level README.md from tools

@philip-paul-mueller

- Removed `vertical_mass_flux_at_cells_on_half_levels`, `tridiagonal_alpha_coeff_at_cells_on_half_levels`, `tridiagonal_beta_coeff_at_cells_on_model_levels`, `exner_explicit_term` and `rho_explicit_term` from intermediate fields passed around - Combined `_vertically_implicit_solver_at_corrector/predictor_step_before/after_solving_w` - Improved initialization of scans (level 0 and 80) (together with @philip-paul-mueller ) - Removed any unnecessary initializations related to them (i.e. `next_w`) - Handle related unit tests - Make sure that the values for the `n_lev` vertical level are passed from `_set_surface_boundary_condtion_for_computation_of_w` to `_vertically_implicit_solver_at_corrector/predictor_step_before_solving_w` properly - Improves DaCe handling of temporaries --------- Co-authored-by: "Ong Chia Rui" <[email protected]>

Fix an issue in velocity advection tests, that caused the dace programs to fail: a cupy array was expected as argument, a numpy array was passed.

This PR enables CI tests on the dace backend. - In the Github Actions, we run the stencil tests on the CPU backend. - In the CSCS CI, we run the tests with serialized data on the GPU backend, only for diffusion and dycore in order to save compute resources.

…into combine_1_13_further

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

edopao · 2026-01-07T12:19:24Z

There are variations in performance, between the version I pushed before vacation and the version from today (see _2026 suffix), but in general the performance is in pair with baseline. I actually cannot explain why there was a speed up of adv_in_vertical_moment_eq in the previous version.

edopao

The performance is in pair with baseline, and additional speedup will be provided by #951. This PR is needed to remove a write-write pattern which is not supported by dace state fusion.

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

nfarabullini · 2026-01-08T14:46:29Z

cscs-ci run default

edopao · 2026-01-08T14:47:13Z

I now see a clear speedup:

The requested change has been applied.

havogt

Mainly questions on how to improve readability, now that boundary conditions close to the interior computations.
Additionally, we can now also combine the last field_operator into this one since we are already using multiple output domains.

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

havogt · 2026-01-09T11:54:29Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

        concat_where(
-            (maximum(1, nflatlev) <= dims.KDim),
+            dims.KDim < surface_level - 1,
            _interpolate_cell_field_to_half_levels_vp(
                wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner
            ),
-            exner_at_cells_on_half_levels,
+            _interpolate_to_surface(
+                wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner
+            ),
        )


I think this should be a named operator (it's repeated below). What should be its name?

we shouldn't change this name

apologies, let me rephrase: it is my personal opinion that this concat_where should stay as is.

havogt · 2026-01-09T11:54:44Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

    perturbed_theta_v_at_cells_on_half_levels = concat_where(
-        dims.KDim >= 1,
+        dims.KDim < surface_level - 1,
        _interpolate_cell_field_to_half_levels_vp(
            wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels
        ),
-        broadcast(0.0, (dims.CellDim, dims.KDim)),
+        _interpolate_to_surface(
+            wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels
+        ),
    )


this is the same operator again

yes, we keep the same name

havogt · 2026-01-09T11:55:40Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

    theta_v_at_cells_on_half_levels = concat_where(
-        dims.KDim >= 1,
+        dims.KDim < surface_level - 1,
        _interpolate_cell_field_to_half_levels_wp(
            wgtfac_c=wgtfac_c_wp, interpolant=current_theta_v
        ),
-        theta_v_at_cells_on_half_levels,
+        reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels,
    )


this is a similar operator, but for wp and with different boundary condition

same answer as above

@OngChia we should check again if this could actually be expressed with the same operator including boundary condition (but on wp instead of vp) or not.

havogt · 2026-01-09T11:57:05Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

+        -vpfloat("0.5")
+        * (
+            (
+                perturbed_theta_v_at_cells_on_half_levels
+                - perturbed_theta_v_at_cells_on_half_levels(Koff[1])
+            )
+            * d2dexdz2_fac1_mc
+            + perturbed_theta_v_at_cells_on_model_levels * d2dexdz2_fac2_mc
+        )


Should this be a generic d2d2z operator? It seems it could go directly from model_levels to model_levels.

since it doesn't seem that this computation is used anywhere else, I don't see the advantage of placing it in a field_operator

I meant it could be directly written as a 3-point stencil from perturbed_theta_v_at_cells_on_model_levels to the derivate on model_levels. Possibly this might be closer to the math, but not sure...

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

havogt · 2026-01-09T12:01:47Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

-        _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho),
-        rho_at_cells_on_half_levels,
-    )
+    rho_at_cells_on_half_levels = _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho)


maybe also here the same operator with boundary conditions could be used (although it's boundary value seems never be used, but optimizations should remove it).

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

havogt · 2026-01-09T12:08:27Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (1, surface_level - 1),


Suggested change

dims.KDim: (1, surface_level - 1),

dims.KDim: (model_top + 1, surface_level - 1),

I am not sure if I like the namings for the vertical ranges in general. Maybe model_top should be a global constant (FrozenNamespace)? Why is surface_level referring to nlev+1? It's actually half_level_surface and surface_level - 1 might be just surface_level?

the concat_where condition was specifically dims.KDim >= 1, hence my assumption that boundary conditions start from 1

Yes, excactly, but model_top==0 always, therefore either we should remove model_top or express it relative to model_top

what if model_top will change in value but the kdim>=1 condition should stay as is?

The 1 refers to 1 above model_top. I didn't check who introduced model_top, but it would have been more consistent to also replace the 1s with model_top+1 when this refactoring was done.

havogt · 2026-01-09T12:12:45Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (maximum(1, nflatlev), surface_level),


Would be interesting to explore if we should actually define all fields that don't start at 0 aka model_top on a smaller domain. Is this field ever read < nflatlev? Similar are the fields that are computed from 1 ever used at level 0?

This is not a priority, we can do this in another PR

…rface

nfarabullini · 2026-01-13T10:08:59Z

cscs-ci run default

github-actions · 2026-01-14T07:15:09Z

Mandatory Tests

Please make sure you run these tests via comment before you merge!

cscs-ci run default

Optional Tests

To run benchmarks you can use:

cscs-ci run benchmark-bencher

To run tests and benchmarks with the DaCe backend you can use:

cscs-ci run dace

To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:

cscs-ci run extra

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

havogt · 2026-01-14T09:05:58Z

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py

+        perturbed_rho_at_cells_on_model_levels,
+        perturbed_theta_v_at_cells_on_model_levels,
+    ) = concat_where(
+        (dims.CellDim >= start_cell_halo_level_2) & (dims.CellDim < end_cell_halo_level_2),


We should either

just drop the concat_where here, or

move the bounds into the output domain.

Probably the second is the better choice.

Reason: Assuming both current_rho and current_theta_v as well as the reference fields are correct in (start_cell_lateral_boundary_level_3, end_cell_halo), then it should be correct to compute it everywhere. Since the now inlined operator _compute_perturbation_of_rho_and_theta computed on a smaller domain, I assume the result is only needed in that domain.

On a second look it seems we are already computing both fields above and there is no change in between. It seems we should delete this call and probably shrink the output domain to level2. (Or keep it the bigger domain as this might be better for fusion in dace, let's check with @edopao.)

Does the bigger domain imply the use of concat_where? In this case it is preferable to shrink the domain.

No concat_where.

Then let's try with the bigger domain.

nfarabullini · 2026-01-16T11:23:00Z

@havogt and @OngChia will take over the rest of the work

DropD and others added 30 commits September 4, 2025 14:42

add first batch of program tests to continuous benchmarking

a460d50

fix static params for apply_diffusion_to_vn

84abc04

Co-authored-by: Hannes Vogt <[email protected]>

improve static params for apply_diffusion_to_w_... test

50e8365

Co-authored-by: Hannes Vogt <[email protected]>

Merge branch 'main' into expand-continuous-benchmarking

ad8d449

add static params and remove indirectly used prog from benchmarking

aacd35b

update dockerfile to trigger a rebuild with newer bencher CLI

d042c98

add compute_derived_horizontal_winds_... to benchmarking

ba94c10

Merge branch 'main' into expand-continuous-benchmarking

4571a34

temporarily remove failing parametrization case.

3dcdd34

add compute_advection_in_vertical... program

367ba1f

add benchmarking for advection_in_horizontal... program

c4b33c3

add benchmarking for compute_perturbed... program

bdc3296

add benchmarking for 'compute_hydrostatic...' program

5836fbf

fix precommit messages

7d041d5

fix a bug in grid-savepoint reading during driver initialization (#866)

a30e5da

Fix bug in the input arguments of `IconSerialDataProvider.from_grid_savepoint` in `initialization_utils.py` of the driver package after this PR #829 is merged.

Fix cupy cfl reduction performance by raveling the array (#872)

ca2f4f2

On the mch-ch1_medium experiment - this is 4% faster on the full timestep compared to the version without `ravel`; - within fluctuations there is no difference between no reduction and this version.

Update to GT4Py v1.0.9 (dace fix) (#840)

36f5583

Co-authored-by: Edoardo Paone <[email protected]>

delete outdated top level README.md from tools (#870)

510dfae

delete top level README.md from tools

Fix memory allocator for zero fields in velocity advection tests (#874)

937128b

Fix an issue in velocity advection tests, that caused the dace programs to fail: a cupy array was expected as argument, a numpy array was passed.

Merge branch 'main' into expand-continuous-benchmarking

02d031f

add benchmarking for 'vertically_implicit_solver...' program

bddb09d

avoid dace backend corner cases in stencil tests

48c9b08

Merge branch 'main' into expand-continuous-benchmarking

4ea9499

add benchmarking for 'interpolate_rho_theta...' program

70e1720

fix typing for static variants in stencil tests

1479e80

add benchmarking for 'apply_divergence_damping...' program

cf730e8

add benchmarking for 'vertically_implicit...corrector...' program

47d647a

Merge branch 'combine_1_13_further' of https://github.com/C2SM/icon4py …

e8c5f26

…into combine_1_13_further

nfarabullini force-pushed the combine_1_13_further branch from f2d42cc to e8c5f26 Compare January 7, 2026 09:03

edopao reviewed Jan 7, 2026

View reviewed changes

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py Show resolved Hide resolved

edopao requested a review from havogt January 7, 2026 12:44

edopao approved these changes Jan 7, 2026

View reviewed changes

havogt reviewed Jan 8, 2026

View reviewed changes

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py Show resolved Hide resolved

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py Outdated Show resolved Hide resolved

edopao previously requested changes Jan 8, 2026

View reviewed changes

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py Show resolved Hide resolved

nfarabullini force-pushed the combine_1_13_further branch from 4839805 to e8c5f26 Compare January 8, 2026 10:21

nfarabullini added 2 commits January 8, 2026 11:21

fixed small mistake

83315ec

extende concat_where

b3f471f

edopao reviewed Jan 8, 2026

View reviewed changes

...e/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py Show resolved Hide resolved

nfarabullini added 3 commits January 8, 2026 14:24

further inlining

33ebde4

further inlining

02f4c70

Merge branch 'main' into combine_1_13_further

c9ba07d

nfarabullini requested review from edopao and havogt January 8, 2026 15:05

havogt reviewed Jan 9, 2026

View reviewed changes

nfarabullini added 4 commits January 12, 2026 09:06

cleanup

a912450

merged last field_operator in

924e774

edit to stencil test

73921e7

renaming of interpolate_to_surface to extrapolate_quadratically_to_su…

05a535e

…rface

Merge branch 'main' into combine_1_13_further

07d5b72

havogt reviewed Jan 14, 2026

View reviewed changes

	dims.KDim: (1, surface_level - 1),
	dims.KDim: (model_top + 1, surface_level - 1),

Combine compute_perturbed_quantities_and_interpolation further #896

Are you sure you want to change the base?

Combine compute_perturbed_quantities_and_interpolation further #896

Uh oh!

Conversation

muellch commented Oct 8, 2025 • edited by nfarabullini Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

edopao commented Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

edopao left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nfarabullini commented Jan 8, 2026

Uh oh!

edopao commented Jan 8, 2026

Uh oh!

havogt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nfarabullini Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nfarabullini commented Jan 13, 2026

Uh oh!

github-actions bot commented Jan 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

muellch commented Oct 8, 2025 •

edited by nfarabullini

Loading

edopao commented Jan 7, 2026 •

edited

Loading

nfarabullini Jan 9, 2026 •

edited

Loading