Global ops reduction

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

@@ -13,6 +13,7 @@
 import logging
 from collections.abc import Sequence
 from enum import Enum
+from types import ModuleType
 from typing import Any, Literal, Protocol, overload, runtime_checkable
 
 import dace  # type: ignore[import-untyped]
@@ -21,6 +22,7 @@
 
 from icon4py.model.common import utils
 from icon4py.model.common.orchestration.halo_exchange import DummyNestedSDFG
+from icon4py.model.common.states import utils as state_utils
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
@@ -323,6 +325,38 @@ class SingleNodeRun(RunType):
     pass
 
 
+class Reductions(Protocol):
+    def min(
+        self, buffer: data_alloc.NDArray, array_ns: ModuleType = np
+    ) -> state_utils.ScalarType: ...
+
+    def max(
+        self, buffer: data_alloc.NDArray, array_ns: ModuleType = np
+    ) -> state_utils.ScalarType: ...
+
+    def sum(
+        self, buffer: data_alloc.NDArray, array_ns: ModuleType = np
+    ) -> state_utils.ScalarType: ...
+
+    def mean(
+        self, buffer: data_alloc.NDArray, array_ns: ModuleType = np
+    ) -> state_utils.ScalarType: ...
+
+
+class SingleNodeReductions(Reductions):
+    def min(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+        return array_ns.min(buffer).item()
+
+    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:
+        return array_ns.sum(buffer).item()
+
+    def mean(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+        return array_ns.sum(buffer).item() / buffer.size
+
+
 @overload
 def get_runtype(with_mpi: Literal[True]) -> MultiNodeRun: ...
 
@@ -365,4 +399,20 @@ def create_single_node_exchange(
     return SingleNodeExchange()
 
 
+@functools.singledispatch
+def create_reduction(props: ProcessProperties) -> Reductions:
+    """
+    Create a Global Reduction depending on the runtime size.
+
+    Depending on the number of processor a SingleNode version is returned or a GHEX context created and a Multinode returned.
+    """
+    raise NotImplementedError(f"Unknown ProcessorProperties type ({type(props)})")
+
+
+@create_reduction.register(SingleNodeProcessProperties)
+def create_single_reduction_exchange(props: SingleNodeProcessProperties) -> Reductions:
+    return SingleNodeReductions()
+
+
 single_node_default = SingleNodeExchange()
+single_node_reductions = SingleNodeReductions()

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

@@ -8,10 +8,12 @@
 
 from __future__ import annotations
 
+import dataclasses
 import functools
 import logging
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Union
 
 import dace  # type: ignore[import-untyped]
@@ -20,8 +22,9 @@
 
 from icon4py.model.common import dimension as dims
 from icon4py.model.common.decomposition import definitions
-from icon4py.model.common.decomposition.definitions import SingleNodeExchange
+from icon4py.model.common.decomposition.definitions import Reductions, SingleNodeExchange
 from icon4py.model.common.orchestration import halo_exchange
+from icon4py.model.common.states import utils as state_utils
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
@@ -423,3 +426,81 @@ def create_multinode_node_exchange(
         return GHexMultiNodeExchange(props, decomp_info)
     else:
         return SingleNodeExchange()
+
+
+@dataclasses.dataclass
+class GlobalReductions(Reductions):
+    props: definitions.ProcessProperties
+
+    def _reduce(
+        self,
+        buffer: data_alloc.NDArray,
+        local_reduction: Callable[[data_alloc.NDArray], data_alloc.ScalarT],
+        global_reduction: mpi4py.MPI.Op,
+        array_ns: ModuleType = np,
+    ) -> state_utils.ScalarType:
+        local_red_val = local_reduction(buffer)
+        recv_buffer = array_ns.empty(1, dtype=buffer.dtype)
+        if hasattr(
+            array_ns, "cuda"
+        ):  # https://mpi4py.readthedocs.io/en/stable/tutorial.html#gpu-aware-mpi-python-gpu-arrays
+            array_ns.cuda.runtime.deviceSynchronize()
+        self.props.comm.Allreduce(local_red_val, recv_buffer, global_reduction)
+        return recv_buffer.item()
+
+    def _calc_buffer_size(
+        self,
+        buffer: data_alloc.NDArray,
+        array_ns: ModuleType = np,
+    ) -> state_utils.ScalarType:
+        return self._reduce(array_ns.asarray([buffer.size]), array_ns.sum, mpi4py.MPI.SUM, array_ns)
+
+    def min(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+        if self._calc_buffer_size(buffer, array_ns) == 0:
+            raise ValueError("global_min requires a non-empty buffer")
+        return self._reduce(
+            buffer if buffer.size != 0 else array_ns.asarray([array_ns.inf]),
+            array_ns.min,
+            mpi4py.MPI.MIN,
+            array_ns,
+        )
+
+    def max(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+        if self._calc_buffer_size(buffer, array_ns) == 0:
+            raise ValueError("global_max requires a non-empty buffer")
+        return self._reduce(
+            buffer if buffer.size != 0 else array_ns.asarray([-array_ns.inf]),
+            array_ns.max,
+            mpi4py.MPI.MAX,
+            array_ns,
+        )
+
+    def sum(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+        if self._calc_buffer_size(buffer, array_ns) == 0:
+            raise ValueError("global_sum requires a non-empty buffer")
+        return self._reduce(
+            buffer if buffer.size != 0 else array_ns.asarray([0]),
+            array_ns.sum,
+            mpi4py.MPI.SUM,
+            array_ns,
+        )
+
+    def mean(self, buffer: data_alloc.NDArray, array_ns: ModuleType = np) -> state_utils.ScalarType:
+        global_buffer_size = self._calc_buffer_size(buffer, array_ns)
+        if global_buffer_size == 0:
+            raise ValueError("global_mean requires a non-empty buffer")
+
+        return (
+            self._reduce(
+                buffer if buffer.size != 0 else array_ns.asarray([0]),
+                array_ns.sum,
+                mpi4py.MPI.SUM,
+                array_ns,
+            )
+            / global_buffer_size
+        )
+
+
+@definitions.create_reduction.register(MPICommProcessProperties)
+def create_global_reduction(props: MPICommProcessProperties) -> Reductions:
+    return GlobalReductions(props)

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

@@ -189,7 +189,7 @@ def _inverse_field_provider(self, field_name: str) -> factory.FieldProvider:
                     self._edge_domain(h_grid.Zone.LOCAL),
                 )
             },
-            do_exchange=False,
+            do_exchange=True,
         )
         return provider
 

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

@@ -8,6 +8,7 @@
 import functools
 import logging
 import pathlib
+from collections.abc import Callable
 from types import ModuleType
 from typing import Literal, Protocol, TypeAlias
 
@@ -87,6 +88,7 @@ def __init__(
         transformation: IndexTransformation,
         grid_file: pathlib.Path | str,
         config: v_grid.VerticalGridConfig,  # TODO(halungge): remove to separate vertical and horizontal grid
+        global_reductions: decomposition.Reductions = decomposition.single_node_reductions,
     ):
         self._transformation = transformation
         self._file_name = str(grid_file)
@@ -96,6 +98,7 @@ def __init__(
         self._geometry: GeometryDict = {}
         self._reader = None
         self._coordinates: CoordinateDict = {}
+        self._global_reductions = global_reductions
 
     def open(self):
         """Open the gridfile resource for reading."""
@@ -130,7 +133,10 @@ def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_va
 
         self._geometry = self._read_geometry_fields(allocator)
         self._grid = self._construct_grid(
-            allocator=allocator, with_skip_values=keep_skip_values, geometry_type=geometry_type
+            allocator=allocator,
+            with_skip_values=keep_skip_values,
+            geometry_type=geometry_type,
+            mean_reduction=self._global_reductions.mean,
         )
         self._coordinates = self._read_coordinates(allocator, geometry_type)
         self.close()
@@ -353,6 +359,9 @@ def _construct_grid(
         allocator: gtx_typing.FieldBufferAllocationUtil | None,
         with_skip_values: bool,
         geometry_type: base.GeometryType,
+        mean_reduction: Callable[
+            [data_alloc.NDArray, ModuleType], data_alloc.ScalarT
+        ] = decomposition.single_node_reductions.mean,
     ) -> icon.IconGrid:
         """Construct the grid topology from the icon grid file.
 
@@ -386,7 +395,6 @@ def _construct_grid(
         mean_dual_cell_area = self._reader.try_attribute(
             gridfile.MPIMPropertyName.MEAN_DUAL_CELL_AREA
         )
-
         edge_lengths = self.geometry_fields[gridfile.GeometryName.EDGE_LENGTH.value].ndarray
         dual_edge_lengths = self.geometry_fields[
             gridfile.GeometryName.DUAL_EDGE_LENGTH.value
@@ -412,6 +420,7 @@ def _construct_grid(
             dual_edge_lengths=dual_edge_lengths,
             cell_areas=cell_areas,
             dual_cell_areas=dual_cell_areas,
+            mean_reduction=mean_reduction,
         )
         grid_size = base.HorizontalGridSize(
             num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells

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

@@ -16,6 +16,7 @@
 from gt4py.next import allocators as gtx_allocators
 
 from icon4py.model.common import constants, dimension as dims
+from icon4py.model.common.decomposition import definitions as decomposition
 from icon4py.model.common.grid import base, horizontal as h_grid
 from icon4py.model.common.utils import data_allocation as data_alloc
 
@@ -97,13 +98,16 @@ def from_fields(
         cell_areas: data_alloc.NDArray | None = None,
         mean_dual_cell_area: float | None = None,
         dual_cell_areas: data_alloc.NDArray | None = None,
+        mean_reduction: Callable[
+            [data_alloc.NDArray, ModuleType], data_alloc.ScalarT
+        ] = decomposition.single_node_reductions.mean,
         **kwargs,
     ) -> _T:
         def init_mean(value: float | None, data: data_alloc.NDArray | None) -> float | None:
             if value is not None:
                 return value
             if data is not None:
-                return array_ns.mean(data).item()
+                return mean_reduction(data, array_ns=array_ns)
             return None
 
         mean_edge_length = init_mean(mean_edge_length, edge_lengths)

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

@@ -5,6 +5,7 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from collections.abc import Callable
 from types import ModuleType
 
 import gt4py.next as gtx
@@ -17,6 +18,7 @@ def compute_coeff_gradekin(
     edge_cell_length: data_alloc.NDArray,
     inv_dual_edge_length: data_alloc.NDArray,
     horizontal_start: gtx.int32,
+    exchange: Callable[[data_alloc.NDArray], None],
     array_ns: ModuleType = np,
 ) -> data_alloc.NDArray:
     """
@@ -41,4 +43,5 @@ def compute_coeff_gradekin(
         * inv_dual_edge_length[horizontal_start:]
     )
     coeff_gradekin_full = array_ns.column_stack((coeff_gradekin_0, coeff_gradekin_1))
+    exchange(coeff_gradekin_full)
     return coeff_gradekin_full

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

@@ -32,6 +32,7 @@
 from gt4py.next.experimental import concat_where
 
 from icon4py.model.common import constants, dimension as dims, field_type_aliases as fa
+from icon4py.model.common.decomposition import definitions as decomposition
 from icon4py.model.common.dimension import C2E, C2E2C, C2E2CO, E2C, C2E2CODim, Koff
 from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import (
     _cell_2_edge_interpolation,
@@ -596,6 +597,9 @@ def compute_nflat_gradp(
     e_owner_mask: data_alloc.NDArray,
     lateral_boundary_level: int,
     nlev: int,
+    min_reduction: Callable[
+        [data_alloc.NDArray, ModuleType], data_alloc.ScalarT
+    ] = decomposition.single_node_reductions.min,
     array_ns: ModuleType = np,
 ) -> int:
     """
@@ -607,8 +611,8 @@ def compute_nflat_gradp(
         flat_idx_max,
         nlev,
     )
-    nflat_gradp = array_ns.min(mask_array)
-    return nflat_gradp.item()
+    nflat_gradp = min_reduction(mask_array, array_ns=array_ns)
+    return nflat_gradp
 
 
 @gtx.field_operator

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

@@ -69,6 +69,7 @@ def __init__(
         exner_expol: float,
         vwind_offctr: float,
         exchange: decomposition.ExchangeRuntime = decomposition.single_node_default,
+        global_reductions: decomposition.Reductions = decomposition.single_node_reductions,
     ):
         self._backend = backend
         self._xp = data_alloc.import_array_ns(backend)
@@ -81,6 +82,7 @@ def __init__(
         self._geometry = geometry_source
         self._exchange = exchange
         self._interpolation_source = interpolation_source
+        self._global_reductions = global_reductions
         log.info(
             f"initialized metrics factory for backend = '{self._backend_name()}' and grid = '{self._grid}'"
         )
@@ -657,7 +659,11 @@ def _register_computed_fields(self) -> None:  # noqa: PLR0915 [too-many-statemen
         self.register_provider(max_flat_index_provider)
 
         nflat_gradp_provider = factory.NumpyDataProvider(
-            func=functools.partial(mf.compute_nflat_gradp, array_ns=self._xp),
+            func=functools.partial(
+                mf.compute_nflat_gradp,
+                array_ns=self._xp,
+                min_reduction=self._global_reductions.min,
+            ),
             domain=(),
             deps={
                 "flat_idx_max": attrs.FLAT_IDX_MAX,
@@ -779,6 +785,7 @@ def _register_computed_fields(self) -> None:  # noqa: PLR0915 [too-many-statemen
             func=functools.partial(
                 compute_coeff_gradekin.compute_coeff_gradekin,
                 array_ns=self._xp,
+                exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim),
             ),
             domain=(dims.EdgeDim, dims.E2CDim),
             fields=(attrs.COEFF_GRADEKIN,),