WIP: Add experimental DSA memory-planner adapter - #2052
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a pluggable standalone DSA (Dynamic Storage Allocation) memory-planning solver and its adapter to PyPTO, adding the MemoryPlanner.DSA option to allow joint lifetime reuse and offset selection. When active, the pipeline skips opportunistic MemoryReuse and uses AllocateMemoryAddr to export a deterministic schema-v1 JSON problem, solve it using the standalone solver, validate the placement, and write back physical addresses. Feedback on the changes highlights two key issues: first, the PYPTO_ENABLE_DSA_SOLVER preprocessor definition must be defined globally in CMakeLists.txt so that conditional compilation blocks in the core C++ sources are correctly compiled; second, the adapter should generate dense, sequential BufferIds instead of casting the raw interval index, as sparse IDs can violate solver assumptions and lead to out-of-bounds indexing.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
670a53b to
8df2ed4
Compare
207be8f to
5302b22
Compare
…w-native-sys#1980) Phase 1: IR-free DsaSolver interface + FirstFitByLifetimeSolver (per-pool, colocation super-buffers, separation adjacency, decreasing-size skyline) + independent Validate + passes.dsa binding + 11 unit tests. Phase 2: ComputeAllocationLifetimes shared entry point (LifetimeInterval moved to utils/lifetime_analysis.h); BuildDsaProblem adapter; MaybeConsultDsaSolver runs the solver on real IR behind PYPTO_DSA_SOLVER=1 and logs DSA-peak vs bump-peak per space, bump stays authoritative. 2159 transform tests pass.
…#1980) ComputeAllocationPlan now returns intervals + separations (index pairs that must stay apart). Pipeline double-buffer clones (same group, different stage via PipelineMembershipsConflict) are emitted as DsaProblem separations so the solver keeps ping-pong operands in distinct buffers instead of collapsing them. Adds dsa::GetSolverMode (PYPTO_DSA_SOLVER: consult/1 vs plan vs off); MemoryReuse skips its pipeline_membership strip when the solver is engaged so pass 31's adapter can see the tags (codegen already tolerates the surviving attr — it does under memory_planner=PTOAS where the strip is skipped anyway). Verified: expert_routed L0A/Left consult goes 4096 -> 8192 (double-buffer kept apart). 2159 transform tests pass.
…ive-sys#1980) ComputeAllocationPlan now also emits the Ascend910B load+tpop_from_aic in-place hazard (backend-gated) and op-semantic forbid-alias (e.g. tile.sel mask/tmp vs output) as separations — the same constraints MemoryReuse honors — so the solver never forms those unsafe aliases. Takes FunctionPtr now (for the hazard guard). Forbid operands resolve to allocations via base_ identity (no reuse chain pre-solve). 2159 transform tests pass; expert_routed consult clean (no validator warnings).
…tive-sys#1980) Adds ApplyDsaSolution (solution offsets -> MemRef byte_offsets, mirroring the bump's slot-base + relative fold) and MaybePlanWithDsaSolver. In 'plan' mode the solver's addresses replace the bump's, but only through a hard gate: feasible AND every space's peak <= the bump's AND validator clean; otherwise it falls back to the bump (never worse). 'consult'/'1' still just logs. Tests use AIV (InCore) kernels — a plain @pl.function is non-InCore so the pass no-ops. plan mode reuses tile_c into tile_a's slot (peak 49152 -> 32768); consult mode leaves the bump untouched. Verified on expert_routed: all 6 kernels apply the plan (gate passed), compile succeeds. 2160 transform tests pass.
…hw-native-sys#1980) Pipeline-clone separations are now capacity-gated (hw-native-sys#1949 core): keep only F_g = max(1, min(depth, floor(cap/slot))) residues apart, so lifetime-disjoint same-residue clones reuse instead of the full-depth spread that overflowed tight L0B kernels. Deferred (surfaced blocker): skipping MemoryReuse's opportunistic merge under =plan (the prerequisite for the hw-native-sys#1908 sub-region fix) is not safe yet — the optimistic F_g cannot replicate MemoryReuse's full mutable shed loop + exact SpaceFootprint refit, so expert_routed's matmuls overflow L0B (3 co-live 32KB Right buffers = 96KB > 64KB) and the gate falls back to an already-overflowing un-merged bump. =plan therefore runs on MERGED groups for now (safe, gate-bounded, all expert_routed kernels apply + compile passes); hw-native-sys#1908 not yet fixed by =plan. 2160 transform tests pass.
Attach deterministic SSA dependency paths to pre-promotion reuse candidates. Add function-keyed PTOAS summary comparison and candidate inspection tools with regression tests and synchronized documentation.
1b95c88 to
bcac01f
Compare
1 participant
Summary
MemoryPlanner.DSApath that exports unmerged physicalallocations, invokes a standalone C++ DSA solver, independently validates
the placement, and writes offsets back to PyPTO
MemRefsseparations are address-disjoint; lifetime-disjoint buffers may reuse and
subdivide freed regions
target/no-alias constraints, and pipeline-stage separations
planner selection, and focused adapter tests
Status
This is a draft/WIP implementation for RFC #1980 and #1908. It does not yet
close either issue.
The external
tonibohnlein/dsa-solverdependency is temporary research scaffolding. PyPTO pins
0af3c08(package version 0.8.0). Once a production heuristic is selected and validated,
it should be ported into PyPTO and the external dependency removed.
Architecture
The DSA path skips opportunistic
MemoryReusecoalescing but reuses itsphi/loop-aware lifetime and constraint analysis. Each physical allocation is
exported with one conservative lifetime hull; unproven gaps between SSA aliases
are not exposed as reusable memory.
Standard DSA geometry
The former global whole-slot restriction is disabled. A freed region may be
subdivided among smaller co-live buffers, which is the geometry required by
#1908.
PyPTO contributes constraints and provenance:
The uncalibrated
pipeline_serializationcost proxy is research-only and is notused by the production first-fit path.
Validation
Current exact pins:
27ceb9f60af3c08a6e897cd9Host and CI gates:
unit tests, system/simulator tests, and PyPTO-Lib models
co-live 32 KiB consumers at a 64 KiB peak
if-phi/loop-carry writebacks
Current device blocker: partial-overlap dependency
The blocking 910B2 gather gate fails deterministically under DSA on four
devices while PyPTO passes:
test_gather.py::TestGatherIndex::test_gather_tile_input_source_unchanged[a2a3]The standard DSA placement reuses part of an expired 4 KiB allocation:
The placement is valid against the exported sequential lifetimes. However,
PTOAS v0.48 does not emit the required cross-pipe WAR synchronization for this
partial overlap. The later 32-byte load can overwrite data still being consumed
by the move, producing exactly eight corrupted FP32 elements. Equal-base reuse
is handled, but arbitrary range overlap is not currently represented safely
through the PyPTO fixed-address path.
This is not the earlier multi-interval-hole defect and not a standalone-solver
feasibility bug. The remaining design work is to preserve standard DSA
subdivision while making partial-overlap reuse dependency-safe, either through
range-aware downstream dependency handling or a narrowly grounded PyPTO
constraint/repair. Reintroducing the former global whole-slot rule would hide
the defect but would also remove the #1908 subdivision capability.
Until this blocker is fixed and rerun, this PR makes no end-to-end correctness
or performance claim for standard partial-overlap placement.
Remaining WIP
output-stationary workaround
objective
Corpus and research
The standalone repository contains 452 unique checked-in PyPTO/PyPTO-Lib
instances after structural deduplication, plus the MiniMalloc benchmark
family. Benchmark reports keep standard DSA peak/runtime separate from
PyPTO-specific research objectives.
Related