diff --git a/CMakeLists.txt b/CMakeLists.txt index 1e565d6e6..03de214be 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -153,6 +153,7 @@ set(PYPTO_SOURCES src/ir/transforms/canonicalize_tile_slice_pass.cpp src/ir/transforms/fold_no_op_reshape_pass.cpp src/ir/transforms/stamp_tfree_split_pass.cpp + src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp src/ir/transforms/materialize_runtime_scopes_pass.cpp src/ir/transforms/fuse_create_assemble_to_slice_pass.cpp src/ir/transforms/memory_reuse_pass.cpp diff --git a/docs/en/dev/passes/00-pass_manager.md b/docs/en/dev/passes/00-pass_manager.md index 5ad322278..cc88113e5 100644 --- a/docs/en/dev/passes/00-pass_manager.md +++ b/docs/en/dev/passes/00-pass_manager.md @@ -388,6 +388,7 @@ The PTO-oriented tile stage shared by `Default` and `DebugTileOptimization` is: 11. [`StampTfreeSplit`](22-stamp_tfree_split.md) (copies each cross-core tpop's split/pipe-id onto its matching tfree op) 12. `NormalizeReturnOrder` 13. [`SkewCrossCorePipeline`](24-skew_cross_core_pipeline.md) (cross-core cube/vector software-pipeline skew; runs immediately before LowerPipelineLoops) + - `ConvertToPtoasMultiBuffer` — **optional, gated by `PassContext.use_ptoas_multi_buffer`** (no-op by default; see [ptoas-multi-buffer.md](ptoas-multi-buffer.md)). Runs in `LowerPipelineLoops`' slot. When enabled the pass manager also **drops `LowerPipelineLoops` and `CanonicalizeIOOrder`**, so this pass owns pipeline lowering: it rewrites a same-core loop's single i-dependent Vec/Mat load into a **same-slot** ptoas multi-buffer access (`t = tile.multi_buffer_load_slot(region, i%N, ...)`, consumed same-slot), and demotes non-eligible loops to `Sequential`. Codegen emits `pto.alloc_multi_tile` + `pto.multi_tile_get %mb[i%N]`; ptoas delivers the cross-iteration overlap. The switch **auto-forces `memory_planner=PtoAS`** (`--pto-level=level2`) — the overlap only materializes there (level3's baked base + dynamic slot defeats ptoas `MemAlias`). 14. [`LowerPipelineLoops`](25-lower_pipeline_loops.md) 15. [`CanonicalizeIOOrder`](26-canonicalize_io_order.md) 16. [`MaterializeTensorStrides`](27-materialize_tensor_strides.md) — wired into the default pipeline starting from RFC #1300 P6 diff --git a/docs/en/dev/passes/ptoas-multi-buffer.md b/docs/en/dev/passes/ptoas-multi-buffer.md new file mode 100644 index 000000000..344a57848 --- /dev/null +++ b/docs/en/dev/passes/ptoas-multi-buffer.md @@ -0,0 +1,142 @@ +# ptoas Multi-Buffer (`use_ptoas_multi_buffer`) + +Updated: 2026-07-08 + +An opt-in switch that lowers a same-core `pl.pipeline(stage=N)` loop's rotating +load to a **ptoas multi-buffer region** instead of pypto's own body-replication +ping-pong. ptoas then delivers the cross-iteration double-buffer overlap, so the +kernel keeps a **single loop body** and a **single N-slot region** (smaller code, +tighter memory) while matching the native pipeline's overlap. + +Gated by `PassContext.use_ptoas_multi_buffer` — a no-op by default; the default +pipeline is byte-identical when off. + +## How to enable + +```python +from pypto.pypto_core import passes + +# Explicit PassContext +with passes.PassContext([], use_ptoas_multi_buffer=True): + ... + +# RunConfig / @pl.jit +cfg = RunConfig(platform="a2a3", use_ptoas_multi_buffer=True) + +# Env fallback +# PYPTO_PTOAS_MULTI_BUFFER=1 +``` + +Requires ptoas ≥ 0.48 (`pto.alloc_multi_tile` support): `PTOAS_ROOT=/usr/local/ptoas/0.48`. + +## Key design decision: same-slot, not explicit prefetch + +ptoas delivers the overlap itself. Given a tile loaded and consumed on the **same +slot `mb[i%N]`, same iteration**, ptoas PlanMemory assigns the N slots concrete +disjoint addresses and its sync pass overlaps **iteration `i`'s load (slot `i%N`) +with iteration `i-1`'s consume (slot `(i-1)%N ≠ i%N`)** via dyn-event +(`set_flag_dyn` / `wait_flag_dyn` + `arith.select`) WAR synchronization. + +A manual prologue + prefetch-next / consume-cur split (writing slot `(i+1)%N` +while reading slot `i%N`) instead **fights** that analysis: ptoas pairs static +events in program order, the consume waits on the current iteration's prefetch of +a *different* slot (a false dependency), and the loop serializes. So the pass +emits the simple same-slot form and lets ptoas pipeline it. + +## Key constraint: requires ptoas memory planning (level2) + +The overlap only materializes at **`--pto-level=level2`**, where ptoas PlanMemory +owns memory and assigns each slot a concrete disjoint address (its `MemAlias` +proves the slots disjoint → dyn-event sync). At `--pto-level=level3` a single +baked base + a *dynamic* slot `i%N` defeats `MemAlias` (it conservatively treats +a dynamic slot as aliasing all slots) → serial. + +Therefore `use_ptoas_multi_buffer=True` **auto-forces `memory_planner=PtoAS`** +(in the `PassContext` constructor; `compile.py` reads the effective planner back +so codegen's level matches). If the caller passed a different planner, a warning +notes the override. At level3 the pass still emits valid (but non-overlapping) IR +as a fallback. + +## The pass + +`ConvertToPtoasMultiBuffer` (`src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp`) +runs in `LowerPipelineLoops`' slot. When the switch is on the pass manager also +**drops `LowerPipelineLoops` and `CanonicalizeIOOrder`** (they would replicate / +reorder), so this pass **owns pipeline lowering** and must leave zero +`ForKind::Pipeline` loops behind. For each same-core pipeline loop it either: + +- **rewrites** it (exactly one i-dependent Vec/Mat load): hoist + `region = tile.multi_buffer_alloc(shape; count=N)` before the loop, and replace + `t = tile.load(args)` **in place** with + `t = tile.multi_buffer_load_slot(region, i%N, args)` (same tile var → consumers + need no rewrite), then demote the loop to `Sequential`; or +- **demotes** it to a plain `Sequential` loop (correct, no double-buffer) when it + is not an eligible shape. + +A load is eligible when it is a single-def Vec/Mat `tile.load` whose offset args +reference the loop variable (an i-*independent* load is a loop invariant — +multi-buffering it is pointless). + +## New IR ops (pass-synthesized, not DSL-exposed) + +Registered under the `tile.` namespace (`src/ir/op/tile_ops/memory.cpp`) so they +reuse the existing printer/parser round-trip; **not** `internal_only` (which would +break reparse): + +| Op | Result | Codegen | +| -- | ------ | ------- | +| `tile.multi_buffer_alloc(shape; count=N, dtype, target_memory)` | region (per-slot `TileType`) | `pto.alloc_multi_tile` (addr only at level3) | +| `tile.multi_buffer_load_slot(region, k, tensor, offsets, shapes, valid_shapes)` | filled slot view | `pto.multi_tile_get %mb[k]` + `pto.tload` | +| `tile.multi_buffer_get_slot(region, k)` | consume view | `pto.multi_tile_get %mb[k]` | + +`multi_buffer_alloc` reuses `DeduceTileCreateTileType` (`count` is an extra int +kwarg). The dynamic slot `k = i%N` is a normal index SSA operand resolved only at +codegen (`%mb[k]`); it never enters the MemRef — pypto MemRef offsets are static, +which is exactly why the rotation must stay a runtime index and be planned by +ptoas. `multi_buffer_get_slot` is registered but currently unused by the pass +(kept for a future multi-consumer view). + +## Memory layer + +- **Region** (`multi_buffer_alloc`) is a normal addressed tile: at level3 + `AllocateMemoryAddr` reserves `N * slotBytes` at its base + (`MultiBufferBaseCollector`) and `MemoryReuse` excludes it from coalescing; at + level2 both passes are skipped and ptoas PlanMemory owns the N-slot region. +- **Slot views** (`get_slot` / `load_slot`) are **buffer-less** (`InitMemRef` + `ProducesBufferLessTile`) so each use gets its own SSA name (no pure-alias + collapse onto the region) and MemoryReuse skips them (no MemRef); codegen emits + their `pto.multi_tile_get` directly. + +## Verification + +- **Codegen**: the emitted `.pto` shows one hoisted `pto.alloc_multi_tile` + (addr-less at level2, `count=N`) and one `pto.multi_tile_get %mb[i%N]` per + iteration (single loop body, no `scf.if` prefetch guard). +- **Overlap (ptoas 0.48, level2)**: the final `.cpp` primes 2 per-slot events and + uses a **variable event id** (`wait_flag(..., v25)` / `set_flag(..., v26)` — the + lowered `wait_flag_dyn` / `set_flag_dyn`), so `load[i]` overlaps `consume[i-1]`. +- **On-device numeric parity** (`tests/st/runtime/test_ptoas_multi_buffer_device.py`): + switch-on == switch-off == torch reference on a2a3. +- **Codegen / round-trip UT**: `tests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.py`. + +## Limitations (M1 scope) + +- One i-dependent Vec/Mat load per pipeline loop; deeper multi-load bodies and + `N > 2` beyond the same-slot generalization are future work. +- Dropping `LowerPipelineLoops` is global: under the switch, **all** non-eligible + pipeline loops (including matmul L0 stage loops) demote to serial (lose + ping-pong). Acceptable because the switch is opt-in / default-off; a future + increment can replicate non-eligible loops in-pass. +- level3 emits valid but non-overlapping code — the overlap needs level2. + +## File map + +| Concern | Path | +| ------- | ---- | +| Pass | `src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp` | +| Ops | `src/ir/op/tile_ops/memory.cpp` (registration), `python/pypto/ir/op/tile_ops.py` (builders) | +| Auto-force planner | `src/ir/transforms/pass_context.cpp` (ctor), `python/pypto/ir/compile.py` | +| Pass-manager skip | `python/pypto/ir/pass_manager.py` | +| Codegen | `src/codegen/pto/pto_codegen.cpp` (region alloc + `EmitMultiTileGet`), `src/backend/common/pto_ops_common.cpp` (op emitters) | +| Memory layer | `src/ir/transforms/init_memref.cpp`, `allocate_memory_addr_pass.cpp`, `memory_reuse_pass.cpp` | +| ptoas design | `~/PTOAS/docs/designs/ptoas-multi-buffer-explicit-design.md` | diff --git a/docs/zh-cn/dev/passes/00-pass_manager.md b/docs/zh-cn/dev/passes/00-pass_manager.md index 3891d2857..f39d092fc 100644 --- a/docs/zh-cn/dev/passes/00-pass_manager.md +++ b/docs/zh-cn/dev/passes/00-pass_manager.md @@ -388,6 +388,7 @@ with passes.PassContext([passes.VerificationInstrument(passes.VerificationMode.A 11. [`StampTfreeSplit`](22-stamp_tfree_split.md)（把每个跨核 tpop 的 split/pipe-id 复制到与之配对的 tfree 算子上） 12. `NormalizeReturnOrder` 13. [`SkewCrossCorePipeline`](24-skew_cross_core_pipeline.md)（cube/vector 跨核软流水 skew；紧接在 LowerPipelineLoops 之前运行） + - `ConvertToPtoasMultiBuffer` —— **可选，由 `PassContext.use_ptoas_multi_buffer` 门控**（默认 no-op；见 [ptoas-multi-buffer.md](ptoas-multi-buffer.md)）。运行在 `LowerPipelineLoops` 的槽位。开启后 pass_manager 还会**摘掉 `LowerPipelineLoops` 和 `CanonicalizeIOOrder`**，于是本 pass 全权接管 pipeline 下降：把同核循环里那条唯一 i-依赖的 Vec/Mat load 改写为**同 slot** 的 ptoas 多缓冲访问（`t = tile.multi_buffer_load_slot(region, i%N, ...)`，同 slot 消费），并把不满足条件的循环降为 `Sequential`。codegen 发射 `pto.alloc_multi_tile` + `pto.multi_tile_get %mb[i%N]`；跨迭代重叠由 ptoas 完成。开关会**自动强制 `memory_planner=PtoAS`**（`--pto-level=level2`）—— 重叠只在那里成立（level3 的烘死基址 + 动态 slot 会击败 ptoas `MemAlias`）。 14. [`LowerPipelineLoops`](25-lower_pipeline_loops.md) 15. [`CanonicalizeIOOrder`](26-canonicalize_io_order.md) 16. [`MaterializeTensorStrides`](27-materialize_tensor_strides.md) —— 自 RFC #1300 P6 起接入默认 pipeline diff --git a/docs/zh-cn/dev/passes/ptoas-multi-buffer.md b/docs/zh-cn/dev/passes/ptoas-multi-buffer.md new file mode 100644 index 000000000..6b73485dd --- /dev/null +++ b/docs/zh-cn/dev/passes/ptoas-multi-buffer.md @@ -0,0 +1,132 @@ +# ptoas 多缓冲（`use_ptoas_multi_buffer`） + +Updated: 2026-07-08 + +一个可选开关，把同核 `pl.pipeline(stage=N)` 循环里轮转的 load 降级为 **ptoas 多 +缓冲 region**,而非 pypto 自己的循环体复制 ping-pong。跨迭代的双缓冲重叠交给 +ptoas 完成,于是 kernel 保持**单体循环 + 单块 N-slot region**(代码更小、内存更 +紧凑),同时达到与原生 pipeline 相同的重叠。 + +由 `PassContext.use_ptoas_multi_buffer` 门控 —— 默认关闭时是 no-op,默认流水字节 +不变。 + +## 如何开启 + +```python +from pypto.pypto_core import passes + +# 显式 PassContext +with passes.PassContext([], use_ptoas_multi_buffer=True): + ... + +# RunConfig / @pl.jit +cfg = RunConfig(platform="a2a3", use_ptoas_multi_buffer=True) + +# 环境变量回退 +# PYPTO_PTOAS_MULTI_BUFFER=1 +``` + +需要 ptoas ≥ 0.48(支持 `pto.alloc_multi_tile`):`PTOAS_ROOT=/usr/local/ptoas/0.48`。 + +## 关键设计:同 slot,而非显式预取 + +重叠由 ptoas 自己完成。只要一个 tile 在**同一 slot `mb[i%N]`、同一迭代**内被 load +并消费,ptoas PlanMemory 就给 N 个 slot 分配**互不相交的具体地址**,其同步 pass 用 +dyn-event(`set_flag_dyn` / `wait_flag_dyn` + `arith.select`)的 WAR 同步,把 +**迭代 `i` 的 load(slot `i%N`)与迭代 `i-1` 的消费(slot `(i-1)%N ≠ i%N`)重叠**。 + +手写「prologue + 预取 next / 消费 cur」的拆分(写 slot `(i+1)%N` 同时读 slot +`i%N`)反而**破坏**这个分析:ptoas 按程序序配静态 event,消费会等到本迭代对**另一 +个 slot** 的预取(假依赖),循环退化为串行。因此本 pass 只发同 slot 形态,让 ptoas +去做流水。 + +## 关键约束:必须配合 ptoas 内存规划(level2) + +重叠只在 **`--pto-level=level2`** 下成立 —— 此时 ptoas PlanMemory 拥有内存,给每个 +slot 分不相交的具体地址(其 `MemAlias` 判定各 slot disjoint → dyn-event 同步)。在 +`--pto-level=level3` 下单一烘死基址 + **动态** slot `i%N` 会击败 `MemAlias`(它保守 +地把动态 slot 当作别名所有 slot)→ 串行。 + +因此 `use_ptoas_multi_buffer=True` 会**自动强制 `memory_planner=PtoAS`**(在 +`PassContext` 构造函数里;`compile.py` 把有效 planner 读回,保证 codegen 的 level +一致)。若调用方传了别的 planner,会 warning 说明被覆盖。level3 下本 pass 仍发射有效 +(但不重叠)的 IR 作为兜底。 + +## Pass 本身 + +`ConvertToPtoasMultiBuffer`(`src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp`) +运行在 `LowerPipelineLoops` 的槽位。开关打开时 pass_manager 还会**摘掉 +`LowerPipelineLoops` 和 `CanonicalizeIOOrder`**(它们会复制/重排),于是本 pass +**全权接管 pipeline 下降**,必须不留任何 `ForKind::Pipeline` 循环。对每个同核 +pipeline 循环,它要么: + +- **改写**(恰好一条 i-依赖的 Vec/Mat load):把 + `region = tile.multi_buffer_alloc(shape; count=N)` 提到循环前,并把 + `t = tile.load(args)` **就地**换成 + `t = tile.multi_buffer_load_slot(region, i%N, args)`(同一 tile 变量 → 消费者 + 无需改动),然后把循环降为 `Sequential`;要么 +- 对不满足条件的循环**降为**普通 `Sequential`(正确、无双缓冲)。 + +一条 load 满足条件当且仅当它是单-def 的 Vec/Mat `tile.load`,且其 offset 参数引用 +了循环变量(i-*无关* 的 load 是循环不变量,多缓冲无意义)。 + +## 新增 IR op(pass 内部合成,不暴露给 DSL) + +注册在 `tile.` 命名空间下(`src/ir/op/tile_ops/memory.cpp`),复用既有的 +printer/parser round-trip;**不**标 `internal_only`(否则无法反解析): + +| Op | 结果 | Codegen | +| -- | ---- | ------- | +| `tile.multi_buffer_alloc(shape; count=N, dtype, target_memory)` | region(每-slot `TileType`) | `pto.alloc_multi_tile`(仅 level3 带 addr) | +| `tile.multi_buffer_load_slot(region, k, tensor, offsets, shapes, valid_shapes)` | 已填充的 slot 视图 | `pto.multi_tile_get %mb[k]` + `pto.tload` | +| `tile.multi_buffer_get_slot(region, k)` | 消费视图 | `pto.multi_tile_get %mb[k]` | + +`multi_buffer_alloc` 复用 `DeduceTileCreateTileType`(`count` 是额外的 int kwarg)。 +动态 slot `k = i%N` 是普通 index SSA 操作数,只在 codegen 落地为 `%mb[k]`;它从不 +进入 MemRef —— pypto 的 MemRef offset 是静态的,这正是轮转必须留作运行期索引、并交 +给 ptoas 规划的原因。`multi_buffer_get_slot` 已注册但当前 pass 未用(留作将来多消费 +视图)。 + +## 内存层 + +- **Region**(`multi_buffer_alloc`)是普通带址 tile:level3 下 + `AllocateMemoryAddr` 在其基址预留 `N * slotBytes`(`MultiBufferBaseCollector`), + `MemoryReuse` 把它排除出合并;level2 下这两个 pass 被跳过,由 ptoas PlanMemory + 拥有整块 N-slot region。 +- **Slot 视图**(`get_slot` / `load_slot`)是 **buffer-less**(`InitMemRef` 的 + `ProducesBufferLessTile`),这样每次使用拿到自己的 SSA 名(不会与 region 纯别名 + 塌缩),MemoryReuse 也跳过它们(无 MemRef);codegen 直接发它们的 + `pto.multi_tile_get`。 + +## 验证 + +- **Codegen**:发射的 `.pto` 有一条提前的 `pto.alloc_multi_tile`(level2 无 addr、 + `count=N`)和每迭代一条 `pto.multi_tile_get %mb[i%N]`(单体循环、无 `scf.if` 预取 + guard)。 +- **重叠(ptoas 0.48,level2)**:最终 `.cpp` 预置 2 个每-slot event,并使用**变量 + event id**(`wait_flag(..., v25)` / `set_flag(..., v26)` —— 即 lowered 的 + `wait_flag_dyn` / `set_flag_dyn`),于是 `load[i]` 与 `consume[i-1]` 重叠。 +- **上板数值 parity**(`tests/st/runtime/test_ptoas_multi_buffer_device.py`): + 开关开 == 开关关 == torch 参考,a2a3 实测通过。 +- **Codegen / round-trip 单测**:`tests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.py`。 + +## 限制(M1 范围) + +- 每个 pipeline 循环恰好一条 i-依赖 Vec/Mat load;更深的多 load 循环体、以及超出同 + slot 泛化的 `N > 2` 是后续工作。 +- 摘掉 `LowerPipelineLoops` 是全局的:开关下**所有**不满足条件的 pipeline 循环 + (包括 matmul L0 stage 循环)都降为串行(丢 ping-pong)。因为开关是 opt-in / 默认 + 关,可接受;后续增量可在 pass 内对非目标循环做复制。 +- level3 发射有效但不重叠的代码 —— 重叠需要 level2。 + +## 文件地图 + +| 关注点 | 路径 | +| ------ | ---- | +| Pass | `src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp` | +| Ops | `src/ir/op/tile_ops/memory.cpp`(注册)、`python/pypto/ir/op/tile_ops.py`(builder) | +| 自动强制 planner | `src/ir/transforms/pass_context.cpp`(构造函数)、`python/pypto/ir/compile.py` | +| pass_manager 摘 pass | `python/pypto/ir/pass_manager.py` | +| Codegen | `src/codegen/pto/pto_codegen.cpp`(region alloc + `EmitMultiTileGet`)、`src/backend/common/pto_ops_common.cpp`(op emitter) | +| 内存层 | `src/ir/transforms/init_memref.cpp`、`allocate_memory_addr_pass.cpp`、`memory_reuse_pass.cpp` | +| ptoas 设计文档 | `~/PTOAS/docs/designs/ptoas-multi-buffer-explicit-design.md` | diff --git a/include/pypto/codegen/pto/pto_codegen.h b/include/pypto/codegen/pto/pto_codegen.h index 361c96ce2..95fe05ef0 100644 --- a/include/pypto/codegen/pto/pto_codegen.h +++ b/include/pypto/codegen/pto/pto_codegen.h @@ -338,6 +338,17 @@ class PTOCodegen : public CodegenBase { */ void EmitAllocTileForVar(const ir::VarPtr& tile_var, const std::shared_ptr& tile_type); + /** + * @brief Emit `pto.multi_tile_get %mb[k]` for a multi-buffer slot view. + * + * Shared by the `tile.multi_buffer_get_slot` (consume) and + * `tile.multi_buffer_load_slot` (prefetch) op emitters. Resolves the region + * (op arg 0) to its `alloc_multi_tile` SSA (`fs_.mb_region`), lowers the slot + * index (op arg 1) to an SSA, and emits the slot pick into the current result + * buffer. load_slot then emits a `pto.tload` filling the selected slot. + */ + void EmitMultiTileGet(const ir::CallPtr& op); + /** * @brief Resolve the DPS element vars of a tuple-returning op call * @@ -748,6 +759,23 @@ class PTOCodegen : public CodegenBase { /// vars share one handle (PTOAS in-place aliasing). std::set emitted_tile_alloc_names; + /// ptoas multi-buffer (route-2 prefetch, tile.multi_buffer_alloc / + /// _get_slot / _load_slot). The region gets one `pto.alloc_multi_tile` + /// emitted by `EmitAllocTileForVar`; buffer-less slot views emit their own + /// `pto.multi_tile_get %mb[k]` (get_slot / load_slot op codegen), with the + /// dynamic slot `k` lowered from the op's index operand. + struct MultiBufferRegion { + std::string mb_ssa; ///< alloc_multi_tile result SSA + std::string mtb_type_str; ///< `!pto.multi_tile_buf<...>` type string + std::string tile_type_str; ///< per-slot `!pto.tile_buf<...>` type string + }; + /// region Var → N (count). Pre-registered in `VisitStmt_(AssignStmt)` before + /// `EmitAllocTileForVar` so the region emits `alloc_multi_tile`, not `alloc_tile`. + std::map mb_region_count; + /// region Var → emitted `alloc_multi_tile` info (filled by EmitAllocTileForVar, + /// read by the get_slot / load_slot op emitters). + std::map mb_region; + ir::FunctionPtr current_function; ir::VarPtr current_result_var; std::string current_result_buf; @@ -808,6 +836,9 @@ class PTOCodegen : public CodegenBase { memref_identity_to_mlir.clear(); emitted_tile_alloc_names.clear(); + mb_region_count.clear(); + mb_region.clear(); + current_function.reset(); current_result_var.reset(); current_result_buf.clear(); diff --git a/include/pypto/ir/transforms/pass_context.h b/include/pypto/ir/transforms/pass_context.h index fc8f87d3c..66f51f10a 100644 --- a/include/pypto/ir/transforms/pass_context.h +++ b/include/pypto/ir/transforms/pass_context.h @@ -262,12 +262,22 @@ class PassContext { * @param memory_planner Who plans on-chip buffer memory (default: PyPTO). * PtoAS makes the pipeline skip the pypto allocation passes so the * ptoas PlanMemory pass owns allocation instead. + * @param use_ptoas_multi_buffer When true, `ConvertToPtoasMultiBuffer` + * rewrites `pl.pipeline` double/multi-buffered loops to emit ptoas + * `pto.alloc_multi_tile` / `pto.multi_tile_get` slot rotation instead + * of pypto's own body-replication ping-pong (default: false). + * **Forces `memory_planner = PtoAS`** (overriding the argument): the + * cross-iteration overlap only materializes at `--pto-level=level2`, + * where ptoas PlanMemory assigns the N slots concrete disjoint + * addresses (level3's baked base + dynamic slot defeats MemAlias and + * serializes). `GetMemoryPlanner()` returns the forced value. */ explicit PassContext(std::vector instruments, VerificationLevel verification_level = VerificationLevel::Basic, DiagnosticPhase diagnostic_phase = DiagnosticPhase::PrePipeline, DiagnosticCheckSet disabled_diagnostics = {DiagnosticCheck::UnusedControlFlowResult}, - MemoryPlanner memory_planner = MemoryPlanner::PyPTO); + MemoryPlanner memory_planner = MemoryPlanner::PyPTO, + bool use_ptoas_multi_buffer = false); /** * @brief Push this context onto the thread-local stack @@ -315,6 +325,16 @@ class PassContext { */ [[nodiscard]] const DiagnosticCheckSet& GetDisabledDiagnostics() const; + /** + * @brief Whether ptoas multi-buffer lowering is enabled for this context. + * + * When true, `ConvertToPtoasMultiBuffer` converts `pl.pipeline` + * double/multi-buffered loops into ptoas `pto.alloc_multi_tile` / + * `pto.multi_tile_get` slot rotation (delegating ping-pong + cross-iteration + * sync to ptoas) instead of pypto's body-replication ping-pong. + */ + [[nodiscard]] bool UsePtoasMultiBuffer() const; + /** * @brief Get the instruments registered on this context */ @@ -350,6 +370,7 @@ class PassContext { DiagnosticPhase diagnostic_phase_; DiagnosticCheckSet disabled_diagnostics_; MemoryPlanner memory_planner_; + bool use_ptoas_multi_buffer_; PassContext* previous_; static thread_local PassContext* current_; diff --git a/include/pypto/ir/transforms/pass_properties.h b/include/pypto/ir/transforms/pass_properties.h index ca10a5fc6..7690941c4 100644 --- a/include/pypto/ir/transforms/pass_properties.h +++ b/include/pypto/ir/transforms/pass_properties.h @@ -333,6 +333,20 @@ inline const PassProperties kFoldNoOpReshapeProperties{ inline const PassProperties kStampTfreeSplitProperties{.required = {IRProperty::SplitIncoreOrch}}; +// -- Convert pl.pipeline ping-pong to ptoas multi-buffer --------------------- +// +// Gated by PassContext::UsePtoasMultiBuffer(); a no-op otherwise. Runs in +// LowerPipelineLoops' slot (which the pass manager drops under the switch, so +// only same-core Pipeline loops remain — cross-core ones were demoted by +// SkewCrossCorePipeline). Rewrites the i-dependent load into a +// tile.multi_buffer_alloc region + tile.multi_buffer_load_slot(region, i%N) and +// demotes the loop to Sequential. Preserves all pipeline-stage properties. +inline const PassProperties kConvertToPtoasMultiBufferProperties{ + .required = {IRProperty::SSAForm, IRProperty::SplitIncoreOrch, IRProperty::IncoreTileOps, + IRProperty::TileOps2D, IRProperty::TileMemoryInferred, IRProperty::NormalizedStmtStructure}, + .produced = {IRProperty::SSAForm, IRProperty::SplitIncoreOrch, IRProperty::IncoreTileOps, + IRProperty::TileOps2D, IRProperty::TileMemoryInferred, IRProperty::NormalizedStmtStructure}}; + } // namespace pass } // namespace ir } // namespace pypto diff --git a/include/pypto/ir/transforms/passes.h b/include/pypto/ir/transforms/passes.h index 3d72efde3..efe011caa 100644 --- a/include/pypto/ir/transforms/passes.h +++ b/include/pypto/ir/transforms/passes.h @@ -776,6 +776,28 @@ Pass MaterializeRuntimeScopes(); */ Pass StampTfreeSplit(); +/** + * @brief Convert `pl.pipeline` ping-pong loops to ptoas multi-buffer slots + * + * Gated by `PassContext::UsePtoasMultiBuffer()` — a no-op when the switch is + * off, so the default pipeline is unchanged. When on, the pass manager also + * drops `LowerPipelineLoops` / `CanonicalizeIOOrder` and the switch forces + * `memory_planner = PtoAS`, so this pass owns pipeline lowering. For each + * same-core `ForKind::Pipeline` loop with `pipeline_stages >= 2` and exactly + * one i-dependent vec/mat load it: + * - hoists `region = tile.multi_buffer_alloc(shape; count=N)` before the loop, and + * - replaces `t = tile.load(args)` in place with + * `t = tile.multi_buffer_load_slot(region, i%N, args)` (same tile var, so + * consumers are unchanged), then demotes the loop to Sequential. + * Non-eligible loops are demoted to Sequential (no Pipeline loop survives). + * + * ptoas delivers the cross-iteration double-buffer overlap itself from this + * same-slot form (dyn-event WAR sync) — only under ptoas PlanMemory + * (`--pto-level=level2`), which is why the switch forces `PtoAS`. Runs in + * `LowerPipelineLoops`' slot. See docs/en/dev/passes/ptoas-multi-buffer.md. + */ +Pass ConvertToPtoasMultiBuffer(); + /** * @brief Verify properties on a program and throw on errors * diff --git a/include/pypto/ir/transforms/utils/attrs.h b/include/pypto/ir/transforms/utils/attrs.h index b4d88601c..3154cf65d 100644 --- a/include/pypto/ir/transforms/utils/attrs.h +++ b/include/pypto/ir/transforms/utils/attrs.h @@ -144,6 +144,12 @@ inline bool PipelineMembershipsConflict(const std::vector> StripAttr( diff --git a/python/bindings/modules/passes.cpp b/python/bindings/modules/passes.cpp index b34f20ad5..3301a8df8 100644 --- a/python/bindings/modules/passes.cpp +++ b/python/bindings/modules/passes.cpp @@ -277,13 +277,14 @@ void BindPass(nb::module_& m) { "verification and the diagnostic channel (warnings + performance\n" "hints) for PassPipeline.") .def(nb::init, VerificationLevel, DiagnosticPhase, DiagnosticCheckSet, - MemoryPlanner>(), + MemoryPlanner, bool>(), nb::arg("instruments"), nb::arg("verification_level") = VerificationLevel::Basic, nb::arg("diagnostic_phase") = DiagnosticPhase::PrePipeline, nb::arg("disabled_diagnostics") = DiagnosticCheckSet{DiagnosticCheck::UnusedControlFlowResult}, - nb::arg("memory_planner") = MemoryPlanner::PyPTO, + nb::arg("memory_planner") = MemoryPlanner::PyPTO, nb::arg("use_ptoas_multi_buffer") = false, "Create a PassContext with instruments, verification level, diagnostic phase gate, " - "optional disabled diagnostic checks, and memory planner selection") + "optional disabled diagnostic checks, memory planner selection, and optional ptoas " + "multi-buffer lowering") .def("__enter__", [](PassContext& self) -> PassContext& { self.EnterContext(); @@ -296,6 +297,8 @@ void BindPass(nb::module_& m) { "Get the diagnostic phase gate for this context") .def("get_disabled_diagnostics", &PassContext::GetDisabledDiagnostics, "Get the diagnostic checks suppressed by this context") + .def("use_ptoas_multi_buffer", &PassContext::UsePtoasMultiBuffer, + "Whether ptoas multi-buffer lowering is enabled for this context") .def("get_instruments", &PassContext::GetInstruments, "Get the instruments registered on this context") .def("get_memory_planner", &PassContext::GetMemoryPlanner, "Get the memory planner selection for this context") @@ -527,6 +530,11 @@ void BindPass(nb::module_& m) { passes.def("stamp_tfree_split", &pass::StampTfreeSplit, "Copy each cross-core tpop's split/pipe-id onto its matching tfree op so codegen\n" "reads them from the op directly. Covers mixed-kernel and explicit AIC/AIV tfrees."); + passes.def("convert_to_ptoas_multi_buffer", &pass::ConvertToPtoasMultiBuffer, + "Convert pl.pipeline ping-pong loops to ptoas multi-buffer slots (gated by\n" + "PassContext.use_ptoas_multi_buffer). Tags vec/mat tile Calls with the slot count\n" + "and downgrades pipeline_stages to 1 so codegen emits pto.alloc_multi_tile /\n" + "pto.multi_tile_get and ptoas owns slot rotation + sync. No-op when the switch is off."); passes.def("materialize_runtime_scopes", &pass::MaterializeRuntimeScopes, "Materialize implicit orchestration scopes as explicit RuntimeScopeStmt nodes.\n\n" "For every Orchestration function, inserts AUTO RuntimeScopeStmt (manual_=false)\n" diff --git a/python/pypto/ir/compile.py b/python/pypto/ir/compile.py index b6badec9a..327e97e9c 100644 --- a/python/pypto/ir/compile.py +++ b/python/pypto/ir/compile.py @@ -53,7 +53,7 @@ def _backend_type_for_platform(platform: str | None, fallback: BackendType) -> B raise ValueError(f"Invalid platform {platform!r}. Expected 'a2a3sim', 'a2a3', 'a5sim', or 'a5'.") -def compile( # noqa: PLR0913 +def compile( # noqa: PLR0913, PLR0912 program: _ir_core.Program, output_dir: str | None = None, strategy: OptimizationStrategy = OptimizationStrategy.Default, @@ -69,6 +69,7 @@ def compile( # noqa: PLR0913 distributed_config: Any = None, block_dim: int | None = None, analyze_auto_scopes_for_deps: bool = False, + use_ptoas_multi_buffer: bool | None = None, ) -> "CompiledProgram | DistributedCompiledProgram": """Compile a Program through passes and codegen. @@ -130,9 +131,16 @@ def compile( # noqa: PLR0913 analyze_auto_scopes_for_deps: If True, let ``AutoDeriveTaskDependencies`` analyze AUTO runtime scopes. The default is False to preserve the existing TensorMap-fallback - behavior unless explicitly enabled. User-written manual scopes are - skipped: they do not get compiler deps or automatic - NoDep/OutputExisting direction rewrites. + behavior unless explicitly enabled. User-written manual scopes do + not get compiler deps, but covered read-only inputs may still be + rewritten to NoDep. + use_ptoas_multi_buffer: When True, ``ConvertToPtoasMultiBuffer`` lowers a + same-core ``pl.pipeline`` loop's rotating load to a ptoas multi-buffer + region (``tile.multi_buffer_alloc`` / ``_load_slot``, slot ``i%N``), + letting ptoas deliver the cross-iteration double-buffer overlap. This + **forces** ``memory_planner=PtoAS`` (the overlap only materializes at + ``--pto-level=level2``). None uses the default (off, or the + ``PYPTO_PTOAS_MULTI_BUFFER`` env var). Returns: A :class:`CompiledProgram` that wraps the output directory and can @@ -176,6 +184,11 @@ def compile( # noqa: PLR0913 "compile() was called with memory_planner while a PassContext is already active. " "Set the memory planner on the existing PassContext instead." ) + if use_ptoas_multi_buffer is not None and outer is not None: + raise RuntimeError( + "compile() was called with use_ptoas_multi_buffer while a PassContext is already active. " + "Set use_ptoas_multi_buffer on the existing PassContext instead." + ) # --- Compile profiling --------------------------------------------------- prof = get_active_profiler() @@ -209,7 +222,31 @@ def compile( # noqa: PLR0913 dphase = diagnostic_phase if diagnostic_phase is not None else _passes.get_default_diagnostic_phase() disabled = disabled_diagnostics if disabled_diagnostics is not None else default_disabled mplan = memory_planner if memory_planner is not None else _passes.MemoryPlanner.PYPTO - ctx = _passes.PassContext(instruments, vlevel, dphase, disabled, mplan) + + # Resolve ptoas multi-buffer lowering. The explicit ``use_ptoas_multi_buffer`` + # argument (from compile() / RunConfig) is the reliable source of truth; the + # PYPTO_PTOAS_MULTI_BUFFER env var is a global fallback that applies whether + # or not an outer PassContext is active. An outer context can only force it + # *on* (its default is off), so it never suppresses the env var. + env_mbuf = os.environ.get("PYPTO_PTOAS_MULTI_BUFFER", "0").strip().lower() in ("1", "true", "yes", "on") + outer_mbuf = outer.use_ptoas_multi_buffer() if outer is not None else False + if use_ptoas_multi_buffer is not None: + mbuf = use_ptoas_multi_buffer + else: + mbuf = outer_mbuf or env_mbuf + ctx = _passes.PassContext(instruments, vlevel, dphase, disabled, mplan, mbuf) + + # ptoas multi-buffer implies the PTOAS memory planner (the overlap only + # materializes at --pto-level=level2). PassContext forces this internally; read + # the effective planner back so codegen's level (generate(memory_planner=...)) + # matches the passes that just ran. Warn if it overrode an explicit request. + if mbuf and memory_planner is not None and memory_planner != _passes.MemoryPlanner.PTOAS: + logger.warning( + "use_ptoas_multi_buffer=True forces memory_planner=PTOAS (requested %s): the " + "multi-buffer double-buffer overlap requires ptoas PlanMemory at --pto-level=level2.", + memory_planner, + ) + mplan = ctx.get_memory_planner() if mplan == _passes.MemoryPlanner.PTOAS: logger.warning( diff --git a/python/pypto/ir/op/tile_ops.py b/python/pypto/ir/op/tile_ops.py index 75b6a5e34..8d1b2b0ec 100644 --- a/python/pypto/ir/op/tile_ops.py +++ b/python/pypto/ir/op/tile_ops.py @@ -160,6 +160,98 @@ def create( create_tile = create +def multi_buffer_alloc( + shape: Sequence[int | Expr] | _ir_core.MakeTuple, + dtype: DataType, + count: int, + target_memory: MemorySpace = MemorySpace.Vec, + span: Span | None = None, +) -> Call: + """Allocate an N-slot ptoas multi-buffer region (route-2 prefetch lowering). + + Pass-synthesized by ``ConvertToPtoasMultiBuffer``; not a user DSL surface. + Structurally a :func:`create` tile plus a ``count`` slot count. Consumed only + by :func:`multi_buffer_get_slot`; codegen lowers it to + ``pto.alloc_multi_tile ... count=N``. + + Args: + shape: Per-slot tile shape, or a MakeTuple + dtype: Per-slot tile dtype + count: Number of slots N (2..16) + target_memory: Memory space of the region (Vec / Mat) + span: Optional source span + + Returns: + Call expression returning the region TileType (per-slot shape) + """ + actual_span = _get_span_or_capture(span) + shape_tuple = _to_make_tuple(shape, actual_span) + kwargs: dict[str, Any] = {"dtype": dtype, "target_memory": target_memory, "count": count} + return _ir_core.create_op_call("tile.multi_buffer_alloc", [shape_tuple], kwargs, actual_span) + + +def multi_buffer_get_slot( + region: Expr, + slot_index: Expr, + span: Span | None = None, +) -> Call: + """Select slot ``slot_index`` of a multi-buffer region (zero-copy view). + + Pass-synthesized by ``ConvertToPtoasMultiBuffer``; not a user DSL surface. + The view inherits the region's per-slot tile and base (offset 0); the dynamic + ``slot_index`` is resolved at codegen (``pto.multi_tile_get %mb[k]``). + + Args: + region: Region tile from :func:`multi_buffer_alloc` + slot_index: Slot to select (index scalar expression) + span: Optional source span + + Returns: + Call expression returning the per-slot tile view + """ + actual_span = _get_span_or_capture(span) + return _ir_core.create_op_call("tile.multi_buffer_get_slot", [region, slot_index], {}, actual_span) + + +def multi_buffer_load_slot( + region: Expr, + slot_index: Expr, + tensor: Expr, + offsets: Sequence[int | Expr] | _ir_core.MakeTuple, + shapes: Sequence[int | Expr] | _ir_core.MakeTuple, + valid_shapes: Sequence[int | Expr] | _ir_core.MakeTuple, + span: Span | None = None, +) -> Call: + """Select slot ``slot_index`` of a multi-buffer region and load into it. + + Pass-synthesized by ``ConvertToPtoasMultiBuffer`` for the prologue/prefetch + fill; not a user DSL surface. Codegen emits ``pto.multi_tile_get %mb[k]`` + followed by a ``pto.tload`` filling that slot (tail args mirror :func:`load`). + + Args: + region: Region tile from :func:`multi_buffer_alloc` + slot_index: Slot to fill (index scalar expression) + tensor: Source tensor + offsets: Load offsets, or a MakeTuple + shapes: Load region shape, or a MakeTuple + valid_shapes: Valid tile shape, or a MakeTuple + span: Optional source span + + Returns: + Call expression returning the filled per-slot tile view + """ + actual_span = _get_span_or_capture(span) + offsets_tuple = _to_make_tuple(offsets, actual_span) + shapes_tuple = _to_make_tuple(shapes, actual_span) + valid_tuple = _to_make_tuple(valid_shapes, actual_span) + return _ir_core.create_op_call( + "tile.multi_buffer_load_slot", + [region, slot_index, tensor, offsets_tuple, shapes_tuple, valid_tuple], + {}, + actual_span, + ) + + def load( tensor: Expr, offsets: Sequence[int | Expr] | _ir_core.MakeTuple, diff --git a/python/pypto/ir/pass_manager.py b/python/pypto/ir/pass_manager.py index 6d970aebe..12ec7a278 100644 --- a/python/pypto/ir/pass_manager.py +++ b/python/pypto/ir/pass_manager.py @@ -162,6 +162,13 @@ def _register_passes(cls): ("StampTfreeSplit", lambda: passes.stamp_tfree_split()), ("NormalizeReturnOrder", lambda: passes.normalize_return_order()), ("SkewCrossCorePipeline", lambda: passes.skew_cross_core_pipeline()), + # Gated by PassContext.use_ptoas_multi_buffer (no-op otherwise). Runs + # right before LowerPipelineLoops so only same-core Pipeline loops + # remain (cross-core ones were demoted by SkewCrossCorePipeline). Tags + # vec/mat tile Calls with the slot count and downgrades pipeline_stages + # to 1 so codegen emits ptoas multi-buffer ops instead of pypto + # replicating the body for ping-pong. + ("ConvertToPtoasMultiBuffer", lambda: passes.convert_to_ptoas_multi_buffer()), ("LowerPipelineLoops", lambda: passes.lower_pipeline_loops()), ("CanonicalizeIOOrder", lambda: passes.canonicalize_io_order()), # MaterializeTensorStrides fills empty stride slots on every @@ -267,10 +274,21 @@ def __init__( skip_mem_planning = ctx is not None and ctx.get_memory_planner() == passes.MemoryPlanner.PTOAS _mem_planning_passes = ("MemoryReuse", "AllocateMemoryAddr") + # When ptoas multi-buffer is on, ConvertToPtoasMultiBuffer owns pipeline + # lowering: it rewrites eligible pipeline loops into a dyn-slot prefetch + # over a multi-buffer region and demotes the rest to Sequential, leaving + # zero Pipeline loops. So the ordinary pipeline-lowering passes are dropped + # (LowerPipelineLoops would double-replicate; CanonicalizeIOOrder would + # reorder the hand-built prefetch order). Mirrors the memory-planner skip. + skip_pipeline_lowering = ctx is not None and ctx.use_ptoas_multi_buffer() + _pipeline_lowering_passes = ("LowerPipelineLoops", "CanonicalizeIOOrder") + # Build pass list for pass_name, pass_factory in self._strategy_passes[strategy]: if skip_mem_planning and pass_name in _mem_planning_passes: continue + if skip_pipeline_lowering and pass_name in _pipeline_lowering_passes: + continue if pass_name == "AutoDeriveTaskDependencies": self.passes.append( passes.auto_derive_task_dependencies(analyze_auto_scopes=analyze_auto_scopes_for_deps) @@ -407,8 +425,17 @@ def after_pass(_pass_obj: passes.Pass, program: core_ir.Program) -> None: inner_phase = passes.DiagnosticPhase.PRE_PIPELINE else: inner_phase = outer_phase - - with passes.PassContext([*outer_instruments, *extra_instruments], level, inner_phase, disabled): + # Preserve ALL outer-context pass config here — this reconstructs the + # context the pipeline runs under, so any PassContext setting the passes + # read (e.g. memory_planner; use_ptoas_multi_buffer gating + # ConvertToPtoasMultiBuffer) must be carried over, or dump mode silently + # changes pass behavior. + 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, *extra_instruments], level, inner_phase, disabled, mplan, mbuf + ): try: return self._pipeline.run(input_ir) finally: diff --git a/python/pypto/jit/cache.py b/python/pypto/jit/cache.py index 902511812..742bac1a4 100644 --- a/python/pypto/jit/cache.py +++ b/python/pypto/jit/cache.py @@ -123,7 +123,7 @@ def compute_source_hash(sources: list[str]) -> str: return h.hexdigest()[:16] -def make_cache_key( +def make_cache_key( # noqa: PLR0913 source_hash: str, param_names: list[str], tensor_shapes: dict[str, tuple[int, ...]], @@ -134,6 +134,7 @@ def make_cache_key( strategy: "OptimizationStrategy | None" = None, distributed_config: Any = None, analyze_auto_scopes_for_deps: bool = False, + use_ptoas_multi_buffer: bool | None = None, ) -> CacheKey: """Build a cache key for a JIT call site. @@ -166,6 +167,10 @@ def make_cache_key( analyze_auto_scopes_for_deps: Compile-side switch for deriving explicit task dependencies from AUTO runtime scopes. Included in the key because it changes generated orchestration dependencies. + use_ptoas_multi_buffer: Compile-side switch for ptoas multi-buffer + lowering. Included in the key because it changes the emitted IR / + kernel (alloc_multi_tile vs replicated alloc_tile); without it the + same kernel compiled with the switch on and off would collide. Returns: Hashable CacheKey tuple. @@ -187,7 +192,10 @@ def make_cache_key( scalar_infos.append(ScalarCacheInfo(name=name, value=scalar_values[name])) dist_key = _freeze(distributed_config) if distributed_config is not None else None - compile_opts = (("analyze_auto_scopes_for_deps", analyze_auto_scopes_for_deps),) + compile_opts = ( + ("analyze_auto_scopes_for_deps", analyze_auto_scopes_for_deps), + ("use_ptoas_multi_buffer", use_ptoas_multi_buffer), + ) return ( source_hash, platform, diff --git a/python/pypto/jit/decorator.py b/python/pypto/jit/decorator.py index 5144627d7..91562c3ec 100644 --- a/python/pypto/jit/decorator.py +++ b/python/pypto/jit/decorator.py @@ -1125,6 +1125,7 @@ def _run_config_compile_kwargs(run_config: Any) -> dict[str, Any]: "diagnostic_phase": run_config.diagnostic_phase, "disabled_diagnostics": run_config.disabled_diagnostics, "analyze_auto_scopes_for_deps": run_config.analyze_auto_scopes_for_deps, + "use_ptoas_multi_buffer": run_config.use_ptoas_multi_buffer, } if run_config.save_kernels_dir is not None: kwargs["output_dir"] = run_config.save_kernels_dir @@ -1458,6 +1459,7 @@ def _resolve_compiled( analyze_auto_scopes_for_deps = ( run_config.analyze_auto_scopes_for_deps if run_config is not None else False ) + use_ptoas_multi_buffer = run_config.use_ptoas_multi_buffer if run_config is not None else None key = make_cache_key( source_hash=self._get_source_hash(), param_names=param_names, @@ -1469,6 +1471,7 @@ def _resolve_compiled( strategy=strategy, distributed_config=distributed_config, analyze_auto_scopes_for_deps=analyze_auto_scopes_for_deps, + use_ptoas_multi_buffer=use_ptoas_multi_buffer, ) # L1 cache lookup diff --git a/python/pypto/pypto_core/passes.pyi b/python/pypto/pypto_core/passes.pyi index f0faa6010..0f3392709 100644 --- a/python/pypto/pypto_core/passes.pyi +++ b/python/pypto/pypto_core/passes.pyi @@ -271,8 +271,15 @@ class PassContext: diagnostic_phase: DiagnosticPhase = DiagnosticPhase.PRE_PIPELINE, disabled_diagnostics: DiagnosticCheckSet = ..., # default: {UnusedControlFlowResult} memory_planner: MemoryPlanner = MemoryPlanner.PYPTO, + use_ptoas_multi_buffer: bool = False, ) -> None: - """Create a PassContext with instruments and pass configuration (incl. memory planner).""" + """Create a PassContext with instruments and pass config (memory planner + ptoas multi-buffer). + + ``use_ptoas_multi_buffer=True`` forces ``memory_planner=PTOAS`` (overriding + the argument): the multi-buffer double-buffer overlap only materializes at + ``--pto-level=level2``, where ptoas PlanMemory assigns the N slots concrete + disjoint addresses. ``get_memory_planner()`` reflects the forced value. + """ ... def __enter__(self) -> PassContext: ... @@ -298,6 +305,10 @@ class PassContext: """Get the memory planner selection for this context.""" ... + def use_ptoas_multi_buffer(self) -> bool: + """Whether ptoas multi-buffer lowering is enabled for this context.""" + ... + def get_instruments(self) -> list[PassInstrument]: """Get the instruments registered on this context.""" ... @@ -640,6 +651,17 @@ def stamp_tfree_split() -> Pass: AIC/AIV tfrees. Runs late, before codegen. """ +def convert_to_ptoas_multi_buffer() -> Pass: + """Convert ``pl.pipeline`` ping-pong loops to ptoas multi-buffer slots. + + Gated by ``PassContext.use_ptoas_multi_buffer`` — a no-op when off. When on, + tags vec/mat tile-producing ``Call`` nodes in each same-core pipeline loop + with the slot count and downgrades ``pipeline_stages`` to 1 so + ``LowerPipelineLoops`` skips replication. Codegen then emits + ``pto.alloc_multi_tile`` / ``pto.multi_tile_get`` and ptoas owns slot + rotation + cross-iteration sync. Runs immediately before ``LowerPipelineLoops``. + """ + def materialize_runtime_scopes() -> Pass: """Materialize implicit orchestration scopes as explicit RuntimeScopeStmt nodes. @@ -818,6 +840,7 @@ __all__ = [ "fuse_create_assemble_to_slice", "fold_no_op_reshape", "stamp_tfree_split", + "convert_to_ptoas_multi_buffer", "VerificationError", "SSAErrorType", "TypeCheckErrorType", diff --git a/python/pypto/runtime/runner.py b/python/pypto/runtime/runner.py index aaebee757..103c3f48e 100644 --- a/python/pypto/runtime/runner.py +++ b/python/pypto/runtime/runner.py @@ -241,6 +241,12 @@ class RunConfig: ring_dep_pool: int | list[int] | tuple[int, ...] | None = None distributed_config: "DistributedConfig | None" = None analyze_auto_scopes_for_deps: bool = False + # Opt into ptoas multi-buffer lowering (ConvertToPtoasMultiBuffer): converts + # pl.pipeline ping-pong loops to pto.alloc_multi_tile / pto.multi_tile_get so + # ptoas owns slot rotation + sync. None defers to the PYPTO_PTOAS_MULTI_BUFFER + # env var; True/False forces it. Threaded explicitly to ir.compile() so it is + # reliable regardless of dump_passes / outer PassContext. + use_ptoas_multi_buffer: bool | None = None def __post_init__(self) -> None: if self.platform not in ("a2a3sim", "a2a3", "a5sim", "a5"): @@ -479,6 +485,7 @@ def run( platform=config.platform, profiling=config.compile_profiling, analyze_auto_scopes_for_deps=config.analyze_auto_scopes_for_deps, + use_ptoas_multi_buffer=config.use_ptoas_multi_buffer, ) if tensors and not config.codegen_only: diff --git a/src/backend/common/pto_ops_common.cpp b/src/backend/common/pto_ops_common.cpp index 306dd2cdb..60185a901 100644 --- a/src/backend/common/pto_ops_common.cpp +++ b/src/backend/common/pto_ops_common.cpp @@ -45,6 +45,7 @@ #include "pypto/ir/comm.h" #include "pypto/ir/expr.h" #include "pypto/ir/kind_traits.h" +#include "pypto/ir/op_registry.h" #include "pypto/ir/scalar_expr.h" #include "pypto/ir/tile_view_semantics.h" #include "pypto/ir/transforms/utils/memref_utils.h" @@ -3388,6 +3389,26 @@ static void EmitTreshapeView(codegen::PTOCodegen& codegen, const ir::ExprPtr& sr codegen.Emit(oss.str()); } +// tile.multi_buffer_get_slot (consume view): emit only the slot pick +// `pto.multi_tile_get %mb[k]` into the current (buffer-less) result buffer. +static std::string MakeMultiBufferGetSlotCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { + auto& codegen = dynamic_cast(codegen_base); + codegen.EmitMultiTileGet(op); + return ""; // EmitMultiTileGet already emitted the line. +} + +// tile.multi_buffer_load_slot (prefetch/prologue): select the slot, then reuse +// tile.load's `pto.tload` emission over the tail args (tensor, offsets, shapes, +// valid_shapes) to fill the selected slot in place. +static std::string MakeMultiBufferLoadSlotCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { + auto& codegen = dynamic_cast(codegen_base); + codegen.EmitMultiTileGet(op); // %slot = pto.multi_tile_get %mb[k] (current result buffer) + std::vector load_args(op->args_.begin() + 2, op->args_.end()); + auto load_op = ir::OpRegistry::GetInstance().GetOp("tile.load"); + auto load_call = std::make_shared(load_op, load_args, op->GetType(), op->span_); + return MakeTileLoadCodegenPTO(load_call, codegen_base); // pto.tload ins(...) outs(%slot) +} + void RegisterPTOOps(Backend& backend, const std::unordered_set& exclude_ops) { // Register simple N-ary ops for (const auto& entry : kSimpleOps) { @@ -3739,6 +3760,20 @@ void RegisterPTOOps(Backend& backend, const std::unordered_set& exc reg("tile.alloc", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) { return MakeTileAllocCodegenPTO(op, codegen); }); + // ptoas multi-buffer (route-2 prefetch). The region alloc emits its + // `pto.alloc_multi_tile` in EmitAllocTileForVar, so this emitter is a no-op; + // the slot views emit `pto.multi_tile_get` (+ `pto.tload` for load_slot). + reg("tile.multi_buffer_alloc", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) { + (void)op; + (void)codegen; + return std::string(""); + }); + reg("tile.multi_buffer_get_slot", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) { + return MakeMultiBufferGetSlotCodegenPTO(op, codegen); + }); + reg("tile.multi_buffer_load_slot", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) { + return MakeMultiBufferLoadSlotCodegenPTO(op, codegen); + }); reg("tile.create", [](const ir::CallPtr& op, codegen::CodegenBase& codegen_base) { (void)op; (void)codegen_base; diff --git a/src/codegen/pto/pto_codegen.cpp b/src/codegen/pto/pto_codegen.cpp index 1baf45b02..bac85db55 100644 --- a/src/codegen/pto/pto_codegen.cpp +++ b/src/codegen/pto/pto_codegen.cpp @@ -43,6 +43,7 @@ #include "pypto/ir/program.h" #include "pypto/ir/scalar_expr.h" #include "pypto/ir/stmt.h" +#include "pypto/ir/transforms/utils/attrs.h" #include "pypto/ir/transforms/utils/core_affinity.h" #include "pypto/ir/transforms/utils/memref_utils.h" #include "pypto/ir/transforms/utils/op_predicates.h" @@ -1138,6 +1139,38 @@ void PTOCodegen::EmitAllocTileForVar(const ir::VarPtr& tile_var, return; } + // ptoas multi-buffer region (route-2 prefetch): emit one `pto.alloc_multi_tile` + // (count=N). The N physical slots are fanned out by ptoas; the buffer-less slot + // views (get_slot / load_slot) emit their own `pto.multi_tile_get %mb[k]` picks. + // + // The cross-iteration double-buffer OVERLAP that this feature exists for is + // only delivered under ptoas memory planning (--pto-level=level2, memory_planner + // =PTOAS): there ptoas PlanMemory assigns each slot a concrete disjoint address, + // so its MemAlias proves the slots disjoint and inserts dyn-event sync. At + // level3 (emit_tile_addr_), a single baked base + dynamic slot defeats MemAlias + // and ptoas serializes — so level3 emits a valid (correct, non-overlapping) + // region while level2 emits the addr-less form ptoas plans. + if (auto rc = fs_.mb_region_count.find(var_key); rc != fs_.mb_region_count.end()) { + int slots = rc->second; + AllocTileFields fields = ComputeAllocTileFields(tile_type); + std::string mtb_type = + "!pto.multi_tile_buf<" + fields.type_str + ", count=" + std::to_string(slots) + ">"; + std::ostringstream line; + line << tile_buf << " = pto.alloc_multi_tile"; + if (emit_tile_addr_) { + INTERNAL_CHECK_SPAN(!fields.addr_ssa.empty(), tile_var->span_) + << "Internal error: level3 multi-buffer region " << tile_var->name_hint_ + << " has no base address"; + line << " addr = " << fields.addr_ssa; + } + if (!fields.valid_row_ssa.empty()) line << " valid_row = " << fields.valid_row_ssa; + if (!fields.valid_col_ssa.empty()) line << " valid_col = " << fields.valid_col_ssa; + line << " : " << mtb_type; + Emit(line.str()); + fs_.mb_region[var_key] = {tile_buf, mtb_type, fields.type_str}; + return; + } + AllocTileFields fields = ComputeAllocTileFields(tile_type); std::ostringstream line; @@ -1151,6 +1184,27 @@ void PTOCodegen::EmitAllocTileForVar(const ir::VarPtr& tile_var, fs_.ssa_to_tile_buf_type[tile_buf] = fields.type_str; } +void PTOCodegen::EmitMultiTileGet(const ir::CallPtr& op) { + // Shared by the get_slot (consume) and load_slot (prefetch) emitters. The slot + // views are buffer-less, so `EmitAllocTileForVar` did not run for them — this + // emits their `pto.multi_tile_get %mb[k]` pick into the current result buffer. + auto region_var = ir::AsVarLike(op->args_[0]); + INTERNAL_CHECK_SPAN(region_var, op->span_) + << "Internal error: multi-buffer slot op first arg must be the region Var"; + auto it = fs_.mb_region.find(region_var.get()); + INTERNAL_CHECK_SPAN(it != fs_.mb_region.end(), op->span_) + << "Internal error: multi-buffer slot op region " << region_var->name_hint_ + << " has no emitted pto.alloc_multi_tile (region alloc must dominate its slots)"; + const auto& region = it->second; + std::string slot_ssa = GetExprAsCode(op->args_[1]); + const std::string& dst = fs_.current_result_buf; + INTERNAL_CHECK_SPAN(!dst.empty(), op->span_) + << "Internal error: multi-buffer slot op has no current result buffer"; + Emit(dst + " = pto.multi_tile_get " + region.mb_ssa + "[" + slot_ssa + "] : " + region.mtb_type_str + + " -> " + region.tile_type_str); + fs_.ssa_to_tile_buf_type[dst] = region.tile_type_str; +} + // ======================================================================== // Private helpers // ======================================================================== @@ -1383,6 +1437,12 @@ void PTOCodegen::VisitStmt_(const AssignStmtPtr& op) { const bool alias_scatter_result_to_input = ShouldAliasScatterResultToInput(op); const bool alias_array_update_to_input = ShouldAliasArrayUpdateResultToInput(op); + // Pre-register a ptoas multi-buffer region so EmitAllocTileForVar emits a + // `pto.alloc_multi_tile` (count=N) instead of an ordinary `pto.alloc_tile`. + if (ir::IsOp(call, "tile.multi_buffer_alloc")) { + fs_.mb_region_count[GetVarKey(op->var_)] = call->GetKwarg("count", 0); + } + if (auto tile_type = ir::GetTileTypeWithMemRef(op->var_->GetType())) { if (!is_set_validshape && !alias_scatter_result_to_input) { EmitAllocTileForVar(op->var_, tile_type); diff --git a/src/codegen/pto/pto_control_flow_codegen.cpp b/src/codegen/pto/pto_control_flow_codegen.cpp index 5289283f5..fa151e088 100644 --- a/src/codegen/pto/pto_control_flow_codegen.cpp +++ b/src/codegen/pto/pto_control_flow_codegen.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include "pypto/codegen/pto/pto_codegen.h" @@ -322,6 +323,12 @@ void PTOCodegen::VisitStmt_(const ForStmtPtr& op) { std::string loop_var_name = NewNamedTemp(op->loop_var_->name_hint_); BindVarToMlir(op->loop_var_, loop_var_name); + // ptoas multi-buffer (route-2): the region `pto.alloc_multi_tile` is emitted + // inline at its own `tile.multi_buffer_alloc` AssignStmt (placed before this + // loop by ConvertToPtoasMultiBuffer), and the buffer-less slot views emit + // their `pto.multi_tile_get` in the get_slot / load_slot op codegen — so no + // per-loop hoist/rotation state is needed here. + // In PTO, only scalar types (index, f32, bool, etc.) need iter_args/yield // for loop-carried value semantics. Non-scalar types (TileType, TensorType) // are mutable references written in-place via outs(), so they are mapped diff --git a/src/ir/op/tile_ops/memory.cpp b/src/ir/op/tile_ops/memory.cpp index 48ca8c09a..ed838ccb5 100644 --- a/src/ir/op/tile_ops/memory.cpp +++ b/src/ir/op/tile_ops/memory.cpp @@ -416,6 +416,27 @@ TypePtr DeduceTileCreateTileType(const std::vector& args, return std::make_shared(tile_shape, dtype, std::nullopt, tile_view, creation_space); } +// tile.multi_buffer_get_slot / tile.multi_buffer_load_slot: reference slot `k` of +// a ptoas multi-buffer region. Both are zero-copy views — the result is the +// per-slot tile, structurally identical to the region's per-slot TileType +// (arg 0). The dynamic slot index (arg 1) is a pure runtime address selector +// resolved at codegen (pto.multi_tile_get %mb[k]), so it never affects the +// result *type*: the view inherits the region type verbatim, and InitMemRef +// shares the region's base (offset 0). load_slot additionally carries a +// tensor/offsets/shapes[/valid_shapes] tail (args 2..) that codegen turns into +// a pto.tload filling the selected slot. +TypePtr DeduceMultiBufferSlotType(const std::vector& args, + const std::vector>& /*kwargs*/, + const std::string& op_name) { + CHECK(args.size() >= 2) << "The operator " << op_name + << " requires at least 2 arguments (region, slot_index), but got " << args.size(); + auto region_type = As(args[0]->GetType()); + CHECK(region_type) << "The operator " << op_name + << " requires its first argument to be a multi-buffer region (TileType), but got " + << args[0]->TypeName(); + return region_type; +} + TypePtr DeduceTileFullType(const std::vector& args, const std::vector>& kwargs, const std::string& op_name) { @@ -773,6 +794,65 @@ REGISTER_OP("tile.create") return DeduceTileCreateTileType(args, kwargs, "tile.create"); }); +// ptoas multi-buffer region: an N-slot rotating buffer for the route-2 prefetch +// lowering (ConvertToPtoasMultiBuffer). Structurally a `tile.create` (same +// per-slot shape/dtype/space, so it reuses DeduceTileCreateTileType) plus a +// `count=N` slot count. AllocateMemoryAddr reserves `N*slotBytes` at its base; +// codegen lowers it to `pto.alloc_multi_tile ... count=N`. Consumed only by +// `tile.multi_buffer_get_slot`; never loaded/stored directly. Pass-synthesized, +// not DSL-exposed. +REGISTER_OP("tile.multi_buffer_alloc") + .set_op_category("TileOp") + .set_description("Allocate an N-slot ptoas multi-buffer region (route-2 prefetch)") + .set_core_affinity(core_affinity::CoreAffinity::SHARED) + .add_argument("shape", "Per-slot shape dimensions (TupleType of ScalarType(INT64))") + .set_attr("dtype") + .set_attr("target_memory") + .set_attr("count") + .set_output_memory_from_kwarg("target_memory") + .f_deduce_type([](const std::vector& args, + const std::vector>& kwargs) { + return DeduceTileCreateTileType(args, kwargs, "tile.multi_buffer_alloc"); + }); + +// Pick slot `k` of a multi-buffer region (consume view): a zero-copy view +// inheriting the region's per-slot tile + base (offset 0). The dynamic +// `slot_index` selects the physical slot at codegen (`pto.multi_tile_get +// %mb[k]`); it does not enter the MemRef (pypto MemRef offsets are static — the +// rotation stays a runtime index). Pass-synthesized, not DSL-exposed. +REGISTER_OP("tile.multi_buffer_get_slot") + .set_op_category("TileOp") + .set_description("Select slot k of a ptoas multi-buffer region (consume view)") + .set_core_affinity(core_affinity::CoreAffinity::SHARED) + .add_argument("region", "Multi-buffer region tile (from tile.multi_buffer_alloc)") + .add_argument("slot_index", "Slot to select (index ScalarType) — resolved at codegen") + .set_output_memory_inherit_input() + .f_deduce_type([](const std::vector& args, + const std::vector>& kwargs) { + return DeduceMultiBufferSlotType(args, kwargs, "tile.multi_buffer_get_slot"); + }); + +// Prefetch/prologue fill: select slot `k` AND load into it from a tensor. The +// tail args (tensor, offsets, shapes[, valid_shapes]) mirror `tile.load`; codegen +// emits `pto.multi_tile_get %mb[k]` then a `pto.tload` filling that slot. Like +// get_slot it is a zero-copy view over the region base (the tload writes the +// region's storage in place). Pass-synthesized, not DSL-exposed. +REGISTER_OP("tile.multi_buffer_load_slot") + .set_op_category("TileOp") + .set_description("Select slot k of a ptoas multi-buffer region and load into it (prefetch)") + .set_core_affinity(core_affinity::CoreAffinity::SHARED) + .add_argument("region", "Multi-buffer region tile (from tile.multi_buffer_alloc)") + .add_argument("slot_index", "Slot to fill (index ScalarType) — resolved at codegen") + .add_argument("tensor", "Source tensor (TensorType)") + .add_argument("offsets", "Load offsets (TupleType of ScalarType)") + .add_argument("shapes", "Load region shape (TupleType of ScalarType)") + .add_argument("valid_shapes", "Valid tile shape (TupleType of ScalarType)") + .set_output_memory_inherit_input() + .f_deduce_type([](const std::vector& args, + const std::vector>& kwargs) { + return DeduceMultiBufferSlotType(args, kwargs, "tile.multi_buffer_load_slot"); + }); + REGISTER_OP("tile.load") .set_op_category("TileOp") .set_description("Copy data from tensor to unified buffer (tile)") diff --git a/src/ir/transforms/allocate_memory_addr_pass.cpp b/src/ir/transforms/allocate_memory_addr_pass.cpp index 8460178fb..d3ad1b314 100644 --- a/src/ir/transforms/allocate_memory_addr_pass.cpp +++ b/src/ir/transforms/allocate_memory_addr_pass.cpp @@ -43,6 +43,7 @@ #include "pypto/ir/transforms/base/visitor.h" #include "pypto/ir/transforms/pass_properties.h" #include "pypto/ir/transforms/passes.h" +#include "pypto/ir/transforms/utils/attrs.h" #include "pypto/ir/transforms/utils/memref_collectors.h" #include "pypto/ir/transforms/utils/memref_utils.h" #include "pypto/ir/transforms/utils/mutable_copy.h" @@ -270,6 +271,41 @@ class MemRefUpdateMutator : public IRMutator { } }; +/** + * @brief Map each ptoas multi-buffer tile's MemRef base_ Ptr to its slot count N. + * + * Walks the function body for tile-producing ``AssignStmt``s whose ``Call`` + * is a ``tile.multi_buffer_alloc`` region (from ConvertToPtoasMultiBuffer). Keyed by + * the tile's MemRef ``base_``, so ``AllocateMemoryAddresses`` can reserve + * ``N * slot_size`` for that base group — the contiguous region ptoas fans the + * hoisted ``pto.alloc_multi_tile`` base address out across (slot k at + * ``base + k*slotBytes``). MemoryReuse already keeps such tiles from sharing a + * base, so each is its own singleton group here. + */ +class MultiBufferBaseCollector : public IRVisitor { + public: + [[nodiscard]] const std::map& base_to_slots() const { return base_to_slots_; } + + void VisitStmt_(const AssignStmtPtr& op) override { + IRVisitor::VisitStmt_(op); + auto call = As(op->value_); + // The multi-buffer region (tile.multi_buffer_alloc, route-2 prefetch) reserves + // `count * slotBytes` at its base so ptoas can fan its N physical slots out + // contiguously from the hoisted `pto.alloc_multi_tile` base. + if (!IsOp(call, "tile.multi_buffer_alloc")) return; + int slots = call->GetKwarg("count", 0); + if (slots < 2) return; + auto tile_type = As(op->var_->GetType()); + if (!tile_type || !tile_type->memref_.has_value()) return; + const auto& memref = tile_type->memref_.value(); + if (!memref || !memref->base_) return; + base_to_slots_[memref->base_.get()] = slots; + } + + private: + std::map base_to_slots_; +}; + /** * @brief Allocate memory addresses using the given allocation policy * @@ -278,10 +314,15 @@ class MemRefUpdateMutator : public IRMutator { * MemRefs (e.g. produced by ``tile.slice``) — every view physically aliases * its parent allocation, so they should share one address slot rather than * each consuming size_ bytes of fresh L1. + * + * A ptoas multi-buffer tile (``base`` in ``multi_buffer_bases``) reserves + * ``N * slot_size`` instead of one slot: ptoas lays its N physical slots + * contiguously from the hoisted ``pto.alloc_multi_tile`` base, so the allocator + * must keep the whole region clear of any following tile. */ std::vector> AllocateMemoryAddresses( const std::vector& memrefs, const ReservedEndBySpace& reserved_end_by_space, - const MemoryAllocatorPolicy& policy) { + const MemoryAllocatorPolicy& policy, const std::map& multi_buffer_bases) { // Group MemRefs by memory space std::unordered_map> space_to_memrefs; for (const auto& [memref, memory_space] : memrefs) { @@ -327,6 +368,16 @@ std::vector> AllocateMemoryAddresses( slot_size = std::max(slot_size, static_cast(ref->size_)); } + // ptoas multi-buffer: reserve N contiguous single-slot regions. ptoas + // fans the hoisted alloc_multi_tile base out to slot k at base + k*slot, + // so the whole [base, base + N*slot) region must be clear of any following + // tile. MemoryReuse keeps such a tile in its own singleton base group, so + // reserving here is exact; the extra (N-1)*slot is charged to this group. + auto mbit = multi_buffer_bases.find(base_key); + const bool is_multi_buffer = mbit != multi_buffer_bases.end(); + const uint64_t reserve_size = + is_multi_buffer ? slot_size * static_cast(mbit->second) : slot_size; + // Bump the whole group to `current_addr`, then preserve each member's // own offset within the slot: new byte_offset = base_addr + old offset. // @@ -359,15 +410,23 @@ std::vector> AllocateMemoryAddresses( // operand; PTO codegen reads this dtype from the ConstInt 1:1. auto member_addr_expr = std::make_shared( static_cast(current_addr) + relative_offset, DataType::INT64, Span::unknown()); + // For a multi-buffer root (offset 0), record the full N-slot footprint + // as size_ so the high-water verifier accounts for the whole reserved + // region (ptoas uses [base, base + N*slot)). Views never occur on a + // multi-buffer tile, so only the root at offset 0 is inflated. + int64_t new_size = old_memref->size_; + if (is_multi_buffer && relative_offset == 0) { + new_size = static_cast(reserve_size); + } // NOTE: MemRef is identity-bearing — each result must get a fresh // unique_id_, so build it via the explicit constructor (MutableCopy is // static_assert-forbidden for Var/MemRef). auto new_memref = std::make_shared(old_memref->name_hint_, old_memref->base_, - member_addr_expr, old_memref->size_, old_memref->span_); + member_addr_expr, new_size, old_memref->span_); memref_pairs.emplace_back(old_memref.get(), new_memref); } - current_addr = policy.AlignAddress(current_addr + slot_size, space); + current_addr = policy.AlignAddress(current_addr + reserve_size, space); } } @@ -411,8 +470,14 @@ FunctionPtr TransformAllocateMemoryAddr(const FunctionPtr& func) { // Step 2: Collect all unique MemRef objects from TileType variables auto memrefs = memref_collectors::CollectMemRefsWithSpace(func->body_); + // ptoas multi-buffer bases (base_ Ptr -> slot count N) so the allocator can + // reserve N contiguous slots for each such tile (use_ptoas_multi_buffer). + MultiBufferBaseCollector mb_collector; + mb_collector.VisitStmt(func->body_); + // Step 3: Allocate memory addresses using the policy - auto memref_pairs = AllocateMemoryAddresses(memrefs, reserve_resolution.reserved_end_by_space, *policy); + auto memref_pairs = AllocateMemoryAddresses(memrefs, reserve_resolution.reserved_end_by_space, *policy, + mb_collector.base_to_slots()); if (memref_pairs.empty() && reserve_resolution.resolved_bases.empty()) { return func; diff --git a/src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp b/src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp new file mode 100644 index 000000000..0a12bc6b3 --- /dev/null +++ b/src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp @@ -0,0 +1,251 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// ConvertToPtoasMultiBuffer — ptoas multi-buffer (same-slot) lowering. +// +// Gated by `PassContext::UsePtoasMultiBuffer()`. When the switch is on, the pass +// manager also drops `LowerPipelineLoops` and `CanonicalizeIOOrder`, so THIS pass +// owns pipeline lowering: it must leave zero `ForKind::Pipeline` loops behind. For +// each same-core pipeline loop it either +// (a) rewrites the i-dependent load into a ptoas multi-buffer region access +// (load slot `i%N`, consumed same-slot same-iteration), or +// (b) demotes it to a plain Sequential loop (correct, no double-buffer) when it +// is not an eligible multi-buffer shape. +// +// Why same-slot, not an explicit prologue/prefetch split: ptoas delivers the +// cross-iteration double-buffer OVERLAP itself. Given `t = mb[i%N]` loaded and +// consumed each iteration, ptoas PlanMemory assigns the N slots disjoint +// addresses and its sync pass overlaps iteration i's load (slot i%N) with +// iteration i-1's consume (slot (i-1)%N ≠ i%N) via dyn-event (`set_flag_dyn` / +// `wait_flag_dyn`) WAR sync. A manual prefetch-next/consume-cur split instead +// fights that analysis and serializes. Verified: same-slot overlaps, explicit +// prefetch does not (see docs/en/dev/passes/ptoas-multi-buffer.md). +// +// **Requires ptoas memory planning (`memory_planner=PTOAS`, --pto-level=level2).** +// The overlap only materializes when ptoas PlanMemory assigns the slots concrete +// disjoint addresses; at level3 (baked base + dynamic slot) ptoas MemAlias is +// conservative and serializes. The pass still emits correct IR at level3, just +// without overlap. +// +// Emits two pass-internal ops (see src/ir/op/tile_ops/memory.cpp): +// tile.multi_buffer_alloc(shape; count=N, ...) -> region (N-slot) +// tile.multi_buffer_load_slot(region, k, tensor, ...) -> load slot k +// The dynamic slot `k = i%N` is a normal index SSA operand resolved only at +// codegen (`pto.multi_tile_get %mb[k]`); it never enters the MemRef. +// +// M1 scope: exactly one i-dependent Vec/Mat load per pipeline loop. Anything else +// falls through to the Sequential demotion. + +#include +#include +#include +#include +#include +#include + +#include "pypto/ir/expr.h" +#include "pypto/ir/function.h" +#include "pypto/ir/kind_traits.h" +#include "pypto/ir/memory_space.h" +#include "pypto/ir/op_registry.h" +#include "pypto/ir/scalar_expr.h" +#include "pypto/ir/stmt.h" +#include "pypto/ir/transforms/base/mutator.h" +#include "pypto/ir/transforms/base/visitor.h" +#include "pypto/ir/transforms/pass_context.h" +#include "pypto/ir/transforms/pass_properties.h" +#include "pypto/ir/transforms/passes.h" +#include "pypto/ir/transforms/utils/attrs.h" +#include "pypto/ir/transforms/utils/mutable_copy.h" +#include "pypto/ir/type.h" + +namespace pypto { +namespace ir { +namespace pass { + +namespace { + +// True when `tile_type` lives in vec/mat local memory — the only spaces ptoas +// multi-buffer supports. +bool IsMultiBufferEligibleSpace(const std::shared_ptr& tile_type) { + if (!tile_type || !tile_type->memory_space_.has_value()) return false; + MemorySpace space = *tile_type->memory_space_; + return space == MemorySpace::Vec || space == MemorySpace::Mat; +} + +// Whether `expr` references `target` anywhere in its subtree. +class VarRefFinder : public IRVisitor { + public: + explicit VarRefFinder(const Var* target) : target_(target) {} + bool found() const { return found_; } + + protected: + void VisitVarLike_(const VarPtr& op) override { + if (op.get() == target_) found_ = true; + IRVisitor::VisitVarLike_(op); + } + + private: + const Var* target_; + bool found_ = false; +}; + +bool ExprReferencesVar(const ExprPtr& expr, const Var* target) { + if (!expr || !target) return false; + VarRefFinder finder(target); + finder.VisitExpr(expr); + return finder.found(); +} + +// A single eligible i-dependent load inside a pipeline loop body. +struct EligibleLoad { + size_t index = 0; // position in the top-level body SeqStmts + AssignStmtPtr assign; // T = tile.load(...) + CallPtr load; // the tile.load Call + VarPtr tile_var; // T + std::shared_ptr tile_type; +}; + +// Return the loop body's top-level statement list (SeqStmts unwrapped). +std::vector BodyStmts(const StmtPtr& body) { + if (auto seq = As(body)) return seq->stmts_; + return {body}; +} + +// Find the sole i-dependent Vec/Mat `tile.load` at the top level of `body`. +// Returns nullopt unless EXACTLY one exists (M1 restriction). +std::optional FindEligibleLoad(const StmtPtr& body, const VarPtr& loop_var) { + auto stmts = BodyStmts(body); + std::optional found; + for (size_t i = 0; i < stmts.size(); ++i) { + auto assign = As(stmts[i]); + if (!assign) continue; + auto call = As(assign->value_); + if (!call || !IsOp(call, "tile.load")) continue; + auto tile_type = As(assign->var_->GetType()); + if (!IsMultiBufferEligibleSpace(tile_type)) continue; + // i-dependence: the load offsets (arg 1) must reference the loop var, else + // the load is loop-invariant and multi-buffering it is pointless. + if (call->args_.size() < 2 || !ExprReferencesVar(call->args_[1], loop_var.get())) continue; + if (found.has_value()) return std::nullopt; // more than one — M1 bails. + found = EligibleLoad{i, assign, call, assign->var_, tile_type}; + } + return found; +} + +// Strip `pipeline_stages` and demote to Sequential, keeping the (kind, attr) +// invariant `kind == Pipeline <=> pipeline_stages` whole (PipelineLoopValid). +StmtPtr DemoteToSequential(const ForStmtPtr& loop) { + auto cleaned = MutableCopy(loop); + cleaned->kind_ = ForKind::Sequential; + cleaned->attrs_ = StripAttr(loop->attrs_, kPipelineStagesAttr); + return cleaned; +} + +ExprPtr MakeIndexConst(int64_t value, const Span& span) { + return std::make_shared(value, DataType::INDEX, span); +} + +// Build `tile.multi_buffer_load_slot(region, slot, tensor, offsets, shapes, +// valid_shapes)` from the original tile.load's arg list. +CallPtr MakeLoadSlot(const VarPtr& region, const ExprPtr& slot, const std::vector& load_args, + const Span& span) { + INTERNAL_CHECK_SPAN(load_args.size() >= 3, span) + << "Internal error: tile.load must have >=3 args (tensor, offsets, shapes)"; + ExprPtr tensor = load_args[0]; + ExprPtr offsets = load_args[1]; + ExprPtr shapes = load_args[2]; + ExprPtr valid_shapes = load_args.size() >= 4 ? load_args[3] : load_args[2]; + std::vector args{region, slot, tensor, offsets, shapes, valid_shapes}; + return OpRegistry::GetInstance().Create("tile.multi_buffer_load_slot", args, span); +} + +// Same-slot multi-buffer rewrite: hoist the region alloc, then replace the +// i-dependent load with a `load_slot(region, i%N, )` on the SAME tile +// var, so its consumers need no rewrite. ptoas delivers the cross-iteration +// overlap; no prologue / prefetch / guard is emitted here. +StmtPtr BuildMultiBuffer(const ForStmtPtr& loop, const EligibleLoad& e, int64_t n) { + const Span& span = loop->span_; + const VarPtr& iv = loop->loop_var_; + ExprPtr n_const = MakeIndexConst(n, span); + + // region alloc (before the loop): multi_buffer_alloc(shape; dtype, mem, count=N). + // Reuses DeduceTileCreateTileType, so shape is a MakeTuple of per-slot ConstInts. + auto shape_tuple = std::make_shared(e.tile_type->shape_, span); + std::vector> alloc_kwargs{ + {"dtype", e.tile_type->dtype_}, + {"target_memory", *e.tile_type->memory_space_}, + {"count", static_cast(n)}}; + auto region_call = + OpRegistry::GetInstance().Create("tile.multi_buffer_alloc", {shape_tuple}, alloc_kwargs, span); + auto region_var = std::make_shared(e.tile_var->name_hint_ + "_mb", region_call->GetType(), span); + auto region_assign = std::make_shared(region_var, region_call, span); + + // Replace the load in place: same tile var, slot `i % N`, original load args. + ExprPtr slot = MakeFloorMod(iv, n_const, span); + auto load_slot = MakeLoadSlot(region_var, slot, e.load->args_, span); + auto new_load_assign = std::make_shared(e.tile_var, load_slot, span); + + std::vector body_stmts = BodyStmts(loop->body_); + body_stmts[e.index] = new_load_assign; + + auto new_loop = MutableCopy(loop); + new_loop->kind_ = ForKind::Sequential; + new_loop->attrs_ = StripAttr(loop->attrs_, kPipelineStagesAttr); + new_loop->body_ = SeqStmts::Flatten(std::move(body_stmts), span); + + std::vector result{region_assign, new_loop}; + return SeqStmts::Flatten(std::move(result), span); +} + +class ConvertMutator : public IRMutator { + protected: + StmtPtr VisitStmt_(const ForStmtPtr& op) override { + // Recurse first so nested pipeline loops convert/demote independently. + auto visited = IRMutator::VisitStmt_(op); + auto loop = As(visited); + if (!loop || loop->kind_ != ForKind::Pipeline || !loop->HasAttr(kPipelineStagesAttr)) return visited; + + int n = loop->GetAttr(kPipelineStagesAttr, 0); + // A single-i-dependent-load stage>=2 loop becomes a same-slot multi-buffer + // access (slot i%N); everything else (stage==1 markers, matmul stage loops, + // non-eligible shapes) demotes to a plain Sequential loop so no Pipeline loop + // survives. + if (n >= 2) { + auto eligible = FindEligibleLoad(loop->body_, loop->loop_var_); + if (eligible.has_value()) return BuildMultiBuffer(loop, *eligible, n); + } + return DemoteToSequential(loop); + } +}; + +} // namespace + +Pass ConvertToPtoasMultiBuffer() { + auto pass_func = [](const FunctionPtr& func) -> FunctionPtr { + if (!func || !func->body_) return func; + // Self-gated: no-op unless the active PassContext enables ptoas multi-buffer. + auto* ctx = PassContext::Current(); + if (ctx == nullptr || !ctx->UsePtoasMultiBuffer()) return func; + + ConvertMutator mutator; + auto new_body = mutator.VisitStmt(func->body_); + if (new_body.get() == func->body_.get()) return func; + return std::make_shared(func->name_, func->params_, func->param_directions_, + func->return_types_, new_body, func->span_, func->func_type_, + func->level_, func->role_, func->attrs_); + }; + return CreateFunctionPass(pass_func, "ConvertToPtoasMultiBuffer", kConvertToPtoasMultiBufferProperties); +} + +} // namespace pass +} // namespace ir +} // namespace pypto diff --git a/src/ir/transforms/init_memref.cpp b/src/ir/transforms/init_memref.cpp index 362b8eca0..f24efd0fb 100644 --- a/src/ir/transforms/init_memref.cpp +++ b/src/ir/transforms/init_memref.cpp @@ -79,6 +79,14 @@ bool ProducesBufferLessTile(const ExprPtr& value, const SourceBufferLess& source if (IsOp(call, "tile.tpop_from_aic") || IsOp(call, "tile.tpop_from_aiv")) { return true; } + // ptoas multi-buffer slot views (route-2 prefetch) own no general-pool + // buffer: their storage is a physical slot of the region, addressed at + // codegen via `pto.multi_tile_get`. Keeping them MemRef-less gives each use + // its own SSA name (no pure-alias collapse onto the region) and keeps them + // out of MemoryReuse's lifetime coalescing. + if (IsOp(call, "tile.multi_buffer_get_slot") || IsOp(call, "tile.multi_buffer_load_slot")) { + return true; + } if (op_predicates::IsBufferAliasingViewOp(call->op_->name_) && !call->args_.empty()) { auto in = AsVarLike(call->args_[0]); return in && source_buffer_less(in.get()); diff --git a/src/ir/transforms/memory_reuse_pass.cpp b/src/ir/transforms/memory_reuse_pass.cpp index 68cd7dfb1..2aa0536ce 100644 --- a/src/ir/transforms/memory_reuse_pass.cpp +++ b/src/ir/transforms/memory_reuse_pass.cpp @@ -86,6 +86,13 @@ struct LifetimeAnalysisResult { /// private for ping-pong); compute↔compute cross-stage reuse is allowed so the /// bulk of intermediates can still coalesce and fit the on-chip budget. std::set pipeline_load_tiles; + /// Reuse-interval representatives defined by ``tile.multi_buffer_alloc`` + /// (ptoas multi-buffer, `use_ptoas_multi_buffer`). + /// Such a tile owns an exclusive N-slot region (rotated slot-by-slot across + /// the loop by ptoas), so it must never share a physical buffer with another + /// tile — ``IdentifyReuseOpportunities`` blocks any coalescing that touches + /// one, and ``AllocateMemoryAddr`` reserves ``N * slot_size`` for it. + std::set multi_buffer_vars; }; /** @@ -1283,6 +1290,7 @@ LifetimeAnalysisResult ComputeLifetimes(const StmtPtr& func_body) { // the subset whose defining op is a load (vs compute). std::map>> pipeline_membership; std::set pipeline_load_tiles; + std::set multi_buffer_vars; for (const auto& var : result.ordered_defs) { if (processed_vars.count(var)) { @@ -1351,8 +1359,15 @@ LifetimeAnalysisResult ComputeLifetimes(const StmtPtr& func_body) { if (IsOp(call, "tile.load") || IsOp(call, "tile.read")) { pipeline_load_tiles.insert(interval.variable.get()); } - break; } + // The ptoas multi-buffer region (route-2 prefetch) owns an exclusive + // N-slot span; mark it so the packer refuses any coalescing that touches + // it. Slot views (get_slot / load_slot) share the region's base and are + // folded into the same sharing group, so excluding the region covers them. + if (IsOp(call, "tile.multi_buffer_alloc")) { + multi_buffer_vars.insert(interval.variable.get()); + } + if (!packed.empty()) break; } for (const auto& group_var : sharing_group) { @@ -1377,7 +1392,8 @@ LifetimeAnalysisResult ComputeLifetimes(const StmtPtr& func_body) { std::move(result.phi_family_ids), std::move(var_liveness), std::move(pipeline_membership), - std::move(pipeline_load_tiles)}; + std::move(pipeline_load_tiles), + std::move(multi_buffer_vars)}; } // NOTE: The former tile-type reuse-compatibility gate (AreTileTypesCompatible) @@ -1635,7 +1651,7 @@ std::map IdentifyReuseOpportunities( const std::map>& sharing_groups, const std::map>& var_liveness, const std::map>>& pipeline_membership, - const std::set& pipeline_load_tiles) { + const std::set& pipeline_load_tiles, const std::set& multi_buffer_vars) { std::map reuse_map; // Members of a sharing group (the vars that already physically share one base). @@ -1801,6 +1817,15 @@ std::map IdentifyReuseOpportunities( // PTO binds a per-var alloc_tile so differing shapes/dtypes legally alias one // base, and largest-first ordering guarantees the buffer is sized to its // representative (no member is ever larger than the buffer it joins). + // ptoas multi-buffer guard: a tile that ptoas rotates through N physical slots + // owns an exclusive [base, base + N*slot) region (AllocateMemoryAddr reserves + // it). Its SSA lifetime looks like a single loop iteration, but slots 1..N-1 + // hold live data across the whole loop, so it must never share storage with + // any other tile. Block symmetrically whenever either side is multi-buffered. + auto multi_buffer_blocks = [&multi_buffer_vars](const LifetimeInterval& a, const LifetimeInterval& b) { + return multi_buffer_vars.count(a.variable.get()) != 0 || multi_buffer_vars.count(b.variable.get()) != 0; + }; + auto can_share = [&](const LifetimeInterval& cand, const LifetimeInterval& member) { // Group-interval overlap is a fast reject; when it fires, fall back to the // precise per-var check so mutually-exclusive / same-value phi-family tiles @@ -1809,7 +1834,8 @@ std::map IdentifyReuseOpportunities( if (LifetimesOverlap(cand, member) && overlap_blocks_sharing(cand, member)) return false; if (hazard_blocks(cand, member) || hazard_blocks(member, cand)) return false; if (forbid_blocks(cand, member) || forbid_blocks(member, cand)) return false; - if (pipeline_blocks(cand, member)) return false; // symmetric — one call suffices + if (pipeline_blocks(cand, member)) return false; // symmetric — one call suffices + if (multi_buffer_blocks(cand, member)) return false; // symmetric return true; }; @@ -2592,7 +2618,7 @@ FunctionPtr TransformMemoryReuse(const FunctionPtr& func) { auto reuse_map = IdentifyReuseOpportunities( analysis_result.lifetimes, hazard, forbid_alias, analysis_result.phi_family_ids, analysis_result.var_sharing_groups, analysis_result.var_liveness, analysis_result.pipeline_membership, - analysis_result.pipeline_load_tiles); + analysis_result.pipeline_load_tiles, analysis_result.multi_buffer_vars); // Step 3: Apply MemRef sharing (skip if no reuse candidates) if (!reuse_map.empty()) { diff --git a/src/ir/transforms/pass_context.cpp b/src/ir/transforms/pass_context.cpp index 50a8ee92f..79690a093 100644 --- a/src/ir/transforms/pass_context.cpp +++ b/src/ir/transforms/pass_context.cpp @@ -304,12 +304,19 @@ std::string DiagnosticInstrument::GetName() const { return "DiagnosticInstrument PassContext::PassContext(std::vector instruments, VerificationLevel verification_level, DiagnosticPhase diagnostic_phase, DiagnosticCheckSet disabled_diagnostics, - MemoryPlanner memory_planner) + MemoryPlanner memory_planner, bool use_ptoas_multi_buffer) : instruments_(std::move(instruments)), verification_level_(verification_level), diagnostic_phase_(diagnostic_phase), disabled_diagnostics_(disabled_diagnostics), - memory_planner_(memory_planner), + // ptoas multi-buffer implies the PTOAS memory planner: the cross-iteration + // double-buffer overlap only materializes at --pto-level=level2, where ptoas + // PlanMemory assigns the N slots concrete disjoint addresses (level3's baked + // base + dynamic slot defeats ptoas MemAlias and serializes). Forcing it here + // keeps every consumer — pass-manager mem-planning skip, codegen level, ptoas + // invocation — consistent no matter how the context was constructed. + memory_planner_(use_ptoas_multi_buffer ? MemoryPlanner::PtoAS : memory_planner), + use_ptoas_multi_buffer_(use_ptoas_multi_buffer), previous_(nullptr) {} VerificationLevel PassContext::GetVerificationLevel() const { return verification_level_; } @@ -320,6 +327,8 @@ DiagnosticPhase PassContext::GetDiagnosticPhase() const { return diagnostic_phas const DiagnosticCheckSet& PassContext::GetDisabledDiagnostics() const { return disabled_diagnostics_; } +bool PassContext::UsePtoasMultiBuffer() const { return use_ptoas_multi_buffer_; } + const std::vector& PassContext::GetInstruments() const { return instruments_; } void PassContext::EnterContext() { diff --git a/src/ir/transforms/python_printer.cpp b/src/ir/transforms/python_printer.cpp index 0be50f441..71aa8bfe6 100644 --- a/src/ir/transforms/python_printer.cpp +++ b/src/ir/transforms/python_printer.cpp @@ -1103,11 +1103,15 @@ void IRPythonPrinter::VisitExpr_(const CallPtr& op) { // tests don't expect. ``pipeline_membership`` (set by LowerPipelineLoops, read // by MemoryReuse) has no such surface and MUST round-trip, else the structural // equality check after those passes fails. The matching reader is - // ``ast_parser`` (``_parse_op_attrs`` -> ``set_call_attrs``). + // ``ast_parser`` (``_parse_op_attrs`` -> ``set_call_attrs``). (ptoas + // multi-buffer carries no such attr — it uses first-class ops + // ``tile.multi_buffer_alloc`` / ``_load_slot`` that round-trip generically.) { std::vector*> serialized_attrs; for (const auto& kv : op->attrs_) { - if (kv.first == kPipelineMembershipAttr) serialized_attrs.push_back(&kv); + if (kv.first == kPipelineMembershipAttr) { + serialized_attrs.push_back(&kv); + } } if (!serialized_attrs.empty()) { stream_ << (need_comma ? ", " : "") << "attrs={"; diff --git a/tests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.py b/tests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.py new file mode 100644 index 000000000..eaf58f4b7 --- /dev/null +++ b/tests/st/codegen/dsl/test_ptoas_multi_buffer_codegen.py @@ -0,0 +1,201 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Codegen checks for the ptoas multi-buffer switch (use_ptoas_multi_buffer). + +A `pl.pipeline(stage=N)` loop over an i-dependent vec load normally lowers to +pypto's own body-replication ping-pong (several `pto.alloc_tile` at disjoint +addresses). With `use_ptoas_multi_buffer=True`, `ConvertToPtoasMultiBuffer` +instead keeps a single loop body and emits ptoas multi-buffer ops: one hoisted +`pto.alloc_multi_tile ... count=N` and a `pto.multi_tile_get %mb[i%N]` per +iteration (same slot for load + consume), delegating slot rotation + the +cross-iteration double-buffer overlap to ptoas. + +The switch **auto-forces `memory_planner=PTOAS`** (--pto-level=level2): the +overlap only materializes when ptoas PlanMemory assigns the N slots concrete +disjoint addresses, so the emitted region is address-less (level2). + +These are codegen-level checks on the emitted `.pto` MLIR text (the JIT compile +may raise post-codegen when the simpler runtime is absent — the assertion is on +the `.pto`, which materializes first). +""" + +import shutil + +import pypto.language as pl +import pypto.pypto_core.ir as _ir +import pytest +import torch +from pypto import backend as _backend +from pypto.backend import BackendType +from pypto.ir.pass_manager import OptimizationStrategy, PassManager +from pypto.pypto_core import passes +from pypto.runtime import RunConfig + +T, N = 256, 64 + + +# Two distinct kernels: `@pl.jit` caches by function + arg signature (not by the +# active PassContext), so a single kernel compiled once would reuse its artifact +# regardless of the switch. Separate functions keep the on/off caches disjoint. +# The load offset is i-DEPENDENT (`[_i*64, 0]`) so consecutive iterations touch +# different data — the only case where double-buffering is observable. +@pl.jit +def mbuf_pipe_add_off( + a: pl.Tensor[[T, N], pl.FP32], + c: pl.Out[pl.Tensor[[T, N], pl.FP32]], +): + """Accumulate a rolling i-dependent [64, N] vec tile across a 2-stage pipeline.""" + with pl.at(level=pl.Level.CORE_GROUP): + acc = pl.load(a, [0, 0], [64, N]) + for _i in pl.pipeline(1, 4, stage=2): # noqa: B007 - pipeline index + t = pl.load(a, [_i * 64, 0], [64, N]) + acc = pl.add(acc, t) + c = pl.store(acc, [0, 0], c) + return c + + +@pl.jit +def mbuf_pipe_add_on( + a: pl.Tensor[[T, N], pl.FP32], + c: pl.Out[pl.Tensor[[T, N], pl.FP32]], +): + """Same kernel as ``mbuf_pipe_add_off``; compiled with the switch on.""" + with pl.at(level=pl.Level.CORE_GROUP): + acc = pl.load(a, [0, 0], [64, N]) + for _i in pl.pipeline(1, 4, stage=2): # noqa: B007 - pipeline index + t = pl.load(a, [_i * 64, 0], [64, N]) + acc = pl.add(acc, t) + c = pl.store(acc, [0, 0], c) + return c + + +def _emit_pto(kernel, dump_dir, use_multi_buffer: bool) -> str: + if dump_dir.exists(): + shutil.rmtree(dump_dir) + cfg = RunConfig( + platform="a2a3", + codegen_only=True, + save_kernels=True, + save_kernels_dir=str(dump_dir), + use_ptoas_multi_buffer=use_multi_buffer, + ) + try: + kernel(torch.randn(T, N), torch.empty(T, N), config=cfg) + except Exception: # noqa: BLE001, S110 - post-codegen (simpler) failures are irrelevant here + pass + ptos = sorted(dump_dir.rglob("*.pto")) + assert ptos, f"codegen emitted no .pto under {dump_dir}" + return ptos[0].read_text() + + +def test_switch_off_uses_ordinary_alloc_tile(tmp_path): + """Default (switch off): no multi-buffer ops; ordinary alloc_tile ping-pong.""" + text = _emit_pto(mbuf_pipe_add_off, tmp_path / "off", use_multi_buffer=False) + assert "pto.alloc_multi_tile" not in text + assert "pto.multi_tile_get" not in text + assert "pto.alloc_tile" in text + + +def test_switch_on_emits_same_slot_multi_buffer(tmp_path): + """Switch on: rotating vec tile lowers to a same-slot multi-buffer region. + + The switch auto-forces memory_planner=PTOAS (level2), so the region is + address-less and ptoas PlanMemory assigns the N slots. The single loop body + keeps one `multi_tile_get` per iteration (slot i%N, no prefetch split). + """ + text = _emit_pto(mbuf_pipe_add_on, tmp_path / "on", use_multi_buffer=True) + assert "pto.alloc_multi_tile" in text + assert "count=2" in text + assert "pto.multi_tile_get" in text + # level2 (ptoas PlanMemory owns addresses): the region carries no baked addr. + assert "pto.alloc_multi_tile addr" not in text + # Same-slot form: no prefetch guard (scf.if) was emitted around the load. + assert "scf.if" not in text + + +@pl.jit +def mbuf_pipe_add_dump( + a: pl.Tensor[[T, N], pl.FP32], + c: pl.Out[pl.Tensor[[T, N], pl.FP32]], +): + """Same kernel; compiled via the explicit RunConfig field with dump_passes on.""" + with pl.at(level=pl.Level.CORE_GROUP): + acc = pl.load(a, [0, 0], [64, N]) + for _i in pl.pipeline(1, 4, stage=2): # noqa: B007 - pipeline index + t = pl.load(a, [_i * 64, 0], [64, N]) + acc = pl.add(acc, t) + c = pl.store(acc, [0, 0], c) + return c + + +def test_runconfig_field_survives_dump_passes(tmp_path): + """Regression: RunConfig.use_ptoas_multi_buffer must reach the pass even with + dump_passes=True. Dump mode reconstructs the PassContext, and previously + dropped use_ptoas_multi_buffer there, silently making the pass a no-op.""" + dump_dir = tmp_path / "dump_on" + cfg = RunConfig( + platform="a2a3", + codegen_only=True, + save_kernels=True, + save_kernels_dir=str(dump_dir), + dump_passes=True, + use_ptoas_multi_buffer=True, + ) + try: + mbuf_pipe_add_dump(torch.randn(T, N), torch.empty(T, N), config=cfg) + except Exception: # noqa: BLE001, S110 - post-codegen (simpler) failures are irrelevant here + pass + ptos = sorted(dump_dir.rglob("*.pto")) + assert ptos, f"codegen emitted no .pto under {dump_dir}" + text = ptos[0].read_text() + assert "pto.alloc_multi_tile" in text + assert "pto.multi_tile_get" in text + + +def test_multi_buffer_ops_round_trip(): + """The pass-synthesized multi_buffer ops must survive print -> parse. + + `tile.multi_buffer_alloc` / `tile.multi_buffer_load_slot` have no user DSL + surface, but IR dumps reparse, so the printer must serialize them and the + parser recover them (structural equality after a round-trip).""" + _backend.reset_for_testing() + _backend.set_backend_type(BackendType.Ascend910B) + try: + + @pl.program + class Prog: + @pl.function + def main( + self, + a: pl.Tensor[[T, N], pl.FP32], + c: pl.Out[pl.Tensor[[T, N], pl.FP32]], + ): + with pl.at(level=pl.Level.CORE_GROUP): + acc = pl.load(a, [0, 0], [64, N]) + for _i in pl.pipeline(1, 4, stage=2): # noqa: B007 + t = pl.load(a, [_i * 64, 0], [64, N]) + acc = pl.add(acc, t) + c = pl.store(acc, [0, 0], c) + return c + + with passes.PassContext([], use_ptoas_multi_buffer=True): + out = PassManager.get_strategy(OptimizationStrategy.Default).run_passes(Prog) + + printed = out.as_python() + assert "multi_buffer_alloc" in printed + assert "multi_buffer_load_slot" in printed + reparsed = pl.parse_program(printed) + _ir.assert_structural_equal(out, reparsed) + finally: + _backend.reset_for_testing() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/st/runtime/test_ptoas_multi_buffer_device.py b/tests/st/runtime/test_ptoas_multi_buffer_device.py new file mode 100644 index 000000000..f8b9e22a0 --- /dev/null +++ b/tests/st/runtime/test_ptoas_multi_buffer_device.py @@ -0,0 +1,79 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""On-device numeric parity for the ptoas multi-buffer switch (use_ptoas_multi_buffer). + +Runs a 2-stage pipeline that accumulates four i-dependent [64, N] vec tiles. +With the switch on, `ConvertToPtoasMultiBuffer` lowers the rotating load to a +ptoas multi-buffer region (slot i%N) and auto-selects memory_planner=PTOAS +(--pto-level=level2), where ptoas PlanMemory delivers the cross-iteration +double-buffer overlap. This verifies the lowered kernel is numerically correct +on real hardware and matches the switch-off (native pipeline) result. +""" + +import dataclasses + +import pypto.language as pl +import pytest +import torch + +T, N = 256, 64 + + +@pl.jit +def mbuf_accum( + a: pl.Tensor[[T, N], pl.FP32], + c: pl.Out[pl.Tensor[[T, N], pl.FP32]], +): + """c[0:64] = a[0:64] + a[64:128] + a[128:192] + a[192:256].""" + with pl.at(level=pl.Level.CORE_GROUP): + acc = pl.load(a, [0, 0], [64, N]) + for _i in pl.pipeline(1, 4, stage=2): # noqa: B007 - pipeline index + t = pl.load(a, [_i * 64, 0], [64, N]) + acc = pl.add(acc, t) + c = pl.store(acc, [0, 0], c) + return c + + +def _reference(a: torch.Tensor) -> torch.Tensor: + return a[0:64] + a[64:128] + a[128:192] + a[192:256] + + +class TestPtoasMultiBufferDevice: + """On-device numeric parity for the ptoas multi-buffer switch.""" + + def test_multi_buffer_numeric_parity(self, test_config): + """Switch-on (multi-buffer, level2) must match the torch reference and + the switch-off (native pipeline) result on real hardware.""" + a = torch.randn(T, N, dtype=torch.float32) + expected = _reference(a) + + # Switch OFF: native pipeline ping-pong (default planner / level3). + mbuf_accum._cache.clear() + c_off = torch.zeros(T, N, dtype=torch.float32) + mbuf_accum(a, c_off, config=dataclasses.replace(test_config, use_ptoas_multi_buffer=False)) + + # Switch ON: ptoas multi-buffer (auto memory_planner=PTOAS / level2). + mbuf_accum._cache.clear() + c_on = torch.zeros(T, N, dtype=torch.float32) + mbuf_accum(a, c_on, config=dataclasses.replace(test_config, use_ptoas_multi_buffer=True)) + + assert torch.allclose(c_off[0:64], expected, rtol=1e-4, atol=1e-4), ( + f"switch-off diff = {(c_off[0:64] - expected).abs().max().item()}" + ) + assert torch.allclose(c_on[0:64], expected, rtol=1e-4, atol=1e-4), ( + f"multi-buffer diff vs reference = {(c_on[0:64] - expected).abs().max().item()}" + ) + assert torch.allclose(c_on[0:64], c_off[0:64], rtol=1e-5, atol=1e-5), ( + f"multi-buffer vs native diff = {(c_on[0:64] - c_off[0:64]).abs().max().item()}" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/ut/ir/transforms/test_pass_manager.py b/tests/ut/ir/transforms/test_pass_manager.py index 1666eadeb..1e5ff5ae6 100644 --- a/tests/ut/ir/transforms/test_pass_manager.py +++ b/tests/ut/ir/transforms/test_pass_manager.py @@ -46,6 +46,7 @@ "StampTfreeSplit", "NormalizeReturnOrder", "SkewCrossCorePipeline", + "ConvertToPtoasMultiBuffer", "LowerPipelineLoops", "CanonicalizeIOOrder", "MaterializeTensorStrides", @@ -86,6 +87,7 @@ "StampTfreeSplit", "NormalizeReturnOrder", "SkewCrossCorePipeline", + "ConvertToPtoasMultiBuffer", "LowerPipelineLoops", "CanonicalizeIOOrder", "MaterializeTensorStrides", diff --git a/tests/ut/jit/test_cache.py b/tests/ut/jit/test_cache.py index b1993c421..01b9c8bc1 100644 --- a/tests/ut/jit/test_cache.py +++ b/tests/ut/jit/test_cache.py @@ -88,7 +88,10 @@ def test_basic_key_structure(self): assert isinstance(tensor_part, tuple) assert isinstance(scalar_part, tuple) assert dist_part is None # single-chip default - assert compile_opts == (("analyze_auto_scopes_for_deps", False),) + assert compile_opts == ( + ("analyze_auto_scopes_for_deps", False), + ("use_ptoas_multi_buffer", None), + ) def test_tensor_shape_in_key(self): key = self._make_key(