Halo construction

model/common/pyproject.toml

@@ -39,7 +39,7 @@ all = ["icon4py-common[dace,distributed,io]"]
 cuda11 = ['cupy-cuda11x>=13.0', 'gt4py[cuda11]']
 cuda12 = ['cupy-cuda12x>=13.0', 'gt4py[cuda12]']
 dace = ["dace<1.0", "gt4py[dace]"] # TODO(egparedes): DaCe max constraint should be transformed to min constraint after updating gt4py
-distributed = ["ghex>=0.3.0", "mpi4py>=3.1.5"]
+distributed = ["ghex>=0.3.0", "mpi4py>=3.1.5","pymetis>2022.1"]
 io = [
   # external dependencies
   "cartopy>=0.22.0",
@@ -107,6 +107,9 @@ filename = "src/icon4py/model/common/__init__.py"
 # parallel = true
 # source_pkgs = ['common']
 
+[tool.pytest]
+log_cli = true
+
 # --ruff --
 [tool.ruff]
 extend = "../../pyproject.toml"

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
@@ -40,6 +41,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):
@@ -91,6 +95,7 @@ def __init__(
         num_vertices: Optional[int] = None,
     ):
         self._global_index = {}
+        self._halo_levels = {}
         self._klevels = klevels
         self._owner_mask = {}
         self._num_vertices = num_vertices
@@ -151,6 +156,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 np.where(self._halo_levels[dim] == level, True, False)
+
 
 class ExchangeResult(Protocol):
     def wait(self):
@@ -392,3 +403,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)
+    """