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,
@@ -51,7 +51,6 @@
 
 from icon4py.model.common.constants import (
     CPD,
-    DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO,
     GAS_CONSTANT_DRY_AIR,
     dbl_eps,
 )
@@ -293,11 +292,11 @@ def __post_init__(self, config):
         ) = self._determine_smagorinski_factor(config)
         object.__setattr__(self, "smagorinski_factor", smagorinski_factor)
         object.__setattr__(self, "smagorinski_height", smagorinski_height)
-        # see mo_interpol_nml.f90:
+        # nudge_max_coeff is already multiplied by factor of 5 in mo_interpol_nml.f90
         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):
@@ -547,9 +546,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,6 +604,7 @@ 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))
 
         # dtime dependent: enh_smag_factor,
         scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={})
@@ -624,9 +624,12 @@ def _do_diffusion_step(
         )
         log.debug("rbf interpolation 1: end")
 
-        # 2.  HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
         log.debug("communication rbf extrapolation of vn - 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 vn - end")
 
         log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start")
@@ -684,7 +687,9 @@ def _do_diffusion_step(
         # TODO (magdalena) move this up and do asynchronous exchange
         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("2nd rbf interpolation: start")
@@ -704,7 +709,11 @@ def _do_diffusion_step(
 
         # 6.  HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields)
         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("running stencils 04 05 06 (apply_diffusion_to_vn): start")
@@ -734,7 +743,9 @@ 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)
+        handle_edge_comm = self._exchange.exchange(
+            EdgeDim, prognostic_state.vn.ndarray[0 : self.grid.num_edges, :]
+        )
 
         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/config.py

@@ -73,3 +73,7 @@ def device(self):
     @cached_property
     def limited_area(self):
         return os.environ.get("ICON4PY_LAM", False)
+
+    @cached_property
+    def parallel_run(self):
+        return os.environ.get("ICON4PY_PARALLEL", False)

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,6 +21,7 @@
 from gt4py.next import Dimension, Field
 
 from icon4py.model.common.decomposition.definitions import SingleNodeExchange
+from icon4py.model.common.settings import device
 
 
 try:
@@ -34,6 +35,7 @@
         make_field_descriptor,
         make_pattern,
     )
+    from ghex.util import Architecture
 
     mpi4py.rc.initialize = False
     mpi4py.rc.finalize = True
@@ -50,6 +52,7 @@
 if TYPE_CHECKING:
     import mpi4py.MPI
 
+ghex_arch = Architecture.GPU if device.name == "GPU" else Architecture.CPU
 
 CommId = Union[int, "mpi4py.MPI.Comm", None]
 log = logging.getLogger(__name__)
@@ -100,8 +103,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 +207,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/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/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,7 @@ 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,

model/common/src/icon4py/model/common/settings.py

@@ -18,3 +18,4 @@
 xp = config.array_ns
 device = config.device
 limited_area = config.limited_area
+parallel_run = config.parallel_run

tools/src/icon4pytools/py2fgen/template.py

@@ -106,7 +106,12 @@ def build_array_size_args() -> dict[str, str]:
     array_size_args = {}
     from icon4py.model.common import dimension
 
-    for var_name, var in vars(dimension).items():
+    from icon4pytools.py2fgen.wrapper_utils import dimension as wrapper_dimensions
+
+    combined_dims = dict(vars(dimension))
+    combined_dims.update(vars(wrapper_dimensions))
+
+    for var_name, var in combined_dims.items():
         if isinstance(var, Dimension):
             dim_name = var_name.replace(
                 "Dim", ""

tools/src/icon4pytools/py2fgen/wrapper_utils/__init__.py

@@ -0,0 +1,13 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+