Enable DaCe external workspace for AMD platform

model/common/src/icon4py/model/common/backend_configuration.py

@@ -0,0 +1,122 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+"""External workspace allocation for the DaCe backend.
+
+The DaCe backend in GT4Py can be configured with an
+:class:`~gt4py.next.program_processors.runners.dace.workflow.common.ExternalWorkspace`
+to provide workspace memory for transient SDFG arrays
+(``transient_memory_mode = EXTERNAL``). When a :class:`BackendConfig` is
+provided, :func:`get_dace_options <icon4py.model.common.model_options.get_dace_options>`
+calls the :class:`IconWorkspaceAllocator` defined here — a process-wide
+singleton that caches a single workspace slab per device and reuses it across
+every compiled program.
+
+The size of the workspace is configurable per experiment via :class:`BackendConfig`
+(see :func:`backend_config_from_env` for an environment-variable based default).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import os
+import typing
+from collections.abc import Iterable
+from typing import ClassVar, Final
+
+import gt4py.next as gtx
+from gt4py.next.program_processors.runners.dace.workflow import common as gtx_wfdcommon
+
+from icon4py.model.common.config import options as common_conf_opt
+from icon4py.model.common.utils import data_allocation
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class BackendConfig:
+    """External DaCe workspace sizing, configurable per experiment."""
+
+    workspace_size: typing.Annotated[
+        int,
+        common_conf_opt.ConfigOption(
+            description=(
+                "Size of the workspace memory (in Bytes) for externally allocated "
+                "temporary fields. This is a performance feature of the DaCe backend "
+                "to avoid runtime allocation of temporary fields, for each program "
+                "call. Note that the memory buffer is allocated once and shared "
+                "across all compiled programs."
+            ),
+            icon_equivalent=None,
+        ),
+    ] = 256 * 1024 * 1024  # 256 MiB
+
+    def __post_init__(self) -> None:
+        if self.workspace_size <= 0:
+            raise ValueError(f"'workspace_size' must be positive, got {self.workspace_size}.")
+
+
+def backend_config_from_env() -> BackendConfig | None:
+    """Build a :class:`BackendConfig` from environment variables.
+
+    Reads ``ICON4PY_BACKEND_WORKSPACE_SIZE`` and returns ``None`` when it is
+    not set.
+    """
+    size = os.environ.get("ICON4PY_BACKEND_WORKSPACE_SIZE")
+    if size is None:
+        return None
+    return BackendConfig(workspace_size=int(size))
+
+
+def _get_slab(nbytes: int, device: gtx.DeviceType) -> data_allocation.NDArray:
+    """Allocate a `nbytes`-byte buffer allocated on ``device``."""
+    xp = data_allocation.array_ns(use_cupy=(device != gtx.DeviceType.CPU))
+    return xp.empty(nbytes, dtype=xp.uint8)
+
+
+class IconWorkspaceAllocator:
+    """Singleton workspace allocator for the DaCe backend.
+
+    Exactly one instance exists per process (enforced by `__new__`); all DaCe
+    backends share it via the module-level `ICON_WORKSPACE_ALLOCATOR`. It keeps
+    a single private workspace slab per device in `_workspace_slabs`, reused
+    across every compiled program. On a cache hit the slab's size is validated
+    against the value passed to `allocate`.
+    """
+
+    _instance: ClassVar[IconWorkspaceAllocator | None] = None
+    _workspace_slabs: ClassVar[dict[gtx.DeviceType, data_allocation.NDArray]] = {}
+
+    def __new__(cls) -> IconWorkspaceAllocator:
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def allocate(
+        self,
+        devices: gtx.DeviceType | Iterable[gtx.DeviceType],
+        *,
+        size: int,
+    ) -> gtx_wfdcommon.ExternalWorkspace:
+        if isinstance(devices, gtx.DeviceType):
+            devices = [devices]
+        wsp: gtx_wfdcommon.ExternalWorkspace = {}
+        for dev in devices:
+            if (cached := self._workspace_slabs.get(dev)) is not None:
+                if cached.nbytes != size:
+                    raise ValueError(
+                        f"Workspace size mismatch for {dev!s}: cached slab has "
+                        f"{cached.nbytes} bytes but 'allocate' was called with "
+                        f"size={size}."
+                    )
+                wsp[dev] = cached
+            else:
+                slab = _get_slab(size, dev)
+                self._workspace_slabs[dev] = slab
+                wsp[dev] = slab
+        return wsp
+
+
+ICON_WORKSPACE_ALLOCATOR: Final[IconWorkspaceAllocator] = IconWorkspaceAllocator()

model/common/src/icon4py/model/common/model_backends.py

@@ -5,12 +5,17 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from __future__ import annotations
+
 from typing import Any, Final, TypeAlias, TypeGuard
 
 import gt4py.next as gtx
+import gt4py.next.custom_layout_allocators as gtx_allocators
 import gt4py.next.typing as gtx_typing
-from gt4py.next import backend as gtx_backend, custom_layout_allocators as gtx_allocators
+from gt4py.next import backend as gtx_backend
 from gt4py.next.program_processors.runners import dace as gtx_dace, gtfn
+from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations
+from gt4py.next.program_processors.runners.dace.workflow import common as gtx_wfdcommon
 
 
 # DeviceType should always be imported from here, as we might replace it by an ICON4Py internal implementation
@@ -81,6 +86,7 @@ def make_custom_dace_backend(
     use_metrics: bool = True,
     use_zero_origin: bool = False,
     use_max_domain_range_on_unstructured_shift: bool | None = None,
+    external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None,
     **_,
 ) -> gtx_typing.Backend:
     """Customize the dace backend with the given configuration parameters.
