feat(ir): Add pl.Buffer for manual memory-reuse control

[lyfne123]

Summary

MemoryReuse coalesces any two tiles whose lifetimes do not overlap. That is the right default for capacity, but it is not free: the pair is then ordered by a WAR dependency the source never asked for, and the hardware serializes work the scheduler could otherwise overlap.

This adds an opt-in, per-tile escape hatch. Declare a buffer, then reference it:

ping = pl.Buffer()   # takes its name from the variable
pong = pl.Buffer()

t0: pl.Tile[[64, 64], pl.FP32, ping, pl.Mem.Vec] = pl.load(a, [0, 0], [64, 64])
t1: pl.Tile[[64, 64], pl.FP32, pong, pl.Mem.Vec] = pl.exp(t0)
t2: pl.Tile[[64, 64], pl.FP32, ping, pl.Mem.Vec] = pl.exp(t1)  # shares t0's buffer

Tiles referencing one buffer share one allocation; nothing else is ever packed into it. The author writes neither a size nor an address — InitMemRef sizes each buffer to its largest member and derives the memory space from the bound tiles. Tiles left unannotated keep the existing automatic reuse.

Measured on a three-tile chain:

Source	On-chip allocations
unannotated (today's behaviour)	1 × 16 KB — the packer coalesces the whole chain
each tile given its own buffer	3 × 16 KB — no coalescing, no false dependencies
t0/t2 on ping, t1 on pong	2 × 16 KB — the author's own sharing

Two things a name should not decide

Which buffer a tile means was a string. A typo in pl.Buffer("scartch") declared a second buffer rather than failing, and a buffer whose name collided with an in-scope variable silently took that variable as its allocation base. A declared pl.Buffer() referenced by variable makes a misspelling a Python NameError, and an unnamed buffer takes the name of the variable it is bound to, so it is named once rather than twice.

The inline pl.Buffer("name") form stays valid, and is not merely legacy: it is what the IR printer emits, because a dumped program has to reparse without a surrounding Python scope. The object form is the source-level spelling; both lower to the same IR.

Whether a MemRef is a binding at all was inferred from size_ == 0. Nothing guarantees a compiler allocation never holds that value — the invariant was implicit — and the printer had to emit a binding as a pl.MemRef(...) with an invented size, so a dump did not say what it was. MemRef::is_user_buffer_ now records it, and a binding prints as pl.Buffer("name"): the form the author wrote, round-tripping without inventing anything.

An explicit field is not free the way the sentinel was. Riding size_ meant every rebuild that copied the size copied the marker with it; a new field is dropped by default. The sites that rebuild the same MemRef now thread it — the deep-clone path (LowerPipelineLoops clones a body per stage; without this every tile in a pipelined region silently unbinds), AllocateMemoryAddr's address rewrite, and both deserializers, where an absent field means the blob predates user buffers and holds only compiler allocations.

How the binding reaches InitMemRef

It rides TileType::memref_, flagged by is_user_buffer_.

The binding lives on the assigned Var, not on the RHS Call. ConvertToSSA merges it into the Var's type, and op type deduction never produces a MemRef, so a rebuilt Call's type has none. Any carry-through that reads call->GetType() therefore matches nothing — silently, since a missing binding just means the buffer is not pinned. Three passes rebuild the type and had to carry it:

• ConvertToSSA merges an LHS MemRef into the RHS-derived authoritative type. Since deduction never produces one, an LHS MemRef is additional information rather than a stale override. This also fixes a latent gap: a re-parsed post-allocation dump previously lost its MemRefs here.
• FlattenTileNdTo2D carries it through all four rebuilds — ND→2D flatten, ≤2D tile.load, rank>2 tile.create/tile.full, and the generic tile-op re-deduction — each sourcing from assign->var_. The latter three fire only once arguments are substituted, which is exactly what happens when tensor params become partition views.
• InferTileMemorySpace syncs the LHS Var type to the rebuilt Call; it already preserved the Var's memory_space_ and now preserves its MemRef for the same reason.

Anything less and the binding survives a plain @pl.function (no substitution, so those statements pass through untouched) but is dropped in any function actually lowered on-core — the only shape that matters in practice.

InitMemRef then consumes the binding: the MemRef it produces is an ordinary one carrying the derived size, with the flag cleared. From there on the allocation's pinned=True kwarg is what marks the buffer as user-owned, so re-running the passes over a dump cannot promote compiler allocations into user buffers.

Rejected at compile time

Each with the offending tile's span and a stated fix: dynamic shape; bound tiles disagreeing on memory space; binding a view / in-place output (it physically is its source's buffer); two independently bound tiles live at the same time. pl.Buffer under memory_planner=PTOAS is rejected outright rather than honored-but-unenforced, since ptoas replaces the pass that provides the isolation.

