Distributed standalone driver

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, common, standalone_driver]
         # 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,12 @@ build_distributed:
     - if: $COMPONENT == 'atmosphere/dycore'
       variables:
         SLURM_TIMELIMIT: '00:15:00'
+    - if: $COMPONENT == 'standalone_driver'
+      variables:
+        SLURM_TIMELIMIT: '00:50:00'
+    - if: $COMPONENT == 'standalone_driver'
+      variables:
+        SLURM_TIMELIMIT: '00:50:00'
     - when: on_success
       variables:
         SLURM_TIMELIMIT: '00:45:00'

model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py

@@ -39,6 +39,7 @@
             "2021-06-20T12:00:10.000",
         ),
         (definitions.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"),
+        (definitions.Experiments.JW, "2008-09-01T00:05:00.000", "2008-09-01T00:05:00.000"),
     ],
 )
 @pytest.mark.parametrize("ndyn_substeps", [2])
@@ -142,6 +143,12 @@ def test_parallel_diffusion(
             prognostic_state=prognostic_state,
             dtime=dtime,
         )
+        if experiment == definitions.Experiments.JW:
+            # TODO (jcanton,ongchia): move this exchange inside
+            # diffusion.run with the proper DiffusionConfig equivalent to
+            # the fortran
+            # IF ( linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN
+            exchange.exchange(dims.CellDim, prognostic_state.w)
     _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
 
     utils.verify_diffusion_fields(

model/common/src/icon4py/model/common/decomposition/definitions.py

@@ -12,10 +12,11 @@
 import dataclasses
 import functools
 import logging
+import warnings
 from collections.abc import Sequence
 from enum import Enum
 from types import ModuleType
-from typing import Any, Literal, Protocol, TypeAlias, overload, runtime_checkable
+from typing import Any, ClassVar, Literal, Protocol, TypeAlias, overload, runtime_checkable
 
 import dace  # type: ignore[import-untyped]
 import gt4py.next as gtx
@@ -160,6 +161,8 @@ class EntryType(int, Enum):
         ALL = 0
         OWNED = 1
         HALO = 2
+        HALO_LEVEL_1 = 11
+        HALO_LEVEL_2 = 12
 
     @utils.chainable
     def set_dimension(
@@ -187,6 +190,14 @@ def local_index(
                 index = self._to_local_index(dim)
                 mask = self._owner_mask[dim]
                 return index[~mask]
+            case DecompositionInfo.EntryType.HALO_LEVEL_1:
+                index = self._to_local_index(dim)
+                mask = self.halo_level_mask(dim=dim, level=DecompositionFlag.FIRST_HALO_LEVEL)
+                return index[mask]
+            case DecompositionInfo.EntryType.HALO_LEVEL_2:
+                index = self._to_local_index(dim)
+                mask = self.halo_level_mask(dim=dim, level=DecompositionFlag.SECOND_HALO_LEVEL)
+                return index[mask]
             case DecompositionInfo.EntryType.OWNED:
                 index = self._to_local_index(dim)
                 mask = self._owner_mask[dim]
@@ -212,6 +223,12 @@ def global_index(
                 return self._global_index[dim][self._owner_mask[dim]]
             case DecompositionInfo.EntryType.HALO:
                 return self._global_index[dim][~self._owner_mask[dim]]
+            case DecompositionInfo.EntryType.HALO_LEVEL_1:
+                mask = self.halo_level_mask(dim=dim, level=DecompositionFlag.FIRST_HALO_LEVEL)
+                return self._global_index[dim][mask]
+            case DecompositionInfo.EntryType.HALO_LEVEL_2:
+                mask = self.halo_level_mask(dim=dim, level=DecompositionFlag.SECOND_HALO_LEVEL)
+                return self._global_index[dim][mask]
             case _:
                 raise NotImplementedError()
 
@@ -408,12 +425,25 @@ def __str__(self) -> str:
 
 @dataclasses.dataclass
 class SingleNodeExchange(ExchangeRuntime):
+    _warning_emitted: ClassVar[bool] = False
+
+    @classmethod
+    def _warn_if_used(cls, *, stacklevel: int = 3) -> None:
+        if not cls._warning_emitted:
+            warnings.warn(
+                "***** SingleNodeExchange is in use; HALO EXCHANGE IS RUNNING IN SINGLE-NODE *****",
+                RuntimeWarning,
+                stacklevel=stacklevel,
+            )
+            cls._warning_emitted = True
+
     def start(
         self,
         dim: gtx.Dimension,
         *fields: gtx.Field | data_alloc.NDArray,
         stream: StreamLike = DEFAULT_STREAM,
     ) -> ExchangeResult:
+        self._warn_if_used()
         return SingleNodeResult()
 
     def my_rank(self) -> int:
@@ -563,11 +593,19 @@ def max(
     ) -> state_utils.ScalarType: ...
 
     def sum(
-        self, buffer: data_alloc.NDArray, array_ns: ModuleType = np
+        self,
+        buffer: data_alloc.NDArray,
+        lower_bound: gtx.int32,
+        upper_bound: gtx.int32,
+        array_ns: ModuleType = np,
     ) -> state_utils.ScalarType: ...
 
     def mean(
-        self, buffer: data_alloc.NDArray, array_ns: ModuleType = np
+        self,
+        buffer: data_alloc.NDArray,
+        lower_bound: gtx.int32,
+        upper_bound: gtx.int32,
+        array_ns: ModuleType = np,
     ) -> state_utils.ScalarType: ...
 
 
@@ -578,10 +616,22 @@ def min(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_ut
     def max(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
         return array_ns.max(buffer).item()
 
-    def sum(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+    def sum(
+        self,
+        buffer: data_alloc.NDArray,
+        lower_bound: gtx.int32,
+        upper_bound: gtx.int32,
+        array_ns: ModuleType = np,
+    ) -> state_utils.ScalarType:
         return array_ns.sum(buffer).item()
 
-    def mean(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+    def mean(
+        self,
+        buffer: data_alloc.NDArray,
+        lower_bound: gtx.int32,
+        upper_bound: gtx.int32,
+        array_ns: ModuleType = np,
+    ) -> state_utils.ScalarType:
         return array_ns.sum(buffer).item() / buffer.size
 
 
@@ -673,5 +723,28 @@ class DecompositionFlag(int, Enum):
     """
 
 
+class ParallelLogger(logging.Filter):
+    def __init__(
+        self,
+        process_properties: ProcessProperties | None = None,
+        print_distributed_debug_msg: bool = False,
+    ) -> None:
+        super().__init__()
+        self._rank_info = ""
+        self._print_distributed_debug_msg = print_distributed_debug_msg
+        self._rank_id = 0
+        if process_properties and process_properties.comm_size > 1:
+            self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] "
+            self._rank_id = process_properties.rank
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        record.rank = self._rank_info
+        return (
+            True
+            if self._rank_id == 0
+            else (self._print_distributed_debug_msg and record.levelno == logging.DEBUG)
+        )
+
+
 single_node_default = SingleNodeExchange()
 single_node_reductions = SingleNodeReductions()

model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py

@@ -85,18 +85,6 @@ def _get_current_comm_or_comm_world(comm_id: CommId) -> mpi4py.MPI.Comm:
         return MPICommProcessProperties(current_comm)
 
 
-class ParallelLogger(logging.Filter):
-    def __init__(self, process_properties: definitions.ProcessProperties | None = None) -> None:
-        super().__init__()
-        self._rank_info = ""
-        if process_properties and process_properties.comm_size > 1:
-            self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] "
-
-    def filter(self, record: logging.LogRecord) -> bool:
-        record.rank = self._rank_info
-        return True
-
-
 @definitions.get_processor_properties.register(definitions.MultiNodeRun)
 def get_multinode_properties(
     s: definitions.MultiNodeRun, comm_id: CommId = None
@@ -489,7 +477,15 @@ def max(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_ut
             array_ns,
         )
 
-    def sum(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+    def sum(
+        self,
+        buffer: data_alloc.NDArray,
+        lower_bound: gtx.int32,
+        upper_bound: gtx.int32,
+        array_ns: ModuleType = np,
+    ) -> state_utils.ScalarType:
+        # TODO (nfarabullini): use owned mask instead of upper lower bound, and move as "internal argument"
+        buffer = buffer[lower_bound:upper_bound]
         if self._calc_buffer_size(buffer, array_ns) == 0:
             raise ValueError("global_sum requires a non-empty buffer")
         return self._reduce(
@@ -499,7 +495,15 @@ def sum(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_ut
             array_ns,
         )
 
-    def mean(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+    def mean(
+        self,
+        buffer: data_alloc.NDArray,
+        lower_bound: gtx.int32,
+        upper_bound: gtx.int32,
+        array_ns: ModuleType = np,
+    ) -> state_utils.ScalarType:
+        # TODO (nfarabullini): use owned mask instead of upper lower bound, and move as "internal argument"
+        buffer = buffer[lower_bound:upper_bound]
         global_buffer_size = self._calc_buffer_size(buffer, array_ns)
         if global_buffer_size == 0:
             raise ValueError("global_mean requires a non-empty buffer")

model/common/src/icon4py/model/common/grid/geometry.py

@@ -110,7 +110,9 @@ def __init__(
         self._decomposition_info = decomposition_info
         self._attrs = metadata
         self._geometry_type: base.GeometryType = grid.global_properties.geometry_type
+        self._cell_domain = h_grid.domain(dims.CellDim)
         self._edge_domain = h_grid.domain(dims.EdgeDim)
+        self._vertex_domain = h_grid.domain(dims.VertexDim)
         self._exchange = exchange
         self._global_reductions = global_reductions
         log.info(
@@ -323,6 +325,10 @@ def _register_computed_fields(self) -> None:
             deps={
                 "buffer": attrs.EDGE_LENGTH,
             },
+            params={
+                "lower_bound": self._grid.start_index(self._edge_domain(h_grid.Zone.LOCAL)),
+                "upper_bound": self._grid.end_index(self._edge_domain(h_grid.Zone.LOCAL)),
+            },
             fields=(attrs.MEAN_EDGE_LENGTH,),
         )
         self.register_provider(mean_edge_length_np)
@@ -336,6 +342,10 @@ def _register_computed_fields(self) -> None:
             deps={
                 "buffer": attrs.DUAL_EDGE_LENGTH,
             },
+            params={
+                "lower_bound": self._grid.start_index(self._edge_domain(h_grid.Zone.LOCAL)),
+                "upper_bound": self._grid.end_index(self._edge_domain(h_grid.Zone.LOCAL)),
+            },
             fields=(attrs.MEAN_DUAL_EDGE_LENGTH,),
         )
         self.register_provider(mean_dual_edge_length_np)
@@ -349,6 +359,10 @@ def _register_computed_fields(self) -> None:
             deps={
                 "buffer": attrs.CELL_AREA,
             },
+            params={
+                "lower_bound": self._grid.start_index(self._cell_domain(h_grid.Zone.LOCAL)),
+                "upper_bound": self._grid.end_index(self._cell_domain(h_grid.Zone.LOCAL)),
+            },
             fields=(attrs.MEAN_CELL_AREA,),
         )
         self.register_provider(mean_cell_area_np)
@@ -362,6 +376,10 @@ def _register_computed_fields(self) -> None:
             deps={
                 "buffer": attrs.DUAL_AREA,
             },
+            params={
+                "lower_bound": self._grid.start_index(self._vertex_domain(h_grid.Zone.LOCAL)),
+                "upper_bound": self._grid.end_index(self._vertex_domain(h_grid.Zone.LOCAL)),
+            },
             fields=(attrs.MEAN_DUAL_AREA,),
         )
         self.register_provider(mean_dual_cell_area_np)

model/common/src/icon4py/model/common/grid/grid_manager.py

@@ -106,8 +106,8 @@ def __call__(
         allocator: gtx_typing.Allocator | None,
         keep_skip_values: bool,
         decomposer: decomp.Decomposer = _single_node_decomposer,
-        run_properties=_single_process_props,
-    ):
+        run_properties: decomposition.ProcessProperties = _single_process_props,
+    ) -> None:
         if not run_properties.is_single_rank() and isinstance(
             decomposer, decomp.SingleNodeDecomposer
         ):

model/common/src/icon4py/model/common/math/stencils/generic_math_operations_array_ns.py

@@ -0,0 +1,57 @@
+# 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
+
+from types import ModuleType
+
+from icon4py.model.common.utils import data_allocation as data_alloc
+
+
+def compute_directional_derivative_on_cells(
+    cell_field: data_alloc.NDArray,
+    e2c: data_alloc.NDArray,
+    inv_dual_edge_length: data_alloc.NDArray,
+    lb_e: int,
+    ub_e: int,
+    num_edges: int,
+    array_ns: ModuleType,
+) -> data_alloc.NDArray:
+    """
+    Compute directional derivative of a cell centered variable with respect to
+    direction normal to triangle edge.
+    """
+    result = array_ns.zeros((num_edges,))
+    result[lb_e:ub_e] = (
+        cell_field[e2c[lb_e:ub_e, 1]] - cell_field[e2c[lb_e:ub_e, 0]]
+    ) * inv_dual_edge_length[lb_e:ub_e]
+    return result
+
+
+def interpolate_edges_to_cell(
+    edge_field: data_alloc.NDArray,
+    c2e: data_alloc.NDArray,
+    e2c: data_alloc.NDArray,
+    edge_cell_length: data_alloc.NDArray,
+    primal_edge_length: data_alloc.NDArray,
+    cell_area: data_alloc.NDArray,
+    ub_c: int,
+    num_cells: int,
+    array_ns: ModuleType,
+) -> data_alloc.NDArray:
+    """
+    Compute interpolation of scalar fields from edge points to cell centers.
+    """
+    e_inn_c = array_ns.zeros((num_cells, 3))
+    jc_indices = array_ns.arange(ub_c)[:, array_ns.newaxis]
+    c2e_local = c2e[:ub_c]
+    idx_ce = (e2c[c2e_local][:, :, 0] != jc_indices).astype(int)
+    e_inn_c[:ub_c] = (
+        edge_cell_length[c2e_local, idx_ce]
+        * primal_edge_length[c2e_local]
+        / cell_area[:ub_c, array_ns.newaxis]
+    )
+    return array_ns.sum(edge_field[c2e] * e_inn_c, axis=1)

model/common/src/icon4py/model/common/metrics/metrics_factory.py

@@ -642,7 +642,7 @@ def _register_computed_fields(self) -> None:  # noqa: PLR0915 [too-many-statemen
                 ),
             },
             fields={"wgtfac_e": attrs.WGTFAC_E},
-            do_exchange=False,
+            do_exchange=True,
         )
         self.register_provider(compute_wgtfac_e)