Enable DaCe external workspace for AMD platform

[edopao]

Summary

Enable DaCe's external workspace feature for the AMD (ROCm) platform. The external workspace avoids the overhead of runtime allocations for transient SDFG arrays, which are expensive in the AMD runtime.

The workspace is sized via workspace_size only; alignment of the workspace slab is handled by the underlying allocator, so no separate alignment configuration is exposed.

Changes

New module: model/common/src/icon4py/model/common/backend_configuration.py

• BackendConfig dataclass for workspace sizing (workspace_size), with validation in __post_init__ (size must be positive).
• backend_config_from_env() reads ICON4PY_BACKEND_WORKSPACE_SIZE from the environment; returns None when it is not set.
• IconWorkspaceAllocator singleton that caches a single workspace slab per device (via _get_slab) and reuses it across every compiled program. On cache hit, the slab size is validated.
• ICON_WORKSPACE_ALLOCATOR module-level instance consumed by get_dace_options.

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

• Added external_workspace parameter to make_custom_dace_backend. When provided, transient_memory_mode is forced to EXTERNAL (or validated if already set).

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

• get_dace_options now takes a backend_config parameter; when set, it allocates the external workspace via ICON_WORKSPACE_ALLOCATOR.allocate(...) and sets transient_memory_mode = EXTERNAL.
• get_options passes backend_config through as a keyword-only argument.
• customize_backend accepts backend_config and falls back to backend_config_from_env() when None.
• setup_program accepts a backend_config parameter.

model/driver/src/icon4py/model/driver/config.py

• Added backend_config field to DriverConfig (defaults to backend_config_from_env()).

model/driver/src/icon4py/model/driver/main.py

• Moved customize_backend / get_allocator to after config = config.with_overrides(...) so the driver config's backend_config is honored.

model/driver/tests/driver/unit_tests/data/test_config.yml

• Added backend_config: (empty/None) to the test config.

pyproject.toml

• Reordered commented-out uv source lines (no functional change).

Configuration

The external workspace is configured via an environment variable:

Variable	Description	Default
ICON4PY_BACKEND_WORKSPACE_SIZE	Workspace size in bytes, per device.	-

When ICON4PY_BACKEND_WORKSPACE_SIZE is not set, no external workspace is allocated.

Backward Compatibility

• The external_workspace parameter defaults to None in make_custom_dace_backend.
• The backend_config parameter defaults to None in customize_backend / setup_program.
• Existing callers that do not pass backend_config are unaffected; the default behavior (no external workspace) is preserved unless ICON4PY_BACKEND_WORKSPACE_SIZE is set in the environment.

Testing

Tests in model/common/tests/common/test_backend_configuration.py and model/common/tests/common/test_model_options.py.

[edopao]

cscs-ci run dace

[msimberg]

Some quick comments. Would appreciate @DropD's comments on config changes.

[msimberg]

Should a BackendConfig be something one can set in an experiment config? Or live purely as a parallel config world which is set differently? If they get merged, then this way of setting env vars to configs should probably be expressible with the config annotations, but I'm not sure that's necessary or a good idea. It would definitely be a new requirement for the rest of the config system.

[edopao]

My idea: the experiment config can be used in granule and driver tests and standalone runs, while the environment variables can be used in stencil test/benchmark.

[msimberg]

Partly answering my own question about whether backend config should be/is in the experiment config. It is added here, and if it should stay then it should probably be aligned a bit better with the other sub-config dataclasses. I'd appreciate @DropD's input on what might make sense here.

[edopao]

Awaiting input from @DropD

[msimberg]

Depending on how important it is to remember to set the BackendConfig, would it make sense to make this non-optional? Or alternatively make the default a BackendConfig() so you don't have to deal with the None case later?

[edopao]

For now, BackendConfig is only used to setup an external workspace, which is a very special case (a workaround for the AMD platform, indeed). I would propose to keep it optional for now. If we extend BackendConfig, it is a good idea to make it non-optional.

[Copilot]

Pull request overview

Enables reusable DaCe external workspaces to reduce ROCm transient-allocation overhead.

Changes:

• Adds validated workspace configuration and aligned per-device allocation.
• Integrates external workspaces into DaCe backend creation and driver configuration.
• Adds allocator/configuration tests and updates test data.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
pyproject.toml	Reorders commented source entries.
model/common/src/icon4py/model/common/backend_configuration.py	Adds workspace configuration and allocation.
model/common/src/icon4py/model/common/model_backends.py	Passes external workspaces to DaCe.
model/common/src/icon4py/model/common/model_options.py	Configures workspaces during backend setup.
model/common/tests/common/test_backend_configuration.py	Tests configuration and allocation.
model/driver/src/icon4py/model/driver/config.py	Adds driver workspace configuration.
model/driver/src/icon4py/model/driver/main.py	Applies configuration before backend creation.
model/driver/tests/driver/unit_tests/data/test_config.yml	Updates serialized test configuration.

Suppressed comments (1)

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

• This allocator is a generic aligned external-workspace allocator, so the Icon prefix does not describe the operation it performs. Per the shared-code generic naming rule, use an operation-based name such as ExternalWorkspaceAllocator; the current name makes reusable common functionality appear caller/project-specific.

class IconWorkspaceAllocator:
    """Singleton workspace allocator for the DaCe backend.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[iomaganaris]

Looks good on my side. Very welcome changes 🚀
Would be nice to have a rule of thumb regarding what is the required size to set for different grids but this can be defined empirically with time

[github-actions]

When developing, you can test your changes on CSCS CI before merge with the default pipeline: cscs-ci run default. This will run a default subset of tests.

You can pass options to override pipeline variables, for example:

• cscs-ci run default;BACKENDS=gtfn_cpu;LEVELS=unit
• cscs-ci run default;MODEL_SUBPACKAGES=common:driver;SESSIONS=model
  Avoid running the pipeline for all tests when you are developing.

Available options are:

• SESSIONS: model, model_mpi, or tools (correspond to nox sessions)
• MODEL_SUBSETS: datatest, basic, or stencils (correspond to nox session selections)
• MODEL_SUBPACKAGES: subpackages for non-MPI tests (last component, e.g. diffusion, driver)
• MODEL_MPI_SUBPACKAGES: subpackages for MPI tests (as above)
• BACKENDS: backends
• GRIDS: grids for stencil tests (simple, icon_regional, or icon_global)
• LEVELS: testing level for non-stencil tests (unit or integration)

For each option, all can be used as a shorthand for all possible values of that variable, e.g. LEVELS=all.

See scripts/python/generate_ci_pipeline.py and noxfile.py for available values for each option.

The all pipeline can be run with cscs-ci run all. This will run all icon4py tests in CSCS CI which can be expensive. This pipeline runs on a schedule on main, and can be run when extensive validation is needed (e.g. before releases).

Merging

Once your PR is approved and ready for merging, add it to the merge queue. The merge CSCS CI pipeline will run automatically on the merge-queue branch and must pass before the PR is merged. A dummy merge check will be triggered on the PR itself since it's required to add a PR to the merge queue.

Optional Tests

To run benchmarks you can use:

• cscs-ci run benchmark-bencher

For more detailed information please look at CI in the EXCLAIM universe.