diff --git a/docs/development/ADRs/next/0027-External_Workspace_Memory.md b/docs/development/ADRs/next/0027-External_Workspace_Memory.md new file mode 100644 index 0000000000..c2d90458ed --- /dev/null +++ b/docs/development/ADRs/next/0027-External_Workspace_Memory.md @@ -0,0 +1,172 @@ +--- +tags: [] +--- + +# External Workspace Memory for DaCe Transients + +- **Status**: valid +- **Authors**: Edoardo Paone (@edopao) +- **Created**: 2026-07-27 +- **Updated**: 2026-08-03 + +In the context of `gt4py.next` DaCe backends that run the same compiled SDFG +many times (e.g. inside a time loop), facing per-call GPU transient allocation +overhead and the desire to bound workspace memory to one SDFG's transient +footprint, we decided to add an explicit `EXTERNAL` transient-memory mode +backed by caller-supplied workspace buffers (a plain `dict` of array-like +objects, one per device) to achieve bounded memory and zero per-call +allocation. + +## Context + +DaCe transients default to `AllocationLifetime.SDFG` (scoped): the generated +code allocates and frees them around every SDFG call. For workloads that run +the same compiled SDFG many times, this per-call allocation/free is overhead, +and for workloads that run many distinct SDFGs in a loop, the natural goal is a +single reusable workspace buffer per storage type, bounded to one SDFG's +transient footprint. + +Two properties are required of such a mode: + +- **Allocate once, reuse across calls.** The workspace is sized once (from the + SDFG's own `get_workspace_sizes()`) and installed via `set_workspace(...)`, + so no per-call allocation appears in the generated runtime path. +- **Explicit lifetime.** The workspace is released when the compiled program + is done — not at GC time, and not per call. + +The existing `PERSISTENT` mode ties one workspace to exactly one SDFG: +transients are allocated once and retained for that SDFG's lifetime, so the +workspace grows without bound across distinct SDFGs and is owned by the +compiled program, not the caller. `EXTERNAL` can reproduce that behaviour — +allocate once, retain — but is more general: it hands ownership of the +workspace memory to the caller. With that ownership the caller can reuse a +single workspace buffer across multiple compiled programs (sized to the +largest, or per storage type), which `PERSISTENT` cannot. Reuse is the +caller's responsibility: `EXTERNAL` issues one workspace per storage type and +the generated runtime path assumes the SDFG runs sequentially on the default +stream, so the caller must ensure no two programs using the same workspace run +concurrently. `POOL` is also not a fit for this reuse-across-programs goal: it +still allocates per SDFG, just amortized through the mempool. + +## Decision + +A single explicit `EXTERNAL` transient-memory mode, backed by **workspace +buffers supplied directly by the caller** (no allocator object). + +- `TransientMemoryMode` enum (`SCOPED`, `PERSISTENT`, `POOL`, `EXTERNAL`), + mutually exclusive, threaded through `gt_auto_optimize()` and the backend + factory. `_gt_auto_post_processing` dispatches per mode: `SCOPED` is a + no-op, `PERSISTENT` sets `AllocationLifetime.Persistent`, `POOL` applies the + GPU mempool pass, and `EXTERNAL` sets `AllocationLifetime.External` on + eligible transients via `gt_configure_transient_lifetime`. +- A public `ExternalWorkspace` `TypeAlias` — + `dict[core_defs.DeviceType, ArrayInterface | CUDAArrayInterface]` — defined in + `workflow/common.py`. The per-device array-like must be accepted by + `dace.dtypes.array_interface_ptr()` as a workspace: a host array exposing + `__array_interface__` or a device array exposing `__cuda_array_interface__`. + No new protocol is introduced. +- The workspace is carried on a dedicated `DaCeBackend` (a `backend.Backend` + subclass, frozen dataclass) through a new `external_workspace` attribute. + `backend.Backend.load_artifact()` is a new overridable hook (defaulting to + `artifact.load()`); `DaCeBackend.load_artifact()` calls the default and then + injects the backend-level workspace onto the wrapped program via + `DaCeDecoratedProgram.set_external_workspace(...)`. `CompiledProgramsPool` + now calls `self.backend.load_artifact(...)` (in both + `_finish_compilation_job` and `_compile_variant`) instead of + `artifact.load()`, so every program loaded through a backend that carries a + workspace gets it installed. +- At runtime, `CompiledDaceProgram.construct_arguments` calls + `sdfg_program.get_workspace_sizes()`, looks up the matching per-device buffer + in `external_workspace` (raising `RuntimeError` if the workspace is unset or + the device is missing), validates it (array interface present; size at least + the required bytes) and installs it via `sdfg_program.set_workspace(...)`. +- **Workspace lifetime is the caller's responsibility.** There is no + `finalize()`, no `__del__`, and no `weakref.finalize` registered on + `CompiledProgramsPool`. The caller owns the buffers and decides when they are + freed; the compiled program simply borrows them for as long as it lives. + +## Consequences + +- A single reusable workspace buffer per storage type can back many sequential + SDFG calls, removing per-call transient allocation/free in the generated + runtime path while keeping memory bounded to one SDFG's transient footprint. +- Workspace lifetime is entirely the caller's: GT4Py never allocates or frees + external workspace, so there is no pool-driven teardown and no resource leak + from GT4Py's side. The flip side is that the caller must keep the buffers + alive for the full lifetime of every compiled program that uses them and free + them when appropriate. +- The public API is a single mode enum plus a plain `dict` of array-like + buffers, threaded through the `DaCeBackend` instance. Incompatible + combinations — a workspace without `EXTERNAL` mode (warns), or `EXTERNAL` + mode without a workspace (raises `ValueError`) — are detected at backend + construction. +- Because the workspace is no longer part of the compilation artifact, there is + no picklability requirement on the buffers. Compilation can be offloaded to a + worker process regardless of the workspace shape; the workspace is injected at + load time, in the main process. +- Multi-stream safety is explicitly **out of scope** for this delivery: + `EXTERNAL` reuses one workspace across sequential SDFG calls on the default + stream. Concurrent use across streams is a follow-up. + +## Alternatives considered + +### An `ExternalMemoryAllocator` protocol (allocate/deallocate) + +- Good, because it gives GT4Py a symmetric allocate/free contract and lets + GT4Py drive teardown at pool finalization, so a forgotten buffer is reclaimed + automatically. +- Bad, because it forces an allocator *object* (a `Protocol` with + `allocate(request) -> wsp` / `deallocate(wsp)` methods, an + `AllocationRequest` dataclass, and a separate `TypeAlias` for the workspace) + into the compilation artifact, which is pickled when compilation is offloaded + to a worker process. That in turn requires a picklability probe + (`pickle.dumps` at `DaCeCompiler.__post_init__`) and a dedicated + `AllocatorNotPicklableError` to avoid silently degrading to in-process + compilation — machinery that exists only to shuttle a buffer from the backend + to the runtime. The actual need is simpler: the caller already has the buffer + and only wants to hand it over. Replaced by passing the workspace buffers + directly, which removes the protocol, the request type, the picklability + probe, the error, and the pool finalizer in one step. The cost is that the + caller now owns the lifetime — acceptable because the caller is the party + that allocated the memory in the first place. + +### `__del__`-based teardown on `CompiledDaceProgram` + +- Good, because it needs no pool integration. +- Bad, because the codebase strongly prefers `weakref.finalize` over `__del__` + (hostile conditions at interpreter shutdown, partial initialization, and an + allocator/buffer that may already be gone). In the direct-buffer design + teardown is the caller's responsibility, so neither `__del__` nor a finalizer + is needed. Rejected. + +### A DLPack-consuming `set_workspace` + +- Good, because DLPack is the cross-framework standard for zero-copy. +- Bad, because `dace.dtypes.array_interface_ptr()` (what DaCe uses for + `set_workspace`) duck-types via `__array_interface__` / + `__cuda_array_interface__` and `hasattr('data_ptr')`, and does not consume + DLPack for this path. Introducing a DLPack requirement would narrow the set + of accepted buffers without buying anything on this code path. Rejected; + `ExternalWorkspace` is kept as a `TypeAlias` over the two array interfaces. + +## References + +- `src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py` + (`TransientMemoryMode`, `_gt_auto_post_processing`). +- `src/gt4py/next/program_processors/runners/dace/workflow/common.py` + (`ExternalWorkspace` `TypeAlias`). +- `src/gt4py/next/program_processors/runners/dace/workflow/backend.py` + (`DaCeBackend`, `make_dace_backend`, the `external_workspace` parameter and + the mode/workspace compatibility checks). +- `src/gt4py/next/program_processors/runners/dace/workflow/compilation.py` + (`CompiledDaceProgram.construct_arguments`/`_configure_external_workspace`, + `_validate_external_workspace`). +- `src/gt4py/next/program_processors/runners/dace/workflow/decoration.py` + (`DaCeDecoratedProgram.set_external_workspace`). +- `src/gt4py/next/backend.py` (`Backend.load_artifact`). +- `src/gt4py/next/otf/compiled_program.py` + (`CompiledProgramsPool._finish_compilation_job`/`_compile_variant` calling + `backend.load_artifact`). +- [ADR 0023](0023-Fingerprinting.md) and + [ADR 0025](0025-Crash_Consistent_Build_Caches.md) for the cache and + build-folder guarantees the external-memory path must not regress. diff --git a/docs/development/ADRs/next/README.md b/docs/development/ADRs/next/README.md index a6d06da00e..24e42da696 100644 --- a/docs/development/ADRs/next/README.md +++ b/docs/development/ADRs/next/README.md @@ -48,6 +48,7 @@ Writing a new ADR is simple: - [0016 - Multiple Backends and Build Systems](0016-Multiple-Backends-and-Build-Systems.md) - [0017 - Toolchain Configuration](0017-Toolchain-Configuration.md) - [0018 - Canonical Form of an SDFG in GT4Py (Especially for Optimizations)](0018-Canonical_SDFG_in_GT4Py_Transformations.md) +- [0027 - External Workspace Memory for DaCe Transients](0027-External_Workspace_Memory.md) ### Python Integration diff --git a/src/gt4py/next/backend.py b/src/gt4py/next/backend.py index ae599ece6d..9063fe09cb 100644 --- a/src/gt4py/next/backend.py +++ b/src/gt4py/next/backend.py @@ -157,6 +157,14 @@ def compile( artifact = self.executor( self.transforms(definitions.ConcreteProgramDef(data=program, args=compile_time_args)) ) + return self.load_artifact(artifact) + + def load_artifact(self, artifact: stages.CompilationArtifact) -> stages.ExecutableProgram: + """Load an artifact into an executable program. + + Backends may override this method to inject backend-specific runtime data + into the loaded program. + """ return artifact.load() @property diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 0784ca0f73..fa89e98dd0 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -583,7 +583,7 @@ def _finish_compilation_job(self, key: CompiledProgramsKey) -> bool: artifact_future = self._compilation_jobs.pop(key) assert isinstance(artifact_future, concurrent.futures.Future) assert key not in self.compiled_programs - self.compiled_programs[key] = artifact_future.result().load() + self.compiled_programs[key] = self.backend.load_artifact(artifact_future.result()) return True def _compile_variant( @@ -660,7 +660,7 @@ def _compile_variant( if future.done(): # Eager so compile() raises now; otherwise the error stays in the # already-resolved future until the next call touches this key. - self.compiled_programs[key] = future.result().load() + self.compiled_programs[key] = self.backend.load_artifact(future.result()) else: self._compilation_jobs[key] = future _ongoing_compilations[future] = ( diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/__init__.py b/src/gt4py/next/program_processors/runners/dace/transformations/__init__.py index 39fb62f1fe..f8cb38334f 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/__init__.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/__init__.py @@ -17,6 +17,7 @@ GT4PyAutoOptHook, GT4PyAutoOptHookFun, GT4PyAutoOptHookStage, + TransientMemoryMode, gt_auto_optimize, ) from .concat_where_mapper import ( @@ -84,7 +85,7 @@ gt_propagate_strides_from_access_node, gt_propagate_strides_of, ) -from .utils import gt_make_transients_persistent +from .utils import gt_configure_transient_lifetime __all__ = [ @@ -119,6 +120,7 @@ "SingleStateGlobalSelfCopyElimination", "SplitAccessNode", "SplitConsumerMemlet", + "TransientMemoryMode", "VerticalMapFusionCallback", "VerticalMapSplitCallback", "constants", @@ -126,13 +128,13 @@ "gt_auto_optimize", "gt_change_strides", "gt_check_if_concat_where_node_is_replaceable", + "gt_configure_transient_lifetime", "gt_create_local_double_buffering", "gt_eliminate_dead_dataflow", "gt_gpu_transform_non_standard_memlet", "gt_gpu_transformation", "gt_horizontal_map_split_fusion", "gt_inline_nested_sdfg", - "gt_make_transients_persistent", "gt_map_strides_to_dst_nested_sdfg", "gt_map_strides_to_src_nested_sdfg", "gt_multi_state_global_self_copy_elimination", diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py b/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py index a7b34f6bbf..d743d8e31e 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py @@ -111,11 +111,42 @@ class GT4PyAutoOptHook(enum.Enum): ] +class TransientMemoryMode(str, enum.Enum): + """ + Policy selecting the lifetime/allocation strategy of transient arrays. + + Supported strategies are: + - `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. + - `PERSISTENT`: Transients are allocated the first time the SDFG is called and + retained through the entire life of the compiled SDFG, and freed only once + it goes out of scope. + - `POOL`: Transients are allocated in a memory pool, associated to the GPU + default stream. These allocations are managed by an asynchronous allocator, + since all memory is allocated and freed in stream order. + - `EXTERNAL`: Transients are backed by workspace memory supplied directly + by the caller (one buffer per storage type, threaded through the DaCe + backend). This strategy allows to reuse a workspace memory across + multiple SDFGs, relying on sequential execution of the programs on the + default stream; the caller owns the workspace lifetime. + Note: + The `EXTERNAL` strategy requires that the `external_workspace` attribute + of the dace backend is set, because it is needed at runtime to install + the memory pointers for transient arrays. + """ + + SCOPED = "SCOPED" + PERSISTENT = "PERSISTENT" + POOL = "POOL" + EXTERNAL = "EXTERNAL" + + def gt_auto_optimize( sdfg: dace.SDFG, gpu: bool, unit_strides_kind: Optional[gtx_common.DimensionKind] = None, - make_persistent: bool = False, + transient_memory_mode: TransientMemoryMode = TransientMemoryMode.POOL, gpu_block_size: Optional[Sequence[int | str] | str] = (32, 8, 1), gpu_block_size_1d: Optional[Sequence[int | str] | str] = (64, 1, 1), gpu_block_size_2d: Optional[Sequence[int | str] | str] = None, @@ -130,7 +161,6 @@ def gt_auto_optimize( reuse_transients: bool = False, gpu_launch_bounds: Optional[int | str] = None, gpu_launch_factor: Optional[int] = None, - gpu_memory_pool: bool = True, constant_symbols: Optional[dict[str, Any]] = None, assume_pointwise: bool = True, optimization_hooks: Optional[dict[GT4PyAutoOptHook, GT4PyAutoOptHookFun]] = None, @@ -174,9 +204,7 @@ def gt_auto_optimize( gpu: Optimize for GPU or CPU. unit_strides_kind: All dimensions of this kind are considered to have unit strides, see `gt_set_iteration_order()` for more. - make_persistent: Turn all transients to persistent lifetime, thus they are - allocated over the whole lifetime of the program, even if the kernel exits. - Thus the SDFG can not be called by different threads. + transient_memory_mode: Lifetime for transient arrays. gpu_block_size: This is used as default thread block size for GPU Maps. See also the `gpu_block_size_*d` arguments gpu_block_size_{1, 2, 3}d: Allows to specify the GPU thread block size for @@ -194,7 +222,6 @@ def gt_auto_optimize( gpu_launch_bounds: Use this value as `__launch_bounds__` for _all_ GPU Maps. gpu_launch_factor: Use the number of threads times this value as `__launch_bounds__` for _all_ GPU Maps. - gpu_memory_pool: Enable CUDA memory pool in gpu codegen. constant_symbols: Symbols listed in this `dict` will be replaced by the respective value inside the SDFG. This might increase performance. assume_pointwise: Assume that the SDFG has no risk for race condition in @@ -400,13 +427,12 @@ def gt_auto_optimize( sdfg = _gt_auto_post_processing( sdfg=sdfg, gpu=gpu, - make_persistent=make_persistent, + transient_memory_mode=transient_memory_mode, # TODO(phimuell): In general `TransientReuse` is a good idea, but the # current implementation also reuses transients scalars inside Map # scopes, which I do not like. Thus we should fix the transformation # to avoid that. reuse_transients=reuse_transients, - gpu_memory_pool=gpu_memory_pool, validate_all=validate_all, ) @@ -918,9 +944,8 @@ def _gt_auto_configure_maps_and_strides( def _gt_auto_post_processing( sdfg: dace.SDFG, gpu: bool, - make_persistent: bool, + transient_memory_mode: TransientMemoryMode, reuse_transients: bool, - gpu_memory_pool: bool, validate_all: bool, ) -> dace.SDFG: """Perform post processing on the SDFG. @@ -939,27 +964,34 @@ def _gt_auto_post_processing( # TODO(phimuell): Fix the bug, it uses the tile value and not the stack array value. dace_aoptimize.move_small_arrays_to_stack(sdfg) - if make_persistent and gpu_memory_pool: - raise ValueError("Cannot set both 'make_persistent' and 'gpu_memory_pool'.") - - if make_persistent: - device = dace.DeviceType.GPU if gpu else dace.DeviceType.CPU - gtx_transformations.gt_make_transients_persistent(sdfg=sdfg, device=device) - - if device == dace.DeviceType.GPU: - # NOTE: For unknown reasons the counterpart of the - # `gt_make_transients_persistent()` function in DaCe, resets the - # `wcr_nonatomic` property of every memlet, i.e. makes it atomic. - # However, it does this only for edges on the top level and on GPU. - # For compatibility with DaCe (and until we found out why) the GT4Py - # auto optimizer will emulate this behaviour. - for state in sdfg.states(): - assert isinstance(state, dace.SDFGState) - for edge in state.edges(): - edge.data.wcr_nonatomic = False - - if gpu and gpu_memory_pool: - gtx_transformations.gpu_utils.gt_gpu_apply_mempool(sdfg) + match transient_memory_mode: + case TransientMemoryMode.PERSISTENT: + gtx_transformations.gt_configure_transient_lifetime( + sdfg=sdfg, lifetime=dace.AllocationLifetime.Persistent + ) + if gpu: + # NOTE: For unknown reasons the counterpart of the + # `gt_make_transients_persistent()` function in DaCe, resets the + # `wcr_nonatomic` property of every memlet, i.e. makes it atomic. + # However, it does this only for edges on the top level and on GPU. + # For compatibility with DaCe (and until we found out why) the GT4Py + # auto optimizer will emulate this behaviour. + for state in sdfg.states(): + assert isinstance(state, dace.SDFGState) + for edge in state.edges(): + edge.data.wcr_nonatomic = False + + case TransientMemoryMode.EXTERNAL: + gtx_transformations.gt_configure_transient_lifetime( + sdfg=sdfg, lifetime=dace.AllocationLifetime.External + ) + + case TransientMemoryMode.POOL: + if gpu: + gtx_transformations.gpu_utils.gt_gpu_apply_mempool(sdfg) + + case TransientMemoryMode.SCOPED: + pass if validate_all: sdfg.validate() diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/utils.py b/src/gt4py/next/program_processors/runners/dace/transformations/utils.py index 80cdda5d06..f80809bec5 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/utils.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/utils.py @@ -10,48 +10,42 @@ from __future__ import annotations -from typing import Optional, Sequence, TypeVar, Union +from typing import Optional, Sequence, Union import dace from dace import data as dace_data, subsets as dace_sbs, symbolic as dace_sym from dace.libraries import standard as dace_stdlib from dace.sdfg import graph as dace_graph, nodes as dace_nodes -from dace.transformation import pass_pipeline as dace_ppl from dace.transformation.passes import analysis as dace_analysis from ordered_set import OrderedSet from gt4py.next.program_processors.runners.dace import library_nodes as gtx_lib -_PassT = TypeVar("_PassT", bound=dace_ppl.Pass) - - -def gt_make_transients_persistent( +def gt_configure_transient_lifetime( sdfg: dace.SDFG, - device: dace.DeviceType, + lifetime: dace.AllocationLifetime, ) -> dict[int, set[str]]: """ - Changes the lifetime of certain transients to `Persistent`. + Configure transient lifetime for eligible data nodes in the given SDFG and all nested SDFGs. - A persistent lifetime means that the transient is allocated only the very first - time the SDFG is executed and only deallocated if the underlying `CompiledSDFG` - object goes out of scope. The main advantage is, that memory must not be - allocated every time the SDFG is run. The downside is that the SDFG can not be - called by different threads. + Eligible data nodes are transient arrays or scalars excluding: + - data nodes with storage type `Register` (relevant for scalars) + - data nodes with lifetime `External` (already externally managed) + - data nodes used inside a scope (e.g., maps) + - data nodes whose size is not fully determined by the SDFG's free symbols (dynamic allocations) Args: sdfg: The SDFG to process. - device: The device type. + lifetime: The desired lifetime to set for eligible transient data nodes. Returns: - A `dict` mapping SDFG IDs to a set of transient arrays that - were made persistent. - - Note: - This function is based on a similar function in DaCe. However, the DaCe - function does, for unknown reasons, also reset the `wcr_nonatomic` property, - but only for GPU. + A dictionary mapping SDFG configuration IDs to the data node names whose + lifetimes were modified. """ + if lifetime not in {dace.AllocationLifetime.Persistent, dace.AllocationLifetime.External}: + raise ValueError(f"Unsupported transient lifetime '{lifetime}'.") + result: dict[int, set[str]] = {} for nsdfg in sdfg.all_sdfgs_recursive(): fsyms: set[str] = nsdfg.free_symbols @@ -70,6 +64,7 @@ def gt_make_transients_persistent( 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 if desc.storage == dace.StorageType.Register: @@ -81,10 +76,8 @@ def gt_make_transients_persistent( continue # If the data is referenced inside a scope, such as a map, it might be possible - # that it is only used inside that scope. If we would make it persistent, then - # it would essentially be allocated outside and be shared among the different - # map iterations. So we can not make it persistent. - # The downside is, that we might have to perform dynamic allocation. + # that it is only used inside that scope. If we would make it global-lifetime, + # it would effectively be shared among map iterations, which is unsafe. if scope_dict[dnode] is not None: not_modify_lifetime.add(dnode.data) continue @@ -100,13 +93,24 @@ def gt_make_transients_persistent( except AttributeError: # total_size is an integer / has no free symbols pass - # Make it persistent. modify_lifetime.add(dnode.data) - # Now setting the lifetime. result[nsdfg.cfg_id] = modify_lifetime - not_modify_lifetime for aname in result[nsdfg.cfg_id]: - nsdfg.arrays[aname].lifetime = dace.AllocationLifetime.Persistent + adesc = nsdfg.arrays[aname] + adesc.lifetime = lifetime + if adesc.storage == dace.StorageType.Default and isinstance(adesc, dace.data.Array): + # GPU transformation have already changed the storage for GPU arrays. + # NOTE: If we do not change the storage during lowering / transformations, + # it will be done in code generation when calling `sdfg.compile()`. + # This side effect is a potential issue, because we do not store + # the program handle, see `sdfg.compile(return_program_handle=False)` + # in `compilation.py`; therefore, we do not have the modified SDFG. + # The original SDFG is deserialized, at call time, and the storage + # type is reset to `Default`. This is a problem for arrays with + # external storage, because `CompiledSDFG.set_workspace()` will + # try to load a symbol from the library for the wrong storage type. + adesc.storage = dace.StorageType.CPU_Heap return result diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py index d86148fb23..34ac8e6f4e 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -8,6 +8,7 @@ from __future__ import annotations +import dataclasses import warnings from typing import Any, Final @@ -16,7 +17,27 @@ import gt4py.next.custom_layout_allocators as next_allocators from gt4py._core import definitions as core_defs from gt4py.next import backend, common, config -from gt4py.next.program_processors.runners.dace.workflow.factory import DaCeWorkflowFactory +from gt4py.next.otf import stages +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, + decoration as gtx_wfddecoration, + factory as gtx_wfdfactory, +) + + +@dataclasses.dataclass(frozen=True) +class DaCeBackend(backend.Backend[Any]): + """DaCe backend with support for injecting an external workspace at load time.""" + + external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None + + 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 {}) + return program class DaCeBackendFactory(factory.Factory): @@ -30,7 +51,7 @@ class DaCeBackendFactory(factory.Factory): """ class Meta: - model = backend.Backend + model = DaCeBackend class Params: name_device = "cpu" @@ -42,7 +63,7 @@ class Params: ) device_type = core_defs.DeviceType.CPU otf_workflow = factory.SubFactory( - DaCeWorkflowFactory, + gtx_wfdfactory.DaCeWorkflowFactory, cached_translation=True, device_type=factory.SelfAttribute("..device_type"), auto_optimize=factory.SelfAttribute("..auto_optimize"), @@ -53,6 +74,7 @@ class Params: executor = factory.LazyAttribute(lambda o: o.otf_workflow) allocator = next_allocators.StandardCPUFieldBufferAllocator() transforms = backend.DEFAULT_TRANSFORMS + external_workspace = None def make_dace_backend( @@ -60,6 +82,7 @@ def make_dace_backend( auto_optimize: bool = True, async_sdfg_call: bool = True, optimization_args: dict[str, Any] | None = None, + external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None, unstructured_horizontal_has_unit_stride: bool = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE, use_metrics: bool = True, use_zero_origin: bool = False, @@ -74,6 +97,8 @@ def make_dace_backend( of GPU kernel execution with the Python driver code. optimization_args: A `dict` containing configuration parameters for the SDFG auto-optimize pipeline, see `gt_auto_optimize()`. + external_workspace: Workspace memory externally allocated, which is used + for SDFG's transient arrays when `transient_memory_mode` is `EXTERNAL`. unstructured_horizontal_has_unit_stride: When the memory layout has unit stride in the horizontal dimension, replace the field stride symbol with '1'. use_metrics: Add SDFG instrumentation to collect the metric for stencil @@ -110,9 +135,29 @@ def make_dace_backend( else None } + if external_workspace is None: + if ( + optimization_args.get("transient_memory_mode") + is gtx_transformations.TransientMemoryMode.EXTERNAL + ): + raise ValueError( + "External memory workspace must be provided when 'transient_memory_mode' is 'EXTERNAL'." + ) + elif transient_memory_mode := optimization_args.get("transient_memory_mode"): + if transient_memory_mode is not gtx_transformations.TransientMemoryMode.EXTERNAL: + warnings.warn( + f"External memory workspace provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{gtx_transformations.TransientMemoryMode.EXTERNAL}'.", + stacklevel=2, + ) + else: + optimization_args["transient_memory_mode"] = ( + gtx_transformations.TransientMemoryMode.EXTERNAL + ) + return DaCeBackendFactory( # type: ignore[return-value] # factory-boy typing not precise enough gpu=gpu, auto_optimize=auto_optimize, + external_workspace=external_workspace, otf_workflow__bare_translation__async_sdfg_call=(async_sdfg_call if gpu else False), otf_workflow__bare_translation__auto_optimize_args=optimization_args, otf_workflow__bare_translation__unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride, diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/common.py b/src/gt4py/next/program_processors/runners/dace/workflow/common.py index 07f053bbc9..883489e90e 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/common.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/common.py @@ -8,11 +8,12 @@ import contextlib import os -from typing import Any, Final, Generator, Optional +from typing import Any, Final, Generator, Optional, TypeAlias import dace from gt4py._core import definitions as core_defs +from gt4py.eve import extended_typing as xtyping from gt4py.next import config as gtx_config from gt4py.next.otf.compilation import common as gtx_compilation_common @@ -33,6 +34,17 @@ """DaCe datatype of `SDFG_ARG_METRIC_COMPUTE_TIME` argument.""" +ExternalWorkspace: TypeAlias = dict[ + core_defs.DeviceType, xtyping.ArrayInterface | xtyping.CUDAArrayInterface +] +""" Mapping from device types to array-like objects. + + The array-like objects must be accepted by `dace.dtypes.array_interface_ptr()` + as a workspace: a host array exposing `gt4py.eve.extended_typing.ArrayInterface` + or a device array exposing `gt4py.eve.extended_typing.CUDAArrayInterface`. +""" + + def set_dace_config( device_type: core_defs.DeviceType, cmake_build_type: Optional[gtx_config.CMakeBuildType] = None, diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py index 7b78942090..87abce5cf6 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -21,6 +21,7 @@ import factory from gt4py._core import definitions as core_defs, locking +from gt4py.eve import extended_typing as xtyping from gt4py.next import common, config, fingerprinting from gt4py.next.otf import code_specs, definitions, stages, workflow from gt4py.next.otf.compilation import cache as gtx_cache @@ -70,6 +71,58 @@ def _add_tx_markers(program_source: SDFGExtensionSource) -> tuple[SDFGExtensionS return new_program_source, sdfg +def _map_storage_to_device(storage: dace.StorageType) -> core_defs.DeviceType: + if storage == dace.StorageType.CPU_Heap: + device = core_defs.DeviceType.CPU + elif storage == dace.StorageType.GPU_Global: + if core_defs.CUPY_DEVICE_TYPE is None: + raise ValueError( + f"Can not map storage type '{storage}' to a device: no GPU device" + " type is configured ('core_defs.CUPY_DEVICE_TYPE' is None)." + ) + device = core_defs.CUPY_DEVICE_TYPE + else: + raise ValueError(f"Unsupported storage type '{storage}' for external workspace allocation.") + + return device + + +def _validate_external_workspace( + wsp: xtyping.ArrayInterface | xtyping.CUDAArrayInterface, storage: dace.StorageType, nbytes: int +) -> None: + """Validate that the provided ``wsp`` workspace satisfies the requirements. + + Args: + wsp: The external workspace to check. + storage: SDFG storage type the workspace buffer is being installed for. + nbytes: Size in bytes required. + + Raises: + TypeError: If ``wsp`` exposes neither ``__array_interface__`` nor + ``__cuda_array_interface__``. + ValueError: If ``wsp`` exposes ``nbytes`` and it is smaller than the + required ``nbytes``. Buffers that do not expose ``nbytes`` are + accepted on a trust basis (their size can not be checked here). + """ + if storage == dace.StorageType.GPU_Global: + if not xtyping.supports_cuda_array_interface(wsp): + raise TypeError( + f"External workspace for storage {storage!r} must expose `__cuda_array_interface__` (got {type(wsp).__name__!r})." + ) + elif storage == dace.StorageType.CPU_Heap: + if not xtyping.supports_array_interface(wsp): + raise TypeError( + f"External workspace for storage {storage!r} must expose `__array_interface__` (got {type(wsp).__name__!r})." + ) + else: + raise ValueError(f"Unsupported storage type {storage!r} for external workspace allocation.") + + if (wsp_nbytes := getattr(wsp, "nbytes", None)) is not None and wsp_nbytes < nbytes: + raise ValueError( + f"External workspace buffer is {wsp_nbytes} bytes for storage {storage!r}, but at least {nbytes} bytes were required." + ) + + class CompiledDaceProgram: sdfg_program: dace.CompiledSDFG @@ -96,6 +149,9 @@ class CompiledDaceProgram: # never updated. csdfg_argv: MutableSequence[Any] | None csdfg_init_argv: Sequence[Any] | None + external_workspace: gtx_wfdcommon.ExternalWorkspace | None = ( + None # This attribute is set at runtime, before the first call. + ) def __init__( self, @@ -122,6 +178,22 @@ def __init__( self.csdfg_argv = None self.csdfg_init_argv = None + def _configure_external_workspace(self, **kwargs: Any) -> None: + self.sdfg_program.initialize(**kwargs) + if workspace_sizes := self.sdfg_program.get_workspace_sizes(): + 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." + ) + for storage, required_nbytes in workspace_sizes.items(): + device = _map_storage_to_device(storage) + workspace = self.external_workspace.get(device) + if workspace is None: + raise RuntimeError(f"External workspace for device {device} not found.") + _validate_external_workspace(workspace, storage, required_nbytes) + self.sdfg_program.set_workspace(storage, workspace) + def construct_arguments(self, **kwargs: Any) -> None: """ This function will process the arguments and store the processed argument @@ -129,6 +201,7 @@ def construct_arguments(self, **kwargs: Any) -> None: """ with dace.config.set_temporary("compiler", "allow_view_arguments", value=True): csdfg_argv, csdfg_init_argv = self.sdfg_program.construct_arguments(**kwargs) + self._configure_external_workspace(**kwargs) # Note we only care about `csdfg_argv` (normal call), since we have to update it, # we ensure that it is a `list`. self.csdfg_argv = [*csdfg_argv] @@ -191,7 +264,7 @@ def load(self) -> stages.ExecutableProgram: sdfg = dace.SDFG.from_json(json.loads(self.sdfg_json)) sdfg_program = dace_compiler.get_program_handle(self.library_path, sdfg) program = CompiledDaceProgram(sdfg_program, self.bind_func_name, self.binding_source_code) - return gtx_wfddecoration.convert_args(program, device=self.device_type) + return gtx_wfddecoration.DaCeDecoratedProgram(program, device_type=self.device_type) @dataclasses.dataclass(frozen=True) diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py index 8c559a6d02..57f45964ce 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py @@ -16,7 +16,6 @@ from gt4py._core import definitions as core_defs from gt4py.next import common as gtx_common, utils as gtx_utils from gt4py.next.instrumentation import metrics -from gt4py.next.otf import stages from gt4py.next.program_processors.runners.dace import sdfg_callable from gt4py.next.program_processors.runners.dace.workflow import common as gtx_wfdcommon @@ -26,21 +25,36 @@ from gt4py.next.program_processors.runners.dace.workflow.compilation import CompiledDaceProgram -def convert_args( - fun: CompiledDaceProgram, - device: core_defs.DeviceType = core_defs.DeviceType.CPU, -) -> stages.ExecutableProgram: - # Retieve metrics level from GT4Py environment variable. - collect_time = metrics.is_level_enabled(metrics.PERFORMANCE) - collect_time_arg = np.array( - [1], dtype=gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME_DTYPE.as_numpy_dtype() - ) - # We use the callback function provided by the compiled program to update the SDFG arglist. - update_sdfg_call_args = functools.partial( - fun.update_sdfg_ctype_arglist, device, fun.sdfg_argtypes - ) - - def decorated_program( +class DaCeDecoratedProgram: + """A compiled DaCe program wrapped as a GT4Py-callable ``ExecutableProgram``. + + On the first call the full SDFG argument vector is constructed via + ``CompiledDaceProgram.construct_arguments``; subsequent calls only update + the argument vector in place through the binding function generated for the + program. External workspace memory (when the SDFG uses + ``TransientMemoryMode.EXTERNAL``) is installed onto the underlying + ``CompiledDaceProgram`` before the first call via `set_external_workspace`; + its lifetime is owned by the caller, not by this wrapper. + """ + + def __init__( + self, + fun: CompiledDaceProgram, + device_type: core_defs.DeviceType = core_defs.DeviceType.CPU, + ) -> None: + self._fun = fun + # Retrieve metrics level from GT4Py environment variable. + self._collect_time = metrics.is_level_enabled(metrics.PERFORMANCE) + self._collect_time_arg = np.array( + [1], dtype=gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME_DTYPE.as_numpy_dtype() + ) + # We use the callback function provided by the compiled program to update the SDFG arglist. + self._update_sdfg_call_args = functools.partial( + fun.update_sdfg_ctype_arglist, device_type, fun.sdfg_argtypes + ) + + def __call__( + self, *args: Any, offset_provider: gtx_common.OffsetProvider, out: Any = None, @@ -55,28 +69,35 @@ def decorated_program( # `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`. - update_sdfg_call_args(args, fun.csdfg_argv, offset_provider) # type: ignore[arg-type] # Will error out in first call. + self._update_sdfg_call_args(args, self._fun.csdfg_argv, offset_provider) # type: ignore[arg-type] # Will error out in first call. except TypeError: # First call. Construct the initial argument vector of the `CompiledDaceProgram`. - assert fun.csdfg_argv is None and fun.csdfg_init_argv is None + assert self._fun.csdfg_argv is None and self._fun.csdfg_init_argv is None flat_args: Sequence[Any] = gtx_utils.flatten_nested_tuple(args) this_call_args = sdfg_callable.get_sdfg_args( - fun.sdfg_program.sdfg, + self._fun.sdfg_program.sdfg, offset_provider, *flat_args, filter_args=False, ) this_call_args |= { gtx_wfdcommon.SDFG_ARG_METRIC_LEVEL: metrics.get_current_level(), - gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME: collect_time_arg, + gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME: self._collect_time_arg, } - fun.construct_arguments(**this_call_args) + self._fun.construct_arguments(**this_call_args) # Perform the call to the SDFG. - fun.fast_call() + self._fun.fast_call() + + if self._collect_time: + metrics.add_sample_to_current_source( + metrics.COMPUTE_METRIC, self._collect_time_arg[0].item() + ) - if collect_time: - metrics.add_sample_to_current_source(metrics.COMPUTE_METRIC, collect_time_arg[0].item()) + def set_external_workspace(self, external_workspace: gtx_wfdcommon.ExternalWorkspace) -> None: + """Set the external workspace for the underlying compiled program. - return decorated_program + This method should be called before the first call to the program. + """ + self._fun.external_workspace = external_workspace diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py index b653dbb40d..e5dfb1262f 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py @@ -8,20 +8,32 @@ """Test the bindings stage of the dace backend workflow.""" -import pytest +import dataclasses +import re import unittest.mock as mock +from typing import Any + +import numpy as np +import pytest + dace = pytest.importorskip("dace") from gt4py import next as gtx from gt4py._core import definitions as core_defs -from gt4py.next.otf import runners +from gt4py.next import config +from gt4py.next.otf import definitions, runners +from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations +from gt4py.next.program_processors.runners.dace.transformations import ( + auto_optimize as gtx_auto_optimize, +) from gt4py.next.program_processors.runners.dace.workflow import ( backend as dace_wf_backend, + common as dace_wf_common, + decoration as dace_wf_decoration, ) -from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations -from next_tests.integration_tests import cases +from next_tests.integration_tests import cases, cases_utils from next_tests.integration_tests.cases_utils import KDim @@ -32,11 +44,13 @@ ], ids=["CPU", "GPU"], ) -def device_type(request) -> str: +def device_type(request) -> gtx.DeviceType: + if request.param == core_defs.CUPY_DEVICE_TYPE: + pytest.importorskip("cupy") return request.param -@pytest.mark.parametrize("auto_optimize", [False, True]) +@pytest.mark.parametrize("auto_optimize", [False, True], ids=["NO_AUTO_OPT", "AUTO_OPT"]) def test_make_backend(auto_optimize, device_type, monkeypatch): on_gpu = device_type == core_defs.CUPY_DEVICE_TYPE @@ -55,20 +69,18 @@ def testee(a: cases.IField, b: cases.IField, out: cases.IField): optimization_args = {} elif on_gpu: optimization_args = { - "make_persistent": False, + "transient_memory_mode": gtx_transformations.TransientMemoryMode.POOL, "gpu_block_size": (32, 8, 1), "gpu_block_size_2d": (20, 20), - "gpu_memory_pool": True, "optimization_hooks": { gtx_transformations.GT4PyAutoOptHook.TopLevelDataFlowPost: mock_top_level_dataflow_hook1, }, } else: optimization_args = { - "make_persistent": True, + "transient_memory_mode": gtx_transformations.TransientMemoryMode.PERSISTENT, "blocking_dim": KDim, "blocking_size": 10, - "gpu_memory_pool": False, "optimization_hooks": { gtx_transformations.GT4PyAutoOptHook.TopLevelDataFlowPost: mock_top_level_dataflow_hook2, }, @@ -124,10 +136,9 @@ def mocked_gpu_transformation(*args, **kwargs) -> dace.SDFG: "__b_IDim_stride": 1, "__out_IDim_stride": 1, }, - make_persistent=optimization_args["make_persistent"], + transient_memory_mode=optimization_args["transient_memory_mode"], gpu_block_size=optimization_args["gpu_block_size"], gpu_block_size_2d=optimization_args["gpu_block_size_2d"], - gpu_memory_pool=optimization_args["gpu_memory_pool"], optimization_hooks=optimization_args["optimization_hooks"], unit_strides_kind=gtx.common.DimensionKind.HORIZONTAL, ) @@ -138,10 +149,9 @@ def mocked_gpu_transformation(*args, **kwargs) -> dace.SDFG: sdfg, gpu=on_gpu, constant_symbols={}, - make_persistent=optimization_args["make_persistent"], + transient_memory_mode=optimization_args["transient_memory_mode"], blocking_dim=optimization_args["blocking_dim"], blocking_size=optimization_args["blocking_size"], - gpu_memory_pool=optimization_args["gpu_memory_pool"], optimization_hooks=optimization_args["optimization_hooks"], unit_strides_kind=None, ) @@ -151,3 +161,284 @@ def mocked_gpu_transformation(*args, **kwargs) -> dace.SDFG: mock_auto_optimize.assert_not_called() mock_top_level_dataflow_hook1.assert_not_called() mock_top_level_dataflow_hook2.assert_not_called() + + +def _make_external_workspace( + device_type: core_defs.DeviceType, *, nbytes: int = 2**20 +) -> dace_wf_common.ExternalWorkspace: + """Return a sufficiently large array-like workspace for ``device_type``.""" + if device_type == core_defs.CUPY_DEVICE_TYPE: + cupy = pytest.importorskip("cupy") + return cupy.empty(nbytes, dtype=cupy.uint8) + return np.empty(nbytes, dtype=np.uint8) + + +class _RecordingWorkspace: + """Minimal picklable array-like workspace for backend-wiring tests. + + Only the identity of the workspace matters here (it is stored on the + backend instance); the array interface is never consumed by these tests. + """ + + nbytes: int = 1024 + __array_interface__: dict[str, Any] = {"shape": (1024,), "typestr": "|u1", "version": 3} + + +def test_make_backend_accepts_external_workspace_with_external_mode(): + workspace = _RecordingWorkspace() + + backend = dace_wf_backend.make_dace_backend( + gpu=False, + auto_optimize=True, + async_sdfg_call=False, + optimization_args={ + "transient_memory_mode": gtx_transformations.TransientMemoryMode.EXTERNAL, + }, + external_workspace={core_defs.DeviceType.CPU: workspace}, + ) + + assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace + + +def test_make_backend_infers_external_mode_when_workspace_is_provided(): + workspace = _RecordingWorkspace() + + backend = dace_wf_backend.make_dace_backend( + gpu=False, + auto_optimize=True, + async_sdfg_call=False, + external_workspace={core_defs.DeviceType.CPU: workspace}, + ) + + assert ( + backend.executor.translation.step.auto_optimize_args["transient_memory_mode"] + == gtx_transformations.TransientMemoryMode.EXTERNAL + ) + assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace + + +def test_make_backend_warns_external_workspace_without_external_mode(): + workspace = _RecordingWorkspace() + + with pytest.warns(UserWarning, match="External memory workspace provided"): + backend = dace_wf_backend.make_dace_backend( + gpu=False, + auto_optimize=True, + async_sdfg_call=False, + optimization_args={ + "transient_memory_mode": gtx_transformations.TransientMemoryMode.POOL, + }, + external_workspace={core_defs.DeviceType.CPU: workspace}, + ) + + # Explicit mode stays as requested by the caller; backend only warns. + assert ( + backend.executor.translation.step.auto_optimize_args["transient_memory_mode"] + == gtx_transformations.TransientMemoryMode.POOL + ) + assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace + + +def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str: + # Helper function to ignore the GPU device initialization code in the generated + # cuda code, which is not relevant to the test. + malloc_re = re.compile( + rf"DACE_GPU_CHECK\(\s*{gpu_api_prefix}Malloc\(\s*\(void \*\*\)\s*&dev_X\s*,\s*1\s*\)\s*\)\s*;" + ) + free_re = re.compile(rf"DACE_GPU_CHECK\(\s*{gpu_api_prefix}Free\(\s*dev_X\s*\)\s*\)\s*;") + + generated_code = "" + device_code_language = ( + "cu" if core_defs.CUPY_DEVICE_TYPE == core_defs.DeviceType.CUDA else "cpp" + ) + for code in sdfg.generate_code(): + if code.name.endswith("_main"): + pass + elif code.language == device_code_language: + clean_code_iter = iter(code.clean_code.splitlines()) + for line in clean_code_iter: + if malloc_re.match(line.strip()): + line = next(clean_code_iter) # not relevant to this test + assert free_re.match(line.strip()) + else: + generated_code += line + "\n" + else: + generated_code += code.clean_code + "\n" + + return generated_code + + +@pytest.mark.parametrize("transient_memory_mode", list(gtx_transformations.TransientMemoryMode)) +def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): + """Each ``TransientMemoryMode`` reaches codegen with the expected memory API. + + The test inspects the *generated* host/device code (via + ``sdfg.generate_code()``) for allocation/free markers such as + ``cudaMalloc``/``cudaFree`` (sync), ``cudaMallocAsync``/``cudaFreeAsync`` + (pool), or ``set_external_memory``/``__dace_get_external_memory_size_`` + (external). These assertions are on DaCe's codegen output and may need to + be updated if a DaCe upgrade changes the emitted strings. + """ + on_gpu = device_type == core_defs.CUPY_DEVICE_TYPE + gpu_api_prefix = "hip" if core_defs.CUPY_DEVICE_TYPE == core_defs.DeviceType.ROCM else "cuda" + gpu_malloc_marker = f"{gpu_api_prefix}Malloc(" + gpu_malloc_async_marker = f"{gpu_api_prefix}MallocAsync(" + gpu_free_marker = f"{gpu_api_prefix}Free(" + gpu_free_async_marker = f"{gpu_api_prefix}FreeAsync(" + # External mode requires a workspace buffer upfront; other modes do not use one. + external_workspace = ( + {device_type: _make_external_workspace(device_type)} + if transient_memory_mode == gtx_transformations.TransientMemoryMode.EXTERNAL + else None + ) + + custom_backend = dace_wf_backend.make_dace_backend( + gpu=on_gpu, + auto_optimize=True, + async_sdfg_call=False, + optimization_args={ + "transient_memory_mode": transient_memory_mode, + }, + external_workspace=external_workspace, + ) + + @gtx.field_operator + def testee_op(a: cases.IField, b: cases.IField) -> cases.IField: + tmp = a + b + return tmp + 1 + + @gtx.program + def testee(a: cases.IField, b: cases.IField, out: cases.IField): + testee_op(a, b, out=out) + + test_case = cases.Case.from_cartesian_grid_descriptor( + cases_utils.simple_cartesian_grid(), + backend=custom_backend, + allocator=custom_backend, + ) + a = cases.allocate(test_case, testee, "a", strategy=cases.UniqueInitializer())() + b = cases.allocate(test_case, testee, "b", strategy=cases.UniqueInitializer())() + out = cases.allocate(test_case, testee, "out")() + + captured_sdfg: dace.SDFG | None = None + translation_step = custom_backend.executor.translation.step + + def mocked_translator(inp: definitions.CompilableProgramDef) -> dace.SDFG: + nonlocal captured_sdfg + result = translation_step(inp) + captured_sdfg = dace.SDFG.from_json(result.source_code) + return result + + custom_backend = dataclasses.replace( + custom_backend, + executor=dataclasses.replace( + custom_backend.executor, + translation=mocked_translator, + ), + ) + + def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: + return sdfg + + monkeypatch.setattr( + gtx_auto_optimize, + "_gt_auto_process_top_level_maps", + no_op_top_level_map_processing, # we need to keep the intermediate transient array + ) + + # The workspace buffer (if any) lives in the main process, so compilation + # is forced in-process so the patched translator is observed. + with mock.patch.object(config, "BUILD_JOBS_MODE", config.BuildJobsMode.SERIAL): + prog = testee.with_backend(custom_backend).compile(offset_provider={}) + + prog(a, b, out=out) + assert len(prog._compiled_programs.compiled_programs) == 1 + _, decorated_program = next(iter(prog._compiled_programs.compiled_programs.items())) + assert isinstance(decorated_program, dace_wf_decoration.DaCeDecoratedProgram) + + assert captured_sdfg is not None + transient_arrays = [ + (aname, adesc) + for aname, adesc in captured_sdfg.arrays.items() + if isinstance(adesc, dace.data.Array) and adesc.transient + ] + assert len(transient_arrays) == 2 + + generated_code = _parse_generated_code_from_sdfg(captured_sdfg, gpu_api_prefix) + + match transient_memory_mode: + case gtx_transformations.TransientMemoryMode.EXTERNAL: + assert all( + tdesc.lifetime == dace.AllocationLifetime.External for _, tdesc in transient_arrays + ) + # load_artifact injected the backend-level workspace onto the program wrapper. + assert ( + decorated_program._fun.external_workspace[device_type] + is external_workspace[device_type] + ) + # External mode wires explicit workspace API calls in generated host code. + assert "set_external_memory" in generated_code + assert "__dace_get_external_memory_size_" in generated_code + if on_gpu: + # Workspace must come from the external mapping, not from runtime GPU alloc/free. + assert not any( + marker in generated_code + for marker in (gpu_malloc_marker, gpu_malloc_async_marker) + ) + assert not any( + marker in generated_code for marker in (gpu_free_marker, gpu_free_async_marker) + ) + else: + # 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")) + + case gtx_transformations.TransientMemoryMode.POOL: + assert all( + tdesc.pool == on_gpu and tdesc.lifetime == dace.AllocationLifetime.Scope + for _, tdesc in transient_arrays + ) + assert "set_external_memory" not in generated_code + assert "__dace_get_external_memory_size_" not in generated_code + if on_gpu: + # Pool mode on GPU should rely on pooled/async allocation APIs. + assert all( + marker in generated_code + for marker in (gpu_malloc_async_marker, gpu_free_async_marker) + ) + assert not any( + marker in generated_code for marker in (gpu_malloc_marker, gpu_free_marker) + ) + else: + # On CPU, pool mode behaves as regular scoped lifetime. + assert any(marker in generated_code for marker in ("new ", "malloc")) + assert any(marker in generated_code for marker in ("delete ", "free")) + + case ( + gtx_transformations.TransientMemoryMode.PERSISTENT, + gtx_transformations.TransientMemoryMode.SCOPED, + ): + expected_lifetime = ( + dace.AllocationLifetime.Persistent + if transient_memory_mode == gtx_transformations.TransientMemoryMode.PERSISTENT + else dace.AllocationLifetime.SDFG + ) + assert all(tdesc.lifetime == expected_lifetime for _, tdesc in transient_arrays) + # `PERSISTENT` and `SCOPED` mode use the same memory APIs, but in different contexts. + assert "set_external_memory" not in generated_code + assert "__dace_get_external_memory_size_" not in generated_code + if on_gpu: + # Persistent and scoped mode on GPU should rely on sync allocation APIs. + assert all( + marker in generated_code for marker in (gpu_malloc_marker, gpu_free_marker) + ) + assert not any( + marker in generated_code + for marker in (gpu_malloc_async_marker, gpu_free_async_marker) + ) + else: + assert any(marker in generated_code for marker in ("new ", "malloc")) + assert any(marker in generated_code for marker in ("delete ", "free")) + + assert np.allclose(out.asnumpy(), a.asnumpy() + b.asnumpy() + 1) diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py index 610567f6b2..bb2ecf6323 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py @@ -8,14 +8,15 @@ """Tests for the compilation stage of the dace backend workflow. -Covers the GPU TX-marker instrumentation and the picklability of -``DaCeCompilationArtifact``. +Covers the GPU TX-marker instrumentation and external workspace handling of +``CompiledDaceProgram`` / ``DaCeCompilationArtifact``. """ import contextlib import pathlib import pickle import unittest.mock as mock +from typing import Any import pytest @@ -27,7 +28,9 @@ from gt4py.next import config from gt4py.next.otf import code_specs, stages from gt4py.next.otf.binding import interface -from gt4py.next.program_processors.runners.dace.workflow import compilation as dace_wf_compilation +from gt4py.next.program_processors.runners.dace.workflow import ( + compilation as dace_wf_compilation, +) _TX = dace.dtypes.InstrumentationType.GPU_TX_MARKERS @@ -165,6 +168,7 @@ def test_compiler_skips_tx_markers_for_non_gpu_device(program_source): def test_dace_compilation_artifact_pickle_round_trip(tmp_path: pathlib.Path): + """The artifact is picklable and does not carry runtime workspace buffers.""" artifact = dace_wf_compilation.DaCeCompilationArtifact( library_path=tmp_path / "build" / "libprogram.so", sdfg_json="{}", @@ -262,3 +266,187 @@ def test_cmake_build_type_changes_artifact(program_source): artifact_debug, _ = _run_compiler(program_source, cmake_build_type=config.CMakeBuildType.DEBUG) assert artifact_release.library_path != artifact_debug.library_path + + +def _make_compiled_program( + *, + external_workspace: dict[core_defs.DeviceType, Any] | None = None, + workspace_sizes: dict[Any, int] | None = None, +): + if workspace_sizes is None: + workspace_sizes = {} + + sdfg = mock.MagicMock() + sdfg.arglist.return_value = {} + sdfg.build_folder = "build-folder" + + sdfg_program = mock.MagicMock() + sdfg_program.sdfg = sdfg + sdfg_program.get_workspace_sizes.return_value = workspace_sizes + sdfg_program.construct_arguments.return_value = ((), ()) + + compiled_program = dace_wf_compilation.CompiledDaceProgram( + program=sdfg_program, + bind_func_name="update_sdfg_args", + binding_source_code="def update_sdfg_args(*a, **k):\n return None\n", + ) + if external_workspace is not None: + compiled_program.external_workspace = external_workspace + return compiled_program + + +def test_construct_arguments_without_external_workspace(): + """If the SDFG does not need a workspace, `construct_arguments` works without it.""" + program = _make_compiled_program( + external_workspace=None, workspace_sizes={} + ) # no workspace needed + + program.construct_arguments(alpha=1) + + program.sdfg_program.initialize.assert_called_once() + program.sdfg_program.get_workspace_sizes.assert_called_once() + program.sdfg_program.set_workspace.assert_not_called() + program.sdfg_program.construct_arguments.assert_called_once() + + +def test_construct_arguments_installs_external_workspace(): + """If the SDFG needs a workspace, `construct_arguments` installs it from the mapping.""" + workspace = _make_array_buffer(nbytes=128) + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: workspace}, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + program.construct_arguments(alpha=1) + + assert program.sdfg_program.initialize.call_count == 1 + assert program.sdfg_program.get_workspace_sizes.call_count == 1 + + set_workspace_call = program.sdfg_program.set_workspace.call_args + assert set_workspace_call.args[0] == dace.StorageType.CPU_Heap + assert set_workspace_call.args[1] is workspace + + +def test_construct_arguments_raises_when_workspace_missing(): + """If the SDFG needs a workspace but none is provided, raise early.""" + program = _make_compiled_program( + external_workspace=None, workspace_sizes={dace.StorageType.CPU_Heap: 128} + ) + + with pytest.raises(RuntimeError, match="External workspace is not set"): + program.construct_arguments(alpha=1) + + +def test_construct_arguments_raises_when_workspace_missing_for_device(): + """If a required device workspace is absent from the mapping, raise early.""" + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: _make_array_buffer(nbytes=128)}, + workspace_sizes={dace.StorageType.GPU_Global: 128}, + ) + with mock.patch.object( + dace_wf_compilation.core_defs, "CUPY_DEVICE_TYPE", core_defs.DeviceType.CUDA + ): + with pytest.raises(RuntimeError, match="External workspace for device .* not found"): + program.construct_arguments(alpha=1) + + +def test_construct_arguments_propagates_validation_error_for_too_small_buffer(): + """A workspace that is too small for the requested storage is rejected.""" + workspace = _make_array_buffer(nbytes=64) + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: workspace}, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + with pytest.raises(ValueError, match="at least 128 bytes were required"): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_not_called() + + +def test_construct_arguments_propagates_validation_error_for_invalid_storage(): + """An unsupported storage type is rejected during device mapping.""" + workspace = _make_array_buffer(nbytes=128) + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: workspace}, + workspace_sizes={dace.StorageType.CPU_Pinned: 128}, + ) + + with pytest.raises(ValueError, match="Unsupported storage type"): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_not_called() + + +def _make_array_buffer(*, nbytes: int, cuda: bool = False) -> mock.MagicMock: + """A minimal array-like buffer accepted by ``_validate_external_workspace``. + + Exposes ``__array_interface__`` (host) or ``__cuda_array_interface__`` + (device); ``nbytes`` matches the requested size. + """ + buffer = mock.MagicMock() + buffer.nbytes = nbytes + interface = {"shape": (nbytes,), "typestr": "|u1", "version": 3} + setattr(buffer, "__cuda_array_interface__" if cuda else "__array_interface__", interface) + return buffer + + +def test_construct_arguments_rejects_buffer_without_array_interface(): + """The workspace must expose an array interface that ``set_workspace`` can consume.""" + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: "not-an-array"}, + workspace_sizes={dace.StorageType.CPU_Heap: 64}, + ) + + with pytest.raises(TypeError, match="must expose `__array_interface__`"): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_not_called() + + +def test_construct_arguments_accepts_cpu_workspace(): + """A host workspace with matching size is accepted and installed.""" + workspace = _make_array_buffer(nbytes=128) + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: workspace}, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_called_once() + assert program.sdfg_program.set_workspace.call_args.args[1] is workspace + + +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}, + ) + with mock.patch.object( + dace_wf_compilation.core_defs, "CUPY_DEVICE_TYPE", core_defs.DeviceType.CUDA + ): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_called_once() + assert program.sdfg_program.set_workspace.call_args.args[0] == dace.StorageType.GPU_Global + assert program.sdfg_program.set_workspace.call_args.args[1] is workspace + + +class _ArrayBufferWithoutNbytes: + __array_interface__ = {"shape": (128,), "typestr": "|u1", "version": 3} + + +def test_construct_arguments_skips_size_check_when_nbytes_missing(): + """When the buffer lacks ``nbytes`` the size contract is trust-based.""" + program = _make_compiled_program( + external_workspace={core_defs.DeviceType.CPU: _ArrayBufferWithoutNbytes()}, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + # Must not raise even though the size cannot be verified. + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_called_once() diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py index d8cf8e33f8..4510a8738e 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py @@ -67,8 +67,8 @@ def test_make_transients_persistent_inner_access(): # Because `b`, the only transient, is used inside a map scope, it is not selected, # although in this situation it would be possible. - change_report: dict[int, set[str]] = gtx_transformations.gt_make_transients_persistent( - sdfg, device=dace.DeviceType.CPU + change_report: dict[int, set[str]] = gtx_transformations.gt_configure_transient_lifetime( + sdfg, lifetime=dace.AllocationLifetime.Persistent ) assert len(change_report) == 1 assert change_report[sdfg.cfg_id] == set() diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_transient_memory_mode.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_transient_memory_mode.py new file mode 100644 index 0000000000..b3eeb91393 --- /dev/null +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_transient_memory_mode.py @@ -0,0 +1,67 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import pytest + +from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations + +dace = pytest.importorskip("dace") + + +def _make_sdfg_with_top_level_transient(storage): + sdfg = dace.SDFG("external_mode_scalar_test") + state = sdfg.add_state("state", is_start_block=True) + sdfg.add_array("a", [10], dace.float64, transient=False, storage=storage) + sdfg.add_array("tmp_arr", [10], dace.float64, transient=True, storage=storage) + sdfg.add_scalar("tmp_scalar", dace.float64, transient=True) + sdfg.add_array("b", [10], dace.float64, transient=False, storage=storage) + + a = state.add_access("a") + tmp_arr = state.add_access("tmp_arr") + tmp_scalar = state.add_access("tmp_scalar") + b = state.add_access("b") + state.add_nedge(a, tmp_arr, dace.Memlet("a[0:10]")) + state.add_nedge(tmp_arr, b, dace.Memlet("tmp_arr[0:10]")) + init_scalar = state.add_tasklet( + "init_scalar", + inputs={}, + outputs={"out"}, + code="out = 1.0", + ) + state.add_edge(init_scalar, "out", tmp_scalar, None, dace.Memlet("tmp_scalar")) + sdfg.validate() + return sdfg + + +@pytest.mark.parametrize( + "lifetime", + [ + dace.AllocationLifetime.Persistent, + dace.AllocationLifetime.External, + ], +) +@pytest.mark.parametrize( + "storage", + [ + dace.StorageType.Default, + dace.StorageType.GPU_Global, + ], +) +def test_configure_transient_lifetime(lifetime, storage): + sdfg = _make_sdfg_with_top_level_transient(storage) + + result = gtx_transformations.gt_configure_transient_lifetime(sdfg, lifetime) + candidates = next(iter(result.values())) + assert candidates == {"tmp_arr", "tmp_scalar"} + + assert sdfg.arrays["tmp_arr"].lifetime == lifetime + assert sdfg.arrays["tmp_arr"].storage == ( + dace.StorageType.CPU_Heap if storage == dace.StorageType.Default else storage + ) + assert sdfg.arrays["tmp_scalar"].lifetime == lifetime + assert sdfg.arrays["tmp_scalar"].storage == dace.StorageType.Default