Rewrite memory planning#913
Conversation
![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 restores the mainline local memory planner (PlanMemory pass) and introduces a fallback address resolution pass (AllocToPointerCast) to handle local allocation handles when automatic planning is absent. It also updates the PTOViewToMemref pass to preserve pto.alloc_tile as a tile-native allocation root, avoiding the intermediate memref.alloc generation. The review feedback highlights critical issues in the fallback address calculation within AllocToPointerCast.cpp that could lead to address collisions, and a missing case in getReserveBufferMemSpec within tools/ptoas/ptoas.cpp that would cause validation to fail for LEFT, RIGHT, and ACC address spaces.
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.
| explicit UnresolvedPointerCastToPointerCastPattern( | ||
| MLIRContext *context, | ||
| DenseMap<Value, SmallVector<uint64_t>> buffer2Offsets) | ||
| : OpRewritePattern<pto::PointerCastOp>(context), | ||
| buffer2Offsets(std::move(buffer2Offsets)) { | ||
| uint64_t maxOff = 0; | ||
| for (const auto &kv : this->buffer2Offsets) { | ||
| for (uint64_t off : kv.second) | ||
| maxOff = std::max(maxOff, off); | ||
| } | ||
| fallbackNextOffset = | ||
| ((maxOff + kDefaultAllocAlignmentBytes - 1) / | ||
| kDefaultAllocAlignmentBytes) * | ||
| kDefaultAllocAlignmentBytes; | ||
| } |
There was a problem hiding this comment.
Recalculating fallbackNextOffset by finding the maximum offset in buffer2Offsets and aligning it up is incorrect and can lead to address collisions. If maxOff is the start offset of the last buffer, setting fallbackNextOffset to alignUp(maxOff) will cause any subsequent fallback allocations to start at the same offset, resulting in overlapping memory ranges. Since the correct fallbackNextOffset is already computed and advanced in runOnOperation(), it should be passed directly to the pattern constructor.
explicit UnresolvedPointerCastToPointerCastPattern(
MLIRContext *context,
DenseMap<Value, SmallVector<uint64_t>> buffer2Offsets,
uint64_t fallbackNextOffset)
: OpRewritePattern<pto::PointerCastOp>(context),
buffer2Offsets(std::move(buffer2Offsets)),
fallbackNextOffset(fallbackNextOffset) {}| patterns.add<UnresolvedPointerCastToPointerCastPattern>( | ||
| patterns.getContext(), std::move(buffer2Offsets)); |
There was a problem hiding this comment.
Pass the advanced fallbackNextOffset from runOnOperation() to the pattern constructor to avoid incorrect recalculation.
| patterns.add<UnresolvedPointerCastToPointerCastPattern>( | |
| patterns.getContext(), std::move(buffer2Offsets)); | |
| patterns.add<UnresolvedPointerCastToPointerCastPattern>( | |
| patterns.getContext(), std::move(buffer2Offsets), fallbackNextOffset); |
| case AddressSpace::LEFT: | ||
| case AddressSpace::RIGHT: | ||
| case AddressSpace::ACC: | ||
| case AddressSpace::BIAS: | ||
| case AddressSpace::SCALING: | ||
| case AddressSpace::GM: | ||
| case AddressSpace::Zero: | ||
| break; |
There was a problem hiding this comment.
The helper getReserveBufferMemSpec is missing cases for the local memory spaces LEFT (L0A), RIGHT (L0B), and ACC (L0C). As a result, any reserve_buffer declared in these spaces will return an empty ReserveBufferMemSpec (with capacity 0), causing validateReserveBufferBase to always fail with a misleading capacity overflow error. Adding the correct capacities and alignments for these spaces prevents false validation failures.
case AddressSpace::LEFT:
return {65536ull, 4096};
case AddressSpace::RIGHT:
return {65536ull, 4096};
case AddressSpace::ACC:
return {arch == PTOArch::A5 ? 262144ull : 131072ull, 4096};
case AddressSpace::BIAS:
case AddressSpace::SCALING:
case AddressSpace::GM:
case AddressSpace::Zero:
break;
Codex Review该评论由 review 机器人自动更新。
SummaryReview failed at stage Findings未生成结构化 findings，因为 review 过程提前失败。 Log Tail |
2809e7f to
72a7f11
Compare
7b42734 to
99270cd
Compare
b9e1bb5 to
ac0493e
Compare
zhangstevenunity
left a comment
zhangstevenunity
left a comment
There was a problem hiding this comment.
Focused correctness review (built this PR head locally, ptoas 0.51, and ran targeted repros). CI is green and the default path stays on legacy; modern is opt-in via --plan-memory-impl=modern.
Findings
- P2 (confirmed miscompile, modern only): the modern planner reuses a loop-invariant buffer's address for a scratch buffer allocated later in the same loop iteration. See inline on
PTOPlanMemoryModern.cpp. Reproduced against this exact build. - P3: redundant double-diagnostic for level1/2
reserve_bufferwith explicit base (inline onptoas.cpp). Same comment notes that the Gemini bot's HIGH finding aboutgetReserveBufferMemSpecmissing LEFT/RIGHT/ACC is actually unreachable (the op verifier restrictsreserve_bufferlocation to vec/mat), so no fix is needed there.
Looks sound
- Legacy liveness changes (GEN = all operands,
tile_buf_addr/TPrintas reads, tile-nativealloc_tileroot collection) broaden liveness in the conservative/safe direction; legacy keeps using MLIR dataflowLiveness, so it is correct on the loop case that breaks modern. - InsertSync cross-root overlap addition (
isLocalBufferOverlapCrossRoot) only ever adds hazard edges when both sides have known absolute addresses and returns the historicalfalsefor unknown addresses, so it cannot introduce a missed-sync. - The two earlier Gemini HIGH findings are stale/non-issues:
UnresolvedPointerCastToPointerCastPatternwas deleted with the standalone pass, and thegetReserveBufferMemSpecgap is unreachable (see above).
Minor consideration (not blocking): broadening legacy GEN to all operands keeps input buffers live longer, which slightly reduces reuse. Worth a sanity check that capacity-tight default kernels don't newly overflow VEC, but the sample/vpto-sim CI passing suggests no widespread regression.
| } | ||
| }; | ||
|
|
||
| static bool lifetimesStrictlyOverlap(const RootInfo &lhs, const RootInfo &rhs) { |
There was a problem hiding this comment.
P2 (confirmed silent miscompile): a loop-invariant buffer read inside a loop can be given the same address as a scratch buffer allocated later in the loop.
Liveness here is a single linear pre-order index: markUse/markWrite set allocIndex..freeIndex, and lifetimesStrictlyOverlap compares those integers. A root defined outside a loop but read inside it gets freeIndex = its last static use index in the body; nothing extends its live range across the loop back-edge to cover the whole scf.for. So a scratch allocated later in the same iteration looks non-overlapping, and canShare() places it at the same offset. The scratch's per-iteration write then clobbers the loop-invariant value, and the next iteration reads garbage.
InsertSync cannot rescue this: it orders ops within one iteration but cannot prevent cross-iteration clobbering of an aliased address; the two buffers are logically distinct yet simultaneously live.
Reproduced against this PR head (ptoas 0.51):
%inv = memref.alloc() : memref<16x16x16xf16, #pto.address_space<vec>>
pto.tload ins(%arg0) outs(%inv) // loaded once, must survive all iterations
scf.for %i = %c0 to %c4 step %c1 {
pto.tstore ins(%inv) outs(%arg1) // loop-invariant read
%tmp = memref.alloc() : memref<16x16x16xf16, #pto.address_space<vec>>
pto.tload ins(%arg2) outs(%tmp) // scratch written every iteration
pto.tstore ins(%tmp) outs(%arg1)
}--mlir-print-ir-after=pto-plan-memory:
- legacy:
%inv -> pointer_cast(0),%tmp -> pointer_cast(8192)(distinct, correct) --plan-memory-impl=modern:%inv -> pointer_cast(0),%tmp -> pointer_cast(0)(same address -> miscompile)
Loop-invariant tiles read inside a loop (weights, scales, KV, ...) are extremely common, so this will hit real kernels once modern is enabled. Existing tests only cover buffers used after the loop (plan_memory_loop_no_reuse_outer_live) and iter_args (seedForIterArgAliases), not loop-invariant reads inside the loop.
Fix direction: when walking a loop region, extend freeIndex of any root defined outside the loop but used inside it to at least the loop op's index (symmetrically, pull allocIndex back for a root written inside a loop but read after it). Legacy is correct precisely because MLIR's dataflow Liveness already models the back-edge.
| } | ||
|
|
||
| if (op.getBaseAttr()) | ||
| (void)validateReserveBufferBase(op, arch); |
There was a problem hiding this comment.
P3: redundant/confusing double diagnostic. For a level1/level2 reserve_buffer with auto = false and a base, this first calls validateReserveBufferBase (which emits its own error, including a bogus reserved range exceeds ... capacity: base N + size M > 0 bytes whenever getReserveBufferMemSpec returns capacity 0), then unconditionally emits the not supported when --pto-level=level1 or level2 error below. That stacks two errors for one condition. Since the op is rejected here regardless, drop the validateReserveBufferBase call and emit only the not supported message.
Related: the Gemini bot flagged (HIGH) that getReserveBufferMemSpec is missing LEFT/RIGHT/ACC/BIAS/SCALING (capacity 0). That path is actually unreachable -- pto.reserve_buffer's verifier restricts location to vec/mat (confirmed: location = #pto.address_space<acc> fails with "expects 'location' to be #pto.address_space or #pto.address_space"), and both vec/mat have correct capacities. So no fix is needed for the missing cases; only this stray call is worth cleaning up.
3 participants
概述
这个 PR 恢复并保留旧版 PTOAS local memory planner，同时新增一套可选的 modern memplan 实现。默认情况下，level1/level2 仍走 legacy memplan，以保持既有行为和兼容性；开发中的新规划器通过 --plan-memory-impl=modern 显式启用。
本次改动继续保留 tile-native pto.alloc_tile lowering 路径：level1/level2 下，pto.alloc_tile(no addr) 不再先降成 memref.alloc，而是透传到 memplan，由 memplan 直接补充常量 addr，后续 pto.t* tile op 继续消费 !pto.tile_buf。
modern memplan 在引入 Largest-First-Fit 和四道禁止复用闸门，用于更系统地约束地址复用安全性。
主要改动
保留旧版 memplan，文件仍为 lib/PTO/Transforms/PTOPlanMemory.cpp / PTOPlanMemory.h。
新增 modern memplan，文件为 lib/PTO/Transforms/PTOPlanMemoryModern.cpp。
pto-plan-memory 仍是 ModuleOp 级别 pass，内部遍历 func::FuncOp 做函数级 local memory planning。
tools/ptoas/ptoas.cpp 负责选择实现：
当显式启用 modern memplan 且用户未指定 --plan-memory-order-by-size 时，modern 默认开启 largest-first 排序；legacy 仍保持原默认行为。
level1/level2：禁止用户显式写 pto.alloc_tile addr。
level1/level2：保留 pto.alloc_tile(no addr) 到 pto-plan-memory，规划完成后直接填充 addr。
level3：要求用户显式提供 local 地址。
pto.alloc_tile(no addr) 不再走 memref.alloc -> pto.pointer_cast -> pto.bind_tile 中间路径。
legacy 和 modern memplan 都支持直接规划 tile-native pto.alloc_tile root。
lowering 路径变化：
非 alloc_tile 的 memref/address root 仍可在 materialize 阶段使用 pto.pointer_cast + pto.bind_tile 表达规划后的地址与 tile metadata。
legacy memplan 增加对 pto.alloc_tile(no addr) 的 root 收集。
BufferInfo 支持从 TileBufType 计算 shape、element type 和 size。
materialize 阶段支持把规划 offset 回写成 pto.alloc_tile addr = ...。
pto.tile_buf_addr 在 legacy liveness 中建模为读取 tile source，避免被未知 local-buffer op 误判。
modern memplan 当前聚焦基础 SPEC_LEVEL_0 复用策略：
modern 的复用判定由四道禁止复用闸门约束：
生命周期闸门：生命周期重叠则禁止复用，touching 场景继续进入后续闸门。
phi family 闸门：为 branch/yield/iter_arg alias 场景预留互斥控制流建模。
target-specific load/tpop hazard：用于 A3 等 target 上 load-derived buffer 与特定 writer touching 复用的硬件/后端风险。
op semantic no-alias：建模 PyPTO 风格 inplace 约束，以及 legacy 中已有的 scratch-output conflict。
level1/level2：
支持 auto = true。
拒绝用户显式指定 base。
auto = true 的 reserve buffer 在普通 local buffer 规划后，通过 aligned first-fit 在剩余空洞中分配 base。
level3：
要求 auto = false。
要求显式提供 base。
测试
恢复之前删除的 test/samples/planmemory 相关 sample。
现有 test/lit/pto/plan_memory_*.pto 增加 modern RUN，使同一批用例同时覆盖 legacy 和 modern。
新增/更新覆盖：
pto.alloc_tile addr level 校验。
tile-native pto.alloc_tile(no addr) 不经过 memref.alloc。
reserve_buffer level 校验与 auto base 规划。
modern Largest-First-Fit。
modern 四道闸门相关用例。
DPS output writer liveness。
target hazard / semantic no-alias 场景。
same-address UB overlap sync 相关回归用例。
限制
默认仍使用 legacy memplan；modern 需要显式通过 --plan-memory-impl=modern 启用。
modern 仍不作为默认生产路径。
modern 没有实现旧版 StorageEntry 合并框架。
modern 没有实现完整 SPEC_LEVEL_1 / SPEC_LEVEL_2。
modern 没有实现旧版回滚式投机规划。
modern 没有实现 bank-conflict-aware planning。
pipeline membership 相关闸门当前不实现；PTOAS 显式 multibuffer 已覆盖 ping-pong slot 分离语义。
aggressive 地址复用后的跨 pipe hazard 需要同步 pass 配合保证正确性。