Halo construction

model/common/pyproject.toml

@@ -30,7 +30,7 @@ requires-python = ">=3.10"
 
 [project.optional-dependencies]
 all = ["icon4py-common[ghex,io]"]
-ghex = ["ghex", "mpi4py"]
+ghex = ["ghex", "mpi4py>=3.1.4", "pymetis>2022.1"]
 io = [
   "icon4py-common[netcdf]",
   "xarray[complete]>=2024.3.0",
@@ -85,9 +85,11 @@ warn_unused_configs = true
 warn_unused_ignores = true
 
 [tool.pytest]
+log_cli = true
 
 [tool.pytest.ini_options]
 addopts = ['-p icon4py.model.common.test_utils.pytest_config']
+log_cli = true
 markers = [
   "datatest: test depending on serialized data generated by a full model run",
   "with_netcdf: test depending on a compatible version of netCDF4"

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

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import enum
 import functools
 import logging
 from dataclasses import dataclass
@@ -18,6 +19,7 @@
 import numpy.ma as ma
 from gt4py.next import Dimension
 
+from icon4py.model.common.settings import xp
 from icon4py.model.common.utils import builder
 
 
@@ -30,6 +32,9 @@ class ProcessProperties(Protocol):
     comm_name: str
     comm_size: int
 
+    def single_node(self) -> bool:
+        return self.comm_size == 1
+
 
 @dataclass(frozen=True, init=False)
 class SingleNodeProcessProperties(ProcessProperties):
@@ -67,12 +72,20 @@ class EntryType(IntEnum):
         HALO = 2
 
     @builder.builder
-    def with_dimension(self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray):
+    def with_dimension(
+        self,
+        dim: Dimension,
+        global_index: np.ndarray,
+        owner_mask: np.ndarray,
+        halo_levels: np.ndarray,
+    ):
         masked_global_index = ma.array(global_index, mask=owner_mask)
         self._global_index[dim] = masked_global_index
+        self._halo_levels[dim] = halo_levels
 
     def __init__(self, klevels: int):
         self._global_index = {}
+        self._halo_levels = {}
         self._klevels = klevels
 
     @property
@@ -113,6 +126,12 @@ def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL):
             case _:
                 raise NotImplementedError()
 
+    def halo_levels(self, dim: Dimension):
+        return self._halo_levels[dim]
+
+    def halo_level_mask(self, dim: Dimension, level: DecompositionFlag):
+        return xp.where(self._halo_levels[dim] == level, True, False)
+
 
 class ExchangeResult(Protocol):
     def wait(self):
@@ -225,3 +244,34 @@ def create_single_node_exchange(
     props: SingleNodeProcessProperties, decomp_info: DecompositionInfo
 ) -> ExchangeRuntime:
     return SingleNodeExchange()
+
+
+class DecompositionFlag(enum.IntEnum):
+    UNDEFINED = -1
+    OWNED = 0
+    """used for locally owned cells, vertices, edges"""
+
+    FIRST_HALO_LINE = 1
+    """
+    used for:
+    - cells that share 1 edge with an OWNED cell
+    - vertices that are on OWNED cell 
+    - edges that are on OWNED cell
+    """
+
+    SECOND_HALO_LINE = 2
+    """
+    used for:
+    - cells that share a 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_LINE = 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)
+    """