Domain decomposition and halo construction

model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py

@@ -32,7 +32,6 @@
 from icon4py.model.common.metrics import metrics_attributes, metrics_factory
 from icon4py.model.common.states import prognostic_state as prognostics
 from icon4py.model.common.utils import data_allocation as data_alloc
-from icon4py.model.testing import grid_utils
 from icon4py.model.testing.fixtures.benchmark import (
     geometry_field_source,
     interpolation_field_source,
@@ -62,8 +61,6 @@ def solve_nonhydro(
 
     nonhydro_params = solve_nh.NonHydrostaticParams(config)
 
-    decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator)
-
     vertical_config = v_grid.VerticalGridConfig(
         mesh.num_levels,
         lowest_layer_thickness=50,
@@ -205,11 +202,7 @@ def solve_nonhydro(
         vertical_params=vertical_grid,
         edge_geometry=edge_geometry,
         cell_geometry=cell_geometry,
-        owner_mask=gtx.as_field(
-            (dims.CellDim,),
-            decomposition_info.owner_mask(dims.CellDim),  # type: ignore[arg-type] # mypy not take the type of owner_mask
-            allocator=allocator,
-        ),
+        owner_mask=geometry_field_source.get("cell_owner_mask"),
         backend=backend_like,
     )
 

model/common/pyproject.toml

@@ -52,6 +52,7 @@ io = [
   "netcdf4>=1.6.1",
   "numpy>=1.23.3",
   "scikit-learn>=1.4.0",
+  "pymetis>2022.1",
   # TODO(halungge): there are failing tests starting from uxarray==2024.4.0: when a data file does not have
   # fields of a given dimension (eg 'edge') then something in uxarray goes wrong with the dimension
   # mapping. It is not yet clear whether this is a uxarray bug or on our side.

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

@@ -0,0 +1,82 @@
+# 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 typing import Protocol, runtime_checkable
+
+import gt4py.next as gtx
+
+from icon4py.model.common.utils import data_allocation as data_alloc
+
+
+@runtime_checkable
+class Decomposer(Protocol):
+    def __call__(
+        self, adjacency_matrix: data_alloc.NDArray, num_partitions: int
+    ) -> data_alloc.NDArray:
+        """
+        Call the decomposition.
+
+        Args:
+            adjacency_matrix: cell-to-cell connectivity matrix on the global (undecomposed) grid. In the Icon4py context this C2E2C
+            n_part: number of nodes
+        """
+        ...
+
+
+class MetisDecomposer(Decomposer):
+    """
+    A simple decomposer using METIS for partitioning a grid topology.
+
+    We use the simple pythonic interface to pymetis: just passing the adjacency matrix, which for ICON is
+    the full grid C2E2C neigbhor table.
+    if more control is needed (for example by using weights we need to switch to the C like interface)
+    https://documen.tician.de/pymetis/functionality.html
+    """
+
+    def __call__(
+        self, adjacency_matrix: data_alloc.NDArray, num_partitions: int
+    ) -> data_alloc.NDArray:
+        """
+        Generate partition labels for this grid topology using METIS:
+        https://github.com/KarypisLab/METIS
+
+        This method utilizes the pymetis Python bindings:
+        https://github.com/inducer/pymetis
+
+        Args:
+            n_part: int, number of partitions to create
+            adjacency_matrix: nd array: neighbor table describing of the main dimension object to be distributed: for example cell -> cell neighbors
+        Returns: data_alloc.NDArray: array with partition label (int, rank number) for each cell
+        """
+
+        import pymetis  # type: ignore [import-untyped]
+
+        # Invalid indices are not allowed here. Metis will segfault or fail
+        # there are any invalid indices in the adjacency matrix.
+        assert (adjacency_matrix >= 0).all()
+
+        # The partitioning is done on all ranks, and this assumes that the
+        # partitioning is deterministic.
+        _, partition_index = pymetis.part_graph(nparts=num_partitions, adjacency=adjacency_matrix)
+        return data_alloc.array_ns_from_array(adjacency_matrix).array(partition_index)
+
+
+class SingleNodeDecomposer(Decomposer):
+    def __call__(
+        self, adjacency_matrix: data_alloc.NDArray, num_partitions: int
+    ) -> data_alloc.NDArray:
+        """Dummy decomposer for single node: assigns all cells to rank = 0"""
+        if num_partitions != 1:
+            raise ValueError(
+                f"SingleNodeDecomposer can only be used for num_partitions=1, but got {num_partitions}"
+            )
+
+        return data_alloc.array_ns_from_array(adjacency_matrix).zeros(
+            adjacency_matrix.shape[0],
+            dtype=gtx.int32,  # type: ignore  [attr-defined]
+        )

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

@@ -20,7 +20,8 @@
 import gt4py.next as gtx
 import numpy as np
 
-from icon4py.model.common import utils
+from icon4py.model.common import dimension as dims, utils
+from icon4py.model.common.grid import base
 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
@@ -35,6 +36,9 @@ class ProcessProperties(Protocol):
     comm_name: str
     comm_size: int
 
+    def is_single_rank(self) -> bool:
+        return self.comm_size == 1
+
 
 @dataclasses.dataclass(frozen=True, init=False)
 class SingleNodeProcessProperties(ProcessProperties):
@@ -71,41 +75,32 @@ def __call__(self) -> int:
 
 
 class DecompositionInfo:
+    def __init__(
+        self,
+    ) -> None:
+        self._global_index: dict[gtx.Dimension, data_alloc.NDArray] = {}
+        self._halo_levels: dict[gtx.Dimension, data_alloc.NDArray] = {}
+        self._owner_mask: dict[gtx.Dimension, data_alloc.NDArray] = {}
+
     class EntryType(int, Enum):
         ALL = 0
         OWNED = 1
         HALO = 2
 
     @utils.chainable
-    def with_dimension(
-        self, dim: gtx.Dimension, global_index: data_alloc.NDArray, owner_mask: data_alloc.NDArray
+    def set_dimension(
+        self,
+        dim: gtx.Dimension,
+        global_index: data_alloc.NDArray,
+        owner_mask: data_alloc.NDArray,
+        halo_levels: data_alloc.NDArray,
     ) -> None:
         self._global_index[dim] = global_index
         self._owner_mask[dim] = owner_mask
+        self._halo_levels[dim] = halo_levels
 
-    def __init__(
-        self,
-        num_cells: int | None = None,
-        num_edges: int | None = None,
-        num_vertices: int | None = None,
-    ):
-        self._global_index: dict[gtx.Dimension, data_alloc.NDArray] = {}
-        self._owner_mask: dict[gtx.Dimension, data_alloc.NDArray] = {}
-        self._num_vertices = num_vertices
-        self._num_cells = num_cells
-        self._num_edges = num_edges
-
-    @property
-    def num_cells(self) -> int | None:
-        return self._num_cells
-
-    @property
-    def num_edges(self) -> int | None:
-        return self._num_edges
-
-    @property
-    def num_vertices(self) -> int | None:
-        return self._num_vertices
+    def is_distributed(self) -> bool:
+        return max(self._halo_levels[dims.CellDim]).item() > DecompositionFlag.OWNED
 
     def local_index(
         self, dim: gtx.Dimension, entry_type: EntryType = EntryType.ALL
@@ -125,19 +120,15 @@ def local_index(
     def _to_local_index(self, dim: gtx.Dimension) -> data_alloc.NDArray:
         data = self._global_index[dim]
         assert data.ndim == 1
-        if isinstance(data, np.ndarray):
-            import numpy as xp
-        else:
-            import cupy as xp  # type: ignore[import-not-found, no-redef]
-
-            xp.arange(data.shape[0])
-        return xp.arange(data.shape[0])
+        return data_alloc.array_ns_from_array(data).arange(data.shape[0])
 
     def owner_mask(self, dim: gtx.Dimension) -> data_alloc.NDArray:
         return self._owner_mask[dim]
 
     def global_index(
-        self, dim: gtx.Dimension, entry_type: EntryType = EntryType.ALL
+        self,
+        dim: gtx.Dimension,
+        entry_type: DecompositionInfo.EntryType = EntryType.ALL,
     ) -> data_alloc.NDArray:
         match entry_type:
             case DecompositionInfo.EntryType.ALL:
@@ -149,6 +140,25 @@ def global_index(
             case _:
                 raise NotImplementedError()
 
+    def get_horizontal_size(self) -> base.HorizontalGridSize:
+        return base.HorizontalGridSize(
+            num_cells=self.global_index(dims.CellDim, self.EntryType.ALL).shape[0],
+            num_edges=self.global_index(dims.EdgeDim, self.EntryType.ALL).shape[0],
+            num_vertices=self.global_index(dims.VertexDim, self.EntryType.ALL).shape[0],
+        )
+
+    def get_halo_size(self, dim: gtx.Dimension, flag: DecompositionFlag) -> int:
+        level_mask = self.halo_level_mask(dim, flag)
+        # TODO(msimberg): Just do array_ns?
+        return data_alloc.array_ns_from_array(level_mask).count_nonzero(level_mask)
+
+    def halo_levels(self, dim: gtx.Dimension) -> data_alloc.NDArray:
+        return self._halo_levels[dim]
+
+    def halo_level_mask(self, dim: gtx.Dimension, level: DecompositionFlag) -> data_alloc.NDArray:
+        levels = self._halo_levels[dim]
+        return data_alloc.array_ns_from_array(levels).where(levels == level.value, True, False)
+
 
 class ExchangeResult(Protocol):
     def wait(self) -> None: ...
@@ -414,5 +424,36 @@ def create_single_reduction_exchange(props: SingleNodeProcessProperties) -> Redu
     return SingleNodeReductions()
 
 
+class DecompositionFlag(int, Enum):
+    UNDEFINED = -1
+    OWNED = 0
+    """used for locally owned cells, vertices, edges"""
+
+    FIRST_HALO_LEVEL = 1
+    """
+    used for:
+    - cells that share 1 edge with an OWNED cell
+    - vertices that are on OWNED cell, but not owned
+    - edges that are on OWNED cell, but not owned
+    """
+
+    SECOND_HALO_LEVEL = 2
+    """
+    used for:
+    - cells that share one vertex with an OWNED cell
+    - vertices that are on a cell(FIRST_HALO_LINE) but not on an owned cell
+    - edges that have _exactly_ one vertex shared with and OWNED Cell
+    """
+
+    THIRD_HALO_LEVEL = 3
+    """
+    This type does not exist in ICON. It denotes the "closing/far" edges of the SECOND_HALO_LINE cells
+    used for:
+    - cells (NOT USED)
+    - vertices (NOT USED)
+    - edges that are only on the cell(SECOND_HALO_LINE)
+    """
+
+
 single_node_default = SingleNodeExchange()
 single_node_reductions = SingleNodeReductions()