Py2f parallel

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -17,10 +17,10 @@
 from dataclasses import InitVar, dataclass, field
 from enum import Enum
 from typing import Final, Optional
-
 from gt4py.next import as_field
 from gt4py.next.common import Dimension
 from gt4py.next.ffront.fbuiltins import Field, int32
+from icon4py.model.common.decomposition.definitions import DecompositionInfo
 
 from icon4py.model.atmosphere.diffusion.diffusion_states import (
     DiffusionDiagnosticState,
@@ -297,7 +297,7 @@ def __post_init__(self, config):
         object.__setattr__(
             self,
             "scaled_nudge_max_coeff",
-            config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO,
+            config.nudge_max_coeff,
         )
 
     def _determine_smagorinski_factor(self, config: DiffusionConfig):
@@ -370,6 +370,9 @@ def __init__(self, exchange: ExchangeRuntime = SingleNodeExchange()):
         self.cell_params: Optional[CellParams] = None
         self._horizontal_start_index_w_diffusion: int32 = 0
 
+    def set_exchange(self, exchange):
+        self._exchange = exchange
+
     def init(
         self,
         grid: IconGrid,
@@ -547,9 +550,9 @@ def _sync_cell_fields(self, prognostic_state):
         log.debug("communication of prognostic cell fields: theta, w, exner - start")
         self._exchange.exchange_and_wait(
             CellDim,
-            prognostic_state.w,
-            prognostic_state.theta_v,
-            prognostic_state.exner,
+            prognostic_state.w.ndarray[0 : self.grid.num_cells, :],
+            prognostic_state.theta_v.ndarray[0 : self.grid.num_cells, :],
+            prognostic_state.exner.ndarray[0 : self.grid.num_cells, :],
         )
         log.debug("communication of prognostic cell fields: theta, w, exner - done")
 
@@ -605,7 +608,21 @@ def _do_diffusion_step(
         vertex_end_local = self.grid.get_end_index(
             VertexDim, HorizontalMarkerIndex.local(VertexDim)
         )
-
+        vertex_end_halo = self.grid.get_end_index(VertexDim, HorizontalMarkerIndex.halo(VertexDim))
+
+        loc_rank = self._exchange.my_rank()
+        # log.debug("cell_start_interior for rank",loc_rank," is ..",cell_start_interior)
+        # log.debug("cell_start_nudging for rank", loc_rank, " is ..", cell_start_nudging)
+        # log.debug("cell_end_local for rank", loc_rank, " is ..", cell_end_local)
+        # log.debug("cell_end_halo for rank", loc_rank, " is ..", cell_end_halo)
+        # log.debug("edge_start_nudging_plus_one for rank", loc_rank, " is ..", edge_start_nudging_plus_one)
+        # log.debug("edge_start_lb_plus4 for rank", loc_rank, " is ..", edge_start_lb_plus4)
+        # log.debug("edge_end_local for rank", loc_rank, " is ..", edge_end_local)
+        # log.debug("edge_end_local_minus2 for rank", loc_rank, " is ..", edge_end_local_minus2)
+        # log.debug("edge_end_halo for rank", loc_rank, " is ..", edge_end_halo)
+        # log.debug("vertex_start_lb_plus1 for rank", loc_rank, " is ..", vertex_start_lb_plus1)
+        # log.debug("vertex_end_local for rank", loc_rank, " is ..", vertex_end_local)
+        # log.debug("vertex_end_halo for rank", loc_rank, " is ..", vertex_end_halo)
         # dtime dependent: enh_smag_factor,
         scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={})
 
@@ -624,10 +641,36 @@ def _do_diffusion_step(
         )
         log.debug("rbf interpolation 1: end")
 
+        # loc_ind_verts=self._exchange._decomposition_info.local_index(VertexDim,DecompositionInfo.EntryType.HALO)
+        # log.debug("loc_ind_verts rank %s", loc_rank, " loc_ind_verts: %s",loc_ind_verts," shape: %s",loc_ind_verts.shape)
+        # log.debug("after rbf rank %s", loc_rank, " u_vert max: %s min: %s",
+        #      xp.max(self.u_vert.ndarray[vertex_start_lb_plus1:vertex_end_local, 0:klevels]),
+        #      xp.min(self.u_vert.ndarray[vertex_start_lb_plus1:vertex_end_local, 0:klevels]))
         # 2.  HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
+        # log.debug("halo....after rbf rank %s", loc_rank, " u_vert max: %s min: %s",
+        #      self.u_vert.ndarray[loc_ind_verts, 0],
+        #      self.u_vert.ndarray[loc_ind_verts, 0])
         log.debug("communication rbf extrapolation of vn - start")
-        self._exchange.exchange_and_wait(VertexDim, self.u_vert, self.v_vert)
+        log.debug(
+            "size of u_vert %s v_vert %s", self.u_vert.ndarray.shape, self.v_vert.ndarray.shape
+        )
+        log.debug(
+            "edge_start_lb_plus4 %s edge_end_local_minus2 %s",
+            edge_start_lb_plus4,
+            edge_end_local_minus2,
+        )
+        self._exchange.exchange_and_wait(
+            VertexDim,
+            self.u_vert.ndarray[0 : self.grid.num_vertices, :],
+            self.v_vert.ndarray[0 : self.grid.num_vertices, :],
+        )
         log.debug("communication rbf extrapolation of vn - end")
+        # log.debug("after exchange rank %s", loc_rank, " u_vert max: %s min: %s",
+        #      xp.max(self.u_vert.ndarray[vertex_start_lb_plus1:vertex_end_local, 0:klevels]),
+        #      xp.min(self.u_vert.ndarray[vertex_start_lb_plus1:vertex_end_local, 0:klevels]))
+        # log.debug("halo....after exchange rank %s", loc_rank, " u_vert max: %s min: %s",
+        #      self.u_vert.ndarray[loc_ind_verts, 0],
+        #      self.u_vert.ndarray[loc_ind_verts, 0])
 
         log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start")
         calculate_nabla2_and_smag_coefficients_for_vn(
@@ -682,10 +725,18 @@ def _do_diffusion_step(
 
         # HALO EXCHANGE  IF (discr_vn > 1) THEN CALL sync_patch_array
         # TODO (magdalena) move this up and do asynchronous exchange
+        # loc_ind_edges=self._exchange._decomposition_info.local_index(EdgeDim,DecompositionInfo.EntryType.HALO)
+        # log.debug("loc_ind_edges rank %s", loc_rank, " loc_ind_edges: %s",loc_ind_edges," shape: %s",loc_ind_edges.shape)
+        # log.debug("halo..z_nabla2_e..before exchange rank %s", loc_rank, " z_nabla2_e: %s",
+        #      self.z_nabla2_e.ndarray[loc_ind_verts, 0])
         if self.config.type_vn_diffu > 1:
             log.debug("communication rbf extrapolation of z_nable2_e - start")
-            self._exchange.exchange_and_wait(EdgeDim, self.z_nabla2_e)
+            self._exchange.exchange_and_wait(
+                EdgeDim, self.z_nabla2_e.ndarray[0 : self.grid.num_edges, :]
+            )
             log.debug("communication rbf extrapolation of z_nable2_e - end")
+        # log.debug("halo..z_nabla2_e..after exchange rank %s", loc_rank, " z_nabla2_e: %s",
+        #      self.z_nabla2_e.ndarray[loc_ind_verts, 0])
 
         log.debug("2nd rbf interpolation: start")
         mo_intp_rbf_rbf_vec_interpol_vertex(
@@ -703,9 +754,19 @@ def _do_diffusion_step(
         log.debug("2nd rbf interpolation: end")
 
         # 6.  HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields)
+        # log.debug("halo....after 2nd... rbf rank %s", loc_rank, " u_vert %s",
+        #      self.u_vert.ndarray[loc_ind_verts, 0], " v_vert:",
+        #      self.u_vert.ndarray[loc_ind_verts, 0])
         log.debug("communication rbf extrapolation of z_nable2_e - start")
-        self._exchange.exchange_and_wait(VertexDim, self.u_vert, self.v_vert)
+        self._exchange.exchange_and_wait(
+            VertexDim,
+            self.u_vert.ndarray[0 : self.grid.num_vertices, :],
+            self.v_vert.ndarray[0 : self.grid.num_vertices, :],
+        )
         log.debug("communication rbf extrapolation of z_nable2_e - end")
+        # log.debug("halo....after 2nd... after exchange rank %s", loc_rank, " u_vert %s",
+        #      self.u_vert.ndarray[loc_ind_verts, 0], " v_vert:",
+        #      self.u_vert.ndarray[loc_ind_verts, 0])
 
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start")
         apply_diffusion_to_vn(
@@ -734,7 +795,14 @@ def _do_diffusion_step(
         )
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end")
         log.debug("communication of prognistic.vn : start")
-        handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn)
+        # log.debug("halo..vn..before exchange rank %s", loc_rank, " vn: %s",
+        #      prognostic_state.vn.ndarray[loc_ind_verts, 0])
+        handle_edge_comm = self._exchange.exchange(
+            EdgeDim, prognostic_state.vn.ndarray[0 : self.grid.num_edges, :]
+        )
+        # handle_edge_comm = self._exchange.exchange_and_wait(EdgeDim, prognostic_state.vn.ndarray[0:self.grid.num_edges,:])
+        # log.debug("halo..vn..after exchange rank %s", loc_rank, " vn: %s",
+        #      prognostic_state.vn.ndarray[loc_ind_verts, 0])
 
         log.debug(
             "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"

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

@@ -19,10 +19,9 @@
 from enum import IntEnum
 from typing import Any, Protocol
 
-import numpy as np
-import numpy.ma as ma
 from gt4py.next import Dimension
 
+from icon4py.model.common.settings import xp
 from icon4py.model.common.utils import builder
 
 
@@ -72,12 +71,13 @@ class EntryType(IntEnum):
         HALO = 2
 
     @builder
-    def with_dimension(self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray):
-        masked_global_index = ma.array(global_index, mask=owner_mask)
-        self._global_index[dim] = masked_global_index
+    def with_dimension(self, dim: Dimension, global_index: xp.ndarray, owner_mask: xp.ndarray):
+        self._global_index[dim] = global_index
+        self._owner_mask[dim] = owner_mask
 
     def __init__(self, klevels: int):
         self._global_index = {}
+        self._owner_mask = {}
         self._klevels = klevels
 
     @property
@@ -90,31 +90,29 @@ def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL):
                 return self._to_local_index(dim)
             case DecompositionInfo.EntryType.HALO:
                 index = self._to_local_index(dim)
-                mask = self._global_index[dim].mask
+                mask = self._owner_mask[dim]
                 return index[~mask]
             case DecompositionInfo.EntryType.OWNED:
                 index = self._to_local_index(dim)
-                mask = self._global_index[dim].mask
+                mask = self._owner_mask[dim]
                 return index[mask]
 
     def _to_local_index(self, dim):
-        data = ma.getdata(self._global_index[dim], subok=False)
+        data = self._global_index[dim]
         assert data.ndim == 1
-        return np.arange(data.shape[0])
+        return xp.arange(data.shape[0])
 
-    def owner_mask(self, dim: Dimension) -> np.ndarray:
-        return self._global_index[dim].mask
+    def owner_mask(self, dim: Dimension) -> xp.ndarray:
+        return self._owner_mask[dim]
 
     def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL):
         match entry_type:
             case DecompositionInfo.EntryType.ALL:
-                return ma.getdata(self._global_index[dim], subok=False)
+                return self._global_index[dim]
             case DecompositionInfo.EntryType.OWNED:
-                global_index = self._global_index[dim]
-                return ma.getdata(global_index[global_index.mask])
+                return self._global_index[dim][self._owner_mask[dim]]
             case DecompositionInfo.EntryType.HALO:
-                global_index = self._global_index[dim]
-                return ma.getdata(global_index[~global_index.mask])
+                return self._global_index[dim][~self._owner_mask[dim]]
             case _:
                 raise NotImplementedError()
 

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

@@ -21,27 +21,29 @@
 from gt4py.next import Dimension, Field
 
 from icon4py.model.common.decomposition.definitions import SingleNodeExchange
-
-
-try:
-    import ghex
-    import mpi4py
-    from ghex.context import make_context
-    from ghex.unstructured import (
-        DomainDescriptor,
-        HaloGenerator,
-        make_communication_object,
-        make_field_descriptor,
-        make_pattern,
-    )
-
-    mpi4py.rc.initialize = False
-    mpi4py.rc.finalize = True
-
-except ImportError:
-    mpi4py = None
-    ghex = None
-    unstructured = None
+from icon4py.model.common.settings import device
+
+
+#try:
+import ghex
+import mpi4py
+from ghex.context import make_context
+from ghex.unstructured import (
+    DomainDescriptor,
+    HaloGenerator,
+    make_communication_object,
+    make_field_descriptor,
+    make_pattern,
+)
+from ghex.util import Architecture
+
+mpi4py.rc.initialize = False
+mpi4py.rc.finalize = True
+
+# except ImportError:
+#     mpi4py = None
+#     ghex = None
+#     unstructured = None
 
 from icon4py.model.common.decomposition import definitions
 from icon4py.model.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim
@@ -51,6 +53,11 @@
     import mpi4py.MPI
 
 
+if device.name == "GPU":
+    ghex_arch = Architecture.GPU
+else:
+    ghex_arch = Architecture.CPU
+
 CommId = Union[int, "mpi4py.MPI.Comm", None]
 log = logging.getLogger(__name__)
 
@@ -100,8 +107,10 @@ def filter(self, record: logging.LogRecord) -> bool:
 
 
 @definitions.get_processor_properties.register(definitions.MultiNodeRun)
-def get_multinode_properties(s: definitions.MultiNodeRun) -> definitions.ProcessProperties:
-    return _get_processor_properties(with_mpi=True)
+def get_multinode_properties(
+    s: definitions.MultiNodeRun, comm_id: CommId = None
+) -> definitions.ProcessProperties:
+    return _get_processor_properties(with_mpi=True, comm_id=comm_id)
 
 
 @dataclass(frozen=True)
@@ -202,15 +211,16 @@ def exchange(self, dim: definitions.Dimension, *fields: Sequence[Field]):
         domain_descriptor = self._domain_descriptors[dim]
         assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found"
         applied_patterns = [
-            pattern(make_field_descriptor(domain_descriptor, f.asnumpy())) for f in fields
+            pattern(make_field_descriptor(domain_descriptor, f, arch=ghex_arch)) for f in fields
         ]
         handle = self._comm.exchange(applied_patterns)
-        log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.")
+        log.debug(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.")
         return MultiNodeResult(handle, applied_patterns)
 
     def exchange_and_wait(self, dim: Dimension, *fields: tuple):
         res = self.exchange(dim, *fields)
         res.wait()
+        log.debug(f"exchange for {len(fields)} fields of dimension ='{dim.value}' done.")
 
 
 @dataclass

model/common/src/icon4py/model/common/dimension.py

@@ -20,6 +20,13 @@
 EdgeDim = Dimension("Edge")
 CellDim = Dimension("Cell")
 VertexDim = Dimension("Vertex")
+SingletonDim = Dimension("Singleton")
+SpecialADim = Dimension("SpecialA")
+SpecialBDim = Dimension("SpecialB")
+SpecialCDim = Dimension("SpecialC")
+CellIndexDim = Dimension("CellIndex")
+EdgeIndexDim = Dimension("EdgeIndex")
+VertexIndexDim = Dimension("VertexIndex")
 CEDim = Dimension("CE")
 CECDim = Dimension("CEC")
 ECDim = Dimension("EC")

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

@@ -163,13 +163,6 @@ def end(cls, dim: Dimension) -> int:
         return cls._end[dim]
 
 
-@dataclass(frozen=True)
-class HorizontalGridSize:
-    num_vertices: int
-    num_edges: int
-    num_cells: int
-
-
 class EdgeParams:
     def __init__(
         self,

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

@@ -14,6 +14,7 @@
 from functools import cached_property
 
 import numpy as np
+from icon4py.model.common.settings import xp
 from gt4py.next.common import Dimension, DimensionKind
 from gt4py.next.ffront.fbuiltins import int32
 
@@ -90,7 +91,8 @@ def __init__(self):
 
     @builder
     def with_start_end_indices(
-        self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray
+        self, dim: Dimension, start_indices: xp.ndarray, end_indices: xp.ndarray
+        #self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray
     ):
         self.start_indices[dim] = start_indices.astype(int32)
         self.end_indices[dim] = end_indices.astype(int32)

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

@@ -11,7 +11,6 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-import numpy as np
 from gt4py.next import Dimension, NeighborTableOffsetProvider
 
 from icon4py.model.common.settings import xp
@@ -23,7 +22,8 @@ def neighbortable_offset_provider_for_1d_sparse_fields(
     neighbor_axis: Dimension,
     has_skip_values: bool,
 ):
-    table = xp.asarray(np.arange(old_shape[0] * old_shape[1]).reshape(old_shape))
+
+    table = xp.asarray(xp.arange(old_shape[0] * old_shape[1]).reshape(old_shape))
     return NeighborTableOffsetProvider(
         table,
         origin_axis,