feat[next-dace]: Support external memory for transient arrays#2719
Open
edopao wants to merge 3 commits into
Open
feat[next-dace]: Support external memory for transient arrays#2719edopao wants to merge 3 commits into
edopao wants to merge 3 commits into
Conversation
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a more flexible and extensible approach to managing transient memory allocation in the DaCe backend. The main improvement is the introduction of the
TransientMemoryModepolicy, which replaces the previous boolean flags for persistent and pooled memory, and enables external memory allocation via a user-provided allocator. The changes are grouped into three main themes:1. Transient Memory Lifetime Refactoring and Extension
TransientMemoryModeenum to specify allocation strategies for transient arrays, supportingSCOPED,PERSISTENT,POOL, andEXTERNALmodes. This replaces the previousmake_persistentandgpu_memory_poolflags.gt_auto_optimizepipeline and related post-processing to useTransientMemoryMode, updating function signatures, documentation, and logic to select the appropriate allocation strategy. [1] [2] [3] [4] [5] [6] [7] [8]__init__.pyto reflect the new mode and utility function. [1] [2] [3]2. Utilities and Backend Support for External Memory Allocation
gt_make_transients_persistentwith the more generalgt_configure_transient_lifetime, which can set the lifetime of eligible transients to eitherPersistentorExternaland adjusts storage as needed. [1] [2] [3] [4]external_memory_allocatorcallable, and ensured that when provided, thetransient_memory_modeis set toEXTERNAL. [1] [2] [3] [4] [5] [6]3. Infrastructure for Device Mapping
_map_workspace_storage_to_deviceto map DaCe storage types to GT4Py device types for use with external memory allocation.These changes make transient memory management more robust and configurable, especially for advanced use cases involving external memory pools or custom allocators.