place exchange in lsq_pseudoinv coefficients

ci/distributed.yml

@@ -91,7 +91,7 @@ build_distributed:
     - ci/scripts/ci-mpi-wrapper.sh pytest -sv -k mpi_tests --with-mpi --backend=$BACKEND model/$COMPONENT --level=$LEVEL
   parallel:
     matrix:
-      - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, common]
+      - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, atmosphere/advection, common]
         # TODO(msimberg): Enable dace_gpu when compilation doesn't take as long
         # or when we can cache across CI jobs.
         BACKEND: [embedded, gtfn_cpu, dace_cpu, gtfn_gpu]
@@ -113,6 +113,9 @@ build_distributed:
     - if: $COMPONENT == 'atmosphere/dycore'
       variables:
         SLURM_TIMELIMIT: '00:15:00'
+    - if: $COMPONENT == 'atmosphere/advection'
+      variables:
+        SLURM_TIMELIMIT: '00:15:00'
     - when: on_success
       variables:
         SLURM_TIMELIMIT: '00:45:00'

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

@@ -205,7 +205,6 @@ def run(
         dtime: ta.wpfloat,
     ) -> None:
         log.debug("advection run - start")
-
         log.debug("communication of prep_adv cell field: mass_flx_ic - start")
         self._exchange.exchange(
             dims.CellDim, prep_adv.mass_flx_ic, stream=decomposition.DEFAULT_STREAM
@@ -445,13 +444,16 @@ def convert_config_to_horizontal_vertical_advection(  # noqa: PLR0912 [too-many-
     horizontal_advection: advection_horizontal.HorizontalAdvection
     match config.horizontal_advection_type:
         case HorizontalAdvectionType.NO_ADVECTION:
-            horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend)
+            horizontal_advection = advection_horizontal.NoAdvection(
+                grid=grid, backend=backend, exchange=exchange
+            )
         case HorizontalAdvectionType.LINEAR_2ND_ORDER:
             tracer_flux = advection_horizontal.SecondOrderMiura(
                 grid=grid,
                 least_squares_state=least_squares_state,
                 horizontal_limiter=horizontal_limiter,
                 backend=backend,
+                exchange=exchange,
             )
             horizontal_advection = advection_horizontal.SemiLagrangian(
                 tracer_flux=tracer_flux,
@@ -478,14 +480,17 @@ def convert_config_to_horizontal_vertical_advection(  # noqa: PLR0912 [too-many-
     vertical_advection: advection_vertical.VerticalAdvection
     match config.vertical_advection_type:
         case VerticalAdvectionType.NO_ADVECTION:
-            vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend)
+            vertical_advection = advection_vertical.NoAdvection(
+                grid=grid, backend=backend, exchange=exchange
+            )
         case VerticalAdvectionType.UPWIND_1ST_ORDER:
             boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend)
             vertical_advection = advection_vertical.FirstOrderUpwind(
                 boundary_conditions=boundary_conditions,
                 grid=grid,
                 metric_state=metric_state,
                 backend=backend,
+                exchange=exchange,
             )
         case VerticalAdvectionType.PPM_3RD_ORDER:
             boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend)
@@ -495,6 +500,7 @@ def convert_config_to_horizontal_vertical_advection(  # noqa: PLR0912 [too-many-
                 grid=grid,
                 metric_state=metric_state,
                 backend=backend,
+                exchange=exchange,
             )
         case _:
             raise NotImplementedError("Unknown vertical advection type.")

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

@@ -215,11 +215,13 @@ def __init__(
         least_squares_state: advection_states.AdvectionLeastSquaresState,
         backend: gtx.typing.Backend | None,
         horizontal_limiter: HorizontalFluxLimiter | None = None,
+        exchange: decomposition.ExchangeRuntime | None = None,
     ):
         self._grid = grid
         self._least_squares_state = least_squares_state
         self._backend = backend
         self._horizontal_limiter = horizontal_limiter or NoLimiter()
+        self._exchange = exchange or decomposition.SingleNodeExchange()
 
         # cell indices
         cell_domain = h_grid.domain(dims.CellDim)
@@ -372,6 +374,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx.typing.Backend | None)
         cell_domain = h_grid.domain(dims.CellDim)
         self._start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING))
         self._end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL))
+        self._exchange = exchange
 
         # stencils
         self._copy_cell_kdim_field = model_options.setup_program(
@@ -408,7 +411,7 @@ def run(
             field_out=p_tracer_new,
         )
         log.debug("running stencil copy_cell_kdim_field - end")
-
+        self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new)
         log.debug("horizontal advection run - end")
 
 
