Merge VelocityAdvection into two programs with per-output domains

bindings/src/icon4py/bindings/dycore_wrapper.py

@@ -374,7 +374,7 @@ def solve_nh_run(  # noqa: PLR0917 [too-many-positional-arguments]
         dynamical_vertical_volumetric_flux_at_cells_on_half_levels=vol_flx_ic,
     )
 
-    # Make `max_vcfl` a 0-d array to avoid cupy synchronization, see `velocity_advection.py`.
+    # Make `max_vcfl` a 0-d array to avoid cupy synchronization, see `_update_max_vertical_cfl`.
     # Note, `max_vcfl` needs to be passed back to Fortran after the timestep.
     max_vcfl = data_alloc.scalar_like_array(max_vcfl_size1_array[0], xp)
 

model/atmosphere/dycore/docs/dycore_numerics.rst

@@ -9,5 +9,4 @@ described in detail.
    :maxdepth: 2
    :caption: Dycore subcomponents:
 
-   dycore_numerics_advection
    dycore_numerics_nonhydro

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

@@ -45,7 +45,12 @@
 from icon4py.model.atmosphere.dycore.stencils.update_theta_and_exner_in_halo import (
     update_theta_and_exner_in_halo,
 )
-from icon4py.model.atmosphere.dycore.velocity_advection import VelocityAdvection
+from icon4py.model.atmosphere.dycore.stencils.velocity_advection_corrector import (
+    compute_velocity_advection_in_corrector_step,
+)
+from icon4py.model.atmosphere.dycore.stencils.velocity_advection_predictor import (
+    compute_velocity_advection_in_predictor_step,
+)
 from icon4py.model.common import (
     constants,
     dimension as dims,
@@ -456,6 +461,22 @@ def __init__(self, config: NonHydrostaticConfig):
         """
 
 
+def _update_max_vertical_cfl(
+    diagnostic_state: nonhydro_states.DiagnosticStateNonHydro,
+    vertical_cfl: fa.CellKHalfField[ta.anyfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+) -> None:
+    # Reductions should be performed on flat, contiguous arrays for best cupy performance
+    # as otherwise cupy won't use cub optimized kernels.
+    max_vertical_cfl = vertical_cfl.array_ns.max(  # type: ignore[attr-defined]
+        vertical_cfl.ndarray[horizontal_start:horizontal_end, :].ravel(order="K")  # type: ignore[attr-defined]
+    )
+    diagnostic_state.max_vertical_cfl = vertical_cfl.array_ns.maximum(  # type: ignore[attr-defined]
+        max_vertical_cfl, diagnostic_state.max_vertical_cfl
+    )
+
+
 class SolveNonhydro:
     def __init__(
         self,
@@ -886,15 +907,82 @@ def __init__(
             },
         )
 
-        self.velocity_advection = VelocityAdvection(
-            grid=grid,
-            metric_state=metric_state_nonhydro,
-            interpolation_state=interpolation_state,
-            vertical_params=vertical_params,
-            edge_params=edge_geometry,
-            owner_mask=owner_mask,
+        cell_horizontal_sizes = {
+            "start_cell_lateral_boundary_level_4": self._start_cell_lateral_boundary_level_4,
+            "end_cell_halo": self._end_cell_halo,
+            "start_edge_nudging_level_2": self._start_edge_nudging_level_2,
+            "end_edge_local": self._end_edge_local,
+        }
+        shared_constant_args: dict[str, gtx.Field | gtx_typing.Scalar] = {
+            "coeff1_dwdz": self._metric_state_nonhydro.coeff1_dwdz,
+            "coeff2_dwdz": self._metric_state_nonhydro.coeff2_dwdz,
+            "c_intp": self._interpolation_state.c_intp,
+            "inv_dual_edge_length": self._edge_geometry.inverse_dual_edge_lengths,
+            "inv_primal_edge_length": self._edge_geometry.inverse_primal_edge_lengths,
+            "tangent_orientation": self._edge_geometry.tangent_orientation,
+            "e_bln_c_s": self._interpolation_state.e_bln_c_s,
+            "ddqz_z_half": self._metric_state_nonhydro.ddqz_z_half,
+            "geofac_n2s": self._interpolation_state.geofac_n2s,
+            "owner_mask": owner_mask,
+            "coriolis_frequency": self._edge_geometry.coriolis_frequency,
+            "geofac_rot": self._interpolation_state.geofac_rot,
+            "coeff_gradekin": self._metric_state_nonhydro.coeff_gradekin,
+            "c_lin_e": self._interpolation_state.c_lin_e,
+            "ddqz_z_full_e": self._metric_state_nonhydro.ddqz_z_full_e,
+            "area_edge": self._edge_geometry.edge_areas,
+            "area": self._cell_params.area,
+            "geofac_grdiv": self._interpolation_state.geofac_grdiv,
+        }
+
+        self._compute_velocity_advection_in_predictor_step = setup_program(
             backend=backend,
+            program=compute_velocity_advection_in_predictor_step,
+            constant_args={
+                "rbf_vec_coeff_e": self._interpolation_state.rbf_vec_coeff_e,
+                "wgtfac_e": self._metric_state_nonhydro.wgtfac_e,
+                "wgtfacq_e": self._metric_state_nonhydro.wgtfacq_e,
+                "ddxn_z_full": self._metric_state_nonhydro.ddxn_z_full,
+                "ddxt_z_full": self._metric_state_nonhydro.ddxt_z_full,
+                "wgtfac_c": self._metric_state_nonhydro.wgtfac_c,
+                **shared_constant_args,
+            },
+            variants={
+                "skip_compute_predictor_vertical_advection": [True, False],
+                # Only True: deriving `apply_extra_diffusion_on_vn` from `max_vertical_cfl` would need a
+                # device synchronization, so the call site fixes it to True (see the TODO there).
+                "apply_extra_diffusion_on_vn": [True],
+            },
+            horizontal_sizes={
+                "start_edge_lateral_boundary_level_5": self._start_edge_lateral_boundary_level_5,
+                "end_edge_halo_level_2": self._end_edge_halo_level_2,
+                **cell_horizontal_sizes,
+            },
+            vertical_sizes={
+                "nflatlev": self._vertical_params.nflatlev,
+                "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer,
+                "vertical_start": gtx.int32(0),
+                "vertical_end": self._grid.num_levels,
+            },
+            offset_provider=self._grid.connectivities,
         )
+
+        self._compute_velocity_advection_in_corrector_step = setup_program(
+            backend=backend,
+            program=compute_velocity_advection_in_corrector_step,
+            constant_args=shared_constant_args,
+            variants={
+                # Only True, as for the predictor step above.
+                "apply_extra_diffusion_on_vn": [True],
+            },
+            horizontal_sizes=cell_horizontal_sizes,
+            vertical_sizes={
+                "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer,
+                "vertical_start": gtx.int32(0),
+                "vertical_end": self._grid.num_levels,
+            },
+            offset_provider=self._grid.connectivities,
+        )
+
         self._allocate_local_fields(model_backends.get_allocator(backend))
 
         self._en_smag_fac_for_zero_nshift(
@@ -1008,6 +1096,9 @@ def _allocate_local_fields(self, allocator: gtx_typing.Allocator | None) -> None
         Declared as enh_divdamp_fac in ICON.
         """
         self.intermediate_fields = IntermediateFields.allocate(grid=self._grid, allocator=allocator)
+        self._vertical_cfl = data_alloc.zero_field(
+            self._grid, dims.CellDim, dims.KHalfDim, allocator=allocator, dtype=ta.vpfloat
+        )
 
     def _determine_local_domains(self) -> None:
         vertex_domain = h_grid.domain(dims.VertexDim)
@@ -1021,6 +1112,9 @@ def _determine_local_domains(self) -> None:
         self._start_cell_lateral_boundary_level_3 = self._grid.start_index(
             cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3)
         )
