feat[next-dace]: Support external memory for transient arrays

[edopao]

This pull request introduces a new "external workspace" memory mode for DaCe transients in the gt4py.next DaCe backend, allowing users to provide their own workspace buffers and thereby avoid per-call GPU memory allocation overhead. The change is documented in a new ADR and is implemented by extending the transient memory mode enum, updating backend and compilation workflows, and making the memory management API more flexible and explicit. The pull request also cleans up and unifies transient memory handling throughout the codebase.

The most important changes are:

External Workspace Memory Mode (EXTERNAL) Implementation:

• Introduced a new TransientMemoryMode enum with an EXTERNAL option, allowing users to supply their own workspace buffers for DaCe transients. This is documented in a detailed ADR (docs/development/ADRs/next/0027-External_Workspace_Memory.md) and referenced in the ADR index. [1] [2] [3]
• Updated the auto-optimization and post-processing pipeline to support the new transient memory mode, including logic for configuring transient lifetimes and workspace installation. [1] [2] [3] [4] [5] [6] [7]

Backend and Compilation Workflow Updates:

• Added a new load_artifact method to the Backend class, allowing backends to inject runtime data (such as external workspaces) into loaded programs. The DaCe backend overrides this to install the user-supplied workspace.
• Modified CompiledProgramsPool to use the backend's load_artifact method instead of directly calling artifact.load(), ensuring that all loaded programs receive the correct workspace configuration. [1] [2]

API and Import Cleanup:

• Updated imports and __all__ lists to expose the new TransientMemoryMode and the refactored utility function for configuring transient lifetimes, replacing the previous persistent-only approach. [1] [2] [3]

These changes collectively provide a more flexible, efficient, and explicit mechanism for managing workspace memory in DaCe-based GT4Py workflows, particularly for repeated or multi-program GPU workloads.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[egparedes]

It looks good. I just some comments and questions, I guess mostly for me to understand some decisions

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated no new comments.

Suppressed comments (10)

src/gt4py/next/program_processors/runners/dace/workflow/backend.py:147

• transient_memory_mode is checked with is not here; use value equality so callers can pass a string value (the enum is str, Enum) without silently bypassing the intended warning path.

    elif transient_memory_mode := optimization_args.get("transient_memory_mode"):
        if transient_memory_mode is not gtx_transformations.TransientMemoryMode.EXTERNAL:

src/gt4py/next/program_processors/runners/dace/workflow/compilation.py:178

• The RuntimeError suggests calling set_external_workspace(), but that method is on DaCeDecoratedProgram / injected via the backend, not on CompiledDaceProgram itself. The message should point callers to the correct API entry point.

            if self.external_workspace is None:
                raise RuntimeError(
                    "External workspace is not set. Please call `set_external_workspace()`"
                    " before the first call to the program."
                )

src/gt4py/next/program_processors/runners/dace/workflow/backend.py:142

• make_dace_backend compares transient_memory_mode using is, which will fail when the value comes in as an equivalent string (e.g. from config/CLI) rather than the exact Enum singleton, causing the EXTERNAL/workspace validation to be skipped.

This issue also appears on line 146 of the same file.

    if external_workspace is None:
        if (
            optimization_args.get("transient_memory_mode")
            is gtx_transformations.TransientMemoryMode.EXTERNAL
        ):

src/gt4py/next/program_processors/runners/dace/workflow/backend.py:40

• DaCeBackend.load_artifact always injects an empty dict when external_workspace is None, which defeats the downstream "workspace is not set" guard (it becomes "device not found") and can make failures harder to understand.

    def load_artifact(self, artifact: stages.CompilationArtifact) -> stages.ExecutableProgram:
        program = super().load_artifact(artifact)
        assert isinstance(program, gtx_wfddecoration.DaCeDecoratedProgram)
        # Inject the backend-level workspace so it is used when arguments are constructed.
        program.set_external_workspace(self.external_workspace or {})

src/gt4py/next/program_processors/runners/dace/workflow/compilation.py:80

• Using assert here turns a missing GPU device type into an AssertionError (and can be stripped with -O). This should be a proper runtime error so external workspace setup fails predictably.

This issue also appears on line 174 of the same file.

    elif storage == dace.StorageType.GPU_Global:
        assert core_defs.CUPY_DEVICE_TYPE is not None
        device = core_defs.CUPY_DEVICE_TYPE
    else:

src/gt4py/next/program_processors/runners/dace/workflow/compilation.py:100

• The _validate_external_workspace docstring claims an alignment check and references request.nbytes/request.alignment, but the function currently only validates array-interface presence and (optionally) nbytes. This documentation should match the actual checks.

        ValueError: If ``wsp`` is smaller than ``request.nbytes`` or its
            base pointer is not aligned to ``request.alignment`` bytes.

src/gt4py/next/program_processors/runners/dace/workflow/compilation.py:107

• The TypeError message still refers to an "External memory allocator" returning the buffer, but the new API passes a workspace buffer directly. This wording is misleading for users debugging invalid workspace objects.

        raise TypeError(
            f"External memory allocator returned {type(wsp).__name__!r} for storage "
            f"{storage!r}, which does not expose `__array_interface__` or "
            f"`__cuda_array_interface__`."
        )

src/gt4py/next/program_processors/runners/dace/workflow/decoration.py:36

• The docstring says teardown is forwarded so the pool can release external workspaces, but external workspace lifetime is explicitly caller-owned (and ExecutableProgram is just a Callable). This sentence is misleading.

    program. Teardown is forwarded to the underlying ``CompiledDaceProgram``
    so that the generic otf pool -- which only sees this callable -- can
    release external-memory workspaces when the pool is finalized.

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py:375

• In the EXTERNAL/CPU branch the comment says external-memory API calls should replace host malloc/free, but the assertions currently require new/malloc and delete/free to appear. This makes the test validate the opposite of what the comment (and EXTERNAL mode intent) describes.

                # CPU external mode should route transient workspace setup via
                # external-memory API calls rather than host malloc/free calls.
                assert any(marker in generated_code for marker in ("new ", "malloc"))
                assert any(marker in generated_code for marker in ("delete ", "free"))

src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py:121

• TransientMemoryMode docs state SCOPED is the default strategy, but gt_auto_optimize currently defaults to TransientMemoryMode.POOL. Either the default or the documentation should be aligned.

    - `SCOPED`: Transients are allocated and deallocated in the scope of the SDFG
        where they are defined, being it the top-level SDFG or a nested one.
        This is the default strategy.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 1 comment.

Suppressed comments (4)

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py:437

• There is no unit test that a host-only workspace (only __array_interface__) is rejected when the SDFG requests GPU_Global workspace. Without this, a regression could silently accept an invalid buffer for GPU storage.

def test_construct_arguments_accepts_gpu_workspace():
    """A device workspace with matching size is accepted and installed."""
    workspace = _make_array_buffer(nbytes=128, cuda=True)
    program = _make_compiled_program(
        external_workspace={core_defs.DeviceType.CUDA: workspace},
        workspace_sizes={dace.StorageType.GPU_Global: 128},
    )

src/gt4py/next/program_processors/runners/dace/workflow/decoration.py:72

• The comment refers to fun.csdfg_args, but the attribute is csdfg_argv. This makes the first-call logic harder to follow when debugging.

            # NOTE: If this is the first time then we will generate an exception because
            #   `fun.csdfg_args` is `None`
            # TODO(phimuell, edopao): Think about refactor the code such that the update
            #   of the argument vector is a Method of the `CompiledDaceProgram`.
            self._update_sdfg_call_args(args, self._fun.csdfg_argv, offset_provider)  # type: ignore[arg-type]  # Will error out in first call.

src/gt4py/next/program_processors/runners/dace/transformations/utils.py:69

• The TODO says “Find out why scalars are processed”, but this branch is taken for data nodes that are not arrays or scalars. This is misleading when reading the lifetime-selection logic.

                desc = dnode.desc(nsdfg)
                if not desc.transient or type(desc) not in {dace.data.Array, dace.data.Scalar}:
                    # TODO(phimuell): Find out why scalars are processed.
                    not_modify_lifetime.add(dnode.data)
                    continue

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py:395

• The comment says CPU external mode should use external-memory APIs “rather than host malloc/free calls”, but the assertions below currently require new/malloc and delete/free markers to be present. Either the comment or the assertions should be adjusted so they describe the same expected behavior.

                # CPU external mode should route transient workspace setup via
                # external-memory API calls rather than host malloc/free calls.
                assert any(marker in generated_code for marker in ("new ", "malloc"))
                assert any(marker in generated_code for marker in ("delete ", "free"))