From 61fbd262576bce4d6eabb64d4ddd3e9dbfdcb661 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Thu, 23 Jul 2026 14:30:00 +0200 Subject: [PATCH 01/25] feat[next]: External memory - step 1 (#2715) This pull request introduces a more flexible and extensible way to control the allocation and lifetime of transient arrays in the DaCe backend, replacing the previous `make_persistent` and `gpu_memory_pool` options with a new `TransientMemoryMode` policy. It also adds support for externally managed memory via an `external_memory_allocator` parameter, and refactors related utility functions to support these changes. **API and Configuration Improvements:** * Introduced the `TransientMemoryMode` enum to specify transient array lifetime/allocation strategies (`SCOPED`, `PERSISTENT`, `POOL`, `EXTERNAL`), and replaced the `make_persistent` and `gpu_memory_pool` parameters throughout the codebase with this unified approach. [[1]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706R114-R127) [[2]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L133) [[3]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L177-R185) [[4]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L197) [[5]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L403-L409) [[6]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L921-L923) [[7]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L942-R963) [[8]](diffhunk://#diff-c7967b0fa16e0de89cb43f3c81e9aee884122772946f031d50fb31d1cb16215bR20) [[9]](diffhunk://#diff-c7967b0fa16e0de89cb43f3c81e9aee884122772946f031d50fb31d1cb16215bR123) * Added support for an `external_memory_allocator` parameter in the DaCe backend and workflow, which allows users to provide custom allocation logic for external memory. This is validated to only be used with `transient_memory_mode=external`. [[1]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbR40) [[2]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbR52) [[3]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbR67) [[4]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbR82-R83) [[5]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbR120-R138) [[6]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR184) [[7]](diffhunk://#diff-91264fda895ecc56c869858ac9724e1ff60ed6514a121d1cd4f323197053e5aaR38) [[8]](diffhunk://#diff-91264fda895ecc56c869858ac9724e1ff60ed6514a121d1cd4f323197053e5aaR77) **Utility Function Refactoring:** * Refactored the persistent transients utility into a more general `_gt_configure_transient_lifetime` function, which now supports both `Persistent` and `External` lifetimes. Exposed two public functions: `gt_make_transients_persistent` and the new `gt_make_transients_external`. [[1]](diffhunk://#diff-6ce6c03a192e5b5b7e735ecc3aa1071dc7ce0ac5d90f206f948b1e3f8e952637L29-R39) [[2]](diffhunk://#diff-6ce6c03a192e5b5b7e735ecc3aa1071dc7ce0ac5d90f206f948b1e3f8e952637L84-R70) [[3]](diffhunk://#diff-c7967b0fa16e0de89cb43f3c81e9aee884122772946f031d50fb31d1cb16215bL87-R88) [[4]](diffhunk://#diff-c7967b0fa16e0de89cb43f3c81e9aee884122772946f031d50fb31d1cb16215bR137) **Documentation and Safety:** * Updated docstrings and comments to reflect the new transient memory modes and clarify the behavior and constraints of each mode, including safety considerations when using global-lifetime transients. [[1]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L177-R185) [[2]](diffhunk://#diff-6ce6c03a192e5b5b7e735ecc3aa1071dc7ce0ac5d90f206f948b1e3f8e952637L84-R70) These changes make the backend's memory management more robust and adaptable, and lay the groundwork for advanced use cases such as externally managed memory. --- .../runners/dace/transformations/__init__.py | 6 +- .../dace/transformations/auto_optimize.py | 94 +++++++++++++------ .../runners/dace/transformations/utils.py | 60 ++++++------ .../runners/dace/workflow/backend.py | 23 +++++ .../runners/dace/workflow/compilation.py | 1 + .../runners/dace/workflow/factory.py | 2 + .../dace_tests/test_dace_backend.py | 67 +++++++++++-- .../test_transient_memory_mode.py | 67 +++++++++++++ 8 files changed, 251 insertions(+), 69 deletions(-) create mode 100644 tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_transient_memory_mode.py 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..6fe296ffa7 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 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 a callable that takes + `(required_nbytes, device_type)` and returns the allocated memory, in the + form of an array object that can handled by `dace.dtypes.array_interface_ptr()`. + The callable is expected to raise an exception if the allocation fails. + """ + + 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..280df88e7b 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -9,6 +9,7 @@ from __future__ import annotations import warnings +from collections.abc import Callable from typing import Any, Final import factory @@ -16,6 +17,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 +37,7 @@ class Meta: class Params: name_device = "cpu" name_postfix = "" + external_memory_allocator = None gpu = factory.Trait( allocator=next_allocators.StandardGPUFieldBufferAllocator(), device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, @@ -46,6 +49,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 +64,7 @@ def make_dace_backend( auto_optimize: bool = True, async_sdfg_call: bool = True, optimization_args: dict[str, Any] | None = None, + external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | 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 +79,9 @@ 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: Callable taking `(required_nbytes, storage_type)` + used later for external-memory workspace allocation. Threaded through + the backend workflow for now. 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 +118,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 7b78942090..f6e0e207a7 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -211,6 +211,7 @@ class DaCeCompiler( bind_func_name: str cache_lifetime: config.BuildCacheLifetime device_type: core_defs.DeviceType + external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None = None add_gpu_trace_markers: bool = dataclasses.field( default_factory=lambda: config.ADD_GPU_TRACE_MARKERS ) 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..c0813c012b 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -35,6 +35,7 @@ class Meta: class Params: auto_optimize: bool = False + external_memory_allocator = 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 +74,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/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..183b07c269 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 @@ -55,20 +55,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 +122,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 +135,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 +147,58 @@ 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 test_make_backend_accepts_external_allocator_with_external_mode(): + external_memory_allocator = lambda size, storage: bytearray(size) + + 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 = lambda size, storage: bytearray(size) + + 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 = lambda size, storage: bytearray(size) + + 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 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 From 449514adbfed7f487f9dba843ba0ddb0fc4d2655 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 24 Jul 2026 15:09:33 +0200 Subject: [PATCH 02/25] feat[next-dace]: External memory - step 2 (#2717) This pull request introduces support for external workspace memory allocation in the DaCe backend, allowing users to provide custom memory allocators for transient workspaces. It updates the compilation workflow to use an external allocator when configured, ensures workspaces are only allocated once, and adds comprehensive tests for the new functionality. The changes also improve device type handling and test coverage for transient memory modes. **External workspace allocator integration:** * Added an `external_memory_allocator` parameter to `CompiledDaceProgram` and `DaCeCompilationArtifact`, allowing external allocation of workspace memory buffers. Workspaces are now tracked and only allocated once per compiled program instance. (`src/gt4py/next/program_processors/runners/dace/workflow/compilation.py` [[1]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR88-R96) [[2]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR116-R137) [[3]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR194-R207) [[4]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR299) * Implemented the `_map_workspace_storage_to_device` helper to map DaCe storage types to device types for allocation. (`src/gt4py/next/program_processors/runners/dace/workflow/compilation.py` [src/gt4py/next/program_processors/runners/dace/workflow/compilation.pyR50-R61](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR50-R61)) * Ensured external workspaces are configured during argument construction and only set up once. (`src/gt4py/next/program_processors/runners/dace/workflow/compilation.py` [src/gt4py/next/program_processors/runners/dace/workflow/compilation.pyR146](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR146)) **Testing and validation:** * Added tests to verify that external workspace allocation works as intended, including correct device routing, error propagation, and single allocation per program instance. (`tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py` [tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.pyR235-R311](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3R235-R311)) * Extended backend tests to check that different transient memory modes (EXTERNAL, POOL, PERSISTENT, SCOPED) use the correct allocation APIs and honor the external allocator, with device-specific checks for CPU and GPU. (`tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py` [tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.pyR213-R378](diffhunk://#diff-19fed38498781ceef82337d097e43458274d2ebe7e0bae4578839e9fbed15194R213-R378)) **Device type handling improvements:** * Updated device type selection in tests to use `core_defs.CUPY_DEVICE_TYPE` for GPU tests, ensuring compatibility with the available GPU backend. (`tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py` [tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.pyR26-R43](diffhunk://#diff-19fed38498781ceef82337d097e43458274d2ebe7e0bae4578839e9fbed15194R26-R43)) These changes make the DaCe backend more flexible and robust by supporting custom workspace allocation, improving test coverage, and ensuring correct behavior across memory modes and device types. --- .../runners/dace/workflow/compilation.py | 47 +++- .../dace_tests/test_dace_backend.py | 203 +++++++++++++++++- .../dace_tests/test_dace_compilation.py | 79 +++++++ 3 files changed, 326 insertions(+), 3 deletions(-) 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 f6e0e207a7..15a659695c 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -70,6 +70,18 @@ def _add_tx_markers(program_source: SDFGExtensionSource) -> tuple[SDFGExtensionS return new_program_source, sdfg +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 + + class CompiledDaceProgram: sdfg_program: dace.CompiledSDFG @@ -96,12 +108,15 @@ class CompiledDaceProgram: # never updated. csdfg_argv: MutableSequence[Any] | None csdfg_init_argv: Sequence[Any] | None + external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None + external_workspaces: dict[dace.StorageType, Any] def __init__( self, program: dace.CompiledSDFG, bind_func_name: str, binding_source_code: str, + external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None = None, ): self.sdfg_program = program @@ -121,6 +136,28 @@ 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) + workspace = self.external_memory_allocator(required_nbytes, device) + 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 construct_arguments(self, **kwargs: Any) -> None: """ @@ -129,6 +166,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] @@ -183,6 +221,7 @@ class DaCeCompilationArtifact: binding_source_code: str bind_func_name: str device_type: core_defs.DeviceType + external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None = None def load(self) -> stages.ExecutableProgram: # TODO(phimuell): Drop ``sdfg_json`` from the artifact once dace @@ -190,7 +229,12 @@ 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) + program = CompiledDaceProgram( + sdfg_program, + self.bind_func_name, + self.binding_source_code, + external_memory_allocator=self.external_memory_allocator, + ) return gtx_wfddecoration.convert_args(program, device=self.device_type) @@ -281,6 +325,7 @@ def __call__(self, inp: SDFGExtensionSource) -> DaCeCompilationArtifact: 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/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 183b07c269..e5388b3b1b 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,6 +8,9 @@ """Test the bindings stage of the dace backend workflow.""" +import re + +import numpy as np import pytest import unittest.mock as mock @@ -16,12 +19,19 @@ from gt4py import next as gtx from gt4py._core import definitions as core_defs 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.workflow import ( backend as dace_wf_backend, ) -from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations +from gt4py.next.program_processors.runners.dace.workflow import ( + translation as gtx_dace_translation, +) +from gt4py.next.program_processors.runners.dace.transformations import ( + auto_optimize as gtx_auto_optimize, +) from next_tests.integration_tests import cases +from next_tests.integration_tests import cases_utils from next_tests.integration_tests.cases_utils import KDim @@ -32,7 +42,9 @@ ], 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 @@ -202,3 +214,190 @@ def test_make_backend_warns_external_allocator_without_external_mode(): == gtx_transformations.TransientMemoryMode.POOL ) assert backend.executor.compilation.external_memory_allocator is external_memory_allocator + + +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 = "" + for code in sdfg.generate_code(): + if code.language == "cu": + 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" + elif not code.name.endswith("_main"): + 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(" + workspace_requests: list[tuple[int, core_defs.DeviceType]] = [] + + def external_memory_allocator(required_nbytes: int, device: core_defs.DeviceType): + workspace_requests.append((required_nbytes, device)) + if device == core_defs.CUPY_DEVICE_TYPE: + import cupy as cp + + return cp.empty((required_nbytes,), dtype=cp.uint8) + return np.empty((required_nbytes,), dtype=np.uint8) + + @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) + + 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 + ) + + 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, + ) + + 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")() + + program = ( + testee.with_grid_type(gtx.common.GridType.CARTESIAN) + .with_backend(custom_backend) + .compile(offset_provider={}) + ) + gtx.wait_for_compilation() + program(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 610567f6b2..58d74e0e10 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 @@ -16,7 +16,9 @@ import pathlib import pickle import unittest.mock as mock +from typing import Any +import numpy as np import pytest dace = pytest.importorskip("dace") @@ -262,3 +264,80 @@ 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_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(side_effect=[np.empty((128,), dtype=np.uint8)]) + 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.call_count == 1 + allocator.assert_called_once_with(128, 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(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.assert_called_once_with(-1, 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(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.assert_called_once_with(16, core_defs.DeviceType.CPU) + program.sdfg_program.set_workspace.assert_not_called() From 90fc3cda000c2db27cf2a139a3b84a39c4e8605b Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 24 Jul 2026 16:55:27 +0200 Subject: [PATCH 03/25] fix tests --- .../runners_tests/dace_tests/test_dace_backend.py | 9 +++++++-- .../test_make_transients_persistent.py | 4 ++-- 2 files changed, 9 insertions(+), 4 deletions(-) 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 e5388b3b1b..746e6174e3 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 @@ -225,8 +225,13 @@ def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str 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.language == "cu": + 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()): @@ -234,7 +239,7 @@ def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str assert free_re.match(line.strip()) else: generated_code += line + "\n" - elif not code.name.endswith("_main"): + else: generated_code += code.clean_code + "\n" return generated_code 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() From 705ba00c1fae92d75c5635d55b46855a7bd3ebf1 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 28 Jul 2026 10:19:24 +0200 Subject: [PATCH 04/25] refactor[next]: Introduce ExternalMemoryAllocator protocol (#2720) This pull request introduces a new "External Memory Allocator" mode for DaCe transients in the `gt4py.next` backend, allowing explicit, caller-managed workspace allocation and teardown. This enables reusing a single workspace buffer across multiple SDFGs, removing per-call allocation overhead and providing more control over memory lifetime. The public API now exposes a typed allocator protocol, and the backend enforces that allocators are picklable to ensure compatibility with process-based compilation. The changes also include explicit pool-driven teardown of resources and improved error handling for allocator pickling. **External Memory Allocator for DaCe Transients** *Design and API additions:* - Added an explicit `EXTERNAL` mode to `TransientMemoryMode`, backed by a caller-supplied `ExternalMemoryAllocator` protocol. This protocol is defined with typed `allocate` and `deallocate` methods and requires picklability. (`[[1]](diffhunk://#diff-b86bd892b20495d184cedc1d9057e4b1ef0e837d723e998628332943d900aa74R1-R147)`, `[[2]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706R117-R178)`, `[[3]](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706L133-R202)`) - Introduced `AllocationRequest` and `Buffer` type alias to formalize the workspace allocation contract. (`[src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.pyR117-R178](diffhunk://#diff-ddbc0e8ab1d5b89929dfbd12477936ae3acad531911fa9afb125066b6fe74706R117-R178)`) - Re-exported `AllocationRequest`, `Buffer`, and `ExternalMemoryAllocator` from the `transformations` package for public use. (`[[1]](diffhunk://#diff-c7967b0fa16e0de89cb43f3c81e9aee884122772946f031d50fb31d1cb16215bL16-R19)`, `[[2]](diffhunk://#diff-c7967b0fa16e0de89cb43f3c81e9aee884122772946f031d50fb31d1cb16215bR95-R99)`) *Backend and resource management:* - Updated the DaCe backend to accept an `ExternalMemoryAllocator` instance, with improved documentation and type safety. (`[[1]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbL67-R66)`, `[[2]](diffhunk://#diff-04258cb5c4aa36c40ab7da9246c27d69866a707135824bba4f76d46e93de72fbL82-R85)`) - Added a check to ensure the allocator is picklable at backend construction, raising a clear error if not. (`[src/gt4py/next/program_processors/runners/dace/workflow/compilation.pyR43-R75](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dR43-R75)`) - Implemented explicit resource teardown: when a `CompiledProgramsPool` is deleted, any compiled programs with a `finalize()` method have it called to release external resources, with warnings on failure. (`[[1]](diffhunk://#diff-cae915aa5b7a62aa4a6eb295978c894fc43270cf6e08e82fca56414bf37d7cdeR88-R113)`, `[[2]](diffhunk://#diff-cae915aa5b7a62aa4a6eb295978c894fc43270cf6e08e82fca56414bf37d7cdeR404-R417)`) *Documentation:* - Added ADR 0026 documenting the rationale, design, and consequences of the external memory allocator feature. (`[docs/development/ADRs/next/0026-External_Memory_Allocator.mdR1-R147](diffhunk://#diff-b86bd892b20495d184cedc1d9057e4b1ef0e837d723e998628332943d900aa74R1-R147)`) These changes provide a robust, explicit, and user-extensible mechanism for managing transient workspace memory in DaCe-based workflows, improving performance and resource control for advanced use cases. --- .../next/0026-External_Memory_Allocator.md | 147 +++++++++ src/gt4py/next/otf/compiled_program.py | 40 +++ .../runners/dace/transformations/__init__.py | 8 +- .../dace/transformations/auto_optimize.py | 74 ++++- .../runners/dace/workflow/backend.py | 11 +- .../runners/dace/workflow/compilation.py | 140 ++++++++- .../runners/dace/workflow/decoration.py | 71 +++-- .../otf_tests/test_compiled_program.py | 69 +++++ .../dace_tests/test_dace_backend.py | 86 ++++-- .../dace_tests/test_dace_compilation.py | 284 +++++++++++++++++- 10 files changed, 860 insertions(+), 70 deletions(-) create mode 100644 docs/development/ADRs/next/0026-External_Memory_Allocator.md 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 f8cb38334f..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,7 +13,10 @@ """ 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, @@ -89,8 +92,11 @@ __all__ = [ + "AllocationRequest", "CopyChainRemover", "DoubleWriteRemover", + "ExternalMemoryAllocator", + "ExternalWorkspace", "FuseHorizontalConditionBlocks", "GPUSetBlockSize", "GT4PyAutoOptHook", 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 6fe296ffa7..aebea60d76 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,10 @@ """Fast access to the auto optimization on DaCe.""" +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, runtime_checkable import dace from dace import data as dace_data @@ -19,6 +20,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,6 +114,66 @@ 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 + + +@runtime_checkable +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. + """ + + 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``. + """ + ... + + 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. @@ -130,10 +193,11 @@ class TransientMemoryMode(str, enum.Enum): 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 a callable that takes - `(required_nbytes, device_type)` and returns the allocated memory, in the - form of an array object that can handled by `dace.dtypes.array_interface_ptr()`. - The callable is expected to raise an exception if the allocation fails. + 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" 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 280df88e7b..765d84ea70 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -9,7 +9,6 @@ from __future__ import annotations import warnings -from collections.abc import Callable from typing import Any, Final import factory @@ -64,7 +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: Callable[[int, core_defs.DeviceType], 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, @@ -79,9 +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: Callable taking `(required_nbytes, storage_type)` - used later for external-memory workspace allocation. Threaded through - the backend workflow for now. + 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 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 15a659695c..1df7f81f62 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, TypeAlias @@ -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, @@ -82,6 +89,52 @@ def _map_workspace_storage_to_device(storage: dace.StorageType) -> core_defs.Dev 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 @@ -108,15 +161,15 @@ class CompiledDaceProgram: # never updated. csdfg_argv: MutableSequence[Any] | None csdfg_init_argv: Sequence[Any] | None - external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None - external_workspaces: dict[dace.StorageType, Any] + 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: Callable[[int, core_defs.DeviceType], Any] | None = None, + external_memory_allocator: ExternalMemoryAllocator | None = None, ): self.sdfg_program = program @@ -154,11 +207,40 @@ def _configure_external_workspaces(self, **kwargs: Any) -> None: ) for storage, required_nbytes in workspace_sizes.items(): device = _map_workspace_storage_to_device(storage) - workspace = self.external_memory_allocator(required_nbytes, device) + 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. + """ + self.sdfg_program.finalize() + 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: """ This function will process the arguments and store the processed argument @@ -221,7 +303,7 @@ class DaCeCompilationArtifact: binding_source_code: str bind_func_name: str device_type: core_defs.DeviceType - external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None = None + external_memory_allocator: ExternalMemoryAllocator | None = None def load(self) -> stages.ExecutableProgram: # TODO(phimuell): Drop ``sdfg_json`` from the artifact once dace @@ -235,7 +317,40 @@ def load(self) -> stages.ExecutableProgram: self.binding_source_code, external_memory_allocator=self.external_memory_allocator, ) - return gtx_wfddecoration.convert_args(program, device=self.device_type) + return gtx_wfddecoration.DaCeDecoratedProgram(program, device_type=self.device_type) + + +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!s}." + " 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 @dataclasses.dataclass(frozen=True) @@ -255,7 +370,11 @@ class DaCeCompiler( bind_func_name: str cache_lifetime: config.BuildCacheLifetime device_type: core_defs.DeviceType - external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None = None + #: 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 ) @@ -266,6 +385,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, 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/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 746e6174e3..202cfd78ad 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 @@ -9,29 +9,28 @@ """Test the bindings stage of the dace backend workflow.""" import re +import unittest.mock as mock import numpy as np import pytest -import unittest.mock as mock + 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.workflow import ( - backend as dace_wf_backend, +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.transformations import ( - auto_optimize as gtx_auto_optimize, -) -from next_tests.integration_tests import cases -from next_tests.integration_tests import cases_utils +from next_tests.integration_tests import cases, cases_utils from next_tests.integration_tests.cases_utils import KDim @@ -161,8 +160,23 @@ def mocked_gpu_transformation(*args, **kwargs) -> dace.SDFG: 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 = lambda size, storage: bytearray(size) + external_memory_allocator = _RecordingAllocator() backend = dace_wf_backend.make_dace_backend( gpu=False, @@ -178,7 +192,7 @@ def test_make_backend_accepts_external_allocator_with_external_mode(): def test_make_backend_infers_external_mode_when_allocator_is_provided(): - external_memory_allocator = lambda size, storage: bytearray(size) + external_memory_allocator = _RecordingAllocator() backend = dace_wf_backend.make_dace_backend( gpu=False, @@ -195,7 +209,7 @@ def test_make_backend_infers_external_mode_when_allocator_is_provided(): def test_make_backend_warns_external_allocator_without_external_mode(): - external_memory_allocator = lambda size, storage: bytearray(size) + external_memory_allocator = _RecordingAllocator() with pytest.warns(UserWarning, match="External memory allocator provided"): backend = dace_wf_backend.make_dace_backend( @@ -216,6 +230,39 @@ def test_make_backend_warns_external_allocator_without_external_mode(): 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. @@ -253,15 +300,14 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): 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(" - workspace_requests: list[tuple[int, core_defs.DeviceType]] = [] - - def external_memory_allocator(required_nbytes: int, device: core_defs.DeviceType): - workspace_requests.append((required_nbytes, device)) - if device == core_defs.CUPY_DEVICE_TYPE: - import cupy as cp - - return cp.empty((required_nbytes,), dtype=cp.uint8) - return np.empty((required_nbytes,), dtype=np.uint8) + external_memory_allocator = _WorkspaceRecordingAllocator() + workspace_requests = external_memory_allocator.requests + + # ``_WorkspaceRecordingAllocator`` is picklable (a module-level class), + # so compilation would otherwise be dispatched to a worker process where + # the ``DaCeTranslator.generate_sdfg`` monkeypatch below does not apply. + # Force in-process compilation so the patched translator is observed. + monkeypatch.setattr(config, "BUILD_JOBS_MODE", config.BuildJobsMode.SERIAL) @gtx.field_operator def testee_op(a: cases.IField, b: cases.IField) -> cases.IField: 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 58d74e0e10..43bd747ddd 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,6 +13,7 @@ """ import contextlib +import dataclasses import pathlib import pickle import unittest.mock as mock @@ -21,6 +22,7 @@ import numpy as np import pytest + dace = pytest.importorskip("dace") from dace.sdfg import nodes as dace_nodes @@ -29,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 @@ -292,7 +297,8 @@ def _make_compiled_program( def test_construct_arguments_installs_external_workspaces_once(): - allocator = mock.MagicMock(side_effect=[np.empty((128,), dtype=np.uint8)]) + 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}, @@ -304,8 +310,11 @@ def test_construct_arguments_installs_external_workspaces_once(): # 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.call_count == 1 - allocator.assert_called_once_with(128, core_defs.DeviceType.CPU) + 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 @@ -316,7 +325,8 @@ def test_construct_arguments_installs_external_workspaces_once(): def test_construct_arguments_propagates_allocator_error_for_invalid_size_request(): - allocator = mock.MagicMock(side_effect=ValueError("invalid workspace 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}, @@ -325,12 +335,15 @@ def test_construct_arguments_propagates_allocator_error_for_invalid_size_request with pytest.raises(ValueError, match="invalid workspace size request"): program.construct_arguments(alpha=1) - allocator.assert_called_once_with(-1, core_defs.DeviceType.CPU) + 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(side_effect=TypeError("invalid storage type 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}, @@ -339,5 +352,262 @@ def test_construct_arguments_propagates_allocator_error_for_invalid_storage_requ with pytest.raises(TypeError, match="invalid storage type request"): program.construct_arguments(alpha=1) - allocator.assert_called_once_with(16, core_defs.DeviceType.CPU) + 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) From cb3745a49b7e45fc200609d54ab05feec6c5ae2f Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 28 Jul 2026 11:33:42 +0200 Subject: [PATCH 05/25] edit --- .../runners/dace/transformations/auto_optimize.py | 3 +++ .../program_processors/runners/dace/workflow/backend.py | 2 +- .../program_processors/runners/dace/workflow/factory.py | 3 ++- .../runners_tests/dace_tests/test_dace_backend.py | 8 +------- 4 files changed, 7 insertions(+), 9 deletions(-) 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 aebea60d76..5f9a9a2a44 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,6 +8,7 @@ """Fast access to the auto optimization on DaCe.""" +import abc import dataclasses import enum import warnings @@ -161,12 +162,14 @@ class ExternalMemoryAllocator(Protocol): 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. 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 765d84ea70..76f50e55b9 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -36,7 +36,7 @@ class Meta: class Params: name_device = "cpu" name_postfix = "" - external_memory_allocator = None + 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, 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 c0813c012b..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,7 +36,7 @@ class Meta: class Params: auto_optimize: bool = False - external_memory_allocator = None + 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 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 202cfd78ad..5222df81b6 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 @@ -356,13 +356,7 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: b = cases.allocate(test_case, testee, "b", strategy=cases.UniqueInitializer())() out = cases.allocate(test_case, testee, "out")() - program = ( - testee.with_grid_type(gtx.common.GridType.CARTESIAN) - .with_backend(custom_backend) - .compile(offset_provider={}) - ) - gtx.wait_for_compilation() - program(a, b, out=out, offset_provider={}) + testee.with_backend(custom_backend)(a, b, out=out, offset_provider={}) assert captured_sdfg is not None transient_arrays = [ From 2ea5ef333cf151135ac8c238c944af1b89367fc4 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 28 Jul 2026 12:50:35 +0200 Subject: [PATCH 06/25] edit --- .../next/program_processors/runners/dace/workflow/compilation.py | 1 - 1 file changed, 1 deletion(-) 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 1df7f81f62..8481657ab3 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -228,7 +228,6 @@ def finalize(self) -> None: raised, so that one failing buffer does not prevent the remaining workspaces from being released. """ - self.sdfg_program.finalize() if self.external_memory_allocator is not None: for wsp in self.external_workspaces.values(): try: From d14777ad63938e41e26ea201df6af403e83d1596 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 28 Jul 2026 14:33:45 +0200 Subject: [PATCH 07/25] undo extra change --- .../runners_tests/dace_tests/test_dace_backend.py | 7 ------- 1 file changed, 7 deletions(-) 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 5222df81b6..d1991c899c 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 @@ -19,7 +19,6 @@ 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 ( @@ -303,12 +302,6 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): external_memory_allocator = _WorkspaceRecordingAllocator() workspace_requests = external_memory_allocator.requests - # ``_WorkspaceRecordingAllocator`` is picklable (a module-level class), - # so compilation would otherwise be dispatched to a worker process where - # the ``DaCeTranslator.generate_sdfg`` monkeypatch below does not apply. - # Force in-process compilation so the patched translator is observed. - monkeypatch.setattr(config, "BUILD_JOBS_MODE", config.BuildJobsMode.SERIAL) - @gtx.field_operator def testee_op(a: cases.IField, b: cases.IField) -> cases.IField: tmp = a + b From b4071528ad62ffbbba5a8982e9bcdd78ff5ffa1f Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 28 Jul 2026 15:05:53 +0200 Subject: [PATCH 08/25] edit --- .../dace_tests/test_dace_backend.py | 46 +++++++++++-------- 1 file changed, 26 insertions(+), 20 deletions(-) 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 d1991c899c..f82bfc37ea 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 @@ -19,6 +19,7 @@ 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 ( @@ -302,6 +303,16 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): 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 @@ -311,6 +322,15 @@ def testee_op(a: cases.IField, b: cases.IField) -> cases.IField: 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 @@ -330,26 +350,12 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: no_op_top_level_map_processing, # we need to keep the intermediate transient array ) - 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, - ) - - 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")() - - testee.with_backend(custom_backend)(a, b, out=out, offset_provider={}) + # ``_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 = [ From 1ff0f9f8bd8c59bfc6fb9b4428fb7a7922eed9e6 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 28 Jul 2026 15:17:02 +0200 Subject: [PATCH 09/25] remove annotation runtime_checkable from ExternalMemoryAllocator --- .../runners/dace/transformations/auto_optimize.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 5f9a9a2a44..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 @@ -12,7 +12,7 @@ import dataclasses import enum import warnings -from typing import Any, Callable, Optional, Protocol, Sequence, TypeAlias, Union, runtime_checkable +from typing import Any, Callable, Optional, Protocol, Sequence, TypeAlias, Union import dace from dace import data as dace_data @@ -145,7 +145,6 @@ class AllocationRequest: ExternalWorkspace: TypeAlias = xtyping.ArrayInterface | xtyping.CUDAArrayInterface -@runtime_checkable class ExternalMemoryAllocator(Protocol): """Allocates and frees workspace memory for ``TransientMemoryMode.EXTERNAL``. From f3010abc960f979876f1d06478ee90bb8f2f7217 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 10:40:09 +0200 Subject: [PATCH 10/25] edit --- .../runners/dace/workflow/backend.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) 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 76f50e55b9..1b5b52ddac 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -119,18 +119,24 @@ 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 + if external_memory_allocator is None: + if ( + optimization_args.get("transient_memory_mode") + is gtx_transformations.TransientMemoryMode.EXTERNAL + ): + raise ValueError( + "External memory allocator 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 allocator 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, @@ -138,7 +144,6 @@ def make_dace_backend( 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, From e1d7b19f12973a16bfe22af907880d9d8e3036cf Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 10:43:12 +0200 Subject: [PATCH 11/25] disable cached translation --- .../program_processors/runners/dace/workflow/backend.py | 6 +++++- .../runners_tests/dace_tests/test_dace_backend.py | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) 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 1b5b52ddac..16da3e5c77 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -42,10 +42,11 @@ class Params: device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, name_device="gpu", ) + cached_translation = True device_type = core_defs.DeviceType.CPU otf_workflow = factory.SubFactory( DaCeWorkflowFactory, - cached_translation=True, + cached_translation=factory.SelfAttribute("..cached_translation"), device_type=factory.SelfAttribute("..device_type"), auto_optimize=factory.SelfAttribute("..auto_optimize"), external_memory_allocator=factory.SelfAttribute("..external_memory_allocator"), @@ -62,6 +63,7 @@ def make_dace_backend( gpu: bool, auto_optimize: bool = True, async_sdfg_call: bool = True, + cached_translation: 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, @@ -76,6 +78,7 @@ def make_dace_backend( auto_optimize: Enable the SDFG auto-optimize pipeline. async_sdfg_call: Make an asynchronous SDFG call on GPU to allow overlapping of GPU kernel execution with the Python driver code. + cached_translation: Enable caching of the SDFG translation step. 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 @@ -141,6 +144,7 @@ def make_dace_backend( return DaCeBackendFactory( # type: ignore[return-value] # factory-boy typing not precise enough gpu=gpu, auto_optimize=auto_optimize, + cached_translation=cached_translation, 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, 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 f82bfc37ea..d5655c1126 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 @@ -303,10 +303,12 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): external_memory_allocator = _WorkspaceRecordingAllocator() workspace_requests = external_memory_allocator.requests + # FIXME(edopao, egparedes): Not clear why the SDFG is cached although the backends are different. custom_backend = dace_wf_backend.make_dace_backend( gpu=on_gpu, auto_optimize=True, async_sdfg_call=False, + cached_translation=False, optimization_args={ "transient_memory_mode": transient_memory_mode, }, From f290603f194e58b22748e08ac942008627c0ff54 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 13:55:10 +0200 Subject: [PATCH 12/25] Revert "disable cached translation" This reverts commit dd47b2e80ddb6ec5d86642b6f73d4169cffd0027. --- .../program_processors/runners/dace/workflow/backend.py | 6 +----- .../runners_tests/dace_tests/test_dace_backend.py | 2 -- 2 files changed, 1 insertion(+), 7 deletions(-) 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 16da3e5c77..1b5b52ddac 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -42,11 +42,10 @@ class Params: device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, name_device="gpu", ) - cached_translation = True device_type = core_defs.DeviceType.CPU otf_workflow = factory.SubFactory( DaCeWorkflowFactory, - cached_translation=factory.SelfAttribute("..cached_translation"), + cached_translation=True, device_type=factory.SelfAttribute("..device_type"), auto_optimize=factory.SelfAttribute("..auto_optimize"), external_memory_allocator=factory.SelfAttribute("..external_memory_allocator"), @@ -63,7 +62,6 @@ def make_dace_backend( gpu: bool, auto_optimize: bool = True, async_sdfg_call: bool = True, - cached_translation: 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, @@ -78,7 +76,6 @@ def make_dace_backend( auto_optimize: Enable the SDFG auto-optimize pipeline. async_sdfg_call: Make an asynchronous SDFG call on GPU to allow overlapping of GPU kernel execution with the Python driver code. - cached_translation: Enable caching of the SDFG translation step. 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 @@ -144,7 +141,6 @@ def make_dace_backend( return DaCeBackendFactory( # type: ignore[return-value] # factory-boy typing not precise enough gpu=gpu, auto_optimize=auto_optimize, - cached_translation=cached_translation, 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, 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 d5655c1126..f82bfc37ea 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 @@ -303,12 +303,10 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): external_memory_allocator = _WorkspaceRecordingAllocator() workspace_requests = external_memory_allocator.requests - # FIXME(edopao, egparedes): Not clear why the SDFG is cached although the backends are different. custom_backend = dace_wf_backend.make_dace_backend( gpu=on_gpu, auto_optimize=True, async_sdfg_call=False, - cached_translation=False, optimization_args={ "transient_memory_mode": transient_memory_mode, }, From 820bd5249a9a99736013f260b9bb6473ea7a6edd Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 14:15:01 +0200 Subject: [PATCH 13/25] serialize backends with subtests --- .../dace_tests/test_dace_backend.py | 47 +++++++++++++------ 1 file changed, 32 insertions(+), 15 deletions(-) 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 f82bfc37ea..4646276d31 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 @@ -292,26 +292,12 @@ def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str 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 +def _test_transient_memory_mode(on_gpu, transient_memory_mode, custom_backend, monkeypatch): 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: @@ -366,6 +352,7 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: assert len(transient_arrays) == 2 generated_code = _parse_generated_code_from_sdfg(captured_sdfg, gpu_api_prefix) + workspace_requests = custom_backend.executor.compilation.external_memory_allocator.requests match transient_memory_mode: case gtx_transformations.TransientMemoryMode.EXTERNAL: @@ -445,3 +432,33 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: assert any(marker in generated_code for marker in ("delete ", "free")) assert np.allclose(out.asnumpy(), a.asnumpy() + b.asnumpy() + 1) + + +def test_transient_memory_mode(device_type, monkeypatch, subtests): + on_gpu = device_type == core_defs.CUPY_DEVICE_TYPE + external_memory_allocator = _WorkspaceRecordingAllocator() + + # Note that the different custom backends are created here, and stored in an + # array, so that they are not garbage collected before all subtests run. + # This is needed to keep the `_compiled_programs` cache in a consistent state. + # Otherwise, it could happen that the same backend id is reused, for a different + # backend object, and the program is loaded from cache instead of being lowered. + configs = [ + ( + mode, + dace_wf_backend.make_dace_backend( + gpu=on_gpu, + async_sdfg_call=False, + optimization_args={ + "transient_memory_mode": mode, + }, + external_memory_allocator=external_memory_allocator, + ), + ) + for mode in gtx_transformations.TransientMemoryMode + ] + + for mode, backend in configs: + with subtests.test(f"transient_memory_mode={mode}"): + _test_transient_memory_mode(on_gpu, mode, backend, monkeypatch) + external_memory_allocator.requests.clear() # reset for next subtest From c8e67b18aa98d43d8825629abedd7da4a349113c Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 15:02:59 +0200 Subject: [PATCH 14/25] Revert "serialize backends with subtests" This reverts commit 696488daa3cf02492b02d0bd6b6632c4a6e984dc. --- .../dace_tests/test_dace_backend.py | 47 ++++++------------- 1 file changed, 15 insertions(+), 32 deletions(-) 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 4646276d31..f82bfc37ea 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 @@ -292,12 +292,26 @@ def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str return generated_code -def _test_transient_memory_mode(on_gpu, transient_memory_mode, custom_backend, monkeypatch): +@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: @@ -352,7 +366,6 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: assert len(transient_arrays) == 2 generated_code = _parse_generated_code_from_sdfg(captured_sdfg, gpu_api_prefix) - workspace_requests = custom_backend.executor.compilation.external_memory_allocator.requests match transient_memory_mode: case gtx_transformations.TransientMemoryMode.EXTERNAL: @@ -432,33 +445,3 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: assert any(marker in generated_code for marker in ("delete ", "free")) assert np.allclose(out.asnumpy(), a.asnumpy() + b.asnumpy() + 1) - - -def test_transient_memory_mode(device_type, monkeypatch, subtests): - on_gpu = device_type == core_defs.CUPY_DEVICE_TYPE - external_memory_allocator = _WorkspaceRecordingAllocator() - - # Note that the different custom backends are created here, and stored in an - # array, so that they are not garbage collected before all subtests run. - # This is needed to keep the `_compiled_programs` cache in a consistent state. - # Otherwise, it could happen that the same backend id is reused, for a different - # backend object, and the program is loaded from cache instead of being lowered. - configs = [ - ( - mode, - dace_wf_backend.make_dace_backend( - gpu=on_gpu, - async_sdfg_call=False, - optimization_args={ - "transient_memory_mode": mode, - }, - external_memory_allocator=external_memory_allocator, - ), - ) - for mode in gtx_transformations.TransientMemoryMode - ] - - for mode, backend in configs: - with subtests.test(f"transient_memory_mode={mode}"): - _test_transient_memory_mode(on_gpu, mode, backend, monkeypatch) - external_memory_allocator.requests.clear() # reset for next subtest From 93744847c4bf88c78b71b537af0c9e6e63acc73e Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 15:10:01 +0200 Subject: [PATCH 15/25] clear compiled_programs cache --- .../runners_tests/dace_tests/test_dace_backend.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 f82bfc37ea..b41c298651 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,6 +8,7 @@ """Test the bindings stage of the dace backend workflow.""" +import copy import re import unittest.mock as mock @@ -350,12 +351,15 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: no_op_top_level_map_processing, # we need to keep the intermediate transient array ) + prog = copy.copy(testee.with_backend(custom_backend)) + prog._compiled_programs.compiled_programs.clear() # clear any cached compiled programs to force recompilation + # ``_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={}) + prog(a, b, out=out, offset_provider={}) assert captured_sdfg is not None transient_arrays = [ From 13ebec7bd23cad8d1500487c2bf004546408a638 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 15:55:44 +0200 Subject: [PATCH 16/25] Revert "clear compiled_programs cache" This reverts commit 63f4bf5f138ccd39dafe0f2514d93fbdf74923c2. --- .../runners_tests/dace_tests/test_dace_backend.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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 b41c298651..f82bfc37ea 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,7 +8,6 @@ """Test the bindings stage of the dace backend workflow.""" -import copy import re import unittest.mock as mock @@ -351,15 +350,12 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: no_op_top_level_map_processing, # we need to keep the intermediate transient array ) - prog = copy.copy(testee.with_backend(custom_backend)) - prog._compiled_programs.compiled_programs.clear() # clear any cached compiled programs to force recompilation - # ``_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): - prog(a, b, out=out, offset_provider={}) + testee.with_backend(custom_backend)(a, b, out=out, offset_provider={}) assert captured_sdfg is not None transient_arrays = [ From ca6debe4c7824e845b0665b4dbc0bbc7ac587c7f Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 29 Jul 2026 16:22:27 +0200 Subject: [PATCH 17/25] edit --- .../dace_tests/test_dace_backend.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) 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 f82bfc37ea..13e1b1cf45 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,6 +8,7 @@ """Test the bindings stage of the dace backend workflow.""" +import dataclasses import re import unittest.mock as mock @@ -20,14 +21,13 @@ 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.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, - translation as gtx_dace_translation, ) from next_tests.integration_tests import cases, cases_utils @@ -332,18 +332,25 @@ def testee(a: cases.IField, b: cases.IField, out: cases.IField): out = cases.allocate(test_case, testee, "out")() captured_sdfg: dace.SDFG | None = None - gt_generate_sdfg = gtx_dace_translation.DaCeTranslator.generate_sdfg + translation_step = custom_backend.executor.translation.step - def mocked_generate_sdfg(self, *args, **kwargs) -> dace.SDFG: + def mocked_translator(inp: definitions.CompilableProgramDef) -> dace.SDFG: nonlocal captured_sdfg - result = gt_generate_sdfg(self, *args, **kwargs) - captured_sdfg = result + 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_dace_translation.DaCeTranslator, "generate_sdfg", mocked_generate_sdfg) monkeypatch.setattr( gtx_auto_optimize, "_gt_auto_process_top_level_maps", From 377492e44130cd4a0d4170ec841e17aff64097a4 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 31 Jul 2026 11:00:30 +0200 Subject: [PATCH 18/25] apply review comments --- src/gt4py/next/otf/compiled_program.py | 41 ++++++++----------- .../runners/dace/workflow/compilation.py | 6 +-- 2 files changed, 19 insertions(+), 28 deletions(-) diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 01d002aab2..e079d58edd 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -85,7 +85,9 @@ def metrics_source_key(pool: CompiledProgramsPool, key: CompiledProgramsKey) -> return source_key -def _finalize_compiled_programs(programs: stages.ExecutableProgram | dict[Any, Any]) -> None: +def _finalize_compiled_programs( + programs: dict[CompiledProgramsKey, stages.ExecutableProgram], +) -> None: """Best-effort cleanup of compiled programs when their pool is deleted. Invoked via :py:func:`weakref.finalize` on a @@ -96,19 +98,16 @@ def _finalize_compiled_programs(programs: stages.ExecutableProgram | dict[Any, A 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, - ) + for program in programs.values(): + if finalize := getattr(program, "finalize", None): + 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 @@ -401,17 +400,9 @@ 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 + # 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. 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) 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 8481657ab3..f825b37a7e 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -77,7 +77,7 @@ def _add_tx_markers(program_source: SDFGExtensionSource) -> tuple[SDFGExtensionS return new_program_source, sdfg -def _map_workspace_storage_to_device(storage: dace.StorageType) -> core_defs.DeviceType: +def workspace_storage_to_device_mapping(storage: dace.StorageType) -> core_defs.DeviceType: if storage == dace.StorageType.CPU_Heap: device = core_defs.DeviceType.CPU elif storage == dace.StorageType.GPU_Global: @@ -206,7 +206,7 @@ def _configure_external_workspaces(self, **kwargs: Any) -> None: "SDFG requires external workspaces, but no allocator was provided." ) for storage, required_nbytes in workspace_sizes.items(): - device = _map_workspace_storage_to_device(storage) + device = workspace_storage_to_device_mapping(storage) request = AllocationRequest(nbytes=required_nbytes, device=device) workspace = self.external_memory_allocator.allocate(request) _validate_external_workspace(storage, request, workspace) @@ -390,7 +390,7 @@ def __post_init__(self) -> None: # 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) + _check_allocator_picklable(self.external_memory_allocator) with gtx_wfdcommon.dace_context( device_type=self.device_type, cmake_build_type=self.cmake_build_type, From 6301c8c51cf98a74f555470f845518d263423b4a Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Thu, 30 Jul 2026 15:09:56 +0200 Subject: [PATCH 19/25] fix[next-dace]: Include modifed SDFG with GPU TX markers in build folder fingerprint (#2739) This pull request refactors and improves the GPU trace marker logic in the DaCe compilation workflow, enhances test coverage, and clarifies test parameterization. The main functional change is to ensure that the fingerprint of the build folder includes the modified SDFG after applying the GPU transaction markers. This way, enabling or not the GPU transaction markers results in different build artifacts. Tests are updated to reflect the new marker application logic and to ensure that build artifacts change when relevant compilation settings differ. **DaCe compilation logic improvements:** * Refactored `_add_tx_markers` to take and return an `ExtensionSource`, only modifying the SDFG if GPU scheduling is detected, and returning a new `ExtensionSource` with updated SDFG JSON. * Updated the main DaCe compiler logic to call `_add_tx_markers` only when GPU trace markers are requested and the device type is GPU, ensuring markers are not redundantly applied. [[1]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dL200-R222) [[2]](diffhunk://#diff-89b749ff9b80be5fdb8a7e0d05411d2b992b6a5892f819c7a3e001c4241a6c0dL218-R235) **Test suite improvements:** * Refactored tests to use a real `program_source` fixture and removed unnecessary mocks and spies, directly verifying the presence or absence of GPU TX markers on the compiled SDFG. [[1]](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3L82-R83) [[2]](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3R100-L109) [[3]](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3L119-L123) [[4]](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3L138-L173) * Added new tests to assert that identical compilation inputs produce the same artifact, and that changing instrumentation or compiler flags results in different artifacts. [[1]](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3R181-R203) [[2]](diffhunk://#diff-57ff279a6de9ce815a6bde8ae126007687da1d07f6ad28626e998541fc9d89b3R213-R248) **Test parameterization and device naming:** * Updated device type parameterization in both compilation and translation tests to use consistent and descriptive IDs (`CPU`, `GPU`, `CUDA`, `ROCM`) and to use the correct device type constants. [[1]](diffhunk://#diff-19fed38498781ceef82337d097e43458274d2ebe7e0bae4578839e9fbed15194L31-R33) [[2]](diffhunk://#diff-8ba502dd24c52188c94e85f0bcd66c31fbe3093077168033f99c14eb19ac1cb4L54-R55) These changes improve the correctness, maintainability, and clarity of both the DaCe compilation workflow and its associated tests. --- .../program_processors/runners/dace/workflow/compilation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 f825b37a7e..2e4ab57304 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -332,7 +332,7 @@ class AllocatorNotPicklableError(TypeError): """ -def _assert_allocator_picklable(allocator: ExternalMemoryAllocator) -> None: +def _check_allocator_picklable(allocator: ExternalMemoryAllocator) -> None: """Fail fast if ``allocator`` is not picklable. Args: @@ -345,7 +345,7 @@ def _assert_allocator_picklable(allocator: ExternalMemoryAllocator) -> None: pickle.dumps(allocator) except Exception as error: # pickle raises arbitrary exceptions raise AllocatorNotPicklableError( - f"external_memory_allocator {allocator!r} is not picklable: {error!s}." + 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 f53fdd886fa335a1019128771a186ab589f7bd6f Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 31 Jul 2026 12:27:12 +0200 Subject: [PATCH 20/25] edit --- .../runners/dace/workflow/compilation.py | 66 +++++++++---------- 1 file changed, 33 insertions(+), 33 deletions(-) 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 2e4ab57304..24f9c3d534 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -280,6 +280,39 @@ def __call__(self, **kwargs: Any) -> None: assert result is None +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 _check_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 + + @dataclasses.dataclass(frozen=True) class DaCeCompilationArtifact: """Result of a DaCe compilation: library path + SDFG bindings + the SDFG itself. @@ -319,39 +352,6 @@ def load(self) -> stages.ExecutableProgram: return gtx_wfddecoration.DaCeDecoratedProgram(program, device_type=self.device_type) -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 _check_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 - - @dataclasses.dataclass(frozen=True) class DaCeCompiler( workflow.ChainableWorkflowMixin[ From 0589f600ddc7dc0fac60c026c63f628e49aaa92c Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 31 Jul 2026 17:11:05 +0200 Subject: [PATCH 21/25] remove external allocator, pass workspace buffers directly --- src/gt4py/next/backend.py | 8 + src/gt4py/next/otf/compiled_program.py | 35 +- .../runners/dace/transformations/__init__.py | 8 +- .../dace/transformations/auto_optimize.py | 76 +--- .../runners/dace/workflow/backend.py | 46 ++- .../runners/dace/workflow/common.py | 14 +- .../runners/dace/workflow/compilation.py | 160 ++------ .../runners/dace/workflow/decoration.py | 9 +- .../runners/dace/workflow/factory.py | 3 - .../otf_tests/test_compiled_program.py | 69 ---- .../dace_tests/test_dace_backend.py | 114 +++--- .../dace_tests/test_dace_compilation.py | 343 +++++------------- 12 files changed, 222 insertions(+), 663 deletions(-) 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 e079d58edd..fa89e98dd0 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -85,31 +85,6 @@ def metrics_source_key(pool: CompiledProgramsPool, key: CompiledProgramsKey) -> return source_key -def _finalize_compiled_programs( - programs: dict[CompiledProgramsKey, stages.ExecutableProgram], -) -> 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. - """ - for program in programs.values(): - if finalize := getattr(program, "finalize", None): - 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, @@ -400,12 +375,6 @@ 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. 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 @@ -614,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( @@ -691,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 e5ea78eb71..f8cb38334f 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,7 @@ """ from . import constants, splitting_tools -from .auto_optimize import ( # re-exported for the external-memory public API - AllocationRequest, - ExternalMemoryAllocator, - ExternalWorkspace, +from .auto_optimize import ( GT4PyAutoOptHook, GT4PyAutoOptHookFun, GT4PyAutoOptHookStage, @@ -92,11 +89,8 @@ __all__ = [ - "AllocationRequest", "CopyChainRemover", "DoubleWriteRemover", - "ExternalMemoryAllocator", - "ExternalWorkspace", "FuseHorizontalConditionBlocks", "GPUSetBlockSize", "GT4PyAutoOptHook", 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 6308442bc8..272621136a 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,11 +8,9 @@ """Fast access to the auto optimization on DaCe.""" -import abc -import dataclasses import enum import warnings -from typing import Any, Callable, Optional, Protocol, Sequence, TypeAlias, Union +from typing import Any, Callable, Optional, Sequence, TypeAlias, Union import dace from dace import data as dace_data @@ -21,8 +19,6 @@ 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, @@ -115,67 +111,6 @@ 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. @@ -194,12 +129,9 @@ class TransientMemoryMode(str, enum.Enum): 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. + The `EXTERNAL` strategy requires that the `external_workspace` attribute + of the dace backend workflow is set, because it is needed at runtime to + install the memory pointers for transient arrays. """ SCOPED = "SCOPED" 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 1b5b52ddac..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,8 +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.otf import stages from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations -from gt4py.next.program_processors.runners.dace.workflow.factory import DaCeWorkflowFactory +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): @@ -31,12 +51,11 @@ class DaCeBackendFactory(factory.Factory): """ class Meta: - model = backend.Backend + model = DaCeBackend 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, @@ -44,11 +63,10 @@ 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"), - external_memory_allocator=factory.SelfAttribute("..external_memory_allocator"), ) auto_optimize = factory.Trait(name_postfix="_opt") @@ -56,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( @@ -63,7 +82,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, + 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, @@ -78,11 +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_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. + 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 @@ -119,18 +135,18 @@ def make_dace_backend( else None } - if external_memory_allocator is None: + if external_workspace is None: if ( optimization_args.get("transient_memory_mode") is gtx_transformations.TransientMemoryMode.EXTERNAL ): raise ValueError( - "External memory allocator must be provided when 'transient_memory_mode' is 'EXTERNAL'." + "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 allocator provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{gtx_transformations.TransientMemoryMode.EXTERNAL}'.", + f"External memory workspace provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{gtx_transformations.TransientMemoryMode.EXTERNAL}'.", stacklevel=2, ) else: @@ -141,7 +157,7 @@ def make_dace_backend( return DaCeBackendFactory( # type: ignore[return-value] # factory-boy typing not precise enough gpu=gpu, auto_optimize=auto_optimize, - external_memory_allocator=external_memory_allocator, + 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 24f9c3d534..de6dce3a34 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -12,7 +12,6 @@ import json import os import pathlib -import pickle import warnings from collections.abc import Callable, MutableSequence, Sequence from typing import Any, Final, TypeAlias @@ -26,11 +25,6 @@ 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, @@ -77,7 +71,7 @@ def _add_tx_markers(program_source: SDFGExtensionSource) -> tuple[SDFGExtensionS return new_program_source, sdfg -def workspace_storage_to_device_mapping(storage: dace.StorageType) -> core_defs.DeviceType: +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: @@ -90,14 +84,14 @@ def workspace_storage_to_device_mapping(storage: dace.StorageType) -> core_defs. def _validate_external_workspace( - storage: dace.StorageType, request: AllocationRequest, wsp: ExternalWorkspace + wsp: xtyping.ArrayInterface | xtyping.CUDAArrayInterface, storage: dace.StorageType, nbytes: int ) -> None: - """Validate that ``wsp`` satisfies ``request`` for ``storage``. + """Validate that the provided ``wsp`` workspace satisfies the requirements. 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. + nbytes: Size in bytes required. + wsp: The external workspace to check. Raises: TypeError: If ``wsp`` exposes neither ``__array_interface__`` nor @@ -111,28 +105,12 @@ def _validate_external_workspace( 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." - ) + if (wsp_nbytes := getattr(wsp, "nbytes", None)) is not None: + if wsp_nbytes < nbytes: + raise ValueError( + f"External workspace buffer is {wsp_nbytes} bytes for storage " + f"{storage!r}, but at least {nbytes} bytes were required." + ) class CompiledDaceProgram: @@ -161,15 +139,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] + external_workspace: gtx_wfdcommon.ExternalWorkspace | None = ( + None # This attribute is set at runtime, before the first call. + ) def __init__( self, program: dace.CompiledSDFG, bind_func_name: str, binding_source_code: str, - external_memory_allocator: ExternalMemoryAllocator | None = None, ): self.sdfg_program = program @@ -189,56 +167,22 @@ 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. + 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_memory_allocator is None: - raise ValueError( - "SDFG requires external workspaces, but no allocator was provided." + 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 = workspace_storage_to_device_mapping(storage) - request = AllocationRequest(nbytes=required_nbytes, device=device) - workspace = self.external_memory_allocator.allocate(request) - _validate_external_workspace(storage, request, workspace) + 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) - # 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: """ @@ -247,7 +191,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) + 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] @@ -280,39 +224,6 @@ def __call__(self, **kwargs: Any) -> None: assert result is None -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 _check_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 - - @dataclasses.dataclass(frozen=True) class DaCeCompilationArtifact: """Result of a DaCe compilation: library path + SDFG bindings + the SDFG itself. @@ -335,7 +246,6 @@ 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 @@ -343,12 +253,7 @@ 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, - external_memory_allocator=self.external_memory_allocator, - ) + program = CompiledDaceProgram(sdfg_program, self.bind_func_name, self.binding_source_code) return gtx_wfddecoration.DaCeDecoratedProgram(program, device_type=self.device_type) @@ -369,11 +274,6 @@ 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 ) @@ -384,13 +284,6 @@ 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: - _check_allocator_picklable(self.external_memory_allocator) with gtx_wfdcommon.dace_context( device_type=self.device_type, cmake_build_type=self.cmake_build_type, @@ -450,7 +343,6 @@ def __call__(self, inp: SDFGExtensionSource) -> DaCeCompilationArtifact: 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 9ed8d14ff6..9611b6e359 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py @@ -94,10 +94,9 @@ def __call__( metrics.COMPUTE_METRIC, self._collect_time_arg[0].item() ) - def finalize(self) -> None: - """Forward teardown to the underlying ``CompiledDaceProgram``. + def set_external_workspace(self, external_workspace: gtx_wfdcommon.ExternalWorkspace) -> None: + """Set the external workspace for the underlying compiled program. - Allows the generic otf pool -- which only sees this callable -- to - release external-memory workspaces when the pool is finalized. + This method should be called before the first call to the program. """ - self._fun.finalize() + self._fun.external_workspace = external_workspace 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 6dacb61885..6238871b8f 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -17,7 +17,6 @@ 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, @@ -36,7 +35,6 @@ 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 @@ -75,6 +73,5 @@ 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 35cb82a355..ed881c9495 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,72 +314,3 @@ 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 13e1b1cf45..39196c2aeb 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 @@ -11,6 +11,7 @@ import dataclasses import re import unittest.mock as mock +from typing import Any import numpy as np import pytest @@ -28,6 +29,7 @@ ) from gt4py.next.program_processors.runners.dace.workflow import ( backend as dace_wf_backend, + common as dace_wf_common, ) from next_tests.integration_tests import cases, cases_utils @@ -47,7 +49,7 @@ def device_type(request) -> gtx.DeviceType: 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 @@ -160,23 +162,29 @@ def mocked_gpu_transformation(*args, **kwargs) -> dace.SDFG: mock_top_level_dataflow_hook2.assert_not_called() -class _RecordingAllocator: - """Minimal `ExternalMemoryAllocator` for backend-wiring tests. +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) - 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") +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. + """ - def deallocate(self, buffer) -> None: - raise AssertionError("backend-wiring tests must not call deallocate") + nbytes: int = 1024 + __array_interface__: dict[str, Any] = {"shape": (1024,), "typestr": "|u1", "version": 3} -def test_make_backend_accepts_external_allocator_with_external_mode(): - external_memory_allocator = _RecordingAllocator() +def test_make_backend_accepts_external_workspace_with_external_mode(): + workspace = _RecordingWorkspace() backend = dace_wf_backend.make_dace_backend( gpu=False, @@ -185,33 +193,33 @@ def test_make_backend_accepts_external_allocator_with_external_mode(): optimization_args={ "transient_memory_mode": gtx_transformations.TransientMemoryMode.EXTERNAL, }, - external_memory_allocator=external_memory_allocator, + external_workspace={core_defs.DeviceType.CPU: workspace}, ) - assert backend.executor.compilation.external_memory_allocator is external_memory_allocator + assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace -def test_make_backend_infers_external_mode_when_allocator_is_provided(): - external_memory_allocator = _RecordingAllocator() +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_memory_allocator=external_memory_allocator, + external_workspace={core_defs.DeviceType.CPU: workspace}, ) 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 + assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace -def test_make_backend_warns_external_allocator_without_external_mode(): - external_memory_allocator = _RecordingAllocator() +def test_make_backend_warns_external_workspace_without_external_mode(): + workspace = _RecordingWorkspace() - with pytest.warns(UserWarning, match="External memory allocator provided"): + with pytest.warns(UserWarning, match="External memory workspace provided"): backend = dace_wf_backend.make_dace_backend( gpu=False, auto_optimize=True, @@ -219,7 +227,7 @@ def test_make_backend_warns_external_allocator_without_external_mode(): optimization_args={ "transient_memory_mode": gtx_transformations.TransientMemoryMode.POOL, }, - external_memory_allocator=external_memory_allocator, + external_workspace={core_defs.DeviceType.CPU: workspace}, ) # Explicit mode stays as requested by the caller; backend only warns. @@ -227,40 +235,7 @@ def test_make_backend_warns_external_allocator_without_external_mode(): 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 + assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str: @@ -300,8 +275,12 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): 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 + # 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, @@ -310,7 +289,7 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): optimization_args={ "transient_memory_mode": transient_memory_mode, }, - external_memory_allocator=external_memory_allocator, + external_workspace=external_workspace, ) @gtx.field_operator @@ -357,10 +336,8 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: 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. + # 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): testee.with_backend(custom_backend)(a, b, out=out, offset_provider={}) @@ -383,7 +360,7 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: 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. + # 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) @@ -391,16 +368,11 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: 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( @@ -409,7 +381,6 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: ) 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( @@ -437,7 +408,6 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: # `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( 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 43bd747ddd..8bd9f33952 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,21 +8,18 @@ """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 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 @@ -31,10 +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.transformations import ( - auto_optimize as gtx_auto_optimize, +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 @@ -172,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="{}", @@ -273,7 +270,7 @@ def test_cmake_build_type_changes_artifact(program_source): def _make_compiled_program( *, - external_memory_allocator=None, + external_workspace: dict[core_defs.DeviceType, Any] | None = None, workspace_sizes: dict[Any, int] | None = None, ): if workspace_sizes is None: @@ -288,97 +285,116 @@ def _make_compiled_program( sdfg_program.get_workspace_sizes.return_value = workspace_sizes sdfg_program.construct_arguments.return_value = ((), ()) - return dace_wf_compilation.CompiledDaceProgram( + 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", - external_memory_allocator=external_memory_allocator, ) + if external_workspace is not None: + compiled_program.external_workspace = external_workspace + return compiled_program -def test_construct_arguments_installs_external_workspaces_once(): - allocator = mock.MagicMock() - allocator.allocate.side_effect = [_make_array_buffer(nbytes=128, address=256)] +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_memory_allocator=allocator, + 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) - 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 + 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_propagates_allocator_error_for_invalid_size_request(): - allocator = mock.MagicMock() - allocator.allocate.side_effect = ValueError("invalid workspace size request") +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_memory_allocator=allocator, - workspace_sizes={dace.StorageType.CPU_Heap: -1}, + 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) + - with pytest.raises(ValueError, match="invalid workspace size request"): +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) - 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") +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_memory_allocator=allocator, - workspace_sizes={dace.StorageType.CPU_Heap: 16}, + external_workspace={core_defs.DeviceType.CPU: workspace}, + workspace_sizes={dace.StorageType.CPU_Pinned: 128}, ) - with pytest.raises(TypeError, match="invalid storage type request"): + with pytest.raises(ValueError, match="Unsupported storage type"): 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. +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) 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. + (device); ``nbytes`` matches the requested size. """ buffer = mock.MagicMock() buffer.nbytes = nbytes - interface = {"data": (address, False), "shape": (nbytes,), "typestr": "|u1", "version": 3} + 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 allocator must return something ``set_workspace`` can consume.""" - allocator = mock.MagicMock() - allocator.allocate.side_effect = ["not-an-array"] # str exposes no array interface + """The workspace must expose an array interface that ``set_workspace`` can consume.""" program = _make_compiled_program( - external_memory_allocator=allocator, + external_workspace={core_defs.DeviceType.CPU: "not-an-array"}, workspace_sizes={dace.StorageType.CPU_Heap: 64}, ) @@ -388,226 +404,49 @@ def test_construct_arguments_rejects_buffer_without_array_interface(): 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 +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_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, + 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.external_workspaces[dace.StorageType.CPU_Heap] is workspace + assert program.sdfg_program.set_workspace.call_args.args[1] 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)] +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_memory_allocator=allocator, + 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 - ), - pytest.raises(ValueError, match="not aligned to the required 256 bytes"), + 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_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}, - ) + 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 - # Must not raise even though alignment can't be verified. - program.construct_arguments(alpha=1) - program.sdfg_program.set_workspace.assert_called_once() +class _ArrayBufferWithoutNbytes: + __array_interface__ = {"shape": (128,), "typestr": "|u1", "version": 3} -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] +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_memory_allocator=allocator, + 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.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) + program.sdfg_program.set_workspace.assert_called_once() From 02009b4ad96215ac3fd51cd5b6d09f82ad139dad Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Mon, 3 Aug 2026 10:00:44 +0200 Subject: [PATCH 22/25] update ADR --- .../next/0026-External_Memory_Allocator.md | 147 --------------- .../next/0027-External_Workspace_Memory.md | 172 ++++++++++++++++++ docs/development/ADRs/next/README.md | 1 + 3 files changed, 173 insertions(+), 147 deletions(-) delete mode 100644 docs/development/ADRs/next/0026-External_Memory_Allocator.md create mode 100644 docs/development/ADRs/next/0027-External_Workspace_Memory.md diff --git a/docs/development/ADRs/next/0026-External_Memory_Allocator.md b/docs/development/ADRs/next/0026-External_Memory_Allocator.md deleted file mode 100644 index c6d5a0e42c..0000000000 --- a/docs/development/ADRs/next/0026-External_Memory_Allocator.md +++ /dev/null @@ -1,147 +0,0 @@ ---- -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/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 From b7eca03dab34757fa5a64e98a6a61d183e4b426d Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Mon, 3 Aug 2026 11:30:03 +0200 Subject: [PATCH 23/25] edit --- .../dace/transformations/auto_optimize.py | 12 +++++----- .../runners/dace/workflow/compilation.py | 15 ++++++++----- .../runners/dace/workflow/decoration.py | 7 +++--- .../dace_tests/test_dace_backend.py | 22 ++++++++++++++++++- 4 files changed, 42 insertions(+), 14 deletions(-) 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 272621136a..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 @@ -125,13 +125,15 @@ class TransientMemoryMode(str, enum.Enum): - `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. + - `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 workflow is set, because it is needed at runtime to - install the memory pointers for transient arrays. + of the dace backend is set, because it is needed at runtime to install + the memory pointers for transient arrays. """ SCOPED = "SCOPED" 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 de6dce3a34..d22ef43ca5 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -75,7 +75,11 @@ 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: - assert core_defs.CUPY_DEVICE_TYPE is not None + 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.") @@ -89,19 +93,20 @@ def _validate_external_workspace( """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. - wsp: The external workspace to check. 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. + 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 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"External workspace is {type(wsp).__name__!r} for storage " f"{storage!r}, which does not expose `__array_interface__` or " f"`__cuda_array_interface__`." ) 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 9611b6e359..57f45964ce 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py @@ -31,9 +31,10 @@ class DaCeDecoratedProgram: 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. + 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__( 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 39196c2aeb..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 @@ -30,6 +30,7 @@ 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 next_tests.integration_tests import cases, cases_utils @@ -269,6 +270,15 @@ def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str @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(" @@ -339,7 +349,12 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: # 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): - testee.with_backend(custom_backend)(a, b, out=out, offset_provider={}) + 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 = [ @@ -356,6 +371,11 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: 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 From a461502764050e74a02f2e718ae5f5bdd4bfbd93 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Mon, 3 Aug 2026 11:40:50 +0200 Subject: [PATCH 24/25] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../runners/dace/workflow/compilation.py | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) 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 d22ef43ca5..87abce5cf6 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -104,18 +104,23 @@ def _validate_external_workspace( required ``nbytes``. Buffers that do not expose ``nbytes`` are accepted on a trust basis (their size can not be checked here). """ - if not (xtyping.supports_array_interface(wsp) or xtyping.supports_cuda_array_interface(wsp)): - raise TypeError( - f"External workspace is {type(wsp).__name__!r} for storage " - f"{storage!r}, which does not expose `__array_interface__` or " - f"`__cuda_array_interface__`." - ) - if (wsp_nbytes := getattr(wsp, "nbytes", None)) is not None: - if wsp_nbytes < nbytes: - raise ValueError( - f"External workspace buffer is {wsp_nbytes} bytes for storage " - f"{storage!r}, but at least {nbytes} bytes were required." + 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: From 817587a5fd2c81154c97f9666d5681911fbcc9a6 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Mon, 3 Aug 2026 11:54:55 +0200 Subject: [PATCH 25/25] fix test --- .../runners_tests/dace_tests/test_dace_compilation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8bd9f33952..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 @@ -398,7 +398,7 @@ def test_construct_arguments_rejects_buffer_without_array_interface(): workspace_sizes={dace.StorageType.CPU_Heap: 64}, ) - with pytest.raises(TypeError, match="does not expose `__array_interface__`"): + with pytest.raises(TypeError, match="must expose `__array_interface__`"): program.construct_arguments(alpha=1) program.sdfg_program.set_workspace.assert_not_called()