Scheduled Halo Exchange (#980)

This PR introduces the [scheduled exchange
feature](ghex-org/GHEX#190) from GHEX into
ICON4Py.

These exchange allows to call the exchange function before all work has
been completed, i.e. the exchange will wait until the previous work is
done. A similar feature is the "scheduled wait", that allows to initiate
the receive without the need to wait on its completion.

In addition to this the function also renamed the functions related to
halo exchange:
- `exchange()` was renamed to `start()`.
- `wait()` was renamed to `finish()` (that might now return before the
transfer has fully concluded).
- `exchange_and_wait()` was renamed to `exchange()`.

All of these functions now accepts the an argument called `stream`,
which defaults to `DEFAULT_STREAM`. It is indicate how synchronization
with the stream should be performed.
In case of `start()` it means that the actual exchange should not start
until all work previously submitted to `stream` has finished. For
`finish()` it means that further work, submitted to `stream`, should not
start until the exchange has ended. For `finish()` it is also possible
to specify `BLOCK`, which means that `finish()` waits until the transfer
has fully finished.

The orchestrator was not updated, but the change were made in such a way
that it continues to work in diffusion, although using the original,
blocking behaviour.


---------

Co-authored-by: Philip Mueller <philip.paul.mueller@bluemain.ch>
Co-authored-by: Hannes Vogt <vogt@hey.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Mikael Simberg <mikael.simberg@iki.fi>
Co-authored-by: Hannes Vogt <hannes@havogt.de>

Author: 5 people

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py

@@ -207,7 +207,9 @@ def run(
         log.debug("advection run - start")
 
         log.debug("communication of prep_adv cell field: mass_flx_ic - start")
-        self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic)
+        self._exchange.exchange(
+            dims.CellDim, prep_adv.mass_flx_ic, stream=decomposition.DEFAULT_STREAM
+        )
         log.debug("communication of prep_adv cell field: mass_flx_ic - end")
 
         log.debug("running stencil copy_cell_kdim_field - start")
@@ -313,7 +315,11 @@ def run(
         log.debug("advection run - start")
 
         log.debug("communication of prep_adv cell field: mass_flx_ic - start")
-        self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic)
+        self._exchange.exchange(
+            dims.CellDim,
+            prep_adv.mass_flx_ic,
+            stream=decomposition.DEFAULT_STREAM,
+        )
         log.debug("communication of prep_adv cell field: mass_flx_ic - end")
 
         # reintegrate density for conservation of mass
@@ -396,7 +402,11 @@ def run(
 
         # exchange updated tracer values, originally happens only if iforcing /= inwp
         log.debug("communication of advection cell field: p_tracer_new - start")
-        self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new)
+        self._exchange.exchange(
+            dims.CellDim,
+            p_tracer_new,
+            stream=decomposition.DEFAULT_STREAM,
+        )
         log.debug("communication of advection cell field: p_tracer_new - end")
 
         # finalize step

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py

@@ -174,7 +174,7 @@ def apply_flux_limiter(
         )
 
         log.debug("communication of advection cell field: r_m - start")
-        self._exchange.exchange_and_wait(dims.CellDim, self._r_m)
+        self._exchange.exchange(dims.CellDim, self._r_m, stream=decomposition.DEFAULT_STREAM)
         log.debug("communication of advection cell field: r_m - end")
 
         # limit outward fluxes

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -417,7 +417,7 @@ def __init__(
         self._cell_params = cell_params
 
         self.halo_exchange_wait = decomposition.create_halo_exchange_wait(
-            self._exchange
+            self._exchange,
         )  # wait on a communication handle
         self.rd_o_cvd: float = constants.GAS_CONSTANT_DRY_AIR / (
             constants.CPD - constants.GAS_CONSTANT_DRY_AIR
@@ -761,11 +761,12 @@ def _sync_cell_fields(self, prognostic_state):
         IF ( linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN
         """
         log.debug("communication of prognostic cell fields: theta, w, exner - start")
-        self._exchange.exchange_and_wait(
+        self._exchange.exchange(
             dims.CellDim,
             prognostic_state.w,
             prognostic_state.theta_v,
             prognostic_state.exner,
+            stream=decomposition.DEFAULT_STREAM,
         )
         log.debug("communication of prognostic cell fields: theta, w, exner - done")
 
@@ -802,12 +803,17 @@ def _do_diffusion_step(
         log.debug("rbf interpolation 1: end")
 
         # 2.  HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
+        # TODO(phimuell, muellch): Is asynchronous mode okay here.
+        # NOTE: We do not specify a stream here but rely on the default argument.
+        #   We do this to ensure that the orchestrator works, but it is not aware
+        #   of the streams.
         log.debug("communication rbf extrapolation of vn - start")
         self._exchange(
             self.u_vert,
             self.v_vert,
             dim=dims.VertexDim,
-            wait=True,
+            full_exchange=True,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
         )
         log.debug("communication rbf extrapolation of vn - end")
 
@@ -850,12 +856,14 @@ def _do_diffusion_step(
         log.debug("2nd rbf interpolation: end")
 
         # 6.  HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields)
+        # TODO(phimuell, muellch): Is asynchronous mode okay here.
         log.debug("communication rbf extrapolation of z_nable2_e - start")
         self._exchange(
             self.u_vert,
             self.v_vert,
             dim=dims.VertexDim,
-            wait=True,
+            full_exchange=True,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
         )
         log.debug("communication rbf extrapolation of z_nable2_e - end")
 
@@ -871,7 +879,12 @@ def _do_diffusion_step(
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end")
 
         log.debug("communication of prognistic.vn : start")
-        handle_edge_comm = self._exchange(prognostic_state.vn, dim=dims.EdgeDim, wait=False)
+        handle_edge_comm = self._exchange(
+            prognostic_state.vn,
+            dim=dims.EdgeDim,
+            full_exchange=False,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
+        )
 
         log.debug(
             "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"
@@ -917,7 +930,8 @@ def _do_diffusion_step(
             log.debug("running stencil 13 to 16 apply_diffusion_to_theta_and_exner: end")
 
         self.halo_exchange_wait(
-            handle_edge_comm
+            handle_edge_comm,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
         )  # need to do this here, since we currently only use 1 communication object.
         log.debug("communication of prognogistic.vn - end")
 

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

@@ -1151,8 +1151,11 @@ def run_predictor_step(
         )
 
         log.debug("exchanging prognostic field 'vn' and local field 'rho_at_edges_on_model_levels'")
-        self._exchange.exchange_and_wait(
-            dims.EdgeDim, prognostic_states.next.vn, z_fields.rho_at_edges_on_model_levels
+        self._exchange.exchange(
+            dims.EdgeDim,
+            prognostic_states.next.vn,
+            z_fields.rho_at_edges_on_model_levels,
+            stream=decomposition.DEFAULT_STREAM,
         )
 
         self._compute_horizontal_velocity_quantities_and_fluxes(
@@ -1223,12 +1226,19 @@ def run_predictor_step(
             log.debug(
                 "exchanging prognostic field 'w' and local field 'dwdz_at_cells_on_model_levels'"
             )
-            self._exchange.exchange_and_wait(
-                dims.CellDim, prognostic_states.next.w, z_fields.dwdz_at_cells_on_model_levels
+            self._exchange.exchange(
+                dims.CellDim,
+                prognostic_states.next.w,
+                z_fields.dwdz_at_cells_on_model_levels,
+                stream=decomposition.DEFAULT_STREAM,
             )
         else:
             log.debug("exchanging prognostic field 'w'")
-            self._exchange.exchange_and_wait(dims.CellDim, prognostic_states.next.w)
+            self._exchange.exchange(
+                dims.CellDim,
+                prognostic_states.next.w,
+                stream=decomposition.DEFAULT_STREAM,
+            )
 
     def run_corrector_step(
         self,
@@ -1319,7 +1329,11 @@ def run_corrector_step(
         )
 
         log.debug("exchanging prognostic field 'vn'")
-        self._exchange.exchange_and_wait(dims.EdgeDim, (prognostic_states.next.vn))
+        self._exchange.exchange(
+            dims.EdgeDim,
+            prognostic_states.next.vn,
+            stream=decomposition.DEFAULT_STREAM,
+        )
 
         self._compute_averaged_vn_and_fluxes(
             spatially_averaged_vn=self.z_vn_avg,
@@ -1389,9 +1403,10 @@ def run_corrector_step(
                 )
 
         log.debug("exchange prognostic fields 'rho' , 'exner', 'w'")
-        self._exchange.exchange_and_wait(
+        self._exchange.exchange(
             dims.CellDim,
             prognostic_states.next.rho,
             prognostic_states.next.exner,
             prognostic_states.next.w,
+            stream=decomposition.DEFAULT_STREAM,
         )