[Feature] Explicit buffer-slot binding for fine-grained multi-level double buffering

[lyfne123]

Background

In pypto-lib, users can create logical tiles in Mat/Left/Right/Acc, but cannot
explicitly bind those tiles to named physical buffer slots. There is no way to
write the CCE-style pattern:

slot = iteration & 1
GM -> l1_buffer[slot]        # 512-token cadence
L1 -> l0b_buffer[sub & 1]    # 128-token cadence
MAD -> l0c_buffer[sub & 1]
TSTORE(l0c_buffer[sub & 1])

pl.pipeline and the allocator own buffer identity, count, reuse, and lifetime.
Consequently a nested stage=2 pipeline happens to allocate two 512-wide L1
buffers and two 128-wide L0B buffers, but only one Acc/L0C buffer. Generated
code drains each result with TSTORE before the next TMATMUL can reuse L0C:

TEXTRACT(B0), TEXTRACT(B1),
TMATMUL(C0, B0), TSTORE(C0),
TMATMUL(C0, B1), TSTORE(C0)

Reproduce with:

PYPTO_BENCH=1 python models/qwen3/14b/repro_two_level_l0c.py -p a2a3 -d 3

The repro case is not committed. Save the inlined script below as
models/qwen3/14b/repro_two_level_l0c.py in pypto-lib.
The only repo-local dependency is the public pypto-lib golden package at the
commit recorded below.

Reproduction environment:

Component	Version
pypto-lib	72ee6a1e2e34ac2f9cad9e89fb6a70f6ee5302eb (branch: main)
pypto	c9b2674c8fe32c738b97d8130a2b47d09daa000f (branch: main)
simpler	8cdb306cb9a81ad1a0561325021105c676a69c1e (detached, matches pin)
ptoas	0.48 (matches pin)
pto-isa	83d01313d9bfc247c4b7c8bcf969d1019f0d106f (detached, matches pin)
CANN	9.0.0 (npu-smi 26.0.rc1)

Diagnosis: pypto — the frontend/IR exposes memory spaces but not explicit
buffer identity. pl.create_tile(..., target_memory=...) and
pl.tile.extract(..., target_memory=...) create logical SSA tiles; users cannot
preallocate buffer[2], select a slot with i & 1, bind an operation to that
slot, or explicitly end the prior lifetime. Fine-grained multi-level double
buffering is therefore left entirely to LowerPipelineLoops,
CanonicalizeIOOrder, and MemoryReuse.

Related: #2040, #1475, #1352.

Summary

PyPTO lacks an explicit buffer-slot API for constructing multi-level
double-buffer schedules with different granularities and reuse cadences.

An outer pl.pipeline(stage=2) over [128, 512] Mat/L1 loads and an inner
pl.pipeline(stage=2) over [128, 128] Right/L0B extracts currently happens to
produce:

• Mat/L1: 2 × 128 KB;
• Right/L0B: 2 × 32 KB.

However, users cannot force or name these two buffers, choose which loop indexes
select them, or request two Acc/L0C slots. Only one 8 KB Acc/L0C buffer is
allocated, so TSTORE remains on the critical path between consecutive
TMATMULs.

The missing capability is not explicit L0 tile creation; it is explicit,
stable physical-buffer assignment and reuse control.

Git Commit ID

c9b2674c8fe32c738b97d8130a2b47d09daa000f

NPU Kind

Ascend 910B

Host Platform

Linux (aarch64)

Reproduction Script

Full runnable repro — repro_two_level_l0c.py

"""Minimal independent L1-512 and L0B-128 pipelines."""

import argparse

import pypto.language as pl
import torch

from golden import TensorSpec, run_jit

STACKS = 4
M = 16
K = 128
STACK_N = 512
L0_N = 128


@pl.jit
def two_level_pingpong(
    q: pl.Tensor[[M, K], pl.BF16],
    b: pl.Tensor[[STACKS * K, STACK_N], pl.BF16],
    out: pl.Out[pl.Tensor[[STACKS * M, STACK_N], pl.FP32]],
) -> pl.Tensor[[STACKS * M, STACK_N], pl.FP32]:
    for _ in pl.spmd(1, name_hint="two_level_pingpong"):
        q_l1: pl.Tile[[M, K], pl.BF16, pl.Mem.Mat] = pl.load(
            q,
            [0, 0],
            [M, K],
            target_memory=pl.MemorySpace.Mat,
        )
        q_l0: pl.Tile[[M, K], pl.BF16, pl.Mem.Left] = pl.tile.extract(
            q_l1,
            0,
            0,
            [M, K],
            target_memory=pl.MemorySpace.Left,
        )
        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(
                b,
                [stack * K, 0],
                [K, STACK_N],
                target_memory=pl.MemorySpace.Mat,
            )
            for col, (out_inner,) in pl.pipeline(
                0,
                STACK_N,
                L0_N,
                stage=2,
                init_values=(out_outer,),
            ):
                b_l0: pl.Tile[[K, L0_N], pl.BF16, pl.Mem.Right] = pl.tile.extract(
                    b_l1,
                    0,
                    col,
                    [K, L0_N],
                    target_memory=pl.MemorySpace.Right,
                )
                acc: pl.Tile[[M, L0_N], pl.FP32, pl.Mem.Acc] = pl.tile.matmul(q_l0, b_l0)
                out_next = pl.store(acc, [stack * M, col], out_inner)
                out_inner_yield = pl.yield_(out_next)
            out_outer_yield = pl.yield_(out_inner_yield)
    return out_outer_yield


