Doubly periodic dycore config

driver/pace/driver/driver.py

@@ -24,7 +24,11 @@
 
 # TODO: move update_atmos_state into pace.driver
 from pace.stencils import update_atmos_state
-from pace.util.communicator import CubedSphereCommunicator
+from pace.util.communicator import (
+    Communicator,
+    CubedSphereCommunicator,
+    TileCommunicator,
+)
 from pace.util.logging import pace_log
 
 from . import diagnostics
@@ -90,6 +94,7 @@ class DriverConfig:
     nz: int
     layout: Tuple[int, int]
     dt_atmos: float
+    grid_type: Optional[int] = 0
     grid_config: GridInitializerSelector = dataclasses.field(
         default_factory=lambda: GridInitializerSelector(
             type="generated", config=GeneratedGridConfig()
@@ -158,7 +163,7 @@ def apply_tendencies(self) -> bool:
 
     def get_grid(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         quantity_factory: Optional[pace.util.QuantityFactory] = None,
     ) -> Tuple[
         pace.util.grid.DampingCoefficients,
@@ -187,7 +192,7 @@ def get_grid(
 
     def get_driver_state(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -213,7 +218,7 @@ def get_driver_state(
             if stencil_factory is None:
                 grid_indexing = (
                     pace.dsl.stencil.GridIndexing.from_sizer_and_communicator(
-                        sizer=sizer, cube=communicator
+                        sizer=sizer, comm=communicator
                     )
                 )
                 stencil_factory = pace.dsl.StencilFactory(
@@ -275,7 +280,7 @@ def from_dict(cls, kwargs: Dict[str, Any]) -> "DriverConfig":
             )
             grid_type = kwargs["grid_config"].config.grid_type
             # Copy grid_type to the DycoreConfig if it's not the default value
-            if grid_type != 0: 
+            if grid_type != 0:
                 kwargs["dycore_config"].grid_type = grid_type
 
         if (
@@ -404,11 +409,19 @@ def __init__(
                 if self.config.performance_config.collect_communication
                 else None
             )
-            communicator = CubedSphereCommunicator.from_layout(
-                comm=self.comm,
-                layout=self.config.layout,
-                timer=comm_timer,
-            )
+            communicator: Communicator
+            if self.config.grid_type <= 3:
+                communicator = CubedSphereCommunicator.from_layout(
+                    comm=self.comm,
+                    layout=self.config.layout,
+                    timer=comm_timer,
+                )
+            else:
+                communicator = TileCommunicator.from_layout(
+                    comm=self.comm,
+                    layout=self.config.layout,
+                    timer=comm_timer,
+                )
             self._update_driver_config_with_communicator(communicator)
 
             if self.config.stencil_config.compilation_config.run_mode == RunMode.Build:
@@ -544,7 +557,7 @@ def exit_instead_of_build(self):
         pace_log.info("initialization of the object done")
 
     def _update_driver_config_with_communicator(
-        self, communicator: CubedSphereCommunicator
+        self, communicator: Communicator
     ) -> None:
         dace_config = DaceConfig(
             communicator=communicator,
@@ -707,7 +720,7 @@ def log_subtile_location(partitioner: pace.util.TilePartitioner, rank: int):
 
 def _setup_factories(
     config: DriverConfig,
-    communicator: pace.util.CubedSphereCommunicator,
+    communicator: pace.util.Communicator,
     stencil_compare_comm,
 ) -> Tuple[pace.util.QuantityFactory, pace.dsl.StencilFactory]:
     """
@@ -735,7 +748,7 @@ def _setup_factories(
     )
 
     grid_indexing = pace.dsl.stencil.GridIndexing.from_sizer_and_communicator(
-        sizer=sizer, cube=communicator
+        sizer=sizer, comm=communicator
     )
     quantity_factory = pace.util.QuantityFactory.from_backend(
         sizer, backend=config.stencil_config.compilation_config.backend

driver/pace/driver/grid.py

@@ -10,7 +10,7 @@
 import pace.stencils
 import pace.util.grid
 from pace.stencils.testing import TranslateGrid
-from pace.util import CubedSphereCommunicator, QuantityFactory
+from pace.util import Communicator, QuantityFactory
 from pace.util.grid import (
     DampingCoefficients,
     DriverGridData,
@@ -35,7 +35,7 @@ class GridInitializer(abc.ABC):
     def get_grid(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
     ) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
         ...
 
@@ -62,7 +62,7 @@ def register(cls, type_name):
     def get_grid(
         self,
         quantity_factory: QuantityFactory,
-        communicator: CubedSphereCommunicator,
+        communicator: Communicator,
     ) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
         return self.config.get_grid(
             quantity_factory=quantity_factory, communicator=communicator
@@ -103,7 +103,7 @@ class GeneratedGridConfig(GridInitializer):
     def get_grid(
         self,
         quantity_factory: QuantityFactory,
-        communicator: CubedSphereCommunicator,
+        communicator: Communicator,
     ) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
 
         metric_terms = MetricTerms(
@@ -158,7 +158,7 @@ def _f90_namelist(self) -> f90nml.Namelist:
     def _namelist(self) -> Namelist:
         return Namelist.from_f90nml(self._f90_namelist)
 
-    def _serializer(self, communicator: pace.util.CubedSphereCommunicator):
+    def _serializer(self, communicator: pace.util.Communicator):
         import serialbox
 
         serializer = serialbox.Serializer(
@@ -170,7 +170,7 @@ def _serializer(self, communicator: pace.util.CubedSphereCommunicator):
 
     def _get_serialized_grid(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         backend: str,
     ) -> pace.stencils.testing.grid.Grid:  # type: ignore
         ser = self._serializer(communicator)
@@ -182,7 +182,7 @@ def _get_serialized_grid(
     def get_grid(
         self,
         quantity_factory: QuantityFactory,
-        communicator: CubedSphereCommunicator,
+        communicator: Communicator,
     ) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
 
         backend = quantity_factory.empty(

driver/pace/driver/initialization.py

@@ -37,7 +37,7 @@ def start_time(self) -> datetime:
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -74,7 +74,7 @@ def start_time(self) -> datetime:
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -105,11 +105,12 @@ class BaroclinicInit(Initializer):
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
     ) -> DriverState:
+        assert isinstance(communicator, pace.util.CubedSphereCommunicator)
         dycore_state = baroclinic_init.init_baroclinic_state(
             grid_data=grid_data,
             quantity_factory=quantity_factory,
@@ -149,12 +150,12 @@ class TropicalCycloneConfig(Initializer):
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
     ) -> DriverState:
-
+        assert isinstance(communicator, pace.util.CubedSphereCommunicator)
         dycore_state = tc_init.init_tc_state(
             grid_data=grid_data,
             quantity_factory=quantity_factory,
@@ -198,7 +199,7 @@ class RestartInit(Initializer):
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -247,7 +248,7 @@ def start_time(self) -> datetime:
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -296,7 +297,7 @@ def _namelist(self) -> Namelist:
 
     def _get_serialized_grid(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         backend: str,
     ) -> pace.stencils.testing.grid.Grid:  # type: ignore
         ser = self._serializer(communicator)
@@ -305,7 +306,7 @@ def _get_serialized_grid(
         ).python_grid()
         return grid
 
-    def _serializer(self, communicator: pace.util.CubedSphereCommunicator):
+    def _serializer(self, communicator: pace.util.Communicator):
         import serialbox
 
         serializer = serialbox.Serializer(
@@ -318,7 +319,7 @@ def _serializer(self, communicator: pace.util.CubedSphereCommunicator):
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -345,7 +346,7 @@ def get_driver_state(
 
     def _initialize_dycore_state(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         backend: str,
     ) -> fv3core.DycoreState:
 
@@ -396,7 +397,7 @@ class PredefinedStateInit(Initializer):
     def get_driver_state(
         self,
         quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,

driver/pace/driver/state.py

@@ -77,7 +77,7 @@ def load_state_from_restart(
         grid_data: pace.util.grid.GridData,
     ) -> "DriverState":
         comm = driver_config.comm_config.get_comm()
-        communicator = pace.util.CubedSphereCommunicator.from_layout(
+        communicator = pace.util.Communicator.from_layout(
             comm=comm, layout=driver_config.layout
         )
         sizer = pace.util.SubtileGridSizer.from_tile_params(
@@ -172,7 +172,7 @@ def _restart_driver_state(
     path: str,
     rank: int,
     quantity_factory: pace.util.QuantityFactory,
-    communicator: pace.util.CubedSphereCommunicator,
+    communicator: pace.util.Communicator,
     damping_coefficients: pace.util.grid.DampingCoefficients,
     driver_grid_data: pace.util.grid.DriverGridData,
     grid_data: pace.util.grid.GridData,

dsl/pace/dsl/dace/build.py

@@ -31,7 +31,7 @@ def unblock_waiting_tiles(comm, sdfg_path: str) -> None:
             comm.send(sdfg_path, dest=tile * tilesize + comm.Get_rank())
 
 
-def get_target_rank(rank: int, partitioner: pace.util.CubedSpherePartitioner):
+def get_target_rank(rank: int, partitioner: pace.util.Partitioner):
     """From my rank & the current partitioner we determine which
     rank we should read from.
     For all layout >= 3,3 this presumes build has been done on a

dsl/pace/dsl/dace/dace_config.py

@@ -7,7 +7,7 @@
 
 from pace.dsl.gt4py_utils import is_gpu_backend
 from pace.util._optional_imports import cupy as cp
-from pace.util.communicator import CubedSphereCommunicator
+from pace.util.communicator import Communicator
 
 
 # This can be turned on to revert compilation for orchestration
@@ -57,7 +57,7 @@ def __call__(self):
 class DaceConfig:
     def __init__(
         self,
-        communicator: Optional[CubedSphereCommunicator],
+        communicator: Optional[Communicator],
         backend: str,
         tile_nx: int = 0,
         tile_nz: int = 0,

dsl/pace/dsl/dace/wrapped_halo_exchange.py

@@ -2,7 +2,7 @@
 from typing import List, Optional
 
 from pace.dsl.dace.orchestration import dace_inhibitor
-from pace.util.communicator import CubedSphereCommunicator
+from pace.util.communicator import Communicator
 from pace.util.halo_updater import HaloUpdater
 
 
@@ -21,7 +21,7 @@ def __init__(
         state,
         qty_x_names: List[str],
         qty_y_names: List[str] = None,
-        comm: Optional[CubedSphereCommunicator] = None,
+        comm: Optional[Communicator] = None,
     ) -> None:
         self._updater = updater
         self._state = state

dsl/pace/dsl/stencil.py

@@ -595,18 +595,18 @@ def domain(self, domain):
 
     @classmethod
     def from_sizer_and_communicator(
-        cls, sizer: pace.util.GridSizer, cube: pace.util.CubedSphereCommunicator
+        cls, sizer: pace.util.GridSizer, comm: pace.util.Communicator
     ) -> "GridIndexing":
         # TODO: if this class is refactored to split off the *_edge booleans,
         # this init routine can be refactored to require only a GridSizer
         domain = cast(
             Tuple[int, int, int],
             sizer.get_extent([pace.util.X_DIM, pace.util.Y_DIM, pace.util.Z_DIM]),
         )
-        south_edge = cube.tile.partitioner.on_tile_bottom(cube.rank)
-        north_edge = cube.tile.partitioner.on_tile_top(cube.rank)
-        west_edge = cube.tile.partitioner.on_tile_left(cube.rank)
-        east_edge = cube.tile.partitioner.on_tile_right(cube.rank)
+        south_edge = comm.tile.partitioner.on_tile_bottom(comm.rank)
+        north_edge = comm.tile.partitioner.on_tile_top(comm.rank)
+        west_edge = comm.tile.partitioner.on_tile_left(comm.rank)
+        east_edge = comm.tile.partitioner.on_tile_right(comm.rank)
         return cls(
             domain=domain,
             n_halo=sizer.n_halo,