feat(codegen): add use_ptoas_multi_buffer switch for ptoas multi-buffer #1953
feat(codegen): add use_ptoas_multi_buffer switch for ptoas multi-buffer #1953lyfne123 wants to merge 6 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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds a new ChangesPtoas Multi-Buffer Feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant RunConfig
participant compile
participant PassManager
participant ConvertToPtoasMultiBuffer
participant PTOCodegen
RunConfig->>compile: use_ptoas_multi_buffer flag
compile->>compile: resolve effective flag (arg/outer/env)
compile->>PassManager: run pipeline with PassContext(mbuf=flag)
PassManager->>ConvertToPtoasMultiBuffer: process pipeline loops
ConvertToPtoasMultiBuffer->>ConvertToPtoasMultiBuffer: tag eligible calls, downgrade pipeline_stages
PassManager->>PTOCodegen: generate function
PTOCodegen->>PTOCodegen: PrescanAndHoistMultiBuffer, emit pto.alloc_multi_tile
PTOCodegen->>PTOCodegen: EmitAllocTileForVar emits pto.multi_tile_get(iv % N)
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 the ConvertToPtoasMultiBuffer pass, which optimizes same-core pipeline loops by converting ping-pong buffers into ptoas multi-buffer slots. When enabled via PassContext.use_ptoas_multi_buffer, the pass tags eligible vector/matrix tile-producing calls with the slot count and downgrades pipeline stages to 1. Consequently, codegen hoists pto.alloc_multi_tile allocations and emits pto.multi_tile_get inside the loop, while the memory allocator reserves contiguous slots and the memory reuse pass ensures exclusive allocation. Feedback on the changes highlights a potential segmentation fault in src/ir/transforms/memory_reuse_pass.cpp where the call pointer is dereferenced without a null check.
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: c906a170a9
ℹ️ 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/pypto/ir/pass_manager.py (1)
454-469: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
_run_with_profilingdropsmemory_planner/use_ptoas_multi_bufferfrom the reconstructed context.This is the same bug class the PR just fixed in
run_passes's dump-mode path (lines 417-427), but it's still present here: whenCompileProfiler.current()is active anddump_ir=False, the rebuiltPassContextomitsmemory_planner/use_ptoas_multi_buffer, silently resetting them to defaults (PyPTO/False). Any caller withuse_ptoas_multi_buffer=Trueormemory_planner=PTOAScombined with an active profiler would get the wrong pipeline behavior without any error.🐛 Proposed fix
ctx = passes.PassContext.current() outer_instruments = list(ctx.get_instruments()) if ctx else [] level = ctx.get_verification_level() if ctx else passes.get_default_verification_level() dphase = ctx.get_diagnostic_phase() if ctx else passes.get_default_diagnostic_phase() if ctx: disabled = ctx.get_disabled_diagnostics() else: disabled = passes.DiagnosticCheckSet() disabled.insert(passes.DiagnosticCheck.UnusedControlFlowResult) + mplan = ctx.get_memory_planner() if ctx else passes.MemoryPlanner.PYPTO + mbuf = ctx.use_ptoas_multi_buffer() if ctx else False - with passes.PassContext([*outer_instruments, timing_instrument], level, dphase, disabled): + with passes.PassContext( + [*outer_instruments, timing_instrument], level, dphase, disabled, mplan, mbuf + ): try: return self._pipeline.run(input_ir) finally: if stage_open: prof._end_stage()🤖 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 `@python/pypto/ir/pass_manager.py` around lines 454 - 469, The reconstructed PassContext in _run_with_profiling is missing the memory_planner and use_ptoas_multi_buffer settings, so profiling mode silently resets them to defaults. Update the PassContext creation in _run_with_profiling to preserve those fields from the current context, matching the fix already applied in run_passes’ dump-mode path, so any active CompileProfiler path keeps the same pipeline configuration.
🧹 Nitpick comments (1)
src/ir/transforms/allocate_memory_addr_pass.cpp (1)
285-304: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider asserting the "singleton multi-buffer group" invariant.
Both the
slotsoverwrite inMultiBufferBaseCollectorand thereserve_sizecomputation inAllocateMemoryAddressesrely on MemoryReuse keeping each multi-buffer tile's base as a singleton group — currently only documented in comments, not enforced. AnINTERNAL_CHECKwould turn a future violation (e.g. a view op sharing the multi-buffer tile's base, or two AssignStmts tagging the same base with differentN) into a loud failure instead of a silently wrong reserved region / high-water mark.🛡️ Suggested defensive check
auto mbit = multi_buffer_bases.find(base_key); const bool is_multi_buffer = mbit != multi_buffer_bases.end(); + if (is_multi_buffer) { + INTERNAL_CHECK_SPAN(group.size() == 1, group.front()->span_) + << "ptoas multi-buffer base '" << base_key->name_hint_ + << "' unexpectedly has " << group.size() << " MemRef members; expected a singleton group"; + } const uint64_t reserve_size = is_multi_buffer ? slot_size * static_cast<uint64_t>(mbit->second) : slot_size;Also applies to: 355-377
🤖 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/allocate_memory_addr_pass.cpp` around lines 285 - 304, Add a defensive INTERNAL_CHECK for the singleton multi-buffer group invariant in the memory allocation pass: both MultiBufferBaseCollector::VisitStmt_ and the reserve_size logic in AllocateMemoryAddresses currently assume each multi-buffer tile base appears in only one group and never gets conflicting slot counts. Enforce this by checking that a base is either unseen or already mapped to the same slots value before writing base_to_slots_, and by asserting the expected singleton ownership when reserve_size is computed so any shared base or conflicting AssignStmt tagging fails loudly instead of producing a wrong reservation.
🤖 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/jit/decorator.py`:
- Line 1128: The JIT cache key is missing the use_ptoas_multi_buffer flag, so
different `@pl.jit` configurations can reuse the same compiled artifact. Update
make_cache_key(...) in decorator.py to include run_config.use_ptoas_multi_buffer
alongside the other codegen-affecting options, using the existing cache-key
construction around the JIT decorator path so distinct values produce distinct
keys.
In `@src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp`:
- Around line 67-86: Tuple-returning tile ops are not getting multi-buffer tags
because VisitStmt_(const AssignStmtPtr&) only checks assign->var_ for TileType
and ignores tuple-typed calls like tile.gather_compare. Update the logic in
ConvertToPtoAsMultiBufferPass::VisitStmt_ to detect eligible tile operations
from the assigned call/result type as well as the destination, and apply
kMultiBufferSlotsAttr to the tensor-producing call even when the assignment is
TupleType. Keep the existing idempotency check and attr stripping, but make sure
tuple-valued tile ops are included in the eligibility path so pipelined
compare-gather lowering still multi-buffers.
In `@tests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.py`:
- Around line 86-87: Update the bare exception handlers in the test helper to
silence both Ruff rules by extending the existing noqa on the `except Exception`
blocks in `test_ptoas_multi_buffer_codegen.py`; the current `# noqa: BLE001`
only covers the blind-except warning, so add S110 to the same comment at both
call sites in the `except Exception` branches. Use the existing `except
Exception` blocks in this test as the anchors and keep the change limited to the
two post-codegen failure swallows.
---
Outside diff comments:
In `@python/pypto/ir/pass_manager.py`:
- Around line 454-469: The reconstructed PassContext in _run_with_profiling is
missing the memory_planner and use_ptoas_multi_buffer settings, so profiling
mode silently resets them to defaults. Update the PassContext creation in
_run_with_profiling to preserve those fields from the current context, matching
the fix already applied in run_passes’ dump-mode path, so any active
CompileProfiler path keeps the same pipeline configuration.
---
Nitpick comments:
In `@src/ir/transforms/allocate_memory_addr_pass.cpp`:
- Around line 285-304: Add a defensive INTERNAL_CHECK for the singleton
multi-buffer group invariant in the memory allocation pass: both
MultiBufferBaseCollector::VisitStmt_ and the reserve_size logic in
AllocateMemoryAddresses currently assume each multi-buffer tile base appears in
only one group and never gets conflicting slot counts. Enforce this by checking
that a base is either unseen or already mapped to the same slots value before
writing base_to_slots_, and by asserting the expected singleton ownership when
reserve_size is computed so any shared base or conflicting AssignStmt tagging
fails loudly instead of producing a wrong reservation.
🪄 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: db460d6a-4378-43dd-a6ce-168f254b33cd
📒 Files selected for processing (23)
CMakeLists.txtdocs/en/dev/passes/00-pass_manager.mddocs/zh-cn/dev/passes/00-pass_manager.mdinclude/pypto/codegen/pto/pto_codegen.hinclude/pypto/ir/transforms/pass_context.hinclude/pypto/ir/transforms/pass_properties.hinclude/pypto/ir/transforms/passes.hinclude/pypto/ir/transforms/utils/attrs.hpython/bindings/modules/passes.cpppython/pypto/ir/compile.pypython/pypto/ir/pass_manager.pypython/pypto/jit/decorator.pypython/pypto/pypto_core/passes.pyipython/pypto/runtime/runner.pysrc/codegen/pto/pto_codegen.cppsrc/codegen/pto/pto_control_flow_codegen.cppsrc/ir/transforms/allocate_memory_addr_pass.cppsrc/ir/transforms/convert_to_ptoas_multi_buffer_pass.cppsrc/ir/transforms/memory_reuse_pass.cppsrc/ir/transforms/pass_context.cppsrc/ir/transforms/python_printer.cpptests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.pytests/ut/ir/transforms/test_pass_manager.py
- cache: include use_ptoas_multi_buffer in the JIT cache key (cache.py + decorator.py) so the same kernel compiled with the switch on vs off no longer collides in the cache (codex/CodeRabbit). - ConvertToPtoasMultiBuffer: stop the body tagger at ANY nested loop (not just Pipeline), matching codegen's PrescanAndHoistMultiBuffer boundary, so a tile in a nested sequential loop is no longer tagged-but-never-hoisted (codex); note that multi-output (TupleType) ops keep ordinary allocation (CodeRabbit). - pre-commit: clang-format the touched C++, add <utility> include for std::move in pto_control_flow_codegen.cpp (cpplint), extend test noqa to S110, add PLR0912/PLR0913 noqa, shorten an over-long stub docstring. memory_reuse null-deref (gemini) is a false positive: `if (!call) continue;` already guards the dereference on the line above.
Add an opt-in path that delegates pl.pipeline ping-pong to ptoas's multi-buffer slots instead of pypto's body-replication double-buffering. Gated by PassContext.use_ptoas_multi_buffer (env PYPTO_PTOAS_MULTI_BUFFER, or compile(use_ptoas_multi_buffer=...)); off by default, so the default pipeline is unchanged. When enabled, the new ConvertToPtoasMultiBuffer pass (runs just before LowerPipelineLoops) tags vec/mat tile-producing Calls in same-core pipeline loops with kMultiBufferSlotsAttr = N and downgrades pipeline_stages to 1 so the loop body is not replicated. Codegen then hoists one pto.alloc_multi_tile addr=<base> count=N per tile and emits pto.multi_tile_get %mb[iv % N] at each use. AllocateMemoryAddr reserves N contiguous slots for such a tile and MemoryReuse keeps it exclusive, so ptoas's base fan-out (slot k at base + k*slotBytes) lands in a clear region; ptoas owns slot rotation and cross-iteration synchronization. Verified end-to-end: emitted .pto compiles under ptoas 0.48 with correct per-slot addresses and auto-inserted set_flag/wait_flag sync.
The switch silently no-op'd on the default path. In dump mode (compile() defaults dump_passes=True), PassManager.run_passes rebuilds the PassContext the pipeline runs under but only carried over verification level / phase / disabled diagnostics — it dropped use_ptoas_multi_buffer, so ConvertToPtoasMultiBuffer read Current()->UsePtoasMultiBuffer() == False and became a no-op. Only dump_passes=False paths (e.g. RunConfig default) ran the transform, which is why effectiveness flipped with dump_passes. - PassManager.run_passes: propagate use_ptoas_multi_buffer into the dump-mode PassContext (with a note that all outer-context config must be carried over). - compile(): resolve the flag as explicit arg > outer-context (on only) > env, so PYPTO_PTOAS_MULTI_BUFFER applies whether or not an outer PassContext is active (previously env was read only when outer was None). - Thread the flag explicitly end-to-end: add RunConfig.use_ptoas_multi_buffer (tri-state; None defers to env), forward it from run() and the @pl.jit RunConfig->compile kwarg builder. Explicit RunConfig/compile is now the reliable source of truth rather than env + outer context. Adds a regression test that the RunConfig field reaches the pass with dump_passes=True.
The multi_buffer_slots attr (set by ConvertToPtoasMultiBuffer, read by
AllocateMemoryAddr + PTO codegen) had no print surface, so it was dropped on
any print -> parse round-trip of the IR — leaving it invisible in dumps and
lost whenever a dumped .pto is reparsed.
Add it to the op-call attr allowlist in the python printer alongside
pipeline_membership; it now emits as attrs={"multi_buffer_slots": N}. The
parser already recovers generic int attrs (_parse_op_attrs -> _parse_attrs_dict
-> set_call_attrs stores Python int as C++ int, matching PrintAttrValue), so no
parser change is needed.
Adds a round-trip test asserting the attr survives print -> parse with
structural equality.
The Default and DebugTileOptimization pass-name inventories assert the exact pipeline; add the gated ConvertToPtoasMultiBuffer entry (between SkewCrossCorePipeline and LowerPipelineLoops) so the equality checks match.
- cache: include use_ptoas_multi_buffer in the JIT cache key (cache.py + decorator.py) so the same kernel compiled with the switch on vs off no longer collides in the cache (codex/CodeRabbit). - ConvertToPtoasMultiBuffer: stop the body tagger at ANY nested loop (not just Pipeline), matching codegen's PrescanAndHoistMultiBuffer boundary, so a tile in a nested sequential loop is no longer tagged-but-never-hoisted (codex); note that multi-output (TupleType) ops keep ordinary allocation (CodeRabbit). - pre-commit: clang-format the touched C++, add <utility> include for std::move in pto_control_flow_codegen.cpp (cpplint), extend test noqa to S110, add PLR0912/PLR0913 noqa, shorten an over-long stub docstring. memory_reuse null-deref (gemini) is a false positive: `if (!call) continue;` already guards the dereference on the line above.
make_cache_key now appends ("use_ptoas_multi_buffer", None) to compile_opts;
update the exact-tuple assertion in test_cache.py to match (unit-tests CI).
lyfne123
force-pushed
the
feat/ptoas-multi-buffer-switch
branch
from
60e1020 to
590c1a1
Compare
|
Took a close look at the multi-buffer lowering and verified a few things on the emitted H1 (high) — non-unit-step pipeline loops silently lose rotationThe slot index is scf.for %i = 0 to 8 step 2 { // %i ∈ {0,2,4,6}
%0 = arith.remui %i, %c2 : index // ≡ 0 every iteration
%t = pto.multi_tile_get %mb[%0] // always slot 0 — ping-pong silently defeated
M1 (medium) — loop-carried accumulators get tagged; accumulator buffer over-reservedThe tagger tags any single-output vec/mat M2 (medium) —
|
2 participants
Add an opt-in path that delegates pl.pipeline ping-pong to ptoas's
multi-buffer slots instead of pypto's body-replication double-buffering.
Gated by PassContext.use_ptoas_multi_buffer (env PYPTO_PTOAS_MULTI_BUFFER,
or compile(use_ptoas_multi_buffer=...)); off by default, so the default
pipeline is unchanged.
When enabled, the new ConvertToPtoasMultiBuffer pass (runs just before
LowerPipelineLoops) tags vec/mat tile-producing Calls in same-core pipeline
loops with kMultiBufferSlotsAttr = N and downgrades pipeline_stages to 1 so
the loop body is not replicated. Codegen then hoists one
pto.alloc_multi_tile addr= count=N per tile and emits
pto.multi_tile_get %mb[iv % N] at each use. AllocateMemoryAddr reserves N
contiguous slots for such a tile and MemoryReuse keeps it exclusive, so
ptoas's base fan-out (slot k at base + k*slotBytes) lands in a clear
region; ptoas owns slot rotation and cross-iteration synchronization.
Verified end-to-end: emitted .pto compiles under ptoas 0.48 with correct
per-slot addresses and auto-inserted set_flag/wait_flag sync.