+        self._start_cell_lateral_boundary_level_4 = self._grid.start_index(
+            cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)
+        )
         self._start_cell_nudging = self._grid.start_index(cell_domain(h_grid.Zone.NUDGING))
         self._start_cell_local = self._grid.start_index(cell_domain(h_grid.Zone.LOCAL))
         self._start_cell_halo = self._grid.start_index(cell_domain(h_grid.Zone.HALO))
@@ -1180,17 +1274,34 @@ def run_predictor_step(
                 and not (at_initial_timestep and at_first_substep)
             )
 
-            assert self._cell_params.area is not None
+            # Note, if we compute `apply_extra_diffusion_on_vn = max_vertical_cfl > VerticalCflConstants.W_LIMIT`
+            # from the reduction below, we would have to synchronize with the device before this call.
+            # TODO (Chia Rui): to decide whether make apply_extra_diffusion_on_vn a config parameter or remove it or always turn on extra diffusion
+            apply_extra_diffusion_on_vn = True
 
-            self.velocity_advection.run_predictor_step(
-                skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection,
-                diagnostic_state=diagnostic_state_nh,
-                prognostic_state=prognostic_states.current,
-                contravariant_correction_at_edges_on_model_levels=self._contravariant_correction_at_edges_on_model_levels,
-                horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels,
+            # TODO(havogt): however, our test data is probably not able to catch cfl_clipping conditions
+            self._compute_velocity_advection_in_predictor_step(
+                tangential_wind=diagnostic_state_nh.tangential_wind,
                 tangential_wind_on_half_levels=z_fields.tangential_wind_on_half_levels,
+                vn_on_half_levels=diagnostic_state_nh.vn_on_half_levels,
+                horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels,
+                contravariant_correction_at_edges_on_model_levels=self._contravariant_correction_at_edges_on_model_levels,
+                contravariant_correction_at_cells_on_half_levels=diagnostic_state_nh.contravariant_correction_at_cells_on_half_levels,
+                vertical_wind_advective_tendency=diagnostic_state_nh.vertical_wind_advective_tendency.predictor,
+                vertical_cfl=self._vertical_cfl,
+                normal_wind_advective_tendency=diagnostic_state_nh.normal_wind_advective_tendency.predictor,
+                vn=prognostic_states.current.vn,
+                w=prognostic_states.current.w,
                 dtime=dtime,
-                cell_areas=self._cell_params.area,
+                skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection,
+                apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn,
+            )
+
+            _update_max_vertical_cfl(
+                diagnostic_state_nh,
+                self._vertical_cfl,
+                self._start_cell_lateral_boundary_level_4,
+                self._end_cell_halo,
             )
 
         self._compute_perturbed_quantities_and_interpolation(
@@ -1356,20 +1467,37 @@ def run_corrector_step(
         # scaling factor for second-order divergence damping: second_order_divdamp_factor_from_sfc_to_divdamp_z*delta_x**2
         # delta_x**2 is approximated by the mean cell area
         # Coefficient for reduced fourth-order divergence d
-        assert self._cell_params.area is not None
         assert self._cell_params.mean_cell_area is not None
         second_order_divdamp_scaling_coeff = (
             second_order_divdamp_factor * self._cell_params.mean_cell_area
         )
 
         log.debug("corrector run velocity advection")
-        self.velocity_advection.run_corrector_step(
-            diagnostic_state=diagnostic_state_nh,
-            prognostic_state=prognostic_states.next,
-            horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels,
+        # Note, if we compute `apply_extra_diffusion_on_vn = max_vertical_cfl > VerticalCflConstants.W_LIMIT`
+        # from the reduction below, we would have to synchronize with the device before this call.
+        # TODO (Chia Rui): to decide whether make apply_extra_diffusion_on_vn a config parameter or remove it or always turn on extra diffusion
+        apply_extra_diffusion_on_vn = True
+
+        self._compute_velocity_advection_in_corrector_step(
+            vertical_wind_advective_tendency=diagnostic_state_nh.vertical_wind_advective_tendency.corrector,
+            vertical_cfl=self._vertical_cfl,
+            normal_wind_advective_tendency=diagnostic_state_nh.normal_wind_advective_tendency.corrector,
+            vn=prognostic_states.next.vn,
+            w=prognostic_states.next.w,
+            tangential_wind=diagnostic_state_nh.tangential_wind,
             tangential_wind_on_half_levels=z_fields.tangential_wind_on_half_levels,
+            vn_on_half_levels=diagnostic_state_nh.vn_on_half_levels,
+            horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels,
+            contravariant_correction_at_cells_on_half_levels=diagnostic_state_nh.contravariant_correction_at_cells_on_half_levels,
             dtime=dtime,
-            cell_areas=self._cell_params.area,
+            apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn,
+        )
+
+        _update_max_vertical_cfl(
+            diagnostic_state_nh,
+            self._vertical_cfl,
+            self._start_cell_lateral_boundary_level_4,
+            self._end_cell_halo,
         )
 
         self._compute_interpolation_and_nonhydro_buoy(

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py

@@ -5,9 +5,15 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+
+# Note: Currently unused, kept for implementing a version faithful to ICON.
+# The ICON version requires the CFL reduction within velocity advection which
+# drains the GPU kernel pipeline.
+
 import gt4py.next as gtx
 from gt4py.next import abs, astype, minimum, neighbor_sum, where  # noqa: A004
 
+from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import VerticalCflConstants
 from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
 from icon4py.model.common.dimension import E2C, E2C2EO, E2V
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -26,25 +32,24 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl(
     geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat],
     vn: fa.EdgeKField[ta.wpfloat],
     ddt_vn_apc: fa.EdgeKField[ta.vpfloat],
-    cfl_w_limit: ta.vpfloat,
-    scalfac_exdiff: ta.wpfloat,
     dtime: ta.wpfloat,
 ) -> fa.EdgeKField[ta.vpfloat]:
     """Formerly known as _mo_velocity_advection_stencil_20."""
-    z_w_con_c_full_wp, ddqz_z_full_e_wp, ddt_vn_apc_wp, cfl_w_limit_wp = astype(
-        (z_w_con_c_full, ddqz_z_full_e, ddt_vn_apc, cfl_w_limit), wpfloat
+    z_w_con_c_full_wp, ddqz_z_full_e_wp, ddt_vn_apc_wp = astype(
+        (z_w_con_c_full, ddqz_z_full_e, ddt_vn_apc), wpfloat
     )
 
     w_con_e = neighbor_sum(c_lin_e * z_w_con_c_full_wp(E2C), axis=dims.E2CDim)
-    difcoef = scalfac_exdiff * minimum(
-        wpfloat("0.85") - cfl_w_limit_wp * dtime,
-        abs(w_con_e) * dtime / ddqz_z_full_e_wp - cfl_w_limit_wp * dtime,
+    vertical_cfl_number_at_edges = abs(w_con_e) * dtime / ddqz_z_full_e_wp
+    difcoef = (VerticalCflConstants.EXTRA_DIFFUSION_SCALING / dtime) * minimum(
+        VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT,
+        vertical_cfl_number_at_edges - VerticalCflConstants.W_LIMIT,
     )
     ddt_vn_apc_wp = where(
         # TODO(havogt): my guess is if the second condition is `True`, then
         # `(levelmask | levelmask(dims.KDim + 1))` is also `True`
         (levelmask | levelmask(dims.KDim + 1))
-        & (abs(w_con_e) > astype(cfl_w_limit * ddqz_z_full_e, wpfloat)),
+        & (vertical_cfl_number_at_edges > VerticalCflConstants.W_LIMIT),
         ddt_vn_apc_wp
         + difcoef
         * area_edge
@@ -72,8 +77,6 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl(
     geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat],
     vn: fa.EdgeKField[ta.wpfloat],
     ddt_vn_apc: fa.EdgeKField[ta.vpfloat],
-    cfl_w_limit: ta.vpfloat,
-    scalfac_exdiff: ta.wpfloat,
     dtime: ta.wpfloat,
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
@@ -92,8 +95,6 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl(
         geofac_grdiv=geofac_grdiv,
         vn=vn,
         ddt_vn_apc=ddt_vn_apc,
-        cfl_w_limit=cfl_w_limit,
-        scalfac_exdiff=scalfac_exdiff,
         dtime=dtime,
         out=ddt_vn_apc,
         domain={