feat(passes): opt-in dbC=2 L0C double-buffer under the PyPTO memory planner (subsumes #1991)#2002
feat(passes): opt-in dbC=2 L0C double-buffer under the PyPTO memory planner (subsumes #1991)#2002tonibohnlein wants to merge 7 commits into
Conversation
…ership LowerPipelineLoops stamps a (group, stage) `pipeline_membership` on every tile a pipeline stage defines so MemoryReuse keeps the stages' clones in separate buffers. That is correct for the operands a stage *loads* — their loads overlap the previous stage's compute, so two stages' operand buffers are genuinely co-live and must stay apart — but wrong for the cube accumulator. A cube accumulator is not loaded; it is written by the single serialized cube (a `tile.matmul*` MAD), which retires one tile's MAD before starting the next regardless of how many stages the scheduler overlaps. So a stage's accumulator is never co-live with the next stage's accumulator. Tagging it anyway makes MemoryReuse's capacity-gate request one L0C buffer per stage and then shed it back onto one — a redundant separation that emits a spurious PH-MR-001, and for an accumulator nested N pipeline loops deep balloons to 2^N requested buffers before shedding. Skip cube-accumulator tiles in the membership tagger, gated on the producer being a `tile.matmul*` MAD (not on `Mem.Acc` alone — a data-movement op that also targets Acc, e.g. `tile.extract(..., target_memory=Acc)`, is a genuine per-stage buffer that overlaps across stages and must keep its membership). The drain-before-next accumulator now coalesces by lifetime alone onto the single L0C buffer it needs, and the spurious perf hint is gone. Operand tagging is unchanged, so the pipelined-operand double-buffer (hw-native-sys#1475) is unaffected. Coalescing accumulators is correctness-safe: a shared L0C buffer forces the same drain-before-next serialization the single-accumulator schedule already relies on. Split-K's in-place `matmul_acc` chain (a RAW dependency, lifetimes touch but don't overlap) also coalesces correctly and no longer gets spuriously separated. Add a regression test asserting a pipelined `tile.matmul` accumulator stays untagged while its operands — and a non-cube Acc producer (`tile.extract` into Acc) — carry distinct-stage membership, and update the pass docs (en + zh-cn) to record the cube-accumulator exception.
…s its ping-pong The cube-accumulator skip added in the previous commit leaves every pipelined accumulator untagged. That is correct for the single-accumulator schedule, but the dbC=2 double-buffered-L0C emit genuinely co-lives two accumulators — tile i's FIXPIPE drain overlaps tile i+1's MAD — so its accumulator DOES need a pipeline_membership to double-buffer under the PyPTO memory planner. Carve out the dbC case: the moving loop that carries kPipelineDoubleBufferCAttr tags its cube accumulator with its own (group, stage); every enclosing pipeline loop still skips it (the outer loop double-buffers operands, but the MADs writing successive outer-stage accumulators still serialize on the one cube). The result is a flat depth-2 membership from the dbC loop alone, so MemoryReuse's capacity-gate (hw-native-sys#1475) allocates exactly the two co-live L0C buffers instead of the per-loop cross-product it previously sheds back to a single serialized buffer. Verified under the PyPTO planner: the four unrolled accumulators carry a flat "0:0"/"0:1" membership and bind to two co-live 64KB L0C buffers (the MM,MM,ST,ST ping-pong), with no spurious PH-MR-001 — matching what the PTOAS planner produces via InitMemRef. dbC=2 activation under PyPTO remains gated off (a later step, pending the opt-in flag and device validation); this only wires the emit so the capacity-gate can honor it.
…t flag) The previous two commits made the PyPTO capacity gate able to honor the dbC=2 accumulator (flat depth-2 membership). This wires the activation: a PassContext flag `enable_pypto_l0c_double_buffer` (default off, experimental) that opens the chooser's dbC=2 design point under the PyPTO planner. cfg.allow_double_buffer_c = ptoas_planner || (pypto_planner && flag); PTOAS keeps dbC=2 unconditionally (device-validated). Under PyPTO the flag is off by default, so default codegen is byte-identical (the "PyPTO stays dbC=1" golden still holds). When on, AutoTileMatmulL0 emits the two co-live accumulators, the tagger gives them a flat depth-2 membership, and MemoryReuse's capacity gate (hw-native-sys#1475) allocates exactly the two co-live L0C buffers — verified end to end (flag off -> 1 Acc buffer; flag on -> 2, the MM,MM,ST,ST ping-pong). Cross-layer: PassContext ctor arg + getter (C++ header/impl), nanobind binding, type stub; threaded through `ir.compile` + `runtime.compile_program`. New unit test asserts the flag flips the PyPTO L0C buffer count 1<->2. Docs (14-auto_tile, en + zh-cn) updated to describe the opt-in and how the pair survives allocation under each planner. Also propagate the flag through PassManager's nested dump/profile PassContexts (the same fix 55301c1 applied to memory_planner) — the chooser reads it *during* pass execution, so a nested context that resets it would silently disable dbC=2 whenever the pipeline dumps IR (i.e. exactly the device-compile path). Kept default-off pending device validation of the numerics + drain-hidden win before any default flip.
Adds `test_direct_store_dbc_pypto` (4/8/16-tile grids) alongside the existing PTOAS cases in test_dbc2_double_buffer.py, run under `memory_planner=PYPTO` with `enable_pypto_l0c_double_buffer=True`. Same shapes and torch golden as the PTOAS sweep, so it is the numerics correctness gate for the new allocation path: under PyPTO the two co-live L0C accumulators are kept apart by MemoryReuse's capacity gate (not skipped as under PTOAS) and the reuse-WAR is enforced by PyPTO codegen sync, so a value check on a >=4-tile grid catches a wrong-order drain/MAD. Threads the opt-in through the harness: `PTOTestCase` gains a `get_enable_pypto_l0c_double_buffer()` hook (ctor override, default off) that the test runner passes to `compile_program`.
…m noqa Rebase cleanup after replaying the PyPTO-planner dbC=2 commits onto main (post-hw-native-sys#1959 merge): - The mat-scratch `xfail(run=False)` had drifted onto `test_direct_store_dbc_pypto` (the PyPTO dbC=2 direct-store test, which passes) instead of `test_mat_scratch_dbc`. Move it back and correct the reason: the mat-scratch failure under PTOAS is the chained-matmul consumer's K-reduction accumulator if-phi handle being split across two L0C buffers by PTO codegen (not the Acc->Mat assemble drain, not dbC-specific), fixed upstream by hw-native-sys#1961. - `compile_program` grew an 11th arg (enable_pypto_l0c_double_buffer); add `# noqa: PLR0913`, matching the sibling `compile()` entry point.
… loop Review nit: kPipelineDoubleBufferCAttr depends only on the loop op, not the clone index, so read it once above the replication loop instead of per clone. Behaviour-identical.
|
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds an experimental PyPTO L0C double-buffer flag, propagates it through compilation and test execution, updates pipeline membership and allocation behavior, and adds unit/system coverage plus English and Chinese documentation. ChangesL0C double-buffering
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Compile
participant PassContext
participant AutoTile
participant LowerPipelineLoops
participant MemoryReuse
Caller->>Compile: enable PyPTO dbC flag
Compile->>PassContext: resolve and store flag
PassContext->>AutoTile: query dbC enablement
AutoTile->>LowerPipelineLoops: emit pipeline double-buffer attribute
LowerPipelineLoops->>MemoryReuse: provide accumulator membership metadata
MemoryReuse->>MemoryReuse: allocate one or two L0C buffers
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 |
![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 an experimental opt-in flag, enable_pypto_l0c_double_buffer, to enable L0C double-buffering (dbC=2) under the PyPTO memory planner. It updates PassContext in C++ and its Python bindings, propagates the flag through the compilation pipeline, and modifies LowerPipelineLoops to avoid tagging cube accumulators unless they belong to a dbC=2 moving loop. Documentation and comprehensive unit/integration tests are also updated to validate this new allocation path. There are no review comments to address, and I have no additional feedback to provide.
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.
![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: 9cb71692d8
ℹ️ 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".
| dbc_flag = ( | ||
| enable_pypto_l0c_double_buffer | ||
| if enable_pypto_l0c_double_buffer is not None | ||
| else outer.get_enable_pypto_l0c_double_buffer() | ||
| ) |
There was a problem hiding this comment.
When this branch inherits enable_pypto_l0c_double_buffer from an outer PassContext, the flag changes the generated AutoTileMatmulL0/MemoryReuse output but the JIT cache does not distinguish it: python/pypto/jit/decorator.py builds the key from platform/strategy/analyze-auto-scopes/memory_planner and has no effective dbC flag. A @pl.jit kernel first called with the PyPTO flag off and later inside PassContext(..., enable_pypto_l0c_double_buffer=True) will reuse the dbC=1 artifact (or keep dbC=2 stuck on in the reverse order), so the opt-in is silently ignored for cached JIT calls; add the effective flag to the JIT compile kwargs/cache key.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 9002ecf — good catch. Added _resolve_enable_pypto_l0c_double_buffer() (reads the active PassContext, mirroring _resolve_memory_planner) and included it in make_cache_key at both call sites (decorator.py). So a JIT kernel compiled with the flag off no longer reuses that dbC=1 artifact when later called under PassContext(enable_pypto_l0c_double_buffer=True) (or vice versa). New UTs pin it: test_dbc_double_buffer_flag_splits_key (the flag splits the key) + test_reads_the_active_pass_context (the resolver reads the ambient context).
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/st/harness/core/harness.py (1)
191-226: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring missing new param.
enable_pypto_l0c_double_bufferis added to the constructor (Line 199) and stored (Line 225), but the docstring'sArgs:section doesn't document it, unlikememory_planner.📝 Proposed docstring addition
memory_planner: Override the on-chip memory planner (PYPTO/PTOAS). If None, falls back to ``get_memory_planner()`` (which returns None, deferring to ir.compile's PYPTO default). + enable_pypto_l0c_double_buffer: Override the experimental PyPTO + dbC=2 opt-in. If None, falls back to + ``get_enable_pypto_l0c_double_buffer()`` (which returns None, + deferring to ir.compile's default of off). """🤖 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 `@tests/st/harness/core/harness.py` around lines 191 - 226, Document the enable_pypto_l0c_double_buffer constructor parameter in the __init__ Args section, describing that it overrides whether PyPTO L0C double buffering is enabled and that None preserves the default behavior.
🤖 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 `@python/pypto/ir/compile.py`:
- Line 67: Document the enable_pypto_l0c_double_buffer parameter in the
compile() docstring’s Args section, alongside memory_planner. Describe that it
is experimental and that None inherits the outer context setting or defaults to
disabled.
---
Nitpick comments:
In `@tests/st/harness/core/harness.py`:
- Around line 191-226: Document the enable_pypto_l0c_double_buffer constructor
parameter in the __init__ Args section, describing that it overrides whether
PyPTO L0C double buffering is enabled and that None preserves the default
behavior.
🪄 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
Run ID: 3a018a0a-be33-4645-a44b-0260729309be
📒 Files selected for processing (18)
docs/en/dev/passes/14-auto_tile_matmul_l0.mddocs/en/dev/passes/25-lower_pipeline_loops.mddocs/zh-cn/dev/passes/14-auto_tile_matmul_l0.mddocs/zh-cn/dev/passes/25-lower_pipeline_loops.mdinclude/pypto/ir/transforms/pass_context.hpython/bindings/modules/passes.cpppython/pypto/ir/compile.pypython/pypto/ir/pass_manager.pypython/pypto/pypto_core/passes.pyipython/pypto/runtime/runner.pysrc/ir/transforms/auto_tile_matmul_l0_pass.cppsrc/ir/transforms/lower_pipeline_loops_pass.cppsrc/ir/transforms/pass_context.cpptests/st/harness/core/harness.pytests/st/harness/core/test_runner.pytests/st/runtime/ops/test_dbc2_double_buffer.pytests/ut/ir/transforms/test_auto_tile_matmul_l0.pytests/ut/ir/transforms/test_lower_pipeline_loops.py
…aram Review (hw-native-sys#2002): - codex (P2): `enable_pypto_l0c_double_buffer` changes the AutoTileMatmulL0 / MemoryReuse artifact, but the @pl.jit cache key did not distinguish it, so a kernel first compiled with the flag off (then called under `PassContext(enable_pypto_l0c_double_buffer=True)`) silently reused the dbC=1 artifact. Add `_resolve_enable_pypto_l0c_double_buffer()` (reads the active PassContext, like `_resolve_memory_planner`) and include it in `make_cache_key` at both call sites. New UTs: the flag splits the key + the resolver reads the ambient context; updated the pinned key-structure test. - coderabbit: document `enable_pypto_l0c_double_buffer` in the `compile()` docstring.
1 participant
Summary
Adds an opt-in dbC=2 (double-buffered L0C) path under the PyPTO memory planner, and the
tagger fix it depends on. dbC=2 keeps two co-live L0C accumulators ping-ponged so tile i's
FIXPIPE drain overlaps tile i+1's MAD. The PTOAS-planner dbC=2 path shipped in #1959; this
adds the PyPTO-planner path behind a flag that is default OFF (experimental).
This PR subsumes #1991 (its tagger fix is the first commit here) — see "Why coupled".
Part 1 — don't tag pipelined cube accumulators (was #1991)
LowerPipelineLoops'sPipelineMembershipTaggertagged every tile in a pipeline-stagebody with
pipeline_membership, including cube-matmul accumulators. An accumulator is writtenby the single serialized cube (drain-before-next) and is never co-live with the next stage's
accumulator, so tagging it made
MemoryReuse's #1949 capacity gate request one L0C buffer perstage and then shed it back — a spurious
PH-MR-001, and for an accumulator nested N loopsdeep, 2^N requested buffers. Fix: skip cube-matmul accumulators (
tile.matmul*writingMem.Acc); data-movement ops that target Acc (e.g.tile.extract(target_memory=Acc)) staytagged. This also removes a spurious
PH-MR-001on the plain non-dbC path.Part 2 — opt-in dbC=2 under the PyPTO planner
kPipelineDoubleBufferCAttr) does tag its cube accumulator → flat depth-2 membership0:0,0:1,0:0,0:1→ the MemoryReuse: capacity-gated packing fixes the pipelined-operand collapse (#1475) #1949 capacity gate allocates exactly two co-live L0C buffers(the ping-pong), reused across the grid.
PassContext(enable_pypto_l0c_double_buffer=…)(default False), threadedthrough
ir.compile/runtime.compile_program/ the s-test harness, and — importantly —through PassManager's two nested dump/profile
PassContexts (the same class of silent-resetbug previously fixed for
memory_planner). The chooser gate becomesallow_double_buffer_c = ptoas_planner || (pypto_planner && flag); the PTOAS path is unchanged.Why coupled
Part 1 is a hard prerequisite for Part 2: without the cube-accumulator skip, the dbC
accumulator gets nested 2×2 membership → the capacity gate sheds it to one buffer → no
ping-pong. Shipping them together keeps the change coherent.
Validation
test_dbc2_pypto_flag_allocates_ping_pongpins flag OFF → 1 acc buffer / ON → 2co-live;
test_pipelined_accumulator_is_not_taggedpins the skip vs. carve-out. Fulltransforms + pass suites green.
match the torch golden, and the two accumulators land in distinct L0C buffers at {0, 32 KB}.
threads through both nested contexts with no silent-reset regression. The one code nit (a
loop-invariant read) is fixed here; the two remaining are test-coverage follow-ups on the
default-off path (a dbC matmul nested in a user
pl.pipeline; the profiling-path forwarding).Notes
test_mat_scratch_dbcisxfail(run=False)— blocked on an unrelated codegen bug(feat(codegen): first-class alloc_tile op fixes PTOAS if/else-yield miscompile (#1956) #1961, a chained-matmul consumer accumulator handle split under PTOAS), not this path.