def _golden(values: dict[str, torch.Tensor]) -> None:
    q = values["q"].float()
    b = values["b"].float()
    out = values["out"]
    for stack in range(STACKS):
        b_stack = b[stack * K : (stack + 1) * K]
        out[stack * M : (stack + 1) * M] = q @ b_stack


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--platform", default="a2a3")
    parser.add_argument("-d", "--device", type=int, default=3)
    args = parser.parse_args()

    torch.manual_seed(2040)
    specs = [
        TensorSpec("q", [M, K], torch.bfloat16, init_value=torch.randn),
        TensorSpec("b", [STACKS * K, STACK_N], torch.bfloat16, init_value=torch.randn),
        TensorSpec("out", [STACKS * M, STACK_N], torch.float32, is_output=True),
    ]
    result = run_jit(
        two_level_pingpong,
        specs,
        golden_fn=_golden,
        compile_cfg={"dump_passes": True},
        runtime_cfg={"platform": args.platform, "device_id": args.device},
        rtol=2e-2,
        atol=2e-2,
    )
    print(f"passed={result.passed} work_dir={result.work_dir}")
    if not result.passed:
        raise RuntimeError(result.error)


if __name__ == "__main__":
    main()

Motivation / Use Case

Hand-written CCE attention uses different ping-pong granularities at different
memory levels: for example, a 512-token L1 stack rotates independently from
128-token L0A/L0B tiles, while two L0C results rotate so FIXPIPE drains one as
the cube computes the other.

PyPTO can describe all relevant memory spaces, but cannot preserve this buffer
topology when automatic liveness/reuse decisions differ from the intended
schedule. A user needs to be able to explicitly construct two fixed slots per
level and select them inside a loop.

Proposed API / Behavior

Provide a public mechanism equivalent in semantics to:

l1_slots = pl.create_tile_buffers(2, [128, 512], pl.BF16, pl.Mem.Mat)
l0b_slots = pl.create_tile_buffers(2, [128, 128], pl.BF16, pl.Mem.Right)
l0c_slots = pl.create_tile_buffers(2, [16, 128], pl.FP32, pl.Mem.Acc)

for stack in pl.range(STACKS):
    l1_slot = stack & 1
    pl.tile.load_into(l1_slots[l1_slot], ...)
    for sub in pl.range(4):
        l0_slot = sub & 1
        pl.tile.extract_into(l0b_slots[l0_slot], l1_slots[l1_slot], ...)
        pl.tile.matmul_into(l0c_slots[l0_slot], q_l0, l0b_slots[l0_slot])
        pl.store(l0c_slots[l0_slot], ...)

The exact API can differ, but it must provide:

• stable buffer identity across loop iterations;
• explicit slot selection by a scalar/modulo expression;
• explicit destination-buffer binding for load/move/matmul;
• a way to delimit/release lifetimes so the compiler generates correct WAR
  fences without coalescing distinct slots; and
• independent rotations at nested L1/L0/L0C granularities.

With two explicit L0C slots, the user should be able to request an order
equivalent to:

TEXTRACT(B0), TEXTRACT(B1),
TMATMUL(C0, B0), TMATMUL(C1, B1),
TSTORE(C0), TSTORE(C1)

This allows FIXPIPE to drain C0 while the cube computes C1, subject to the
normal WAR fence before the next reuse of C0.

Observed Behavior

The case passes golden validation and runs at 13.8 us median, but the allocation
report shows:

Mat   256 KB: 2 x 128 KB buffers
Left    4 KB: 1 x   4 KB buffer
Right  64 KB: 2 x  32 KB buffers
Acc     8 KB: 1 x   8 KB buffer

The generated instruction order is:

TEXTRACT(B0), TEXTRACT(B1),
TMATMUL(C0, B0), TSTORE(C0),
TMATMUL(C0, B1), TSTORE(C0)

Thus automatic operand buffering works for this isolated shape, while L0C/FIX
remains serialized. There is no public API to override the allocation and force
the desired physical slot topology.

Profiling Data

Relevant generated-code sequence:

115: TEXTRACT(v36, ...)
125: TEXTRACT(v38, ...)
134: TMATMUL(v41, ..., v36)
146: TSTORE(..., v41)
155: TMATMUL(v47, ..., v38)
165: TSTORE(..., v47)

The generated addresses confirm two Mat buffers and two Right buffers, but only
one Acc allocation at [0, 8192).

Additional Context

This issue intentionally does not use the PTOAS memory planner. The requested
behavior is a frontend/IR capability usable with the PyPTO planner.

The result also corrects an overly broad formulation of #2040: PyPTO can express
logical tiles at different L1 and L0 granularities. It cannot explicitly assign
and rotate physical buffers to enforce the intended fine-grained reuse and event
schedule.