@@ -97,15 +103,34 @@ def make_custom_dace_backend(
         use_max_domain_range_on_unstructured_shift: When True, compute `as_fieldop`
             expressions everywhere. Otherwise, when all connectivities are given
             at compile time, infer the minimal domain of all `as_fieldop` statically.
+        external_workspace: The external workspace memory to use as storage for
+            the transient arrays. If `None`, the transient arrays will be allocated
+            inside the SDFG with scope lifetime.
 
     Returns:
         A dace backend with custom configuration for the target device.
     """
+    if external_workspace is not None:
+        if optimization_args is None:
+            optimization_args = {
+                "transient_memory_mode": gtx_transformations.TransientMemoryMode.EXTERNAL,
+            }
+        elif transient_memory_mode := optimization_args.get("transient_memory_mode"):
+            if transient_memory_mode != gtx_transformations.TransientMemoryMode.EXTERNAL:
+                raise ValueError(
+                    f"Cannot use external workspace with transient_memory_mode={transient_memory_mode}."
+                )
+        else:
+            optimization_args["transient_memory_mode"] = (
+                gtx_transformations.TransientMemoryMode.EXTERNAL
+            )
+
     on_gpu = device == GPU
     return gtx_dace.make_dace_backend(
         gpu=on_gpu,
         auto_optimize=auto_optimize,
         async_sdfg_call=async_sdfg_call,
+        external_workspace=external_workspace,
         optimization_args=optimization_args,
         unstructured_horizontal_has_unit_stride=True,
         use_metrics=use_metrics,

model/common/src/icon4py/model/common/model_options.py

@@ -16,7 +16,7 @@
 from gt4py.next import backend as gtx_backend
 from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations
 
-from icon4py.model.common import model_backends
+from icon4py.model.common import backend_configuration as backend_cfg, model_backends
 
 
 log = logging.getLogger(__name__)
@@ -35,11 +35,25 @@ def _dace_remove_access_node_copies(sdfg: dace.SDFG) -> None:
 
 
 def get_dace_options(
-    program_name: str, **backend_descriptor: Any
+    program_name: str,
+    backend_config: backend_cfg.BackendConfig | None,
+    **backend_descriptor: Any,
 ) -> model_backends.BackendDescriptor:
-    is_rocm_device = backend_descriptor.get("device") == model_backends.DeviceType.ROCM
+    device = backend_descriptor.get("device") or model_backends.CPU
     optimization_args = backend_descriptor.get("optimization_args", {})
     optimization_hooks = optimization_args.get("optimization_hooks", {})
+
+    if backend_config is not None:
+        # The workspace memory allows to avoid the overhead of runtime allocations,
+        # which are expensive in the AMD runtime.
+        backend_descriptor["external_workspace"] = backend_cfg.ICON_WORKSPACE_ALLOCATOR.allocate(
+            device,
+            size=backend_config.workspace_size,
+        )
+        optimization_args["transient_memory_mode"] = (
+            gtx_transformations.TransientMemoryMode.EXTERNAL
+        )
+
     if program_name in [
         "vertically_implicit_solver_at_corrector_step",
         "vertically_implicit_solver_at_predictor_step",
@@ -60,7 +74,7 @@ def get_dace_options(
         backend_descriptor["use_zero_origin"] = True
     if program_name == "graupel_run":
         optimization_args["fuse_tasklets"] = True
-        if not is_rocm_device:
+        if device != model_backends.DeviceType.ROCM:
             optimization_args["gpu_maxnreg"] = 80
             optimization_args["gpu_block_size_2d"] = (64, 6)
         optimization_args["gpu_memory_pool"] = False
@@ -78,12 +92,17 @@ def get_gtfn_options(
     return backend_descriptor
 
 
-def get_options(program_name: str, **backend_descriptor: Any) -> model_backends.BackendDescriptor:
+def get_options(
+    program_name: str,
+    *,
+    backend_config: backend_cfg.BackendConfig | None,
+    **backend_descriptor: Any,
+) -> model_backends.BackendDescriptor:
     if "backend_factory" not in backend_descriptor:
         # here we could set a backend_factory per program
         backend_descriptor["backend_factory"] = model_backends.make_custom_dace_backend
     if backend_descriptor["backend_factory"] == model_backends.make_custom_dace_backend:
-        backend_descriptor = get_dace_options(program_name, **backend_descriptor)
+        backend_descriptor = get_dace_options(program_name, backend_config, **backend_descriptor)
     if backend_descriptor["backend_factory"] == model_backends.make_custom_gtfn_backend:
         backend_descriptor = get_gtfn_options(program_name, **backend_descriptor)
 
@@ -96,7 +115,9 @@ def customize_backend(
     | model_backends.DeviceType
     | model_backends.BackendDescriptor
     | None,
+    backend_config: backend_cfg.BackendConfig | None = None,
 ) -> gtx_typing.Backend | None:
+    backend_config = backend_config or backend_cfg.backend_config_from_env()
     program_name = program.__name__ if program is not None else ""
     if backend is None or isinstance(backend, gtx_backend.Backend):
         backend_name = backend.name if backend is not None else "embedded"
@@ -106,7 +127,9 @@ def customize_backend(
     backend_descriptor = (
         {"device": backend} if isinstance(backend, model_backends.DeviceType) else backend
     )
-    backend_descriptor = get_options(program_name, **backend_descriptor)
+    backend_descriptor = get_options(
+        program_name, backend_config=backend_config, **backend_descriptor
+    )
     backend_descriptor["device"] = backend_descriptor.get(
         "device", model_backends.CPU
     )  # set default device
@@ -132,10 +155,11 @@ def setup_program(
     horizontal_sizes: dict[str, gtx.int32] | None = None,
     vertical_sizes: dict[str, gtx.int32] | None = None,
     offset_provider: gtx_typing.OffsetProvider | None = None,
+    backend_config: backend_cfg.BackendConfig | None = None,
 ) -> Callable[..., None]:
     """
     This function processes arguments to the GT4Py program. It
-    - binds arguments that don't change during model run ('constant_args', 'horizontal_sizes', "vertical_sizes');
+    - binds arguments that don't change during model run ('constant_args', 'horizontal_sizes', 'vertical_sizes');
     - inlines scalar arguments into the GT4Py program at compile-time (via GT4Py's 'compile').
     Args:
         - backend: GT4Py backend,
@@ -145,14 +169,16 @@ def setup_program(
         - horizontal_sizes: horizontal domain bounds,
         - vertical_sizes: vertical domain bounds,
         - offset_provider: GT4Py offset_provider,
+        - backend_config: external DaCe workspace sizing, or `None` to fall back
+          to the 'ICON4PY_BACKEND_WORKSPACE_SIZE' environment variable.
     """
     constant_args = {} if constant_args is None else constant_args
     variants = {} if variants is None else variants
     horizontal_sizes = {} if horizontal_sizes is None else horizontal_sizes
     vertical_sizes = {} if vertical_sizes is None else vertical_sizes
     offset_provider = {} if offset_provider is None else offset_provider
 
-    backend = customize_backend(program, backend)
+    backend = customize_backend(program, backend, backend_config=backend_config)
 
     bound_static_args = {k: v for k, v in constant_args.items() if gtx.is_scalar_type(v)}
     static_args_program = program.with_backend(backend)

model/common/src/icon4py/model/common/utils/data_allocation.py

@@ -8,7 +8,6 @@
 
 from __future__ import annotations
 
-import logging as log
 import math
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, TypeAlias, TypeGuard, TypeVar
@@ -71,23 +70,25 @@ def as_numpy(array: NDArrayInterface) -> np.ndarray:
         return cp.asnumpy(array)
 
 
-def _array_ns(try_cupy: bool) -> ModuleType:
-    """CuPy if requested and installed, NumPy otherwise."""
-    if try_cupy:
+def array_ns(use_cupy: bool) -> ModuleType:
+    """CuPy if requested, NumPy otherwise.
+
+    Raises RuntimeError if CuPy is requested but not available.
+    """
+    if use_cupy:
         try:
             import cupy as cp  # noqa: PLC0415 [import-outside-top-level]
-
-            return cp
-        except ImportError:
-            log.warning("No cupy installed, falling back to numpy for array_ns")
+        except ImportError as err:
+            raise RuntimeError(f"cupy is not available: {err!r}.") from err
+        return cp
     import numpy as np  # noqa: PLC0415 [import-outside-top-level]
 
     return np
 
 
 def import_array_ns(allocator: gtx_typing.Allocator | None) -> ModuleType:
     """Import cupy or numpy depending on a chosen GT4Py backend DevicType."""
-    return _array_ns(device_utils.is_cupy_device(allocator))
+    return array_ns(device_utils.is_cupy_device(allocator))
 
 
 def scalar_like_array[ScalarT: gtx_typing.Scalar](