@@ -443,7 +446,6 @@ def run(
             p_mflx_tracer_h=p_mflx_tracer_h,
             dtime=dtime,
         )
-
         log.debug("horizontal advection run - end")
 
     @abstractmethod

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

@@ -65,6 +65,7 @@
     model_options,
     type_alias as ta,
 )
+from icon4py.model.common.decomposition import definitions as decomposition
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid
 from icon4py.model.common.utils import data_allocation as data_alloc
 
@@ -405,12 +406,13 @@ def run(
 class NoAdvection(VerticalAdvection):
     """Class that implements disabled vertical advection."""
 
-    def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None):
+    def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None, exchange):
         log.debug("vertical advection class init - start")
 
         # input arguments
         self._grid = grid
         self._backend = backend
+        self._exchange = exchange
 
         # cell indices
         cell_domain = h_grid.domain(dims.CellDim)
@@ -469,7 +471,6 @@ def run(
             horizontal_end=horizontal_end,
         )
         log.debug("running stencil copy_cell_kdim_field - end")
-
         log.debug("vertical advection run - end")
 
 
@@ -543,6 +544,7 @@ def __init__(
         grid: icon_grid.IconGrid,
         metric_state: advection_states.AdvectionMetricState,
         backend: gtx_typing.Backend | None,
+        exchange: decomposition.ExchangeRuntime | None = None,
     ):
         log.debug("vertical advection class init - start")
 
@@ -551,6 +553,7 @@ def __init__(
         self._grid = grid
         self._metric_state = metric_state
         self._backend = backend
+        self._exchange = exchange
 
         # cell indices
         cell_domain = h_grid.domain(dims.CellDim)
@@ -629,7 +632,6 @@ def _compute_numerical_flux(
         horizontal_start, horizontal_end = self._get_horizontal_start_end(
             even_timestep=even_timestep
         )
-
         log.debug("running stencil compute_vertical_tracer_flux_upwind - start")
         self._compute_vertical_tracer_flux_upwind(
             p_cc=p_tracer_now,
@@ -678,7 +680,7 @@ def _update_unknowns(
             horizontal_end=horizontal_end,
         )
         log.debug("running stencil integrate_tracer_vertically - end")
-
+        self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new)
         log.debug("vertical unknowns update - end")
 
 
@@ -692,6 +694,7 @@ def __init__(
         grid: icon_grid.IconGrid,
         metric_state: advection_states.AdvectionMetricState,
         backend: gtx_typing.Backend | None,
+        exchange: decomposition.ExchangeRuntime | None = None,
     ):
         log.debug("vertical advection class init - start")
 
@@ -701,6 +704,7 @@ def __init__(
         self._grid = grid
         self._metric_state = metric_state
         self._backend = backend
+        self._exchange = exchange
 
         # cell indices
         cell_domain = h_grid.domain(dims.CellDim)
@@ -1055,15 +1059,13 @@ def _compute_numerical_flux(
         log.debug("running stencil compute_ppm4gpu_integer_flux - end")
 
         ## set boundary conditions
-
         self._boundary_conditions.run(
             p_mflx_tracer_v=p_mflx_tracer_v,
             horizontal_start=horizontal_start,
             horizontal_end=horizontal_end,
         )
 
         ## apply flux limiter
-
         self._vertical_limiter.limit_fluxes(
             horizontal_start=horizontal_start, horizontal_end=horizontal_end
         )
@@ -1098,6 +1100,6 @@ def _update_unknowns(
             horizontal_start=horizontal_start,
             horizontal_end=horizontal_end,
         )
+        self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new)
         log.debug("running stencil integrate_tracer_vertically - end")
-
         log.debug("vertical unknowns update - end")

model/atmosphere/advection/tests/advection/fixtures.py

@@ -46,3 +46,13 @@ def advection_exit_savepoint(data_provider, date):
     fixture, passing 'date=<iso_string>'.
     """
     return data_provider.from_advection_exit_savepoint(size=data_provider.grid_size, date=date)
+
+
+@pytest.fixture
+def advection_lsq_state(
+    interpolation_savepoint: serialbox.InterpolationSavepoint,
+) -> advection_states.AdvectionLeastSquaresState:
+    return advection_states.AdvectionLeastSquaresState(
+        lsq_pseudoinv_1=interpolation_savepoint.lsq_pseudoinv_1(),
+        lsq_pseudoinv_2=interpolation_savepoint.lsq_pseudoinv_2(),
+    )

model/atmosphere/advection/tests/advection/mpi_tests/__init__.py

@@ -0,0 +1,7 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause