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 |
FangRui0
force-pushed
the
feature_memplan
branch
3 times, most recently
from
3078197 to
88f8522
Compare
2 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。
主要改动
保留旧版 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 负责选择实现：--plan-memory-impl=legacy：默认，走旧版 memplan。
--plan-memory-impl=modern：显式启用新 memplan。
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。
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 复用策略：
收集 local allocation root。
做 alias-aware liveness 分析。
记录 semantic conflict。
按 AddressSpace 独立规划。
使用 first-fit 候选地址选择。
将规划出的 offset materialize 回 PTO IR。
address space 相同。
容量足够。
生命周期不重叠。
无 semantic conflict。
暂不检查 element type / tile type 是否一致。
level1/level2：支持 auto = true。
level1/level2：拒绝用户显式指定 base。
level3：要求 auto = false 且显式提供 base。
auto = true 的 reserve buffer 在普通 local buffer 规划后，通过 aligned first-fit 在剩余空洞中分配 base。
测试
恢复之前删除的 test/samples/planmemory 相关 sample。
现有 test/lit/pto/plan_memory_*.pto 增加 modern RUN，使同一批用例同时覆盖 legacy 和 modern。
新增/更新 alloc_tile addr、reserve_buffer level 校验、modern SPEC_LEVEL_0 等 lit。
增加设计说明文档：docs/designs/ptoas-dual-memory-planner-design.md。
限制
默认仍使用 legacy memplan；modern memplan 需要显式通过 --plan-memory-impl=modern 启用。
modern 当前只实现基础 SPEC_LEVEL_0 复用。
modern 没有实现旧版 StorageEntry 合并框架。
modern 没有实现完整 SPEC_LEVEL_1 / SPEC_LEVEL_2。
modern 没有实现 pipe-conflict-aware / bank-conflict-aware planning。
modern 没有实现回滚式投机规划。
modern 当前控制流支持是 alias-aware 的保守传播，不等价于旧版强 inplace merge 模型。