[Refactor | Feature] Introducing NDSL backend

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_PERFORMANCE_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_PERFORMANCE_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,8 @@
+from .backend import Backend, BackendFramework, BackendTargetDevice
+
+
+__all__ = [
+    "Backend",
+    "BackendFramework",
+    "BackendTargetDevice",
+]

ndsl/config/backend.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+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"""
+
+_NDSL_SHORT_TO_LONG_NAME = {
+    "numpy": "st:python:cpu:numpy",
+    "debug": "st:python:cpu:debug",
+    "performance_cpu": "orch:dace:cpu:IJK",
+    "hybrid_fortran_cpu": "orch:dace:cpu:KJI",
+    "performance_gpu": "orch:dace:gpu:KJI",
+}
+"""Internal: match short name to long name form"""
+
+
+class Backend:
+    """Backend for NDSL.
+
+    The backend is a string concatenating information on the intent of the user
+    for a given execution seperated by a ':'.
+
+    It describes to NDSL the strategy, device and framwork to be used
+    on the frontend code. Additionaly, 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 (see internal
+    `_NDSL_SHORT_TO_LONG_NAME`).
+    """
+
+    def __init__(self, ndsl_backend: str) -> None:
+        # Swap short form to long form
+        if ndsl_backend in _NDSL_SHORT_TO_LONG_NAME.keys():
+            ndsl_backend = _NDSL_SHORT_TO_LONG_NAME[ndsl_backend]
+
+        # Checks for existence and form
+        if ndsl_backend not in _NDSL_TO_GT4PY_BACKEND_NAMING.keys():
+            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"Coding error: NDSL backend requested {self._humanly_readable} "
+                f"translate to non-GPU {self.as_gt4py()} GT4Py backend"
+            )
+
+    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 debug() -> Backend:
+        return Backend("debug")
+
+    @staticmethod
+    def python() -> Backend:
+        return Backend("numpy")
+
+    @staticmethod
+    def performance_cpu() -> Backend:
+        return Backend("performance_cpu")
+
+    @staticmethod
+    def hybrid_fortran_cpu() -> Backend:
+        return Backend("hybrid_fortran_cpu")
+
+    @staticmethod
+    def performance_gpu() -> Backend:
+        return Backend("performance_gpu")
+
+    @property
+    def device(self) -> BackendTargetDevice:
+        return self._device
+
+    @property
+    def framework(self) -> BackendFramework:
+        return self._framework
+
+    def as_gt4py(self) -> str:
+        """Given a NDSL backend, give back a GT4Py equivalent"""
+        if self._humanly_readable in _NDSL_TO_GT4PY_BACKEND_NAMING.keys():
+            return _NDSL_TO_GT4PY_BACKEND_NAMING[self._humanly_readable]
+        raise ValueError(
+            f"Backend {self._humanly_readable} cannot be translate to GT4Py"
+        )
+
+    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
+
+
+# Those two internal values are used for default parameters values in functions/methods
+# as it is bad practice to call a function in default argument value
+_BACKEND_PERFORMANCE_CPU = Backend.performance_cpu()
+"""Internal: cache performance CPU. Please use Backend.performance_cpu()."""
+_BACKEND_PYTHON = Backend.python()
+"""Internal: cache performance CPU. Please use Backend.python()."""

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