[Refactor | Feature] Introducing NDSL backend

.github/workflows/fv3_translate_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   fv3_translate_tests:
-    uses: NOAA-GFDL/pyFV3/.github/workflows/translate.yaml@develop
+    uses: floriandeconinck/pyFV3/.github/workflows/translate.yaml@update/2026.02.00
     with:
       component_trigger: true
       component_name: NDSL

.github/workflows/pace_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   pace_main_tests:
-    uses: NOAA-GFDL/pace/.github/workflows/main_unit_tests.yaml@develop
+    uses: floriandeconinck/pace/.github/workflows/main_unit_tests.yaml@update/2026.02.00
     with:
       component_trigger: true
       component_name: NDSL

.github/workflows/shield_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   shield_translate_tests:
-    uses: NOAA-GFDL/pySHiELD/.github/workflows/translate.yaml@develop
+    uses: floriandeconinck/pySHiELD/.github/workflows/translate.yaml@update/2026.02.00
     with:
       component_trigger: true
       component_name: NDSL

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,15 +114,17 @@ 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."""
+    if not isinstance(backend, Backend):
+        raise RuntimeError(f"Backend {backend} is not of class Backend")
     return _get_factories(
         nx=nx,
         ny=ny,
         nz=nz,
         nhalo=nhalo,
         backend=backend,
-        orchestration=DaCeOrchestration.Python,
+        orchestration=DaCeOrchestration.BuildAndRun,
         topology="tile",
     )

ndsl/config/__init__.py

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

ndsl/config/backend.py

@@ -0,0 +1,209 @@
+from enum import Enum
+from typing import 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"
+    NUMPY = "numpy"
+
+
+class BackendLoopOrder(Enum):
+    """Cartesian loop order generated by the Backend"""
+
+    IJK = "IJK"
+    IKJ = "IKJ"
+    JKI = "JKI"
+    JIK = "JIK"
+    KIJ = "KIJ"
+    KJI = "KJI"
+
+
+_NDSL_TO_GT4PY_BACKEND_NAMING = {
+    "st:python:cpu:IJK": "debug",
+    "st:numpy:cpu:IJK": "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 = BackendLoopOrder(parts[3].upper())
+
+        # Check it exists in GT4Py
+        try:
+            gt4py_backend = gt_backend.from_name(self.as_gt4py())
+        except ValueError:
+            raise ValueError(
+                f"NDSL backend {ndsl_backend} does not have a working GT4Py version. "
+                f"GT4Py backend {self.as_gt4py()} is not registered. "
+                f"Contact the team."
+            )
+
+        # Check GPU capacity
+        if (
+            self._device == BackendTargetDevice.GPU
+            and gt4py_backend.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:
+        if isinstance(other, str):
+            other = Backend(other)
+        if not isinstance(other, Backend):
+            raise NotImplementedError(
+                f"Backend equality operator for {type(other)} is not implemented."
+            )
+        return self._humanly_readable == other._humanly_readable
+
+    def __hash__(self) -> int:
+        return hash(self._humanly_readable)
+
+    @staticmethod
+    def python() -> "Backend":
+        """Default backend for quick iterative work."""
+        return backend_python
+
+    @staticmethod
+    def cpu() -> "Backend":
+        """Default performance backend targeting CPU devices."""
+        return backend_cpu
+
+    @staticmethod
+    def gpu() -> "Backend":
+        """Default performance backend targeting GPU devices."""
+        return backend_gpu
+
+    @property
+    def device(self) -> BackendTargetDevice:
+        return self._device
+
+    @property
+    def framework(self) -> BackendFramework:
+        return self._framework
+
+    @property
+    def loop_order(self) -> BackendLoopOrder:
+        return self._loop_order
+
+    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 as_layout_map(self) -> tuple[int, ...]:
+        loop_order_as_string = self._loop_order.value
+        return tuple(
+            len(loop_order_as_string) - 1 - loop_order_as_string.index(axis)
+            for axis in "IJK"
+        )
+
+    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."""
+
+        # Dev NOTE: this probably should live as an accessor directly on the
+        # storage_info or layout_info of GT4Py, rather than stacked up on NDSL
+        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:IJK")
+"""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."""