diff --git a/docs/development/ADRs/next/0026-External_Memory_Allocator.md b/docs/development/ADRs/next/0026-External_Memory_Allocator.md new file mode 100644 index 0000000000..c6d5a0e42c --- /dev/null +++ b/docs/development/ADRs/next/0026-External_Memory_Allocator.md @@ -0,0 +1,147 @@ +--- +tags: [] +--- + +# External Memory Allocator for DaCe Transients + +- **Status**: valid +- **Authors**: Edoardo Paone (@edopao) +- **Created**: 2026-07-27 +- **Updated**: 2026-07-27 + +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 a caller-supplied `ExternalMemoryAllocator` protocol (allocate once +per SDFG storage type, release when the compiled program is finalized). + +## 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, pool-driven 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, release at finalize -- 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 a caller-supplied +allocator. + +- `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 `ExternalMemoryAllocator` Protocol + (`allocate(request: AllocationRequest) -> ExternalWorkspace` / + `deallocate(wsp) -> None`), defined in `transformations/auto_optimize.py` + alongside `AllocationRequest` (`nbytes`, `device`, `alignment=256`) and + `ExternalWorkspace` (a `TypeAlias` over the existing `ArrayInterface` / + `CUDAArrayInterface` from `gt4py.eve.extended_typing`, not a new protocol). +- At runtime, `CompiledDaceProgram.construct_arguments` calls + `sdfg_program.get_workspace_sizes()`, invokes `allocate` once per storage + type, validates the returned buffer (array interface, size, alignment) and + installs it via `sdfg_program.set_workspace(...)`. The buffers are kept on + the `CompiledDaceProgram` for its lifetime. +- Teardown is explicit and pool-driven, not `__del__`-based: + `CompiledDaceProgram.finalize()` finalizes the underlying SDFG and calls + `deallocate` once per storage type (and is idempotent and resilient: a + failing `deallocate` is warned, not raised, so one bad buffer does not + strand the rest). `DaCeDecoratedProgram` in `workflow/decoration.py` + forwards `finalize()` to the underlying `CompiledDaceProgram`, and + `CompiledProgramsPool.__post_init__` registers a `weakref.finalize` that + walks `compiled_programs` and calls `finalize()` on each when the pool is + collected. This mirrors the existing `metrics_source_key` finalizer and its + "avoid id reuse once a pool dies" rationale. +- The allocator is part of `DaCeCompilationArtifact` and is therefore + **picklable**: when the OTF runner offloads compilation to a + `ProcessPoolExecutor` it pickles the executor chain, which carries the + allocator. `DaCeCompiler.__post_init__` probes the allocator with + `pickle.dumps` and raises `AllocatorNotPicklableError` (a `TypeError`) at + backend construction if it cannot be pickled, rather than letting a closure + or lambda silently degrade to in-process compilation via the generic runner + warning. The error chains the original pickle failure and names the + recommended shape (module-level class or `functools.partial` of picklable + callables). + +## 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 explicit and tied to the compiled program's lifetime + through the pool finalizer, not to GC timing. Backends owning external + resources (this one) get `finalize()` called at pool teardown. +- The public API has a typed allocator protocol and a single mode enum; + incompatible combinations (e.g. an allocator with a non-`EXTERNAL` mode) + are detected and warned at backend construction. +- A non-picklable allocator fails loudly at construction instead of silently + serializing compilation, at the cost of probing every allocator once with + `pickle.dumps` (cheap for the common module-level-class shape). +- 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 (the plan's Phase 6). + +## Alternatives considered + +### `__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 that may already be gone). Rejected in favor of the explicit + `finalize()` forwarded through the callable and driven by the pool finalizer. + +### 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` + (`ExternalMemoryAllocator`, `AllocationRequest`, `ExternalWorkspace`, + `TransientMemoryMode`, `_gt_auto_post_processing`). +- `src/gt4py/next/program_processors/runners/dace/workflow/compilation.py` + (`CompiledDaceProgram.construct_arguments`/`finalize`, + `DaCeCompilationArtifact`, `DaCeCompiler`, `AllocatorNotPicklableError`). +- `src/gt4py/next/program_processors/runners/dace/workflow/decoration.py` + (`DaCeDecoratedProgram.finalize` forwarding). +- `src/gt4py/next/otf/compiled_program.py` + (`_finalize_compiled_programs`, `CompiledProgramsPool.__post_init__` + finalizer). +- [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/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 0784ca0f73..01d002aab2 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -85,6 +85,32 @@ def metrics_source_key(pool: CompiledProgramsPool, key: CompiledProgramsKey) -> return source_key +def _finalize_compiled_programs(programs: stages.ExecutableProgram | dict[Any, Any]) -> None: + """Best-effort cleanup of compiled programs when their pool is deleted. + + Invoked via :py:func:`weakref.finalize` on a + :py:class:`CompiledProgramsPool`. The pool holds each compiled program + as a generic :py:data:`stages.ExecutableProgram` (a backend-specific + callable, e.g. the DaCe ``DaCeDecoratedProgram``); backends that + own external resources expose a ``finalize()`` method, which is forwarded + to the underlying compiled object. Failures are surfaced as warnings + rather than raised, because finalizers cannot propagate exceptions. + """ + values = programs.values() if isinstance(programs, dict) else (programs,) + for program in values: + finalize = getattr(program, "finalize", None) + if finalize is None: + continue + try: + finalize() + except Exception: + warnings.warn( + f"Compiled program {type(program).__name__!r} raised during " + f"pool teardown; its resources may be leaked.", + stacklevel=1, + ) + + @hook_machinery.event_hook def compile_variant_hook( program_pool: CompiledProgramsPool, @@ -375,6 +401,20 @@ def definition(self) -> types.FunctionType: return self.definition_stage.definition def __post_init__(self) -> None: + # Best-effort teardown: when this pool is deleted (and its + # ``compiled_programs`` dict goes with it), finalize any compiled + # program that exposes a ``finalize()`` method so backends that own + # external resources -- e.g. the DaCe external-memory allocator -- + # can release them. The dict is passed by reference so the + # finalizer walks the live collection (programs may be added after + # registration); the pool is held weakly so this does not extend + # its lifetime. Mirrors the finalizer registered in + # ``metrics_source_key`` (which avoids id reuse once a pool dies). + # + # Registered first so a ``__post_init__`` that fails validation + # still installs teardown for whatever is already cached. + weakref.finalize(self, _finalize_compiled_programs, self.compiled_programs) + # TODO(havogt): We currently don't support pos_only or kw_only args at the program level. # This check makes sure we don't miss updating this code if we add support for them in the future. assert not self.program_type.definition.kw_only_args 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..e5ea78eb71 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/__init__.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/__init__.py @@ -13,10 +13,14 @@ """ from . import constants, splitting_tools -from .auto_optimize import ( +from .auto_optimize import ( # re-exported for the external-memory public API + AllocationRequest, + ExternalMemoryAllocator, + ExternalWorkspace, GT4PyAutoOptHook, GT4PyAutoOptHookFun, GT4PyAutoOptHookStage, + TransientMemoryMode, gt_auto_optimize, ) from .concat_where_mapper import ( @@ -84,12 +88,15 @@ 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__ = [ + "AllocationRequest", "CopyChainRemover", "DoubleWriteRemover", + "ExternalMemoryAllocator", + "ExternalWorkspace", "FuseHorizontalConditionBlocks", "GPUSetBlockSize", "GT4PyAutoOptHook", @@ -119,6 +126,7 @@ "SingleStateGlobalSelfCopyElimination", "SplitAccessNode", "SplitConsumerMemlet", + "TransientMemoryMode", "VerticalMapFusionCallback", "VerticalMapSplitCallback", "constants", @@ -126,13 +134,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..6308442bc8 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 @@ -8,9 +8,11 @@ """Fast access to the auto optimization on DaCe.""" +import abc +import dataclasses import enum import warnings -from typing import Any, Callable, Optional, Sequence, TypeAlias, Union +from typing import Any, Callable, Optional, Protocol, Sequence, TypeAlias, Union import dace from dace import data as dace_data @@ -19,6 +21,8 @@ from dace.transformation.auto import auto_optimize as dace_aoptimize from dace.transformation.passes import analysis as dace_analysis +from gt4py._core import definitions as core_defs +from gt4py.eve import extended_typing as xtyping from gt4py.next import common as gtx_common, utils as gtx_utils from gt4py.next.program_processors.runners.dace import ( library_nodes as gtx_library_nodes, @@ -111,11 +115,104 @@ class GT4PyAutoOptHook(enum.Enum): ] +@dataclasses.dataclass(frozen=True) +class AllocationRequest: + """A single workspace allocation request issued by the DaCe runtime. + + Attributes: + nbytes: Number of bytes the compiled SDFG requires for this workspace. + device: Target device for the workspace, derived from the SDFG + transient storage type (``CPU`` for ``CPU_Heap``, the configured + GPU device for ``GPU_Global``). + alignment: Minimum byte alignment required for the returned buffer. + Allocators may return more strictly aligned memory; the default + matches the alignment DaCe assumes for transient storage on the + target device. + """ + + nbytes: int + device: core_defs.DeviceType + alignment: int = 256 + + +#: Array-like object that ``dace.dtypes.array_interface_ptr()`` accepts as a +#: workspace: a host array exposing :class:`~gt4py.eve.extended_typing.ArrayInterface` +#: or a device array exposing :class:`~gt4py.eve.extended_typing.CUDAArrayInterface`. +#: The returned object must be at least as large as the requested number of +#: bytes; allocators that return a larger slab are acceptable. Reuses the +#: existing array-interface protocols from :mod:`gt4py.eve.extended_typing` +#: rather than introducing a new one. +ExternalWorkspace: TypeAlias = xtyping.ArrayInterface | xtyping.CUDAArrayInterface + + +class ExternalMemoryAllocator(Protocol): + """Allocates and frees workspace memory for ``TransientMemoryMode.EXTERNAL``. + + The allocator owns the lifetime of the memory it hands out. For each SDFG + that requires external workspaces, ``allocate`` is called once per + SDFG storage type during `CompiledDaceProgram.construct_arguments`, + and ``deallocate`` is called once per storage type when the `CompiledDaceProgram` + is finalized. Allocators that wish to reuse a single slab across many programs + must keep the slab alive after ``deallocate`` returns: ``deallocate`` + signals "this program is done with the workspace", not "destroy the memory". + + Implementations must be picklable (module-level classes or + :py:func:`functools.partial` of picklable callables are recommended), + because the allocator is part of the compilation artifact. + """ + + @abc.abstractmethod + def allocate(self, request: AllocationRequest) -> ExternalWorkspace: + """Return a buffer of at least ``request.nbytes`` bytes on + ``request.device``. Must raise on failure; never return ``None``. + """ + ... + + @abc.abstractmethod + def deallocate(self, wsp: ExternalWorkspace) -> None: + """Release or reclaim ``wsp``. Called once per workspace when the + owning compiled program is finalized. + """ + ... + + +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 allocated and deallocated by an external allocator. + This strategy allows to reuse a workspace memory for multiple SDFGs, relying + on sequential execution of the programs on the default stream. + Note: + The `EXTERNAL` strategy requires that the `external_memory_allocator` + attribute of the dace backend workflow is set to an + `ExternalMemoryAllocator`. Workspaces are allocated once per compiled + program (one `allocate` call per SDFG storage type) and freed when the + compiled program is finalized (one `deallocate` call per storage type). + The allocator must be picklable. + """ + + 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 +227,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 +270,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 +288,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 +493,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 +1010,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 +1030,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..76f50e55b9 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -16,6 +16,7 @@ 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 import transformations as gtx_transformations from gt4py.next.program_processors.runners.dace.workflow.factory import DaCeWorkflowFactory @@ -35,6 +36,7 @@ class Meta: class Params: name_device = "cpu" name_postfix = "" + external_memory_allocator: gtx_transformations.ExternalMemoryAllocator | None = None gpu = factory.Trait( allocator=next_allocators.StandardGPUFieldBufferAllocator(), device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, @@ -46,6 +48,7 @@ class Params: cached_translation=True, device_type=factory.SelfAttribute("..device_type"), auto_optimize=factory.SelfAttribute("..auto_optimize"), + external_memory_allocator=factory.SelfAttribute("..external_memory_allocator"), ) auto_optimize = factory.Trait(name_postfix="_opt") @@ -60,6 +63,7 @@ def make_dace_backend( auto_optimize: bool = True, async_sdfg_call: bool = True, optimization_args: dict[str, Any] | None = None, + external_memory_allocator: gtx_transformations.ExternalMemoryAllocator | 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 +78,11 @@ 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_memory_allocator: Allocator used to provide workspace memory + when `transient_memory_mode` is `EXTERNAL`. Threaded through the + backend workflow and called once per SDFG storage type when + arguments are constructed; see + `gtx_transformations.ExternalMemoryAllocator` for the contract. 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,11 +119,26 @@ def make_dace_backend( else None } + if external_memory_allocator is not None: + expected_mode = gtx_transformations.TransientMemoryMode.EXTERNAL + if "transient_memory_mode" in optimization_args: + if ( + transient_memory_mode := optimization_args["transient_memory_mode"] + ) is not expected_mode: + warnings.warn( + f"External memory allocator provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{expected_mode}'.", + stacklevel=2, + ) + else: + optimization_args["transient_memory_mode"] = expected_mode + return DaCeBackendFactory( # type: ignore[return-value] # factory-boy typing not precise enough gpu=gpu, auto_optimize=auto_optimize, + external_memory_allocator=external_memory_allocator, 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__compilation__external_memory_allocator=external_memory_allocator, otf_workflow__bare_translation__unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride, otf_workflow__bare_translation__use_metrics=use_metrics, otf_workflow__bare_translation__disable_field_origin_on_program_arguments=use_zero_origin, 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 c018bb13a1..d86a538172 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -12,6 +12,7 @@ import json import os import pathlib +import pickle import warnings from collections.abc import Callable, MutableSequence, Sequence from typing import Any, Final @@ -21,9 +22,15 @@ 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 +from gt4py.next.program_processors.runners.dace.transformations.auto_optimize import ( + AllocationRequest, + ExternalMemoryAllocator, + ExternalWorkspace, +) from gt4py.next.program_processors.runners.dace.workflow import ( common as gtx_wfdcommon, decoration as gtx_wfddecoration, @@ -33,6 +40,39 @@ _COMPILE_COMPLETE_MARKER: Final = ".gt4py_compile_complete" +class AllocatorNotPicklableError(TypeError): + """Raised when an ``external_memory_allocator`` cannot be pickled. + + The allocator is part of the compilation artifact and is pickled when + compilation is offloaded to a worker process. Allocators that can not be + pickled -- typically closures, lambdas, or classes defined inside a + function -- would otherwise degrade silently to in-process compilation + via a generic runner warning. This error surfaces the contract failure + early, at backend construction, with the original :mod:`pickle` error + chained as ``__cause__``. + """ + + +def _assert_allocator_picklable(allocator: ExternalMemoryAllocator) -> None: + """Fail fast if ``allocator`` is not picklable. + + Args: + allocator: The allocator to probe; must not be ``None``. + + Raises: + AllocatorNotPicklableError: If ``pickle.dumps(allocator)`` raises. + """ + try: + pickle.dumps(allocator) + except Exception as error: # pickle raises arbitrary exceptions + raise AllocatorNotPicklableError( + f"external_memory_allocator {allocator!r} is not picklable: {error!r}." + " The allocator is part of the compilation artifact and is pickled" + " when compilation is offloaded to a worker process. Use a" + " module-level class or functools.partial of picklable callables." + ) from error + + def _add_tx_markers(sdfg: dace.SDFG) -> None: has_gpu_schedule = any( getattr(node, "schedule", dace.dtypes.ScheduleType.Default) in dace.dtypes.GPU_SCHEDULES @@ -47,6 +87,64 @@ def _add_tx_markers(sdfg: dace.SDFG) -> None: node.instrument = dace.dtypes.InstrumentationType.GPU_TX_MARKERS +def _map_workspace_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: + assert core_defs.CUPY_DEVICE_TYPE is not 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( + storage: dace.StorageType, request: AllocationRequest, wsp: ExternalWorkspace +) -> None: + """Validate that ``wsp`` satisfies ``request`` for ``storage``. + + Args: + storage: SDFG storage type the workspace buffer is being installed for. + request: Allocation request that was issued. + wsp: External workspace returned by the external allocator. + + Raises: + TypeError: If ``wsp`` exposes neither ``__array_interface__`` nor + ``__cuda_array_interface__``. + ValueError: If ``wsp`` is smaller than ``request.nbytes`` or its + base pointer is not aligned to ``request.alignment`` bytes. + """ + if not (xtyping.supports_array_interface(wsp) or xtyping.supports_cuda_array_interface(wsp)): + 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__`." + ) + nbytes = getattr(wsp, "nbytes", None) + if nbytes is not None and nbytes < request.nbytes: + raise ValueError( + f"External memory allocator returned a buffer of {nbytes} bytes for storage " + f"{storage!r}, but at least {request.nbytes} were required." + ) + # Validate alignment against the base pointer DaCe will hand to the SDFG + # (see ``dace.dtypes.array_interface_ptr``). The ``data`` field is + # optional on the host array interface; if it is missing the alignment + # contract is trust-based and the check is skipped, mirroring ``nbytes``. + interface = ( + getattr(wsp, "__cuda_array_interface__", None) + if storage == dace.StorageType.GPU_Global + else getattr(wsp, "__array_interface__", None) + ) + data = interface.get("data") if interface is not None else None + if data is not None and request.alignment > 1 and data[0] % request.alignment != 0: + raise ValueError( + f"External memory allocator returned a buffer for storage {storage!r} " + f"whose base pointer ({data[0]}) is not aligned to the required " + f"{request.alignment} bytes." + ) + + class CompiledDaceProgram: sdfg_program: dace.CompiledSDFG @@ -73,12 +171,15 @@ class CompiledDaceProgram: # never updated. csdfg_argv: MutableSequence[Any] | None csdfg_init_argv: Sequence[Any] | None + external_memory_allocator: ExternalMemoryAllocator | None + external_workspaces: dict[dace.StorageType, ExternalWorkspace] def __init__( self, program: dace.CompiledSDFG, bind_func_name: str, binding_source_code: str, + external_memory_allocator: ExternalMemoryAllocator | None = None, ): self.sdfg_program = program @@ -98,6 +199,56 @@ def __init__( # Since the SDFG hasn't been called yet. self.csdfg_argv = None self.csdfg_init_argv = None + self.external_memory_allocator = external_memory_allocator + self.external_workspaces = {} + + def _configure_external_workspaces(self, **kwargs: Any) -> None: + if self.external_workspaces: + # We already allocated the external workspaces, no need to do it again. + return + + # DaCe computes workspace sizes during ``initialize`` and stores them + # for subsequent ``get_workspace_sizes``/``set_workspace`` calls. + self.sdfg_program.initialize(**kwargs) + if workspace_sizes := self.sdfg_program.get_workspace_sizes(): + if self.external_memory_allocator is None: + raise ValueError( + "SDFG requires external workspaces, but no allocator was provided." + ) + for storage, required_nbytes in workspace_sizes.items(): + device = _map_workspace_storage_to_device(storage) + request = AllocationRequest(nbytes=required_nbytes, device=device) + workspace = self.external_memory_allocator.allocate(request) + _validate_external_workspace(storage, request, workspace) + self.sdfg_program.set_workspace(storage, workspace) + # Keep the workspace buffers alive as long as the compiled program lives. + self.external_workspaces[storage] = workspace + + def finalize(self) -> None: + """Release external workspaces. + + Finalizes the underlying ``sdfg_program`` and calls ``deallocate`` + once per allocated storage type. Safe to call multiple times: after + the first call the per-storage workspace buffers are dropped from + ``external_workspaces`` and subsequent calls are no-ops. A ``None`` + allocator performs no work but still clears any externally-installed + workspaces. + + Failures during deallocation are surfaced as warnings rather than + raised, so that one failing buffer does not prevent the remaining + workspaces from being released. + """ + if self.external_memory_allocator is not None: + for wsp in self.external_workspaces.values(): + try: + self.external_memory_allocator.deallocate(wsp) + except Exception: + warnings.warn( + f"Failed to deallocate external workspace " + f"({type(wsp).__name__!r}); it may be leaked.", + stacklevel=1, + ) + self.external_workspaces = {} def construct_arguments(self, **kwargs: Any) -> None: """ @@ -106,6 +257,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_workspaces(**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] @@ -153,6 +305,7 @@ class DaCeCompilationArtifact: binding_source_code: str bind_func_name: str device_type: core_defs.DeviceType + external_memory_allocator: ExternalMemoryAllocator | None = None def load(self) -> stages.ExecutableProgram: # TODO(phimuell): Drop ``sdfg_json`` from the artifact once dace @@ -160,8 +313,13 @@ def load(self) -> stages.ExecutableProgram: # into the returned ``CompiledSDFG``. 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) + program = CompiledDaceProgram( + sdfg_program, + self.bind_func_name, + self.binding_source_code, + external_memory_allocator=self.external_memory_allocator, + ) + return gtx_wfddecoration.DaCeDecoratedProgram(program, device_type=self.device_type) @dataclasses.dataclass(frozen=True) @@ -181,6 +339,11 @@ class DaCeCompiler( bind_func_name: str cache_lifetime: config.BuildCacheLifetime device_type: core_defs.DeviceType + #: Allocator providing external workspace memory when + #: ``transient_memory_mode`` is ``EXTERNAL``. Must be picklable (a + #: module-level class or :py:func:`functools.partial` of picklable + #: callables is recommended); probed at construction time. + external_memory_allocator: ExternalMemoryAllocator | None = None add_gpu_trace_markers: bool = dataclasses.field( default_factory=lambda: config.ADD_GPU_TRACE_MARKERS ) @@ -191,6 +354,13 @@ class DaCeCompiler( dace_config_nondefaults: dict[str, Any] = dataclasses.field(init=False) def __post_init__(self) -> None: + # The allocator is part of the compilation artifact and is pickled + # when compilation is offloaded to a worker process. Probe it here, + # at backend construction, so a non-picklable allocator (closure, + # lambda, local class) fails fast with an actionable error instead + # of silently degrading to in-process compilation. + if self.external_memory_allocator is not None: + _assert_allocator_picklable(self.external_memory_allocator) with gtx_wfdcommon.dace_context( device_type=self.device_type, cmake_build_type=self.cmake_build_type, @@ -251,6 +421,7 @@ def __call__( binding_source_code=inp.binding_source.source_code, bind_func_name=self.bind_func_name, device_type=self.device_type, + external_memory_allocator=self.external_memory_allocator, ) 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..9ed8d14ff6 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,35 @@ 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. 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. + """ + + 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 +68,36 @@ 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 finalize(self) -> None: + """Forward teardown to the underlying ``CompiledDaceProgram``. - return decorated_program + Allows the generic otf pool -- which only sees this callable -- to + release external-memory workspaces when the pool is finalized. + """ + self._fun.finalize() diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py index 6238871b8f..6dacb61885 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -17,6 +17,7 @@ from gt4py.next import config from gt4py.next.otf import recipes, stages, workflow from gt4py.next.otf.compilation import cache +from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations from gt4py.next.program_processors.runners.dace.workflow import bindings as bindings_step from gt4py.next.program_processors.runners.dace.workflow.compilation import ( DaCeCompilationStepFactory, @@ -35,6 +36,7 @@ class Meta: class Params: auto_optimize: bool = False + external_memory_allocator: gtx_transformations.ExternalMemoryAllocator | None = None device_type: core_defs.DeviceType = core_defs.DeviceType.CPU cmake_build_type: config.CMakeBuildType = factory.LazyFunction( # type: ignore[assignment] # factory-boy typing not precise enough lambda: config.CMAKE_BUILD_TYPE @@ -73,5 +75,6 @@ class Params: bind_func_name=_GT_DACE_BINDING_FUNCTION_NAME, cache_lifetime=factory.LazyFunction(lambda: config.BUILD_CACHE_LIFETIME), device_type=factory.SelfAttribute("..device_type"), + external_memory_allocator=factory.SelfAttribute("..external_memory_allocator"), cmake_build_type=factory.SelfAttribute("..cmake_build_type"), ) diff --git a/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py b/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py index ed881c9495..35cb82a355 100644 --- a/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py +++ b/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py @@ -314,3 +314,72 @@ def test_f(): compiled_program._pools_per_root = _pools_per_root ctx.run(test_f) + + +class _FinalizableProgram: + """Minimal stand-in for a backend compiled program exposing ``finalize()``.""" + + def __init__(self) -> None: + self.finalize_count = 0 + + def finalize(self) -> None: + self.finalize_count += 1 + + +class _FailingFinalizeProgram: + def finalize(self) -> None: + raise RuntimeError("teardown blew up") + + +def test_finalize_compiled_programs_calls_finalize_on_each_value(): + a = _FinalizableProgram() + b = _FinalizableProgram() + programs = {("a",): a, ("b",): b} + + compiled_program._finalize_compiled_programs(programs) + + assert a.finalize_count == 1 + assert b.finalize_count == 1 + + +def test_finalize_compiled_programs_skips_programs_without_finalize(): + class _NoFinalize: + pass + + programs = {("a",): _NoFinalize(), ("b",): _FinalizableProgram()} + compiled_program._finalize_compiled_programs(programs) # must not raise on _NoFinalize + + +def test_finalize_compiled_programs_surfaces_finalize_failures_as_warnings(): + programs = {("a",): _FailingFinalizeProgram(), ("b",): _FinalizableProgram()} + ok = programs[("b",)] + + with pytest.warns(UserWarning, match="raised during pool teardown"): + compiled_program._finalize_compiled_programs(programs) + + # A failing program does not stop teardown of the rest. + assert ok.finalize_count == 1 + + +def test_pool_finalizer_finalizes_compiled_programs_when_pool_is_deleted(): + """A program held in ``CompiledProgramsPool.compiled_programs`` is finalized + when the pool is garbage-collected, so backends that own external + resources release them.""" + # ``CompiledProgramsPool.__post_init__`` validates its (heavy) + # constructor arguments, which is irrelevant to teardown. Install the + # live ``compiled_programs`` dict by hand and register the same + # finalizer ``__post_init__`` would -- this is exactly the field and + # helper production uses. + pool = compiled_program.CompiledProgramsPool.__new__(compiled_program.CompiledProgramsPool) + pool.compiled_programs = {} + weakref.finalize(pool, compiled_program._finalize_compiled_programs, pool.compiled_programs) + + program = _FinalizableProgram() + pool.compiled_programs[("only",)] = program + + pool_ref = weakref.ref(pool) + del pool + gc.collect() + + assert pool_ref() is None + assert program.finalize_count == 1 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 b7bbe4216c..7d5e97fdb5 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,30 +8,41 @@ """Test the bindings stage of the dace backend workflow.""" -import pytest +import re import unittest.mock as mock +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 import config from gt4py.next.otf import 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, + translation as gtx_dace_translation, ) -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 @pytest.fixture( params=[ pytest.param(core_defs.DeviceType.CPU), - pytest.param(core_defs.DeviceType.CUDA, marks=pytest.mark.requires_gpu), + pytest.param(core_defs.CUPY_DEVICE_TYPE, marks=pytest.mark.requires_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 @@ -54,20 +65,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, }, @@ -123,10 +132,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, ) @@ -137,10 +145,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, ) @@ -150,3 +157,290 @@ 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() + + +class _RecordingAllocator: + """Minimal `ExternalMemoryAllocator` for backend-wiring tests. + + Only the identity of the allocator matters here (it is threaded through + to ``executor.compilation.external_memory_allocator``); ``allocate`` is + never called by these tests. + """ + + def allocate(self, request: gtx_auto_optimize.AllocationRequest): + raise AssertionError("backend-wiring tests must not call allocate") + + def deallocate(self, buffer) -> None: + raise AssertionError("backend-wiring tests must not call deallocate") + + +def test_make_backend_accepts_external_allocator_with_external_mode(): + external_memory_allocator = _RecordingAllocator() + + 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_memory_allocator=external_memory_allocator, + ) + + assert backend.executor.compilation.external_memory_allocator is external_memory_allocator + + +def test_make_backend_infers_external_mode_when_allocator_is_provided(): + external_memory_allocator = _RecordingAllocator() + + backend = dace_wf_backend.make_dace_backend( + gpu=False, + auto_optimize=True, + async_sdfg_call=False, + external_memory_allocator=external_memory_allocator, + ) + + assert ( + backend.executor.translation.step.auto_optimize_args["transient_memory_mode"] + == gtx_transformations.TransientMemoryMode.EXTERNAL + ) + assert backend.executor.compilation.external_memory_allocator is external_memory_allocator + + +def test_make_backend_warns_external_allocator_without_external_mode(): + external_memory_allocator = _RecordingAllocator() + + with pytest.warns(UserWarning, match="External memory allocator 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_memory_allocator=external_memory_allocator, + ) + + # 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.executor.compilation.external_memory_allocator is external_memory_allocator + + +class _WorkspaceRecordingAllocator: + """Minimal picklable `ExternalMemoryAllocator` that records every request. + + Allocations are recorded as ``(nbytes, device)`` tuples in ``requests``; + ``deallocate`` is a no-op. Defined at module scope so the allocator can + be pickled when compilation is dispatched to a worker process. + """ + + def __init__(self) -> None: + self.requests: list[tuple[int, core_defs.DeviceType]] = [] + + def allocate(self, request: gtx_auto_optimize.AllocationRequest): + # Overallocate by `request.alignment - 1` bytes and slices forward to the + # nearest aligned boundary, using `request.alignment` directly. This makes + # the returned buffer deterministically aligned to the requested value + # (256 by default) for any workspace size — both host (`__array_interface__`) + # and device (`__cuda_array_interface__`) paths. + self.requests.append((request.nbytes, request.device)) + if request.device == core_defs.CUPY_DEVICE_TYPE: + import cupy as cp + + raw = cp.empty(request.nbytes + request.alignment - 1, dtype=cp.uint8) + offset = (-raw.__cuda_array_interface__["data"][0]) % request.alignment + return raw[offset : offset + request.nbytes] + + raw = np.empty(request.nbytes + request.alignment - 1, dtype=np.uint8) + offset = (-raw.__array_interface__["data"][0]) % request.alignment + return raw[offset : offset + request.nbytes] + + def deallocate(self, buffer) -> None: + pass + + +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): + 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_memory_allocator = _WorkspaceRecordingAllocator() + workspace_requests = external_memory_allocator.requests + + 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_memory_allocator=external_memory_allocator, + ) + + @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 + gt_generate_sdfg = gtx_dace_translation.DaCeTranslator.generate_sdfg + + def mocked_generate_sdfg(self, *args, **kwargs) -> dace.SDFG: + nonlocal captured_sdfg + result = gt_generate_sdfg(self, *args, **kwargs) + captured_sdfg = result + return result + + def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: + return sdfg + + monkeypatch.setattr(gtx_dace_translation.DaCeTranslator, "generate_sdfg", mocked_generate_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 + ) + + # ``_WorkspaceRecordingAllocator`` is picklable (a module-level class), + # so compilation would otherwise be dispatched to a worker process where + # the ``DaCeTranslator.generate_sdfg`` monkeypatch above does not apply. + # Force in-process compilation so the patched translator is observed. + with mock.patch.object(config, "BUILD_JOBS_MODE", config.BuildJobsMode.SERIAL): + testee.with_backend(custom_backend)(a, b, out=out, offset_provider={}) + + 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 + ) + # 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 allocator, 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) + ) + expected_device = core_defs.CUPY_DEVICE_TYPE + 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")) + expected_device = core_defs.DeviceType.CPU + + assert workspace_requests + assert all(device == expected_device for _, device in workspace_requests) + + 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 + assert not workspace_requests + 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 + assert not workspace_requests + 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 f85e2bdea6..d5fcb63375 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 @@ -13,12 +13,16 @@ """ import contextlib +import dataclasses import pathlib import pickle import unittest.mock as mock +from typing import Any +import numpy as np import pytest + dace = pytest.importorskip("dace") from dace.sdfg import nodes as dace_nodes @@ -27,6 +31,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.transformations import ( + auto_optimize as gtx_auto_optimize, +) from gt4py.next.program_processors.runners.dace.workflow import compilation as dace_wf_compilation @@ -230,3 +237,345 @@ def test_cmake_build_type_changes_build_folder(): artifact_debug, _, _ = _run_compiler(cmake_build_type=config.CMakeBuildType.DEBUG) assert artifact_release.library_path != artifact_debug.library_path + + +def _make_compiled_program( + *, + external_memory_allocator=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 = ((), ()) + + return 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", + external_memory_allocator=external_memory_allocator, + ) + + +def test_construct_arguments_installs_external_workspaces_once(): + allocator = mock.MagicMock() + allocator.allocate.side_effect = [_make_array_buffer(nbytes=128, address=256)] + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + program.construct_arguments(alpha=1) + program.construct_arguments(alpha=2) + + # Workspace configuration is done exactly once and reused afterwards. + assert program.sdfg_program.initialize.call_count == 1 + assert program.sdfg_program.get_workspace_sizes.call_count == 1 + assert allocator.allocate.call_count == 1 + allocate_request = allocator.allocate.call_args.args[0] + assert isinstance(allocate_request, gtx_auto_optimize.AllocationRequest) + assert allocate_request.nbytes == 128 + assert allocate_request.device == core_defs.DeviceType.CPU + assert program.sdfg_program.set_workspace.call_count == 1 + assert program.sdfg_program.construct_arguments.call_count == 2 + + set_workspace_call = program.sdfg_program.set_workspace.call_args + assert set_workspace_call.args[0] == dace.StorageType.CPU_Heap + configured_workspace = set_workspace_call.args[1] + assert program.external_workspaces[dace.StorageType.CPU_Heap] is configured_workspace + + +def test_construct_arguments_propagates_allocator_error_for_invalid_size_request(): + allocator = mock.MagicMock() + allocator.allocate.side_effect = ValueError("invalid workspace size request") + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: -1}, + ) + + with pytest.raises(ValueError, match="invalid workspace size request"): + program.construct_arguments(alpha=1) + + allocator.allocate.assert_called_once() + assert allocator.allocate.call_args.args[0].nbytes == -1 + assert allocator.allocate.call_args.args[0].device == core_defs.DeviceType.CPU + program.sdfg_program.set_workspace.assert_not_called() + + +def test_construct_arguments_propagates_allocator_error_for_invalid_storage_request(): + allocator = mock.MagicMock() + allocator.allocate.side_effect = TypeError("invalid storage type request") + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 16}, + ) + + with pytest.raises(TypeError, match="invalid storage type request"): + program.construct_arguments(alpha=1) + + allocator.allocate.assert_called_once() + assert allocator.allocate.call_args.args[0].nbytes == 16 + assert allocator.allocate.call_args.args[0].device == core_defs.DeviceType.CPU + program.sdfg_program.set_workspace.assert_not_called() + + +def _make_array_buffer(*, nbytes: int, address: int, cuda: bool = False) -> mock.MagicMock: + """A minimal array-like buffer with a configurable base pointer. + + Exposes ``__array_interface__`` (host) or ``__cuda_array_interface__`` + (device) so it is accepted by ``_validate_external_workspace``; the + ``data`` tuple carries the address that DaCe's ``array_interface_ptr`` + would hand to the SDFG. + """ + buffer = mock.MagicMock() + buffer.nbytes = nbytes + interface = {"data": (address, False), "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 allocator must return something ``set_workspace`` can consume.""" + allocator = mock.MagicMock() + allocator.allocate.side_effect = ["not-an-array"] # str exposes no array interface + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 64}, + ) + + with pytest.raises(TypeError, match="does not expose `__array_interface__`"): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_not_called() + + +def test_construct_arguments_rejects_misaligned_buffer(): + """A host buffer whose base pointer is not aligned is rejected.""" + allocator = mock.MagicMock() + allocator.allocate.side_effect = [_make_array_buffer(nbytes=128, address=100)] # 100 % 256 + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + with pytest.raises(ValueError, match="not aligned to the required 256 bytes"): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_not_called() + + +def test_construct_arguments_accepts_aligned_buffer(): + """A host buffer whose base pointer is aligned is accepted.""" + allocator = mock.MagicMock() + workspace = _make_array_buffer(nbytes=128, address=1024) # 1024 % 256 == 0 + allocator.allocate.side_effect = [workspace] + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_called_once() + assert program.external_workspaces[dace.StorageType.CPU_Heap] is workspace + + +def test_construct_arguments_rejects_misaligned_gpu_buffer(): + """A device buffer whose base pointer is not aligned is rejected.""" + allocator = mock.MagicMock() + allocator.allocate.side_effect = [_make_array_buffer(nbytes=128, address=100, cuda=True)] + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.GPU_Global: 128}, + ) + + with ( + mock.patch.object( + dace_wf_compilation.core_defs, "CUPY_DEVICE_TYPE", core_defs.DeviceType.CUDA + ), + pytest.raises(ValueError, match="not aligned to the required 256 bytes"), + ): + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_not_called() + + +def test_construct_arguments_skips_alignment_when_data_missing(): + """When the array interface omits ``data`` alignment is trust-based.""" + allocator = mock.MagicMock() + buffer = mock.MagicMock() + buffer.nbytes = 64 + # ``data`` is optional on the host array interface. + buffer.__array_interface__ = {"shape": (64,), "typestr": "|u1", "version": 3} + allocator.allocate.side_effect = [buffer] + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 64}, + ) + + # Must not raise even though alignment can't be verified. + program.construct_arguments(alpha=1) + + program.sdfg_program.set_workspace.assert_called_once() + + +def test_finalize_calls_deallocate_once_per_storage(): + """``finalize()`` releases each workspace exactly once and is idempotent.""" + allocator = mock.MagicMock() + workspace = _make_array_buffer(nbytes=128, address=256) + allocator.allocate.side_effect = [workspace] + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={dace.StorageType.CPU_Heap: 128}, + ) + + program.construct_arguments(alpha=1) + + program.finalize() + assert allocator.deallocate.call_count == 1 + assert allocator.deallocate.call_args.args[0] is workspace + assert program.external_workspaces == {} + # finalize() is idempotent. + program.finalize() + assert allocator.deallocate.call_count == 1 + + +def test_finalize_continues_and_is_idempotent_when_deallocate_fails(): + """If one ``deallocate`` raises, the remaining workspaces are still released. + + A failing buffer must not prevent the others from being deallocated, and + ``external_workspaces`` must still be cleared so a subsequent ``finalize()`` + (e.g. the pool finalizer) is a no-op rather than re-deallocating the + buffers that already succeeded. + """ + allocator = mock.MagicMock() + workspace_a = _make_array_buffer(nbytes=16, address=256) # aligned for CPU + workspace_b = _make_array_buffer(nbytes=32, address=512, cuda=True) # aligned for GPU + allocator.allocate.side_effect = [workspace_a, workspace_b] + program = _make_compiled_program( + external_memory_allocator=allocator, + workspace_sizes={ + dace.StorageType.CPU_Heap: 16, + dace.StorageType.GPU_Global: 32, + }, + ) + with mock.patch.object( + dace_wf_compilation.core_defs, "CUPY_DEVICE_TYPE", core_defs.DeviceType.CUDA + ): + program.construct_arguments(alpha=1) + # The first deallocate raises; the second must still be called. + allocator.deallocate.side_effect = [RuntimeError("boom"), None] + + with pytest.warns(UserWarning, match="Failed to deallocate"): + program.finalize() + + assert allocator.deallocate.call_count == 2 + assert program.external_workspaces == {} + # finalize() is idempotent even after partial failures. + program.finalize() + assert allocator.deallocate.call_count == 2 + + +def test_finalize_with_no_allocator_is_a_safe_noop(): + """A ``None`` allocator performs no work but still clears workspaces.""" + program = _make_compiled_program(external_memory_allocator=None) + + # finalize() must not raise even though no allocator is configured. + program.finalize() + assert program.external_workspaces == {} + + +# --- Phase 5: allocator pickleability ------------------------------------- +# +# ``DaCeCompiler`` is the step that gets pickled when the OTF runner offloads +# compilation to a ``ProcessPoolExecutor`` (``otf/runners.py``), and it carries +# the ``external_memory_allocator``. A non-picklable allocator (closure, +# lambda, local class) must fail fast at construction with +# ``AllocatorNotPicklableError`` rather than silently degrading to in-process +# compilation via a generic runner warning. + + +@dataclasses.dataclass(frozen=True) +class _ModuleLevelPicklableAllocator: + """A picklable allocator defined at module scope. + + ``allocate``/``deallocate`` are never called by the tests below; only the + type's picklability and identity through a round-trip matter. Defined at + module scope (not inside a test) so ``pickle`` can locate it by qualname. + Frozen with no fields so two instances are structurally equal, mirroring + a stateless allocator and the frozenness of ``DaCeCompilationArtifact``. + """ + + def allocate(self, request: gtx_auto_optimize.AllocationRequest): + raise AssertionError("pickleability tests must not call allocate") + + def deallocate(self, buffer) -> None: + raise AssertionError("pickleability tests must not call deallocate") + + +def _make_compiler(allocator=None) -> dace_wf_compilation.DaCeCompiler: + return dace_wf_compilation.DaCeCompiler( + bind_func_name="bind", + cache_lifetime=config.BuildCacheLifetime.SESSION, + device_type=core_defs.DeviceType.CPU, + external_memory_allocator=allocator, + ) + + +def test_dace_compiler_rejects_non_picklable_allocator(): + """An allocator that can not be pickled fails fast at construction.""" + + class _LocalAllocator: # local class -> not picklable by qualname + def allocate(self, request): ... + + def deallocate(self, buffer) -> None: ... + + with pytest.raises( + dace_wf_compilation.AllocatorNotPicklableError, + match="external_memory_allocator .* is not picklable", + ) as excinfo: + _make_compiler(allocator=_LocalAllocator()) + + # The original pickle error is chained so the user can see *why*. + assert isinstance(excinfo.value.__cause__, Exception) + assert "Can't pickle" in str(excinfo.value.__cause__) + + +def test_dace_compiler_accepts_picklable_allocator(): + """A module-level allocator (and the ``None`` default) pass the gate.""" + # ``None`` default: no probe, no raise. + _make_compiler(allocator=None) + + # Module-level class: picklable, no raise. + _make_compiler(allocator=_ModuleLevelPicklableAllocator()) + + +def test_dace_compilation_artifact_pickle_round_trip_with_allocator(tmp_path: pathlib.Path): + """The allocator round-trips through the pickled compilation artifact. + + The existing ``test_dace_compilation_artifact_pickle_round_trip`` covers the + no-allocator default; this ensures a real allocator is carried through + serialization with identity of intent preserved (structural equality, + since the allocator class defines no per-instance state). + """ + allocator = _ModuleLevelPicklableAllocator() + artifact = dace_wf_compilation.DaCeCompilationArtifact( + library_path=tmp_path / "build" / "libprogram.so", + sdfg_json="{}", + binding_source_code="def update_sdfg_args(*a, **k): ...", + bind_func_name="update_sdfg_args", + device_type=core_defs.DeviceType.CPU, + external_memory_allocator=allocator, + ) + + restored = pickle.loads(pickle.dumps(artifact)) + + assert restored == artifact + assert isinstance(restored.external_memory_allocator, _ModuleLevelPicklableAllocator) 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