feat(codegen): Lower a multi-slot MemRef to a ptoas multi-buffer region - #2257
feat(codegen): Lower a multi-slot MemRef to a ptoas multi-buffer region#2257lyfne123 wants to merge 1 commit 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:
📝 WalkthroughWalkthroughThe PR adds PTOAS support for declarative multi-slot ChangesPTOAS multi-slot MemRef support
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
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: dea0372885
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| } else if (type_str != candidate.slot_type_str) { | ||
| candidate.blocker = "its slots hold differently shaped tiles, and ptoas slots are uniform"; |
There was a problem hiding this comment.
When slots have the same physical tile type but different TileView.valid_shape values—or a later slot has a runtime valid shape—this comparison still considers them uniform because GetTileBufTypeStringFromTileType() deliberately renders v_row=? and v_col=?. TryComputeStaticValidShape() is then applied only to the first tile, and its extents are shared by the entire alloc_multi_tile region, so other slots silently receive the wrong valid extent and can truncate or overrun computation. Compare every slot's effective valid shape and reject mismatches or dynamic extents before creating the region.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (3)
include/pypto/codegen/pto/pto_codegen.h (1)
873-878: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClear the new region state in
FunctionState::Reset().
Reset()at Lines 940-990 does not clearmulti_buffer_regionsormulti_buffer_region_order. Every other member ofFunctionStateis cleared there. Today this is safe becausePlanMultiBufferRegionsclears both at its start, andGenerateFunctionalways calls it. The state is therefore correct only because of the call order betweenfs_.Reset()(Line 621 ofsrc/codegen/pto/pto_codegen.cpp) andPlanMultiBufferRegions(Line 870). If a later change reads a region before planning runs, the previous function's%mbhandle is returned. Add the clears soReset()matches its documented contract.♻️ Proposed change in
FunctionState::Reset()extra_alloc_tiles.clear(); ssa_to_tile_buf_type.clear(); subview_materializations.clear(); + + multi_buffer_regions.clear(); + multi_buffer_region_order.clear();🤖 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 `@include/pypto/codegen/pto/pto_codegen.h` around lines 873 - 878, Update FunctionState::Reset() to clear both multi_buffer_regions and multi_buffer_region_order alongside the other per-function state. Ensure reset removes all multi-buffer region entries and discovery-order data before the next function is processed.src/codegen/pto/pto_codegen.cpp (1)
1249-1260: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the shared
dimsselection.Lines 1253-1259 repeat the
dimsselection ofComputeAllocTileFields(Lines 1215-1221) exactly. The comment states that the two must agree. A small file-local helper that returns theconst std::vector<ir::ExprPtr>*makes that agreement structural instead of documented, so a later change to one rule cannot skip the other.🤖 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/codegen/pto/pto_codegen.cpp` around lines 1249 - 1260, Extract the repeated valid-dimension selection logic from PTOCodegen::TryComputeStaticValidShape and ComputeAllocTileFields into a shared file-local helper returning const std::vector<ir::ExprPtr>*. Update both callers to use it, preserving the priority of non-empty tile_view_->valid_shape over tile_type->shape_ and returning nullptr when neither is available.tests/ut/codegen/test_multi_buffer_codegen.py (1)
212-214: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGuard the index into the
multi_tile_getlines.
_lines(mlir, "pto.multi_tile_get")[0]raisesIndexErrorwhen nopto.multi_tile_getis emitted. The reported failure then names the list index, not the missing operation. Assert the list is non-empty first, so a regression in region lowering reports the actual cause with the MLIR attached.♻️ Proposed change
mlir = _codegen(RotatingSlot, passes.MemoryPlanner.PTOAS) - get = _lines(mlir, "pto.multi_tile_get")[0] + gets = _lines(mlir, "pto.multi_tile_get") + assert gets, f"expected the slot to be read from the region:\n{mlir}" + get = gets[0] slot_ssa = get.split("[")[1].split("]")[0]🤖 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/ut/codegen/test_multi_buffer_codegen.py` around lines 212 - 214, Update the test around _codegen(RotatingSlot, passes.MemoryPlanner.PTOAS) to store the matching pto.multi_tile_get lines, assert that the collection is non-empty with the generated MLIR included in the failure message, then extract the first line for slot_ssa. Preserve the existing parsing behavior when the operation is present.
🤖 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 `@include/pypto/codegen/pto/pto_type_utils.h`:
- Around line 49-58: Update the documentation around kMinMultiTileBufSlots and
kMaxMultiTileBufSlots to state that PTOAS rejects declarations outside [2, 16].
Clarify that ordinary one-allocation-per-slot lowering applies only within the
PyPTO planner, and do not imply unsupported counts are lowered normally by
PTOAS.
In `@python/bindings/modules/ir.cpp`:
- Around line 823-835: Update CheckSlotCount, used by the MemRef Python
constructors, to throw pypto::ValueError when slots is zero instead of
triggering CHECK. Ensure invalid pl.MemRef(..., slots=0) calls surface as Python
ValueError while valid slot counts retain their current behavior.
- Around line 823-842: Update the resolved MemRef overloads in the DSL typing
wrapper at memref.py to accept both slots and slot, matching the constructor
signatures exposed by the C++ bindings and core stub. Ensure the printed form
MemRef(base, offset, size, slots=N)[k] type-checks and round-trips, and run
pyright for the affected Python typing files.
In `@src/codegen/pto/pto_codegen.cpp`:
- Around line 1308-1324: Update the pass-1 candidate discovery loop to require
both a multi-slot allocation and a valid memref slot index before assigning
geometry and reference_tile, using memref->slot_index_. Keep candidates seeded
for every eligible slotted base so pass 2 still reports unsubscripted-only
bindings, while ensuring slot_type_str, slot_tile_type, and reference_tile come
from the first tile that actually selects a slot.
In `@tests/ut/codegen/test_multi_buffer_codegen.py`:
- Around line 320-321: Update the assertion over slots in the ConstantSlots test
to narrow each MemRef byte_offset_ to ConstInt before accessing value, asserting
the expected constant type for both folded offsets. Preserve the existing offset
set comparison while using the appropriate ConstInt type or established
narrowing mechanism.
---
Nitpick comments:
In `@include/pypto/codegen/pto/pto_codegen.h`:
- Around line 873-878: Update FunctionState::Reset() to clear both
multi_buffer_regions and multi_buffer_region_order alongside the other
per-function state. Ensure reset removes all multi-buffer region entries and
discovery-order data before the next function is processed.
In `@src/codegen/pto/pto_codegen.cpp`:
- Around line 1249-1260: Extract the repeated valid-dimension selection logic
from PTOCodegen::TryComputeStaticValidShape and ComputeAllocTileFields into a
shared file-local helper returning const std::vector<ir::ExprPtr>*. Update both
callers to use it, preserving the priority of non-empty tile_view_->valid_shape
over tile_type->shape_ and returning nullptr when neither is available.
In `@tests/ut/codegen/test_multi_buffer_codegen.py`:
- Around line 212-214: Update the test around _codegen(RotatingSlot,
passes.MemoryPlanner.PTOAS) to store the matching pto.multi_tile_get lines,
assert that the collection is non-empty with the generated MLIR included in the
failure message, then extract the first line for slot_ssa. Preserve the existing
parsing behavior when the operation is present.
🪄 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: cf608035-82e6-4912-ad0d-7d0974c4f4da
📒 Files selected for processing (18)
docs/en/dev/codegen/00-pto_codegen.mddocs/en/dev/language/00-python_syntax.mddocs/en/dev/passes/29-init_memref.mddocs/zh/dev/codegen/00-pto_codegen.mddocs/zh/dev/language/00-python_syntax.mddocs/zh/dev/passes/29-init_memref.mdinclude/pypto/codegen/pto/pto_codegen.hinclude/pypto/codegen/pto/pto_type_utils.hinclude/pypto/ir/memref.hpython/bindings/modules/ir.cpppython/pypto/language/parser/type_resolver.pypython/pypto/pypto_core/ir.pyisrc/codegen/pto/pto_codegen.cppsrc/codegen/pto/pto_type_utils.cppsrc/ir/transforms/init_memref.cppsrc/ir/transforms/python_printer.cppsrc/ir/transforms/utils/deep_clone_utils.cpptests/ut/codegen/test_multi_buffer_codegen.py
| /// The slot-count bounds ptoas's `!pto.multi_tile_buf` verifier enforces | ||
| /// (`MAX_MULTI_BUFFER_NUM = 16`); a declaration outside them keeps the ordinary | ||
| /// one-alloc-per-slot lowering. | ||
| inline constexpr uint64_t kMinMultiTileBufSlots = 2; | ||
| inline constexpr uint64_t kMaxMultiTileBufSlots = 16; | ||
|
|
||
| /// Wrap a single-slot `!pto.tile_buf<...>` into ptoas's N-slot container type, | ||
| /// `!pto.multi_tile_buf<<slot>, count=N>` — the result of `pto.alloc_multi_tile` | ||
| /// and the operand of `pto.multi_tile_get`. `count` must be within ptoas's | ||
| /// `[2, 16]`; callers gate on that before reaching here. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the unsupported-slot behavior.
The comment says a declaration outside [2, 16] uses ordinary per-slot allocations. PlanMultiBufferRegions instead rejects a PTOAS declaration above kMaxMultiTileBufSlots. State that PTOAS rejects unsupported counts. Keep the ordinary allocation behavior limited to the PyPTO planner.
🤖 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 `@include/pypto/codegen/pto/pto_type_utils.h` around lines 49 - 58, Update the
documentation around kMinMultiTileBufSlots and kMaxMultiTileBufSlots to state
that PTOAS rejects declarations outside [2, 16]. Clarify that ordinary
one-allocation-per-slot lowering applies only within the PyPTO planner, and do
not imply unsupported counts are lowered normally by PTOAS.
| CheckSlotCount(slots); | ||
| new (self) MemRef(base, byte_offset, size, span, is_pinned, slots, slot); | ||
| }, | ||
| nb::arg("base"), nb::arg("byte_offset"), nb::arg("size"), nb::arg("span") = Span::unknown(), | ||
| nb::arg("is_pinned") = false, nb::arg("slots") = 1, nb::arg("slot") = nb::none(), | ||
| "Create a memory reference with base Ptr, integer byte_offset, and size. Set is_pinned for an " | ||
| "author-declared allocation whose size the parser leaves for InitMemRef to derive; slots/slot " | ||
| "select one slot of a multi-slot declaration") | ||
| .def(nb::init<VarPtr, ExprPtr, uint64_t, Span>(), nb::arg("base"), nb::arg("byte_offset"), | ||
| nb::arg("size"), nb::arg("span") = Span::unknown(), | ||
| "Create a memory reference with base Ptr, byte_offset expression, and size") | ||
| .def( | ||
| "__init__", | ||
| [](MemRef* self, const VarPtr& base, const ExprPtr& byte_offset, uint64_t size, const Span& span, | ||
| uint64_t slots, const ExprPtr& slot) { | ||
| CheckSlotCount(slots); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Raise ValueError for slots=0.
CheckSlotCount uses CHECK for public Python input. Replace it with throw pypto::ValueError(...). An invalid pl.MemRef(..., slots=0) call must raise a Python exception instead of triggering an internal assertion.
🤖 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/bindings/modules/ir.cpp` around lines 823 - 835, Update
CheckSlotCount, used by the MemRef Python constructors, to throw
pypto::ValueError when slots is zero instead of triggering CHECK. Ensure invalid
pl.MemRef(..., slots=0) calls surface as Python ValueError while valid slot
counts retain their current behavior.
| CheckSlotCount(slots); | ||
| new (self) MemRef(base, byte_offset, size, span, is_pinned, slots, slot); | ||
| }, | ||
| nb::arg("base"), nb::arg("byte_offset"), nb::arg("size"), nb::arg("span") = Span::unknown(), | ||
| nb::arg("is_pinned") = false, nb::arg("slots") = 1, nb::arg("slot") = nb::none(), | ||
| "Create a memory reference with base Ptr, integer byte_offset, and size. Set is_pinned for an " | ||
| "author-declared allocation whose size the parser leaves for InitMemRef to derive; slots/slot " | ||
| "select one slot of a multi-slot declaration") | ||
| .def(nb::init<VarPtr, ExprPtr, uint64_t, Span>(), nb::arg("base"), nb::arg("byte_offset"), | ||
| nb::arg("size"), nb::arg("span") = Span::unknown(), | ||
| "Create a memory reference with base Ptr, byte_offset expression, and size") | ||
| .def( | ||
| "__init__", | ||
| [](MemRef* self, const VarPtr& base, const ExprPtr& byte_offset, uint64_t size, const Span& span, | ||
| uint64_t slots, const ExprPtr& slot) { | ||
| CheckSlotCount(slots); | ||
| new (self) MemRef(base, byte_offset, size, span, /*is_pinned=*/false, slots, slot); | ||
| }, | ||
| nb::arg("base"), nb::arg("byte_offset"), nb::arg("size"), nb::arg("span") = Span::unknown(), | ||
| nb::arg("slots") = 1, nb::arg("slot") = nb::none(), | ||
| "Create a memory reference with base Ptr, byte_offset expression, and size. slots/slot record " | ||
| "which slot of a multi-slot allocation this is, which outlives InitMemRef resolving the index " | ||
| "into the offset") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Synchronize the DSL MemRef overloads.
PrintMemRef emits pl.MemRef(base, offset, size, slots=N)[k]. The overloads in python/pypto/language/typing/memref.py accept slots only for declaration forms. Pyright rejects this valid printed form before it can round-trip.
Add slots and slot to the resolved MemRef overloads in the DSL typing wrapper. Match the core stub and these bindings.
As per coding guidelines, **/*.{py,pyi} must use pyright.
🤖 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/bindings/modules/ir.cpp` around lines 823 - 842, Update the resolved
MemRef overloads in the DSL typing wrapper at memref.py to accept both slots and
slot, matching the constructor signatures exposed by the C++ bindings and core
stub. Ensure the printed form MemRef(base, offset, size, slots=N)[k] type-checks
and round-trips, and run pyright for the affected Python typing files.
Source: Coding guidelines
| // Pass 1: find the slotted allocations and take each one's geometry from the | ||
| // first tile that actually selects a slot. Reading it from whichever tile came | ||
| // first instead would make the outcome depend on discovery order — an unslotted | ||
| // binding seen first would fix the count at 1 and skip the checks below, which | ||
| // is the silent degradation they exist to prevent. | ||
| for (const auto& [tile_var, tile_type] : fs_.tile_var_allocs) { | ||
| auto memref = ir::GetDefinedMemRef(tile_type); | ||
| if (memref->slot_count_ <= 1) continue; | ||
| const ir::Var* base = memref->base_.get(); | ||
| auto [it, fresh] = candidates.try_emplace(base); | ||
| if (!fresh) continue; | ||
| discovery_order.push_back(base); | ||
| it->second.count = memref->slot_count_; | ||
| it->second.slot_type_str = GetTileBufTypeStringFromTileType(tile_type); | ||
| it->second.slot_tile_type = tile_type; | ||
| it->second.reference_tile = tile_var; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The pass 1 filter does not match its comment.
The comment states that the geometry is taken from "the first tile that actually selects a slot". The loop filters only on memref->slot_count_ <= 1. It does not test memref->slot_index_.
InitMemRef::UserBoundMemRef forwards binding->slot_count_ to every tile bound to the allocation, including a tile that binds it whole with no subscript (the UnsubscriptedBinding case in tests/ut/codegen/test_multi_buffer_codegen.py). Such a tile can therefore become the reference_tile and supply slot_type_str and slot_tile_type.
The allocation is still rejected, because pass 2 sets the "binds it without selecting a slot" blocker. The effect is on the diagnostic: the reported blocker and the anchoring span can name a different tile than intended, and the uniform-type baseline is taken from a tile that selects no slot. Add the slot-index test so the code states what the comment claims.
🐛 Proposed fix
for (const auto& [tile_var, tile_type] : fs_.tile_var_allocs) {
auto memref = ir::GetDefinedMemRef(tile_type);
if (memref->slot_count_ <= 1) continue;
+ if (!memref->slot_index_.has_value() || !*memref->slot_index_) continue;
const ir::Var* base = memref->base_.get();
auto [it, fresh] = candidates.try_emplace(base);
if (!fresh) continue;Pass 2 then needs a candidate for every slotted base, including one whose only bindings are unsubscripted. Seed the map in a separate walk, or keep the try_emplace for discovery and set the geometry on the first slot-selecting tile only.
🤖 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/codegen/pto/pto_codegen.cpp` around lines 1308 - 1324, Update the pass-1
candidate discovery loop to require both a multi-slot allocation and a valid
memref slot index before assigning geometry and reference_tile, using
memref->slot_index_. Keep candidates seeded for every eligible slotted base so
pass 2 still reports unsubscripted-only bindings, while ensuring slot_type_str,
slot_tile_type, and reference_tile come from the first tile that actually
selects a slot.
| # ...and the offset is still resolved, as before. | ||
| assert {mr.byte_offset_.value for mr in slots} == {0, 64 * 64 * 4} |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Fix the pyright failure on byte_offset_.value.
The pre-commit pyright hook fails here: Cannot access attribute "value" for class "Expr". MemRef.byte_offset_ is declared as Expr in python/pypto/pypto_core/ir.pyi. value exists only on ConstInt. Narrow the type before reading it.
Both slots of ConstantSlots carry a constant index, so InitMemRef folds each offset to a ConstInt. Asserting that explicitly also documents the folding the test relies on.
🐛 Proposed fix
# ...and the offset is still resolved, as before.
- assert {mr.byte_offset_.value for mr in slots} == {0, 64 * 64 * 4}
+ offsets = [mr.byte_offset_ for mr in slots]
+ assert all(isinstance(off, ir.ConstInt) for off in offsets), (
+ "a constant slot index must fold to a constant offset"
+ )
+ assert {off.value for off in offsets} == {0, 64 * 64 * 4}🧰 Tools
🪛 GitHub Actions: CI / 8_pre-commit.txt
[error] 321-321: Pyright: Cannot access attribute "value" for class "Expr"; attribute "value" is unknown (reportAttributeAccessIssue).
🪛 GitHub Actions: CI / pre-commit
[error] 321-321: Pyright: Cannot access attribute "value" for class "Expr"; attribute "value" is unknown (reportAttributeAccessIssue). The pyright pre-commit hook failed with exit code 1.
🤖 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/ut/codegen/test_multi_buffer_codegen.py` around lines 320 - 321, Update
the assertion over slots in the ConstantSlots test to narrow each MemRef
byte_offset_ to ConstInt before accessing value, asserting the expected constant
type for both folded offsets. Preserve the existing offset set comparison while
using the appropriate ConstInt type or established narrowing mechanism.
Sources: Coding guidelines, Pipeline failures
`pl.MemRef(slots=N)` says "one allocation, N uniform slots, this use takes slot k" — which is exactly what ptoas's `pto.alloc_multi_tile` / `pto.multi_tile_get` describe. Under `memory_planner=PTOAS` codegen now hands ptoas the declaration whole instead of emitting N unrelated `alloc_tile`s, so ptoas plans the slots as one region it may not merge, and derives per-slot synchronization from the slot expression. The slot **index** is what reaches ptoas, not the byte offset InitMemRef resolves it into: ptoas matches the index's affine form to decide which accesses can share a slot, and that is what earns the rotation dynamic event ids. Verified with ptoas 0.54 `--pto-level=level2`: the emitted kernel primes two event ids and uses `wait_flag(..., v24)` / `set_flag(..., v25)` in the loop, overlapping iteration i's load with iteration i-1's compute. - InitMemRef keeps `slot_count_` / `slot_index_` on the resolved MemRef. Resolving the index into `byte_offset_` says *where* the slot lands; it does not stop the MemRef from being slot k of an N-slot allocation. Both fields print, so a post-InitMemRef dump round-trips as `pl.MemRef(base, offset, size, slots=N)[k]`. - InitMemRef's blanket rejection of declared allocations under PTOAS is narrowed to single-slot ones. That rejection exists because ptoas could coalesce what the author separated — which a region provably prevents for its own slots. A single slot still has no ptoas counterpart. - A slotted allocation ptoas cannot describe (non-uniform slot types, a space other than Vec/Mat/Acc, a runtime valid shape, a slot carried out as a phi, a count outside [2,16]) is a ValueError naming the shape, not a silent fallback: falling back to per-slot allocs would undo the separation. - DeepClone now remaps `slot_index_` as it does `byte_offset_` — the index outlives InitMemRef and names SSA values, so a clone must follow it. The PyPTO planner is untouched: at `--pto-level=level3` ptoas does not fold its per-slot address fan-out, so the region form there loses the slot analysis it exists for (and is measurably worse than the baked-address path — an extra false WAR pair between two constant slots). Filed as hw-native-sys/PTOAS#1106; when it lands, widening the gate is one condition with no API change.
dea0372 to
a576434
Compare
1 participant
Summary
pl.MemRef(slots=N)(#2210) says "one allocation, N uniform slots, this use takes slot k" — which is exactly what ptoas'spto.alloc_multi_tile/pto.multi_tile_getdescribe. Undermemory_planner=PTOAScodegen now hands ptoas the declaration whole instead of emitting N unrelatedalloc_tiles.That is what buys the double-buffering: ptoas plans the N slots as one region it is forbidden to merge, and derives per-slot (dynamic event id) synchronization from the slot expression, overlapping iteration i's load with iteration i-1's compute.
No new IR op, no new pass, no new user-facing switch — the declaration already carries the intent, and the memory planner already decides who owns addresses.
Generated
.pto(PTOAS planner)The operand is the slot index, not the byte offset
InitMemRefresolved it into — ptoas matches the index's affine form to decide which accesses can share a slot, and that is what earns the dynamic event ids. Handing iti % 2 * 16384would defeat the analysis.What changed
InitMemRefslot_count_/slot_index_on the resolved MemRef. Resolving the index intobyte_offset_says where the slot lands; it does not stop the MemRef from being slot k of an N-slot allocationInitMemRefPlanMultiBufferRegionsdecides eligibility per allocation; onepto.alloc_multi_tilein the function head, onepto.multi_tile_getper use. Gated onemit_tile_addr_ == falsepl.MemRef(base, offset, size, slots=N)[k]round-trips, so a post-InitMemRefdump reparses as the same slotDeepCloneslot_index_as it doesbyte_offset_— the index outlivesInitMemRefand names SSA values, so a clone must follow itWhy the PTOAS rejection could be narrowed
InitMemRefrejected every declared allocation undermemory_planner=PTOAS, because the isolation is enforced byMemoryReuse, which ptoas replaces wholesale — "honoring the allocation but not its isolation would hand back exactly the coalescing the author declared it to prevent".A multi-slot declaration is the exception: it becomes a ptoas region whose slots ptoas plans into disjoint physical segments it is explicitly forbidden to alias-merge. The separation the author asked for is carried into ptoas rather than lost. A single-slot declaration still has no counterpart and stays rejected — now pointing at
slots=Nas the supported spelling.Unsupported shapes are loud, not silent
A slotted allocation ptoas cannot describe — non-uniform slot types, a space other than Vec/Mat/Acc, a runtime valid shape, a slot carried out of an
if/loop as a phi, a count outside[2, 16]— raises aValueErrornaming the shape. Falling back to per-slotalloc_tilewould silently undo the separation, which is the one thing the declaration exists to state.Why PTOAS-planner only
Measured with ptoas 0.54 (
--enable-insert-sync, a3) on the same double-buffer loop:alloc_tile(today)alloc_multi_tile+multi_tile_getwait_flag(..., v24)/set_flag(..., v25)in-loop → real overlapEVENT_ID0, serializedEVENT_ID0— no differenceRoot cause of the level3 rows: its address fan-out emits an unfolded
arith.addi %base, %c512, and apto.pointer_castwith a non-constant address falls back to conservative aliasing, so slot narrowing is lost. Filed as hw-native-sys/PTOAS#1106;--enable-graph-sync-solverbehaves the same. When that lands, widening the gate is one condition inPlanMultiBufferRegionswith no API change.So the PyPTO planner path is untouched — it keeps baking addresses at level3.
Testing
tests/ut/codegen/test_multi_buffer_codegen.py(15): region emission, slot operand is the index (not a byte offset), constant slots share one region, every blocker branch errors with its reason, PyPTO planner unaffected, slot geometry + print/parse round-trip surviveInitMemRef..ptocompiles under ptoas 0.54--pto-level=level2 --enable-insert-sync, emitting two primed event ids plus in-loop dynamicset_flag_dyn/wait_flag_dyn.tests/ut+tests/lint— 8584 passed, 2 skipped. clang-tidy clean on the 13 changed sources.language/00-python_syntax,passes/29-init_memref,codegen/00-pto_codegen) in bothenandzh.Follow-ups (not in this PR)
multi_tile_get. Measured cost: one extraTileobject +TASSIGN; ptoas CSEs the duplicate index and the sync output is byte-identical to a shared handle — so this is cleanup, not a correctness or performance issue.🤖 Generated with Claude Code