[Refactor | Feature] Introducing NDSL backend

examples/mpi/zarr_monitor.py

@@ -12,6 +12,7 @@
     TilePartitioner,
     ZarrMonitor,
 )
+from ndsl.config import backend_python
 from ndsl.constants import I_DIM, J_DIM, K_DIM
 
 
@@ -20,9 +21,9 @@
 
 def get_example_state(time):
     sizer = SubtileGridSizer(
-        nx=48, ny=48, nz=70, n_halo=3, data_dimensions={}, backend="debug"
+        nx=48, ny=48, nz=70, n_halo=3, data_dimensions={}, backend=backend_python
     )
-    allocator = QuantityFactory(sizer, np)
+    allocator = QuantityFactory(sizer, backend=backend_python)
     air_temperature = allocator.zeros([I_DIM, J_DIM, K_DIM], units="degK")
     air_temperature.view[:] = np.random.randn(*air_temperature.extent)
     return {"time": time, "air_temperature": air_temperature}

ndsl/__init__.py

@@ -1,20 +1,14 @@
+# isort:skip_file
 from . import dsl  # isort:skip
 from .logging import ndsl_log  # isort:skip
 from .comm.communicator import CubedSphereCommunicator, TileCommunicator
 from .comm.local_comm import LocalComm
 from .comm.mpi import MPIComm
 from .comm.partitioner import CubedSpherePartitioner, TilePartitioner
+from .config.backend import Backend
 from .constants import ConstantVersions
 from .dsl.caches.codepath import FV3CodePath
-from .dsl.dace.dace_config import DaceConfig, DaCeOrchestration
-from .dsl.dace.orchestration import orchestrate, orchestrate_function
-from .dsl.dace.utils import (
-    ArrayReport,
-    DaCeProgress,
-    MaxBandwidthBenchmarkProgram,
-    StorageReport,
-)
-from .dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
+from .quantity import Quantity
 from .dsl.ndsl_runtime import NDSLRuntime
 from .dsl.stencil import FrozenStencil, GridIndexing, StencilFactory, TimingCollector
 from .dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
@@ -25,14 +19,25 @@
 from .performance.collector import NullPerformanceCollector, PerformanceCollector
 from .performance.profiler import NullProfiler, Profiler
 from .performance.report import Experiment, Report, TimeReport
-from .quantity import Local, LocalState, Quantity, State
+from .quantity import Local, LocalState, State
 from .quantity.field_bundle import FieldBundle, FieldBundleType  # Break circular import
 from .types import Allocator
 from .utils import MetaEnumStr
 
