Scheduled Halo Exchange

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

@@ -188,7 +188,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")
@@ -272,7 +274,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
@@ -365,7 +371,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

@@ -155,7 +155,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

@@ -386,7 +386,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
@@ -731,11 +731,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")
 
@@ -772,12 +773,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")
 
@@ -817,7 +823,13 @@ def _do_diffusion_step(
         # TODO(halungge): move this up and do asynchronous exchange
         if self.config.type_vn_diffu > 1:
             log.debug("communication rbf extrapolation of z_nable2_e - start")
-            self._exchange(self.z_nabla2_e, dim=dims.EdgeDim, wait=True)
+            # TODO(phimuell, muellch): Is asynchronous mode okay here?
+            self._exchange(
+                self.z_nabla2_e,
+                dim=dims.EdgeDim,
+                full_exchange=True,
+                # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
+            )
             log.debug("communication rbf extrapolation of z_nable2_e - end")
 
         log.debug("2nd rbf interpolation: start")
@@ -827,12 +839,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")
 
@@ -848,7 +862,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"
@@ -894,7 +913,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

@@ -1169,8 +1169,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(
@@ -1241,12 +1244,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,
@@ -1334,7 +1344,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,
@@ -1401,9 +1415,10 @@ def run_corrector_step(
                 r_nsubsteps=r_nsubsteps,
             )
             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,
             )