tile.transpose is not in that set: it inherits its input's memory space but permutes into a fresh buffer (pto.ttrans is registered not_inplace_safe()), so it owns its storage and may be bound.

That last rule also covers pl.pipeline: naming a buffer inside a stage=2 body is rejected, because the cloned stages make the tile co-live with itself on one allocation. Explicit slots replace compiler pipelining at a given level rather than stacking on top of it — an author taking over that level writes pl.range plus their own named slots.

What it enables

An explicit multi-level double-buffer schedule (#2131). An outer pl.pipeline(stage=2) keeps compiler-managed L1 buffers while an unrolled inner pl.range drives author-declared L0B/L0C slots:

l0b_ping, l0b_pong = pl.Buffer(), pl.Buffer()
l0c_ping, l0c_pong = pl.Buffer(), pl.Buffer()

for stack, (out_outer,) in pl.pipeline(STACKS, stage=2, init_values=(out,)):
    b_l1: pl.Tile[[K, STACK_N], pl.BF16, pl.Mem.Mat] = pl.load(...)          # compiler-managed
    for col, (out_inner,) in pl.range(0, STACK_N, 2 * L0_N, init_values=[out_outer]):
        b_ping: pl.Tile[[K, L0_N], pl.BF16, l0b_ping, pl.Mem.Right] = ...
        a_ping: pl.Tile[[M, L0_N], pl.FP32, l0c_ping, pl.Mem.Acc]   = ...
        b_pong: pl.Tile[[K, L0_N], pl.BF16, l0b_pong, pl.Mem.Right] = ...
        a_pong: pl.Tile[[M, L0_N], pl.FP32, l0c_pong, pl.Mem.Acc]   = ...

Space	Before	After
Mat / L1	2 × 128 KB (pipeline)	2 × 128 KB (pipeline, unchanged)
Right / L0B	2 × 32 KB (packer)	2 × 32 KB, pinned and named
Acc / L0C	1 × 8 KB	2 × 8 KB, pinned and named

With one accumulator both TMATMULs emit addr = %c0_i64, so a TSTORE has to drain L0C between them. With two slots they land on %c0_i64 and %c8192_i64, and the store is no longer a WAR dependency on the critical path.

Cost to kernels that do not use it

Isolation is a per-slot flag inside the packing loop rather than another can_share gate, and the co-liveness check is a sorted sweep rather than a pairwise scan. Both are short-circuited when no buffer is declared, so the O(M²) pack is unchanged for existing kernels.

Testing

• tests/ut/ir/transforms/test_user_buffer.py — 27 tests: binding via both spellings, sharing, size/space derivation, packer exemption, survival through the full pipeline (incl. ND flattening, tile.create flattening, and a pl.spmd Mat/Left/Right/Acc kernel), the rejections, and the negative cases — a zero-sized ordinary MemRef is not a binding, and re-parsed dumps are not treated as user buffers.
• The pl.spmd cube-kernel test is the regression guard for the Var-vs-Call sourcing fix — verified to fail without it.
• Full tests/ut/ suite: 7999 passed, 2 skipped.
• clang-tidy, clang-format, cpplint, ruff check/format, pyright, doc EN↔ZH parity: clean.
• Docs updated in both languages: 29-init_memref.md, 31-memory_reuse.md, 00-python_syntax.md.

Known limitation

Buffers are function-scoped and do not clone per pipeline stage, so a binding inside a pl.pipeline body is rejected rather than silently shared across stages (see above). Per-stage privacy for load buffers is already handled automatically by the existing pipeline-stage guard in MemoryReuse; letting user buffers clone per stage would need LowerPipelineLoops to freshen the identity, and is left out of this PR.

There is no slot subscript — no buffer[i & 1] — so an N-slot rotation is spelled as an N× unrolled body referencing N declared buffers.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 193bc631-e7b7-43f7-a7db-c6fb34109cb4

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds pl.Buffer("name") annotations for explicit tile buffer sharing, carries bindings through IR transformations, emits pinned allocations, isolates them during memory reuse, validates lifetimes and unsupported cases, and adds English/Chinese documentation and comprehensive tests.

Changes

User buffer support

Layer / File(s)	Summary
Buffer annotation and allocation contracts python/pypto/language/typing/*, python/pypto/language/parser/*, python/pypto/language/op/tile_ops.py, src/ir/op/tile_ops/memory.cpp, docs/{en,zh-cn}/dev/language/*	Adds the public Buffer annotation, parses named bindings into MemRefs, supports pinned allocation syntax, and documents usage.
MemRef propagation and pinned allocation representation include/pypto/ir/transforms/utils/*, src/ir/transforms/convert_to_ssa_pass.cpp, src/ir/transforms/flatten_tile_nd_to_2d/*, src/ir/transforms/infer_tile_memory_space_pass.cpp	Preserves MemRef metadata through rewritten and rebuilt tile operations and adds pinned-allocation and source-buffer predicates.
User buffer initialization src/ir/transforms/init_memref.cpp, docs/{en,zh-cn}/dev/passes/29-init_memref.md	Collects bindings, derives sizes and memory spaces, validates unsupported bindings, and emits pinned allocations.
Memory reuse isolation and validation src/ir/transforms/memory_reuse_pass.cpp, tests/ut/ir/transforms/test_user_buffer.py, docs/{en,zh-cn}/dev/passes/31-memory_reuse.md	Prevents pinned buffers from being coalesced with unrelated intervals, checks overlapping lifetimes, and validates pipeline behavior and rejection cases.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant TileAnnotation
  participant TypeResolver
  participant InitMemRef
  participant MemoryReuse
  participant AllocateMemoryAddr
  TileAnnotation->>TypeResolver: Resolve pl.Buffer("name")
  TypeResolver->>InitMemRef: Attach named MemRef binding
  InitMemRef->>InitMemRef: Derive size and memory space
  InitMemRef->>MemoryReuse: Emit pinned tile.alloc
  MemoryReuse->>MemoryReuse: Validate lifetimes and isolate slots
  MemoryReuse->>AllocateMemoryAddr: Pass distinct allocation bases

Loading

Poem

I’m a rabbit with buffers, named neat as can be,
Ping shares with ping, while private stays free.
Pinned tiles hop past the reuse parade,
Lifetimes checked where allocations are made.
Fluffy tests cheer: the bindings won’t fade!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.47% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly states the main change: adding pl.Buffer for manual memory-reuse control.
Description check	✅ Passed	The description is detailed and directly describes the user-buffer memory-reuse changes in the PR.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b10d95e9d0

ℹ️ 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".

[coderabbitai]

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/ir/transforms/init_memref.cpp (1)

444-524: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

User-buffer rejection check misses two paths that also discard the binding silently.

The new CHECK_SPAN (Lines 464-472) only guards the Call-based inherit/reuse-input path. Two structurally identical paths in this same function bypass it entirely, silently dropping a pl.Buffer(...) binding with no diagnostic:

• Lines 453-457: ProducesBufferLessTile → MakeMemRefLessAssign unconditionally clones with std::nullopt MemRef, discarding any user binding on a cross-core tpop result (or a buffer-aliasing view chained off one).
• Lines 511-519: the plain tile-alias branch (a = b) calls ShareMemRefFrom unconditionally — t1: pl.Tile[..., pl.Buffer("mine"), ...] = t0 has no Call on the RHS, so HasUserBinding is never consulted, and t1's binding is silently replaced by t0's MemRef.

Both are exactly the "output physically is its source's buffer, so it cannot be bound to a buffer of its own" scenario the new check exists to reject (per the accompanying docs), just reached via a Var RHS or a buffer-less producer instead of a Call.

🐛 Suggested fix: extend the guard to both paths

     if (As<TileType>(op->var_->GetType()) && ProducesBufferLessTile(new_value, [](const Var* v) {
           return !GetTypeMemRef(v->GetType()).has_value();
         })) {
+      CHECK_SPAN(!HasUserBinding(op->var_->GetType()), op->var_->span_)
+          << "Tile '" << op->var_->name_hint_ << "' has no general-pool buffer of its own "
+             "(cross-core tpop result or a view over one), so it cannot be bound to a buffer.";
       return MakeMemRefLessAssign(op, new_value);
     }
@@
     if (auto value_var = As<Var>(new_value)) {
       if (As<TileType>(op->var_->GetType())) {
+        CHECK_SPAN(!HasUserBinding(op->var_->GetType()), op->var_->span_)
+            << "Tile '" << op->var_->name_hint_ << "' is an alias of '" << value_var->name_hint_
+            << "', which reuses its source tile's buffer, so it cannot be bound to a buffer of "
+               "its own. Bind the source tile instead.";
         auto result = ShareMemRefFrom(value_var, op, new_value);
         if (result) 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/init_memref.cpp` around lines 444 - 524, Extend the
user-buffer rejection in VisitStmt_ to cover both buffer-less tile assignments
and plain tile aliases. Before MakeMemRefLessAssign in the
ProducesBufferLessTile branch, and before ShareMemRefFrom in the Var RHS alias
branch, validate that HasUserBinding(op->var_->GetType()) is false; otherwise
emit the same span-aware diagnostic used for call-based buffer reuse, while
preserving valid unbound assignments.

🤖 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/ir/transforms/utils/op_predicates.h`:
- Around line 49-59: Update the documentation for the predicate covering
OutputMemoryInheritsInput and GetOutputReusesInputArg to replace “same data as
their source” with wording that describes using a source-provided allocation.
Preserve the distinction from IsBufferAliasingViewOp so tile.transpose is
recognized as sharing storage without implying data aliasing.

In `@python/pypto/language/parser/type_resolver.py`:
- Around line 1589-1642: Update _resolve_buffer so Buffer names bypass
scope_lookup and resolve only through resolver-local base-pointer interning. Add
a dedicated helper or an explicit _intern_base_ptr option for this behavior,
while preserving scope_lookup for ordinary MemRef base resolution and ensuring
repeated Buffer names share the same interned Ptr.

In `@python/pypto/language/typing/buffer.py`:
- Around line 10-17: Document the restriction that Buffer declarations are
rejected inside pl.pipeline(stage=2) bodies, explaining that cloned stages would
otherwise share one function-scoped allocation; update the Buffer docstring in
python/pypto/language/typing/buffer.py (lines 10-17), the English guide in
docs/en/dev/language/00-python_syntax.md (lines 72-92), and the corresponding
Chinese guide in docs/zh-cn/dev/language/00-python_syntax.md (lines 72-89).

In `@src/ir/transforms/flatten_tile_nd_to_2d/rewrite.cpp`:
- Around line 68-86: Update the rank>2 tile.create/tile.full rewrite to pass
new_call->GetType() through WithCarriedMemRef using the relevant assignment
before constructing flat_var. Use this carried type for both the recreated call
and replacement variable, preserving existing shape/type metadata while
retaining the assigned variable’s MemRef binding.

---

Outside diff comments:
In `@src/ir/transforms/init_memref.cpp`:
- Around line 444-524: Extend the user-buffer rejection in VisitStmt_ to cover
both buffer-less tile assignments and plain tile aliases. Before
MakeMemRefLessAssign in the ProducesBufferLessTile branch, and before
ShareMemRefFrom in the Var RHS alias branch, validate that
HasUserBinding(op->var_->GetType()) is false; otherwise emit the same span-aware
diagnostic used for call-based buffer reuse, while preserving valid unbound
assignments.

🪄 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: e2a9702c-411b-40b4-ae4b-fe13b96fcbf7

📥 Commits

Reviewing files that changed from the base of the PR and between 6b47c80 and 4bf4a98.

📒 Files selected for processing (22)

• docs/en/dev/language/00-python_syntax.md
• docs/en/dev/passes/29-init_memref.md
• docs/en/dev/passes/31-memory_reuse.md
• docs/zh-cn/dev/language/00-python_syntax.md
• docs/zh-cn/dev/passes/29-init_memref.md
• docs/zh-cn/dev/passes/31-memory_reuse.md
• include/pypto/ir/transforms/utils/memref_utils.h
• include/pypto/ir/transforms/utils/op_predicates.h
• python/pypto/language/__init__.py
• python/pypto/language/op/tile_ops.py
• python/pypto/language/parser/ast_parser.py
• python/pypto/language/parser/type_resolver.py
• python/pypto/language/typing/__init__.py
• python/pypto/language/typing/buffer.py
• src/ir/op/tile_ops/memory.cpp
• src/ir/transforms/convert_to_ssa_pass.cpp
• src/ir/transforms/flatten_tile_nd_to_2d/rewrite.cpp
• src/ir/transforms/infer_tile_memory_space_pass.cpp
• src/ir/transforms/init_memref.cpp
• src/ir/transforms/memory_reuse_pass.cpp
• src/ir/transforms/utils/op_predicates.cpp
• tests/ut/ir/transforms/test_user_buffer.py