+from .dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
+from .dsl.dace.utils import (
+    ArrayReport,
+    DaCeProgress,
+    MaxBandwidthBenchmarkProgram,
+    StorageReport,
+)
+from .dsl.dace.dace_config import DaceConfig, DaCeOrchestration
+from .dsl.dace.orchestration import orchestrate, orchestrate_function
+
 
 __all__ = [
     "dsl",
+    "Backend",
     "CubedSphereCommunicator",
     "TileCommunicator",
     "LocalComm",

ndsl/boilerplate.py

@@ -1,6 +1,7 @@
 import warnings
 
 from ndsl import (
+    Backend,
     CompilationConfig,
     DaceConfig,
     DaCeOrchestration,
@@ -14,14 +15,15 @@
     TileCommunicator,
     TilePartitioner,
 )
+from ndsl.config.backend import backend_cpu, backend_python
 
 
 def _get_factories(
     nx: int,
     ny: int,
     nz: int,
     nhalo: int,
-    backend: str,
+    backend: Backend,
     orchestration: DaCeOrchestration,
     topology: str,
 ) -> tuple[StencilFactory, QuantityFactory]:
@@ -91,14 +93,14 @@ def get_factories_single_tile_orchestrated(
     ny: int,
     nz: int,
     nhalo: int,
-    backend: str = "dace:cpu",
+    backend: Backend = backend_cpu,
     *,
     orchestration_mode: DaCeOrchestration | None = None,
 ) -> tuple[StencilFactory, QuantityFactory]:
     """Build the pair of (StencilFactory, QuantityFactory) for orchestrated code on a single tile topology."""
 
-    if backend is not None and not backend.startswith("dace"):
-        raise ValueError("Only `dace:*` backends can be orchestrated.")
+    if backend is not None and not backend.is_orchestrated():
+        raise ValueError(f"Only `orch:*` backends can be orchestrated, got {backend}.")
 
     return _get_factories(
         nx=nx,
@@ -112,7 +114,7 @@ def get_factories_single_tile_orchestrated(
 
 
 def get_factories_single_tile(
-    nx: int, ny: int, nz: int, nhalo: int, backend: str = "numpy"
+    nx: int, ny: int, nz: int, nhalo: int, backend: Backend = backend_python
 ) -> tuple[StencilFactory, QuantityFactory]:
     """Build the pair of (StencilFactory, QuantityFactory) for stencils on a single tile topology."""
     return _get_factories(
@@ -121,6 +123,6 @@ def get_factories_single_tile(
         nz=nz,
         nhalo=nhalo,
         backend=backend,
-        orchestration=DaCeOrchestration.Python,
+        orchestration=DaCeOrchestration.BuildAndRun,
         topology="tile",
     )

ndsl/config/__init__.py

@@ -0,0 +1,20 @@
+from .backend import (
+    Backend,
+    BackendFramework,
+    BackendStrategy,
+    BackendTargetDevice,
+    backend_cpu,
+    backend_gpu,
+    backend_python,
+)
+
+
+__all__ = [
+    "Backend",
+    "BackendFramework",
+    "BackendStrategy",
+    "BackendTargetDevice",
+    "backend_python",
+    "backend_cpu",
+    "backend_gpu",
+]

ndsl/config/backend.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Final
+
+import gt4py.cartesian.backend as gt_backend
+
+
+class BackendStrategy(Enum):
+    """Strategy for the code execution"""
+
+    STENCIL = "st"
+    ORCHESTRATION = "orch"
+
+
+class BackendTargetDevice(Enum):
+    """Target device"""
+
+    CPU = "cpu"
+    GPU = "gpu"
+
+
+class BackendFramework(Enum):
+    """Main lower-level framework (or language) backend relies on"""
+
+    GRIDTOOLS = "gt"
+    DACE = "dace"
+    PYTHON = "python"
+
+
+_NDSL_TO_GT4PY_BACKEND_NAMING = {
+    "st:python:cpu:debug": "debug",
+    "st:python:cpu:numpy": "numpy",
+    "st:gt:cpu:IJK": "gt:cpu_kfirst",
+    "st:gt:cpu:KJI": "gt:cpu_ifirst",
+    "st:gt:gpu:KJI": "gt:gpu",
+    "st:dace:cpu:IJK": "dace:cpu_kfirst",
+    "orch:dace:cpu:IJK": "dace:cpu_kfirst",
+    "st:dace:cpu:KIJ": "dace:cpu",
+    "orch:dace:cpu:KIJ": "dace:cpu",
+    "st:dace:cpu:KJI": "dace:cpu_KJI",
+    "orch:dace:cpu:KJI": "dace:cpu_KJI",
+    "st:dace:gpu:KJI": "dace:gpu",
+    "orch:dace:gpu:KJI": "dace:gpu",
+}
+"""Internal: match the NDSL backend names with the GT4Py names"""
+
+_FORTRAN_LOOP_LAYOUT = (2, 1, 0)
+"""Fortran is a column-first (or stride-first) memory system,
+which in the internal gt4py loop layout means I (or axis[0]) has
+the higher value, e.g. "higher importance to run first":
+
+for k # Layout=0
+    for j # Layout=1
+        for i # Layout=2
+
+"""
+
+
+class Backend:
+    """Backend for NDSL.
+
+    The backend is a string concatenating information on the intent of the user
+    for a given execution separated by a ':'.
+
+    It describes to NDSL the strategy, device and framework to be used
+    on the frontend code. Additionally, it gives a hint toward the macro-strategy
+    for loop ordering (IJK, KJI, etc.) or a more broad intent (debug, numpy).
+
+    For convenience, shorcuts are given to the most common needs (
+    `backend_python`, `backend_cpu`, `backend_gpu`).
+    """
+
+    def __init__(self, ndsl_backend: str) -> None:
+        # Checks for existence and form
+        if ndsl_backend not in _NDSL_TO_GT4PY_BACKEND_NAMING:
+            raise ValueError(
+                f"Unknown {ndsl_backend}, options are {list(_NDSL_TO_GT4PY_BACKEND_NAMING.keys())}"
+            )
+        parts = ndsl_backend.split(":")
+        if len(parts) != 4:
+            raise ValueError(f"Backend {ndsl_backend} is ill-formed.")
+
+        # Breakdown and save into internal parameters
+        self._humanly_readable = ndsl_backend
+        self._strategy = BackendStrategy(parts[0].lower())
+        self._framework = BackendFramework(parts[1].lower())
+        self._device = BackendTargetDevice(parts[2].lower())
+        self._loop_order = parts[3]
+
+        # Check GPU capacity
+        if (
+            self._device == BackendTargetDevice.GPU
+            and gt_backend.from_name(self.as_gt4py()).storage_info["device"] != "gpu"
+        ):
+            raise ValueError(
+                f"NDSL backend requested ({self._humanly_readable}) tagets GPU,"
+                f"but requests a non-GPU backend from GT4Py ({self.as_gt4py()})."
+            )
+
+    def __str__(self) -> str:
+        return self.as_humanly_readable()
+
+    def __repr__(self) -> str:
+        return self.as_humanly_readable()
+
+    def __eq__(self, other: object) -> bool:
+        return self._humanly_readable == self._humanly_readable
+
+    def __hash__(self) -> int:
+        return hash(self._humanly_readable)
+
+    def __add__(self, other: str) -> Any:
+        """Concatenation operators"""
+        if isinstance(other, Backend):
+            raise TypeError("OperationError: Backend cannot add to another Backend")
+        return str(self) + other
+
+    def __radd__(self, other: str) -> Any:
+        """Concatenation operators"""
+        if isinstance(other, Backend):
+            raise TypeError("OperationError: Backend cannot add to another Backend")
+        return other + str(self)
+
+    @staticmethod
+    def python() -> Backend:
+        """Default backend for quick iterative work."""
+        return backend_python
+
+    @staticmethod
+    def cpu() -> Backend:
+        """Default performance backend targeting CPU device"""
+        return backend_cpu
+
+    @staticmethod
+    def gpu() -> Backend:
+        """Default performance backend targeting GPU device"""
+        return backend_gpu
+
+    @property
+    def device(self) -> BackendTargetDevice:
+        return self._device
+
+    @property
+    def framework(self) -> BackendFramework:
+        return self._framework
+
+    def as_gt4py(self) -> str:
+        """Given an NDSL backend, give back a GT4Py equivalent"""
+        return _NDSL_TO_GT4PY_BACKEND_NAMING[self._humanly_readable]
+
+    def as_humanly_readable(self) -> str:
+        return self._humanly_readable
+
+    def as_safe_for_path(self) -> str:
+        return self._humanly_readable.replace(":", "_")
+
+    def is_orchestrated(self) -> bool:
+        return self._strategy == BackendStrategy.ORCHESTRATION
+
+    def is_stencil(self) -> bool:
+        return self._strategy == BackendStrategy.STENCIL
+
+    def is_gpu_backend(self) -> bool:
+        return self._device == BackendTargetDevice.GPU
+
+    def is_fortran_aligned(self) -> bool:
+        """Check that the standard 3D field on cartesian axis is memory-aligned with Fortran
+        striding."""
+        return _FORTRAN_LOOP_LAYOUT == gt_backend.from_name(
+            self.as_gt4py()
+        ).storage_info["layout_map"](("I", "J", "K"))
+
+
+backend_python: Final[Backend] = Backend("st:python:cpu:debug")
+"""Default backend for quick iterative work."""
+
+backend_cpu: Final[Backend] = Backend("orch:dace:cpu:IJK")
+"""Default performance backend targeting CPU device"""
+
+backend_gpu: Final[Backend] = Backend("orch:dace:gpu:KJI")
+"""Default performance backend targeting GPU device"""

ndsl/dsl/dace/build.py

@@ -2,6 +2,7 @@
 from dace.sdfg import SDFG
 from gt4py.cartesian import config as gt_config
 
+from ndsl.config import Backend
 from ndsl.dsl.caches.cache_location import get_cache_directory, get_cache_fullpath
 from ndsl.dsl.dace.dace_config import DaceConfig, DaCeOrchestration
 from ndsl.logging import ndsl_log
@@ -23,7 +24,10 @@ def build_info_filepath() -> str:
 
 
 def write_build_info(
-    sdfg: SDFG, layout: tuple[int, int], resolution_per_tile: list[int], backend: str
+    sdfg: SDFG,
+    layout: tuple[int, int],
+    resolution_per_tile: list[int],
+    backend: Backend,
 ) -> None:
     """Write down all relevant information on the build to identify
     it at load time."""
@@ -86,7 +90,7 @@ def get_sdfg_path(
         build_info_file.readline()
         # Read in
         build_backend = build_info_file.readline().rstrip()
-        if config.get_backend() != build_backend:
+        if config.get_backend() != Backend(build_backend):
             raise RuntimeError(
                 f"SDFG build for {build_backend}, {config._backend} has been asked"
             )
@@ -110,12 +114,12 @@ def set_distributed_caches(config: DaceConfig) -> None:
     """In Run mode, check required file then point current rank cache to source cache"""
 
     # Execute specific initialization per orchestration state
-    orchestration_mode = config.get_orchestrate()
-    if orchestration_mode == DaCeOrchestration.Python:
+    if not config.get_backend().is_orchestrated():
         return
 
     # Check that we have all the file we need to early out in case
     # of issues.
+    orchestration_mode = config.get_orchestrate()
     if orchestration_mode == DaCeOrchestration.Run:
         import os