perf(autotile): Double-buffer L0C in eligible pipelines - #2178
perf(autotile): Double-buffer L0C in eligible pipelines#2178tonibohnlein wants to merge 7 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. 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:
📝 WalkthroughWalkthroughAutoTileMatmulL0 now recognizes eligible existing PyPTO pipelines and enables capacity-checked L0C double buffering. CanonicalizeIOOrder schedules depth-two compute/drain groups and rotates Acc memberships, with documentation and comprehensive tests covering supported and rejected pipeline shapes. ChangesExisting Pipeline L0C Double Buffering
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant AutoTileMatmulL0
participant LowerPipelineLoops
participant CanonicalizeIOOrder
participant MemoryReuse
AutoTileMatmulL0->>LowerPipelineLoops: set pipeline_double_buffer_c
LowerPipelineLoops->>CanonicalizeIOOrder: emit staged memberships
CanonicalizeIOOrder->>MemoryReuse: rotate Acc membership modulo 2
MemoryReuse->>MemoryReuse: allocate two L0C residues
Possibly related issues
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48407afe5c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/ir/transforms/canonicalize_io_order_pass.cpp (1)
55-66: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse
AppendPipelineMembershipfor the repack.
attrs.halready owns the"group:stage"/;encoding; hand-rolling it here means two places must stay in sync if the format changes.♻️ Suggested repack
std::string result; for (const auto& [group, member_stage] : memberships) { - if (!result.empty()) result += ";"; - result += std::to_string(group) + ":" + std::to_string(member_stage); + result = AppendPipelineMembership(result, group, member_stage); } return result;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ir/transforms/canonicalize_io_order_pass.cpp` around lines 55 - 66, Update RotateInnermostPipelineStage to repack memberships using the existing AppendPipelineMembership helper instead of manually concatenating group, stage, and semicolon separators. Preserve the current rotation logic and returned encoding while delegating format ownership to AppendPipelineMembership.src/ir/transforms/auto_tile_matmul_l0_pass.cpp (1)
1724-1795: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename
empty: the same set is used as the "no plan" sentinel and as the accepted-plan accumulator.Lines 1790-1794 insert the selected loops into a variable named
emptyand return it, which reads as a bug at a glance. Aplanlocal (returning{}for the early bails) is clearer, and the structured bindings then no longer need(void)discards.♻️ Suggested rename
std::unordered_set<const ForStmt*> BuildPipelineDbCPlan(const FunctionPtr& func) { - std::unordered_set<const ForStmt*> empty; const auto* ctx = PassContext::Current(); const MemoryPlanner planner = ctx ? ctx->GetMemoryPlanner() : MemoryPlanner::PyPTO; - if (planner != MemoryPlanner::PyPTO) return empty; + if (planner != MemoryPlanner::PyPTO) return {};...and at the end:
- for (const auto& [loop, candidate] : candidates.candidates) { - (void)candidate; - empty.insert(loop); - } - return empty; + std::unordered_set<const ForStmt*> plan; + for (const auto& entry : candidates.candidates) plan.insert(entry.first); + return plan;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ir/transforms/auto_tile_matmul_l0_pass.cpp` around lines 1724 - 1795, Rename the result set in BuildPipelineDbCPlan from empty to plan, returning an empty set literal or equivalent empty result on all early exits. Insert accepted loops into plan at the end, and remove the unnecessary (void) discards from the final structured binding since the candidate value is no longer unused.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/en/dev/passes/15-auto_tile_matmul_l0.md`:
- Around line 9-10: Update the AutoTile documentation and public API comment to
include tile.write wherever the accepted direct-drain operations are listed,
alongside tile.store. Apply this consistently in
docs/en/dev/passes/15-auto_tile_matmul_l0.md at lines 9-10 and 62-63,
docs/zh-cn/dev/passes/15-auto_tile_matmul_l0.md at lines 9-10 and 62-63, and
include/pypto/ir/transforms/passes.h at lines 475-477; no recognizer code change
is needed.
- Around line 9-10: Update the AutoTile transfer-count contract to match the
recognizer: describe only the selected moving operand’s direct Mat→L0 producer,
rather than requiring exactly one transfer. Apply this correction in
docs/en/dev/passes/15-auto_tile_matmul_l0.md at lines 9-10 and 62-63,
docs/zh-cn/dev/passes/15-auto_tile_matmul_l0.md at lines 9-10 and 62-63, and the
public API comment in include/pypto/ir/transforms/passes.h at lines 471-474; no
implementation change is requested.
In `@tests/ut/ir/transforms/test_auto_tile_matmul_l0.py`:
- Around line 2174-2208: Rename the intentionally unused loop-carried value q_r
in test_rejects_loop_carried_matmul_operand to use the underscore-prefixed
convention, matching nearby unused pipeline outputs and satisfying RUF059. Keep
the yielded q_i value and all test behavior unchanged.
---
Nitpick comments:
In `@src/ir/transforms/auto_tile_matmul_l0_pass.cpp`:
- Around line 1724-1795: Rename the result set in BuildPipelineDbCPlan from
empty to plan, returning an empty set literal or equivalent empty result on all
early exits. Insert accepted loops into plan at the end, and remove the
unnecessary (void) discards from the final structured binding since the
candidate value is no longer unused.
In `@src/ir/transforms/canonicalize_io_order_pass.cpp`:
- Around line 55-66: Update RotateInnermostPipelineStage to repack memberships
using the existing AppendPipelineMembership helper instead of manually
concatenating group, stage, and semicolon separators. Preserve the current
rotation logic and returned encoding while delegating format ownership to
AppendPipelineMembership.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 50259a56-797c-4e37-a9b7-01f1f9fa078a
📒 Files selected for processing (14)
docs/en/dev/passes/15-auto_tile_matmul_l0.mddocs/en/dev/passes/26-lower_pipeline_loops.mddocs/en/dev/passes/27-canonicalize_io_order.mddocs/zh-cn/dev/passes/15-auto_tile_matmul_l0.mddocs/zh-cn/dev/passes/26-lower_pipeline_loops.mddocs/zh-cn/dev/passes/27-canonicalize_io_order.mdinclude/pypto/ir/transforms/passes.hinclude/pypto/ir/transforms/utils/attrs.hpython/bindings/modules/passes.cpppython/pypto/pypto_core/passes.pyisrc/ir/transforms/auto_tile_matmul_l0_pass.cppsrc/ir/transforms/canonicalize_io_order_pass.cppsrc/ir/transforms/lower_pipeline_loops_pass.cpptests/ut/ir/transforms/test_auto_tile_matmul_l0.py
1 participant
Summary
matmul, matmul, drain, drain, while preserving the full user-selected pipeline depth for L0A/L0B and rotating only L0C over two slotsRelationship to #2131 and #2150
Related to #2131 and complementary to #2150.
Issue #2131 asks for a general public mechanism for stable physical buffer identity, runtime slot selection, destination binding, and explicit lifetime release. PR #2150 provides that explicit user-facing capability through tile buffer sets and destination-form operations, including independently selected L1/L0B/L0C rotations.
This PR does not replace or depend on #2150. It addresses the common canonical form automatically: a normal
pl.pipelinecontaining one already-L0 matmul, one stationary operand, one per-iteration Mat-to-L0 operand transfer, and one loop-carried drain. The source program does not need explicit buffers or annotations. For the #2131 reproduction, AutoTile recognizes the inner 128-wide L0 loop and requests a second accumulator slot automatically.The two changes therefore serve different layers:
Eligibility and profitability
The recognizer requires:
pl.pipeline(stage=F)withF >= 2and a trip count divisible byFtile.matmulwith static Left/Right operandstile.extractortile.move; the other operand is loop-invariantPath-specific profitability gates are intentionally conservative:
tile.store: trip count >= 4tile.assemble: trip count >= 8 and aligned Acc footprint >=ceil(L0C / 4)Before marking any loop, AutoTile computes a conservative whole-function L0C inventory, including pipeline replication of non-cube Acc producers, and adds one extra aligned accumulator slot for every accepted loop. The complete plan is declined unless it fits the backend L0C capacity.
The automatic recognizer is PyPTO-planner-only. PTOAS output remains unchanged: the reproduced pipeline already received four physical Acc placements there, and source marking was performance-neutral in the device study.
Lowering
AutoTile attaches
pipeline_double_buffer_c=trueand disables the weaker store-overlap policy on accepted loops.LowerPipelineLoopsthen tags the replicated cube accumulators with pipeline membership.CanonicalizeIOOrderemits depth-two compute/drain chunks and rotates only the innermost Acc membership modulo two beforeMemoryReuseconsumes it. This creates two genuinely co-live L0C accumulators without multiplying L0C usage by a deeper operand pipeline.Validation
Exact final rebased head:
cmake --build build --parallel 2PYTHONPATH=$(pwd)/python python -m pytest tests/ut/ir/transforms/test_auto_tile_matmul_l0.py tests/ut/ir/transforms/test_lower_pipeline_loops.py tests/ut/ir/transforms/test_canonicalize_io_order.py -qAscend 910B2 device campaigns on logic-equivalent pre-rebase commits:
The final rebases only incorporated unrelated upstream changes and documentation renumbering; no device-validated implementation behavior changed. The exact final head was rebuilt and host-tested as above.
Limitations
L0C / 4size threshold; broader trip-count/compute-drain modeling is left for follow-up rather than fitting another rule to one point