diff --git a/docs/en/dev/ir/05-operators.md b/docs/en/dev/ir/05-operators.md index 424a866dfc..0ea1d0ce13 100644 --- a/docs/en/dev/ir/05-operators.md +++ b/docs/en/dev/ir/05-operators.md @@ -196,6 +196,16 @@ Two limitations, both diagnosed rather than silently dropped: accumulate is therefore the author's responsibility, exactly as an oversized unpredicated `tile.matmul_acc` already is. +`AutoTileMatmulL0` *emits* this form: the K-loop it builds for a plain +`tile.matmul` is a single `tile.matmul_acc` predicated on `init_cond=(ko == 0)`, +not an `if ko == 0` over a fresh `tile.matmul` and an in-place accumulate. That +matters for more than tidiness. The branchy form has two producers of one logical +value on two different L0C buffers, and every pass downstream must then agree on +which buffer the phi lives in — with no `Acc`->`Acc` copy available when they +disagree. A predicated accumulate is in place on both paths, so one buffer is the +only possibility. `tile.matmul_bias` has no `init_cond` operand, so the bias form +still branches: its first K step must apply the bias exactly once. + At the tile layer, `tile.batch_matmul` provides batched semantics for `TileType` operands. It accepts rank >= 2 tiles, broadcasts the leading batch dimensions, and keeps the same operand-only interface style as `tile.matmul`. diff --git a/docs/en/dev/passes/15-auto_tile_matmul_l0.md b/docs/en/dev/passes/15-auto_tile_matmul_l0.md index eaa63c16a5..de9e616b2d 100644 --- a/docs/en/dev/passes/15-auto_tile_matmul_l0.md +++ b/docs/en/dev/passes/15-auto_tile_matmul_l0.md @@ -50,7 +50,7 @@ For each `tile.matmul`, `tile.matmul_acc`, or `tile.matmul_bias` in an InCore-ty - Sub-byte dtypes (cube path doesn't support them) — `PH-AT-003`. - `ChooseL0Tile` rejects the configuration — `PH-AT-005`. 5. **Build the K-loop** (per output sub-tile — the whole output when K-only, or each `[m, n]` sub-tile when M/N tiling): - - `tile.matmul` — iter-arg init is an Acc-resident `tile.create([m, n], dtype, target_memory=Acc)` placeholder; the loop body branches on `ko == 0` between `tile.matmul` (fresh Acc) and `tile.matmul_acc` (accumulating into the iter-arg). The `IfStmt` materializes a phi return_var that the outer yield carries back to the iter-arg. + - `tile.matmul` — iter-arg init is an Acc-resident `tile.create([m, n], dtype, target_memory=Acc)` placeholder; the loop body is a single `tile.matmul_acc` predicated on `init_cond=(ko == 0)`, so the first K step overwrites the accumulator and later steps accumulate into it. The backend expands the predicate into the `pto.tmatmul` / `pto.tmatmul.acc` pair, so the selection costs no IR-level control flow. This keeps the body in place on both paths and therefore on one L0C buffer — a branch between a *fresh* `tile.matmul` and an in-place accumulate would put one logical value on two Acc buffers, which nothing downstream can reconcile (there is no `Acc`→`Acc` copy). - `tile.matmul_acc` — iter-arg init is the caller's accumulator directly (its type already matches the per-iter `tile.matmul_acc` output); every iteration is uniform `tile.matmul_acc`, so no if-else. - `tile.matmul_bias` — uses the same fresh-accumulator loop as `tile.matmul`, but the first K block is `tile.matmul_bias` and every later block is `tile.matmul_acc`, so bias is added exactly once. For N tiling, the matching window is reconstructed as an independent tensor→Mat `tile.load`, then transferred with `tile.move` to Bias (`pto.tload` + `pto.tmov`). A zero-copy one-row Mat `tile.slice` is not used because its boxed `pto.subview` is not PTOAS-legal. - Per-iter operand extracts use `tile.extract(src, idx_row, idx_col, [shape], target_memory=Left|Right)` — the SSA-form fusion of the older `tile.slice` (Mat-resident result) + `tile.mov` (Mat→Left/Right) pair. This eliminates the intermediate Mat-resident slice tile and lowers to `pto.textract` rather than `pto.subview`, sidestepping the latter's `valid_row` codegen mismatch. For an output sub-tile at origin `(mi, ni)` the extracts slice `lhs[mi:mi+m, ko:ko+k]` and `rhs[ko:ko+k, ni:ni+n]`; the K-only case is `mi == ni == 0`, `m == M`, `n == N`. @@ -167,13 +167,8 @@ class After: for ko, (c_iter,) in pl.pipeline(0, 256, 64, init_values=(c_l0_init,), stage=2): sa = pl.tile.extract(a_mat, 0, ko, [128, 64], target_memory=Left) sb = pl.tile.extract(b_mat, ko, 0, [64, 128], target_memory=Right) - if ko == 0: - c_first = pl.tile.matmul(sa, sb) - c_phi = pl.yield_(c_first) - else: - c_acc = pl.tile.matmul_acc(c_iter, sa, sb) - c_phi = pl.yield_(c_acc) - c = pl.yield_(c_phi) + c_acc = pl.tile.matmul_acc(c_iter, sa, sb, init_cond=(ko == 0)) + c = pl.yield_(c_acc) # c (the yield-LHS) holds the accumulated Acc-typed result. ... ``` @@ -208,13 +203,8 @@ c_t1_init = pl.tile.create([256, 256], dtype=pl.FP32, target_memory=Acc) for ko, (c_iter,) in pl.pipeline(0, 512, 32, init_values=(c_t1_init,), stage=2): sa = pl.tile.extract(lhs_mat, 256, ko, [256, 32], target_memory=Left) sb = pl.tile.extract(rhs_mat, ko, 0, [32, 256], target_memory=Right) - if ko == 0: - c_first = pl.tile.matmul(sa, sb) - c_phi = pl.yield_(c_first) - else: - c_acc = pl.tile.matmul_acc(c_iter, sa, sb) - c_phi = pl.yield_(c_acc) - c_t1 = pl.yield_(c_phi) + c_acc = pl.tile.matmul_acc(c_iter, sa, sb, init_cond=(ko == 0)) + c_t1 = pl.yield_(c_acc) out_t1 = pl.store(c_t1, [256, 0], out_t0) # store sub-tile to out[256:512, 0:256] ``` diff --git a/docs/en/dev/passes/33-memory_reuse.md b/docs/en/dev/passes/33-memory_reuse.md index ad5331614e..2f77afd5f1 100644 --- a/docs/en/dev/passes/33-memory_reuse.md +++ b/docs/en/dev/passes/33-memory_reuse.md @@ -59,13 +59,12 @@ program_optimized = reuse_pass(program) 1. **Lifetime Analysis**: Walk the full IR tree (including nested control flow bodies) to compute variable lifetimes via def-use analysis. Variables defined outside a loop but used inside have their lifetime extended to the end of the loop (loop-aware extension) 2. **Interference Check**: Identify variables with overlapping lifetimes 3. **MemRef Sharing** (global first-fit-decreasing packing, `IdentifyReuseOpportunities`): Within each memory space, intervals are packed **largest-first**; every later interval joins the first existing buffer all of whose members it can share with (non-overlapping lifetime + hazard/no-alias safe — see `can_share`). A buffer's allocated size is fixed by its first (largest) member, so admitting a smaller member afterwards costs nothing — and a *later, larger* interval can now host an *earlier, smaller* one. (The previous definition-order greedy had a one-directional size gate `source.size >= target.size`, so two lifetime-disjoint tiles whose small one was defined first were never coalesced.) The representative each member is rebased onto is the buffer's largest member; its `tile.alloc` dominates the whole function because InitMemRef hoists every alloc to the function-body head, so a representative defined after some of its members is safe. Because the packer no longer runs in program order, every pairwise gate (hazard, no-alias) is checked in both directions. -4. **Loop-carry re-alignment** (`AlignLoopCarriesToInitMutator`): Sharing (step 3) only retypes `AssignStmt`-defined vars (producers/init); loop-carried `iter_arg`/`return_var` nodes are excluded from the lifetime/sharing maps and keep their original MemRef. This step walks `ForStmt`s **top-down** and retypes each loop's `iter_arg`/`return_var` to its (now-reused) `initValue` MemRef, seeding `var_remap_` before recursing so a nested loop observes the corrected outer `iter_arg` as its init. Without it, a reused **nested pipelined `matmul_acc`** accumulator splits across two Acc buffers and step 5 emits invalid `acc→acc tile.move` ops that ptoas rejects on Ascend 910B ([#1352](https://github.com/hw-native-sys/pypto/issues/1352)) -5. **Accumulator if-phi coalescing** (`TopDownRetargeter::CoalesceAccumulatorIfPhis`): `LowerPipelineLoops` peels a stage-2 K-loop into `if`-phis whose live branch is an in-place `matmul_acc` (on the accumulator buffer) and whose dead `if k==0` branch is a fresh `matmul` seed on a *different* Acc buffer. Left alone, step 6 would try to reconcile them with an `acc→acc tile.move` — a second co-live L0C buffer (overflow) that ptoas also rejects (no legal Acc→Acc `tmov`). This step identifies the in-place-accumulator branch by its `reuses_input` producer and retargets the *other* branch's seed onto the accumulator buffer, so both branches share it and no move is emitted (matching `mad_acc`'s shared-`%dst` semantics). Scoped to `Acc`; the retarget is **mandatory** (a declined retarget is an `INTERNAL_CHECK`, never a move — no legal Acc→Acc move exists). It bypasses the *global* dead-at-assign liveness check (which would false-decline on the legitimate post-if phi consumer), but only after verifying the two preconditions branch exclusivity actually needs: (a) the seed producer is a `Call` lexically **inside** the branch (a pre-if value yielded through the branch runs unconditionally and would clobber the accumulator the sibling in-place branch reads), and (b) a **branch-scoped** `IsTargetDeadAtAssign` (bounded to stop at the enclosing `if`) finds no same-branch tail read of the accumulator buffer. When either fails, the phi remains uncoalesced and step 6 fails loudly rather than manufacturing unsupported Acc→Acc IR -6. **Yield fixup**: Fix MemRef mismatches in control flow return variables. Acc→Acc copies are not legal, so a residual mismatched Acc carry is an internal error rather than a `tile.move`: +4. **Loop-carry re-alignment** (`AlignLoopCarriesToInitMutator`): Sharing (step 3) only retypes `AssignStmt`-defined vars (producers/init); loop-carried `iter_arg`/`return_var` nodes are excluded from the lifetime/sharing maps and keep their original MemRef. This step walks `ForStmt`s **top-down** and retypes each loop's `iter_arg`/`return_var` to its (now-reused) `initValue` MemRef, seeding `var_remap_` before recursing so a nested loop observes the corrected outer `iter_arg` as its init. Without it, a reused **nested pipelined `matmul_acc`** accumulator splits across two Acc buffers and step 5 rejects the carry as unreconcilable ([#1352](https://github.com/hw-native-sys/pypto/issues/1352)) +5. **Yield fixup**: Fix MemRef mismatches in control flow return variables. There is no Acc→Acc copy on any target (nothing reads L0C but the FIXPIPE drain, whose destinations are L1 and GM), so a mismatched Acc carry is reported to the author — naming both buffers — rather than reconciled with a `tile.move`: - **ForStmt**: Ensure all 4 loop-carry variables (initValue, iter_arg, yield value, return_var) share the same MemRef. Inserts `tile.move` before yield if MemRefs differ - **IfStmt**: Patch return_vars to match yield value's MemRef -7. **Identity-copy buffer normalization** (`NormalizeIdentityCopyBuffersMutator`): after step 5 retargets an accumulator if-phi, a downstream bare-`Var` SSA identity copy of the (now-moved) return_var can still carry the pre-coalesce buffer (e.g. `c: …mem_acc_17 = c_phi` after `c_phi` moved to `mem_acc_5`). An `x = y` copy (value a bare `Var`, not a `Call`) is a pure rename and must alias `y`'s buffer, so this single forward pass retypes such a copy's LHS to the RHS's MemRef and substitutes the LHS's downstream uses. No-op when no mismatch exists -8. **Remove redundant allocs**: Collect all MemRefs still referenced by TileType variables, then remove `tile.alloc` statements whose MemRef is no longer in use +6. **Identity-copy buffer normalization** (`NormalizeIdentityCopyBuffersMutator`): when an earlier step retargets a value, a downstream bare-`Var` SSA identity copy of it can still carry the pre-retarget buffer (e.g. `c: …mem_acc_17 = c_phi` after `c_phi` moved to `mem_acc_5`). An `x = y` copy (value a bare `Var`, not a `Call`) is a pure rename and must alias `y`'s buffer, so this single forward pass retypes such a copy's LHS to the RHS's MemRef and substitutes the LHS's downstream uses. No-op when no mismatch exists +7. **Remove redundant allocs**: Collect all MemRefs still referenced by TileType variables, then remove `tile.alloc` statements whose MemRef is no longer in use **Reuse conditions**: @@ -247,9 +246,8 @@ Pass MemoryReuse(); - `ComputeLifetimes` builds MemRef sharing groups and lifetime intervals - `IdentifyReuseOpportunities` finds reuse candidates - `ApplyMemRefSharing` updates MemRef pointers via `MemRefSharingMutator` -- `TopDownRetargeter::CoalesceAccumulatorIfPhis` coalesces peeled loop-carried accumulator `if`-phis by retargeting the dead-branch seed onto the in-place accumulator buffer, so `YieldFixupMutator` never emits an illegal `acc→acc tile.move` (see Algorithm step 5) -- `YieldFixupMutator` fixes ForStmt/IfStmt yield/return_var MemRef mismatches after reuse (inserts `tile.move` when legal; rejects residual Acc→Acc mismatches) -- `NormalizeIdentityCopyBuffersMutator` reconciles bare-`Var` SSA identity copies whose LHS/RHS buffers diverged after accumulator if-phi coalescing (see Algorithm step 7) +- `YieldFixupMutator` fixes ForStmt/IfStmt yield/return_var MemRef mismatches after reuse (inserts `tile.move` when legal; reports Acc→Acc mismatches, naming both buffers) +- `NormalizeIdentityCopyBuffersMutator` reconciles bare-`Var` SSA identity copies whose LHS/RHS buffers diverged after an earlier retarget (see Algorithm step 6) - `UsedMemRefCollector` gathers still-referenced MemRef pointers after sharing - `RemoveUnusedAllocStatements` filters out redundant `tile.alloc` statements from `SeqStmts` diff --git a/docs/zh/dev/ir/05-operators.md b/docs/zh/dev/ir/05-operators.md index 8d5ea2511f..d873c5c493 100644 --- a/docs/zh/dev/ir/05-operators.md +++ b/docs/zh/dev/ir/05-operators.md @@ -179,6 +179,14 @@ init 操作数。由于 `matmul_acc` 是原地操作（`set_output_reuses_input( split-K 所期望的。超尺寸的带谓词累加因而由编写者负责，这与超尺寸的无谓词 `tile.matmul_acc` 现有行为一致。 +`AutoTileMatmulL0` 会*生成*这种形式：它为普通 `tile.matmul` 构造的 K 循环体是一条以 +`init_cond=(ko == 0)` 为谓词的 `tile.matmul_acc`，而不是在全新 `tile.matmul` 与原地累加 +之间做 `if ko == 0` 分支。这不只是写法上的整洁：分支形式让同一个逻辑值有两个生产者、 +落在两块不同的 L0C 缓冲上，之后每一个 pass 都必须就 phi 属于哪一块达成一致 —— 而一旦 +不一致，并不存在可用于协调的 `Acc`->`Acc` 拷贝。带谓词的累加在两条路径上都是原地写， +因此只可能有一块缓冲。`tile.matmul_bias` 没有 `init_cond` 操作数，所以 bias 形式仍然分支： +它的第一个 K 步必须恰好施加一次 bias。 + 在 tile 层，`tile.batch_matmul` 为 `TileType` 操作数提供批量语义。它接受 rank >= 2 的 tile，广播前导批量维度，并保持与 `tile.matmul` 相同的纯操作数接口风格。如果批量操作数 需要转置语义，可以通过两种等价方式表达：在输入上显式使用 `tile.transpose(...)`，或在 diff --git a/docs/zh/dev/passes/15-auto_tile_matmul_l0.md b/docs/zh/dev/passes/15-auto_tile_matmul_l0.md index 55b63114ff..3af6cd6327 100644 --- a/docs/zh/dev/passes/15-auto_tile_matmul_l0.md +++ b/docs/zh/dev/passes/15-auto_tile_matmul_l0.md @@ -50,7 +50,7 @@ program_tiled = l0_tile_pass(program) - 子字节 dtype（cube path 不支持）—— `PH-AT-003`。 - `ChooseL0Tile` 拒绝该配置 —— `PH-AT-005`。 5. **构造 K-loop**（针对一个输出子块——K 切分时即整个输出，M/N 切分时为每个 `[m, n]` 子块）： - - `tile.matmul` —— iter-arg 初值为 Acc-resident 的 `tile.create([m, n], dtype, target_memory=Acc)` 占位；循环体用 `IfStmt` 在 `ko == 0` 时走 `tile.matmul`（产生新的 Acc），其它迭代走 `tile.matmul_acc`（向 iter-arg 上累加）。`IfStmt` 物化一个 phi 形式的 `return_var`，由外层 yield 写回 iter-arg。 + - `tile.matmul` —— iter-arg 初值为 Acc-resident 的 `tile.create([m, n], dtype, target_memory=Acc)` 占位；循环体是一条以 `init_cond=(ko == 0)` 为谓词的 `tile.matmul_acc`：第一个 K 步覆盖写累加器，其余迭代在其上累加。后端把该谓词展开为 `pto.tmatmul` / `pto.tmatmul.acc` 二选一，因此这一选择不产生任何 IR 层面的控制流。这样循环体在两条路径上都是原地写、始终只用一块 L0C；若改为在全新 `tile.matmul` 与原地累加之间分支，同一个逻辑值就会落在两块 Acc 缓冲上，而下游无法协调（不存在 `Acc`→`Acc` 拷贝）。 - `tile.matmul_acc` —— iter-arg 初值就是调用方传入的累加器（其类型已经与每次迭代的 `tile.matmul_acc` 输出一致）；每次迭代统一是 `tile.matmul_acc`，无需 if-else。 - `tile.matmul_bias` —— 使用与 `tile.matmul` 相同的新累加器循环，但第一个 K block 使用 `tile.matmul_bias`，之后都使用 `tile.matmul_acc`，因此 bias 恰好只加一次。N 切分时，对应窗口从原始 tensor 重新 `tile.load` 到独立 Mat tile，再通过 `tile.move` 传到 Bias（`pto.tload` + `pto.tmov`）。不会使用单行 Mat `tile.slice`，因为其 boxed `pto.subview` 不满足 PTOAS 合法性。 - 每次迭代的操作数抽取使用 `tile.extract(src, idx_row, idx_col, [shape], target_memory=Left|Right)` —— 这是旧版 `tile.slice`（Mat-resident 中间 tile）+ `tile.mov`（Mat→Left/Right）的 SSA 化合并。这样既消除了 Mat-resident 中间 slice tile，也使得 lower 后是 `pto.textract` 而不是 `pto.subview`，从而绕开后者的 `valid_row` codegen 不一致问题。对于原点为 `(mi, ni)` 的输出子块，抽取的是 `lhs[mi:mi+m, ko:ko+k]` 与 `rhs[ko:ko+k, ni:ni+n]`；K 切分情形即 `mi == ni == 0`、`m == M`、`n == N`。 @@ -155,13 +155,8 @@ class After: for ko, (c_iter,) in pl.pipeline(0, 256, 64, init_values=(c_l0_init,), stage=2): sa = pl.tile.extract(a_mat, 0, ko, [128, 64], target_memory=Left) sb = pl.tile.extract(b_mat, ko, 0, [64, 128], target_memory=Right) - if ko == 0: - c_first = pl.tile.matmul(sa, sb) - c_phi = pl.yield_(c_first) - else: - c_acc = pl.tile.matmul_acc(c_iter, sa, sb) - c_phi = pl.yield_(c_acc) - c = pl.yield_(c_phi) + c_acc = pl.tile.matmul_acc(c_iter, sa, sb, init_cond=(ko == 0)) + c = pl.yield_(c_acc) # c（即 yield-LHS）持有累加得到的 Acc 类型结果。 ... ``` @@ -196,13 +191,8 @@ c_t1_init = pl.tile.create([256, 256], dtype=pl.FP32, target_memory=Acc) for ko, (c_iter,) in pl.pipeline(0, 512, 32, init_values=(c_t1_init,), stage=2): sa = pl.tile.extract(lhs_mat, 256, ko, [256, 32], target_memory=Left) sb = pl.tile.extract(rhs_mat, ko, 0, [32, 256], target_memory=Right) - if ko == 0: - c_first = pl.tile.matmul(sa, sb) - c_phi = pl.yield_(c_first) - else: - c_acc = pl.tile.matmul_acc(c_iter, sa, sb) - c_phi = pl.yield_(c_acc) - c_t1 = pl.yield_(c_phi) + c_acc = pl.tile.matmul_acc(c_iter, sa, sb, init_cond=(ko == 0)) + c_t1 = pl.yield_(c_acc) out_t1 = pl.store(c_t1, [256, 0], out_t0) # 子块 store 到 out[256:512, 0:256] ``` diff --git a/docs/zh/dev/passes/33-memory_reuse.md b/docs/zh/dev/passes/33-memory_reuse.md index 77776641fd..7c5aa03448 100644 --- a/docs/zh/dev/passes/33-memory_reuse.md +++ b/docs/zh/dev/passes/33-memory_reuse.md @@ -55,13 +55,12 @@ program_optimized = reuse_pass(program) 1. **生命周期分析**：遍历完整 IR 树（包括嵌套控制流体内的语句）通过 def-use 分析计算变量生命周期。在循环外定义但在循环内使用的变量，其生命周期会延展到循环结束（循环感知延展） 2. **干涉检查**：识别生命周期重叠的变量 3. **MemRef 共享**（全局「最大优先 + first-fit」装箱，`IdentifyReuseOpportunities`）：在每个内存空间内，按 **大小从大到小** 装箱；后续每个区间加入第一个其全部成员都能与之共享的缓冲区（生命周期不重叠 + hazard / no-alias 安全，见 `can_share`）。缓冲区的分配大小由其首个（最大）成员固定，因此之后纳入更小的成员是「免费」的 —— 且 *后定义的较大区间* 现在可以承载 *先定义的较小区间*。（此前的定义序贪心带有单向的大小门槛 `source.size >= target.size`，因此两个生命周期不相交、但较小者先定义的 tile 永远无法合并。）每个成员被重定位到的「代表」是该缓冲区的最大成员；由于 InitMemRef 会把所有 `tile.alloc` 提升到函数体头部，代表的 alloc 支配整个函数，因此代表即使定义在其部分成员之后也是安全的。由于装箱器不再按程序序处理，每个成对门槛（hazard、no-alias）都会在两个方向上检查。 -4. **循环携带变量重对齐**（`AlignLoopCarriesToInitMutator`）：共享（步骤 3）只会重写由 `AssignStmt` 定义的变量（producer/init），而循环携带的 `iter_arg`/`return_var` 节点被排除在生命周期/共享映射之外、仍保留原始 MemRef。本步骤**自外向内**遍历 `ForStmt`，将每个循环的 `iter_arg`/`return_var` 重对齐到其（已复用的）`initValue` 的 MemRef，并在递归前写入 `var_remap_`，使嵌套循环能观察到已修正的外层 `iter_arg` 作为其 init。若缺少本步骤，被复用的**嵌套流水化 `matmul_acc`** 累加器会分裂到两个 Acc 缓冲区，导致步骤 5 插入非法的 `acc→acc tile.move`，被 Ascend 910B 的 ptoas 拒绝（[#1352](https://github.com/hw-native-sys/pypto/issues/1352)） -5. **累加器 if-phi 合并**（`TopDownRetargeter::CoalesceAccumulatorIfPhis`）：`LowerPipelineLoops` 会把 stage-2 的 K 循环剥离成 `if`-phi，其活跃分支是就地累加的 `matmul_acc`（位于累加器缓冲区），而失效的 `if k==0` 分支是位于*不同* Acc 缓冲区上的全新 `matmul` seed。若不处理，步骤 6 会尝试用 `acc→acc tile.move` 协调二者 —— 产生第二个同时存活的 L0C 缓冲区（溢出），且 ptoas 也会拒绝（不存在合法的 Acc→Acc `tmov`）。本步骤通过 `reuses_input` 的 producer 识别就地累加分支，并把*另一*分支的 seed 重定向到累加器缓冲区，使两个分支共享同一缓冲区、不再产生 move（符合 `mad_acc` 共享 `%dst` 的语义）。仅作用于 `Acc`；重定向是**强制的**（被拒绝的重定向会触发 `INTERNAL_CHECK`，绝不退化为 move —— 因为不存在合法的 Acc→Acc move）。它会跳过*全局* dead-at-assign 活跃性检查（否则会因 if 之后合法的 phi 消费者而误判拒绝），但仅在验证分支互斥真正所需的两个前提之后：(a) seed 的 producer 是词法上位于该分支**内部**的 `Call`（经由分支透传的 if 前值会无条件执行，从而破坏 sibling 就地分支所读取的累加器），以及 (b) **限定分支范围** 的 `IsTargetDeadAtAssign`（在所属 `if` 处停止）确认分支内 seed 之后没有对累加器缓冲区的尾部读取。任一前提不满足时，该 phi 保持未合并，步骤 6 会明确失败，而不会生成不受支持的 Acc→Acc IR -6. **Yield 修复**：修复控制流返回变量的 MemRef 不一致。Acc→Acc 拷贝不合法，因此残留的不匹配 Acc 循环携带值会触发内部错误，而不是生成 `tile.move`： +4. **循环携带变量重对齐**（`AlignLoopCarriesToInitMutator`）：共享（步骤 3）只会重写由 `AssignStmt` 定义的变量（producer/init），而循环携带的 `iter_arg`/`return_var` 节点被排除在生命周期/共享映射之外、仍保留原始 MemRef。本步骤**自外向内**遍历 `ForStmt`，将每个循环的 `iter_arg`/`return_var` 重对齐到其（已复用的）`initValue` 的 MemRef，并在递归前写入 `var_remap_`，使嵌套循环能观察到已修正的外层 `iter_arg` 作为其 init。若缺少本步骤，被复用的**嵌套流水化 `matmul_acc`** 累加器会分裂到两个 Acc 缓冲区，导致步骤 5 判定该循环携带值无法协调而报错（[#1352](https://github.com/hw-native-sys/pypto/issues/1352)） +5. **Yield 修复**：修复控制流返回变量的 MemRef 不一致。任何目标上都不存在 Acc→Acc 拷贝（除 FIXPIPE 排空外没有任何单元读 L0C，而它的目的地只有 L1 与 GM），因此不匹配的 Acc 循环携带值会连同两块缓冲区的名字一起报告给编写者，而不是用 `tile.move` 协调： - **ForStmt**：确保 4 个循环携带变量（initValue、iter_arg、yield value、return_var）共享同一个 MemRef。若 MemRef 不同则在 yield 前插入 `tile.move` - **IfStmt**：修补 return_vars 使其 MemRef 与 yield value 一致 -7. **恒等拷贝缓冲区归一化**（`NormalizeIdentityCopyBuffersMutator`）：在步骤 5 重定向累加器 if-phi 后，对（已被移动的）return_var 的下游裸 `Var` SSA 恒等拷贝可能仍携带合并前的缓冲区（例如 `c_phi` 移到 `mem_acc_5` 后，`c: …mem_acc_17 = c_phi`）。`x = y` 拷贝（值为裸 `Var` 而非 `Call`）是纯重命名、必须与 `y` 共用缓冲区，因此本次单向前向遍历把这类拷贝的 LHS 重定型到 RHS 的 MemRef，并替换 LHS 的下游使用。无不一致时为空操作 -8. **移除冗余 alloc**：收集仍被 TileType 变量引用的所有 MemRef，然后移除不再使用的 `tile.alloc` 语句 +6. **恒等拷贝缓冲区归一化**（`NormalizeIdentityCopyBuffersMutator`）：当此前某步重定向了某个值后，对（已被移动的）值的下游裸 `Var` SSA 恒等拷贝可能仍携带重定向前的缓冲区（例如 `c_phi` 移到 `mem_acc_5` 后，`c: …mem_acc_17 = c_phi`）。`x = y` 拷贝（值为裸 `Var` 而非 `Call`）是纯重命名、必须与 `y` 共用缓冲区，因此本次单向前向遍历把这类拷贝的 LHS 重定型到 RHS 的 MemRef，并替换 LHS 的下游使用。无不一致时为空操作 +7. **移除冗余 alloc**：收集仍被 TileType 变量引用的所有 MemRef，然后移除不再使用的 `tile.alloc` 语句 **复用条件**： @@ -235,9 +234,8 @@ Pass MemoryReuse(); - `ComputeLifetimes` 构建 MemRef 共享组和生命周期区间 - `IdentifyReuseOpportunities` 查找复用候选 - `ApplyMemRefSharing` 通过 `MemRefSharingMutator` 更新 MemRef 指针 -- `TopDownRetargeter::CoalesceAccumulatorIfPhis` 通过把失效分支的 seed 重定向到就地累加器缓冲区，合并被剥离的循环携带累加器 `if`-phi，使 `YieldFixupMutator` 不再产生非法的 `acc→acc tile.move`（见算法步骤 5） -- `YieldFixupMutator` 修复 ForStmt/IfStmt 在复用后的 yield/return_var MemRef 不一致（合法时插入 `tile.move`；拒绝残留的 Acc→Acc 不一致） -- `NormalizeIdentityCopyBuffersMutator` 协调累加器 if-phi 合并后 LHS/RHS 缓冲区不一致的裸 `Var` SSA 恒等拷贝（见算法步骤 7） +- `YieldFixupMutator` 修复 ForStmt/IfStmt 在复用后的 yield/return_var MemRef 不一致（合法时插入 `tile.move`；对残留的 Acc→Acc 不一致给出报错） +- `NormalizeIdentityCopyBuffersMutator` 协调此前某步重定向后 LHS/RHS 缓冲区不一致的裸 `Var` SSA 恒等拷贝（见算法步骤 6） - `UsedMemRefCollector` 收集共享后仍被引用的 MemRef 指针 - `RemoveUnusedAllocStatements` 从 `SeqStmts` 中过滤掉冗余的 `tile.alloc` 语句 diff --git a/src/ir/transforms/auto_tile_matmul_l0_pass.cpp b/src/ir/transforms/auto_tile_matmul_l0_pass.cpp index 1e90167cc1..19d1d6f537 100644 --- a/src/ir/transforms/auto_tile_matmul_l0_pass.cpp +++ b/src/ir/transforms/auto_tile_matmul_l0_pass.cpp @@ -22,12 +22,13 @@ /// can alias it onto QK's freed L0B (peak L0B = ``max(QK, PV)`` instead of the /// sum). The K-loop has the shape: /// -/// * ``tile.matmul`` — the loop body branches on the iteration index -/// (``ko == 0``) so the first iteration uses ``tile.matmul`` (fresh -/// accumulator) and subsequent iterations use ``tile.matmul_acc`` -/// (accumulating into the iter-arg). The iter-arg init is an Acc- -/// resident ``tile.create`` placeholder so the iter-arg / yield / -/// return_var chain is Acc-typed end-to-end. +/// * ``tile.matmul`` — every iteration is a ``tile.matmul_acc`` predicated on +/// ``init_cond=(ko == 0)``: the first iteration overwrites the accumulator +/// with ``lhs @ rhs``, later ones accumulate into it. The iter-arg init is +/// an Acc-resident ``tile.create`` placeholder so the iter-arg / yield / +/// return_var chain is Acc-typed end-to-end. The predicate keeps the body +/// in place on both paths, so the chain never splits across two L0C buffers +/// (there is no ``Acc``->``Acc`` copy to reconcile one if it did). /// * ``tile.matmul_acc`` — every iteration is ``tile.matmul_acc``; the /// iter-arg init is the caller-provided accumulator directly, so the /// chain is uniform and no if-else is needed. @@ -51,13 +52,13 @@ /// for ko in pl.pipeline(0, K, k, init_values=(c_init,), stage=2): /// sa = tile.extract(x_mat, 0, ko, [m, k], target_memory=Left) /// sb = tile.extract(y_mat, ko, 0, [k, n], target_memory=Right) -/// if ko == 0: -/// c1 = tile.matmul(sa, sb) // fresh Acc -/// c_phi = pl.yield_(c1) // if's return_var -/// else: -/// c2 = tile.matmul_acc(c_iter, sa, sb) // accumulate -/// c_phi = pl.yield_(c2) -/// yield c_phi +/// c_new = tile.matmul_acc(c_iter, sa, sb, init_cond=(ko == 0)) +/// yield c_new +/// +/// The backend expands ``init_cond`` into the ``pto.tmatmul`` / +/// ``pto.tmatmul.acc`` pair (a literal predicate picks one arm outright), so the +/// selection costs no IR-level control flow — see the ``init_cond`` section of +/// ``docs/en/dev/ir/05-operators.md``. /// /// Layout for ``tile.matmul_acc`` (acc_init is the caller's accumulator): /// for ko in pl.pipeline(0, K, k, init_values=(acc_init,), stage=2): @@ -479,18 +480,41 @@ VarPtr BuildBiasOperand(std::vector& stmts, const VarPtr& bias_src, int return bias->var_; } -/// Body of the K-loop for a fresh ``tile.matmul`` or ``tile.matmul_bias``: -/// branches on ``ko == 0`` between a fresh Acc (with the optional bias) and -/// ``tile.matmul_acc``. The ``IfStmt`` materializes a phi return_var that the -/// outer yield carries back to the iter-arg. +/// Body of the K-loop for a fresh ``tile.matmul`` or ``tile.matmul_bias``. +/// +/// Without a bias this is a single predicated ``tile.matmul_acc`` whose +/// ``init_cond`` is ``ko == 0``: on the first K step the accumulator is +/// overwritten with ``lhs @ rhs``, afterwards it is accumulated into. The +/// backend expands the predicate into the ``pto.tmatmul`` / ``pto.tmatmul.acc`` +/// pair (a literal predicate picks one arm outright), so the branch exists only +/// in the emitted code — never in the IR. +/// +/// This matters beyond tidiness. The former shape branched on ``ko == 0`` +/// between a *fresh* Acc tile and an in-place accumulate, which are two +/// producers of one logical value on two different L0C buffers. Every pass +/// downstream then had to agree on which buffer the phi lives in, and there is +/// no ``Acc``->``Acc`` copy to fall back on when they disagree. A predicated +/// accumulate is in place on both paths (``set_output_reuses_input(0)``), so one +/// buffer is the only possibility and no phi is materialized at all. +/// +/// ``tile.matmul_bias`` has no ``init_cond`` operand, so the bias form keeps the +/// ``IfStmt``: its first K step must apply the bias exactly once. StmtPtr BuildMatmulBody(const VarPtr& ko_var, const IterArgPtr& c_iter, const AssignStmtPtr& sa, const AssignStmtPtr& sb, const VarPtr& bias, const std::string& base, const Span& sp) { auto& reg = OpRegistry::GetInstance(); + auto init_cond = MakeEq(ko_var, MakeIndex(0, sp), sp); + + if (!bias) { + auto c_call = reg.Create("tile.matmul_acc", {ExprPtr(c_iter), sa->var_, sb->var_, init_cond}, sp); + auto c_var = std::make_shared(base + "_l0_c_acc", c_call->GetType(), sp); + auto c_assign = std::make_shared(c_var, c_call, sp); + auto body_yield = std::make_shared(std::vector{c_var}, sp); + return SeqStmts::Flatten(std::vector{sa, sb, c_assign, body_yield}, sp); + } - // Then-branch: fresh Acc tile, applying the optional bias exactly once. - auto c_then_call = bias ? reg.Create("tile.matmul_bias", {sa->var_, sb->var_, bias}, sp) - : reg.Create("tile.matmul", {sa->var_, sb->var_}, sp); + // Then-branch: fresh Acc tile, applying the bias exactly once. + auto c_then_call = reg.Create("tile.matmul_bias", {sa->var_, sb->var_, bias}, sp); auto c_then_var = std::make_shared(base + "_l0_c_first", c_then_call->GetType(), sp); auto c_then_assign = std::make_shared(c_then_var, c_then_call, sp); auto then_yield = std::make_shared(std::vector{c_then_var}, sp); @@ -504,8 +528,7 @@ StmtPtr BuildMatmulBody(const VarPtr& ko_var, const IterArgPtr& c_iter, const As StmtPtr else_body = SeqStmts::Flatten(std::vector{c_else_assign, else_yield}, sp); auto c_phi = std::make_shared(base + "_l0_c_phi", c_then_call->GetType(), sp); - auto cond = MakeEq(ko_var, MakeIndex(0, sp), sp); - auto if_stmt = std::make_shared(cond, then_body, std::optional(else_body), + auto if_stmt = std::make_shared(init_cond, then_body, std::optional(else_body), std::vector{c_phi}, sp); auto outer_yield = std::make_shared(std::vector{c_phi}, sp); return SeqStmts::Flatten(std::vector{sa, sb, if_stmt, outer_yield}, sp); diff --git a/src/ir/transforms/memory_reuse_pass.cpp b/src/ir/transforms/memory_reuse_pass.cpp index 643afe9bd3..b40e573c20 100644 --- a/src/ir/transforms/memory_reuse_pass.cpp +++ b/src/ir/transforms/memory_reuse_pass.cpp @@ -283,6 +283,17 @@ bool IsA5Target() { bool IsA5Prelu(const CallPtr& call) { return IsOp(call, "tile.prelu") && IsA5Target(); } +/// Human-readable identity of a MemRef for diagnostics: the allocation's name plus +/// its byte extent, e.g. `mem_acc_3[131072 B]`. Two MemRefs print differently exactly +/// when they name different allocations, which is what a mismatch report needs to show. +std::string MemRefLabel(const MemRefPtr& memref) { + if (!memref) return ""; + std::ostringstream os; + os << (memref->base_ ? memref->base_->name_hint_ : std::string("")) << "[" << memref->size_ + << " B]"; + return os.str(); +} + /// Plans top-down retypes. Produces (old Var -> new Type) map. class TopDownRetargeter { public: @@ -295,36 +306,6 @@ class TopDownRetargeter { return std::move(rewrites_); } - /// Coalesce peeled loop-carried accumulator if-phis. - /// - /// LowerPipelineLoops peels a stage=2 K-loop into an epilogue IfStmt whose live - /// branch is an in-place accumulator (matmul_acc, output aliasing input on the - /// accumulator buffer) and whose dead `if k==0` branch is a fresh matmul seed on a - /// *different* buffer. Left alone, YieldFixupMutator reconciles the two by copying - /// the accumulator onto the seed's buffer via an Acc->Acc tile.move — both a second - /// co-live L0C buffer (overflow) and an op TMovOp::verify rejects on every target - /// (there is no legal Acc->Acc tmov pair). We instead retarget the seed producer - /// onto the accumulator buffer so both branches share it and no move is emitted, - /// matching mad_acc's shared-%dst in-place semantics. - /// - /// The seed retype bypasses the *global* dead-at-assign check (the accumulator - /// buffer is legitimately live at the post-if phi consumer, which the global - /// check would treat as a conflict), but only after `TryCoalesceAccIfPhi` - /// verifies the two preconditions branch exclusivity actually needs: (a) the - /// seed producer is lexically inside the branch, and (b) a branch-scoped - /// liveness scan (`IsTargetDeadAtAssign(..., stop_at=if)`) finds no same-branch - /// tail read of the accumulator buffer. When either fails, that phi is left to - /// YieldFixup instead of being coalesced. Returns the rewrite map (apply via - /// RetypeApplier). A needed-but-declined retarget (after the preconditions - /// hold) is a hard error — no legal Acc->Acc move exists to fall back to. - std::map CoalesceAccumulatorIfPhis(const StmtPtr& func_body) { - DefMapVisitor def_v; - def_v.Run(func_body); - defs_ = std::move(def_v.defs); - VisitIfPhisForAccumulator(func_body); - return std::move(rewrites_); - } - private: std::map defs_; std::map rewrites_; @@ -371,116 +352,6 @@ class TopDownRetargeter { } } - // Walk the IR; coalesce every accumulator if-phi we encounter. - void VisitIfPhisForAccumulator(const StmtPtr& stmt) { - if (!stmt) return; - if (auto seq = As(stmt)) { - for (const auto& s : seq->stmts_) VisitIfPhisForAccumulator(s); - } else if (auto for_stmt = As(stmt)) { - VisitIfPhisForAccumulator(for_stmt->body_); - } else if (auto if_stmt = As(stmt)) { - TryCoalesceAccIfPhi(if_stmt); - VisitIfPhisForAccumulator(if_stmt->then_body_); - if (if_stmt->else_body_.has_value()) VisitIfPhisForAccumulator(if_stmt->else_body_.value()); - } else if (auto scope = As(stmt)) { - VisitIfPhisForAccumulator(scope->body_); - } - } - - // True when `var` is produced by an in-place accumulator op: a Call whose op - // reuses input `k` (matmul_acc) and whose output MemRef aliases input `k`'s — - // i.e. mad_acc's shared %dst. This branch's buffer is the one we keep; the - // other branch's producer is the seed we retarget onto it. - bool IsInplaceAccumulatorProducer(const VarPtr& var) { - auto it = defs_.find(var); - if (it == defs_.end() || it->second.kind != VarDef::kAssign) return false; - auto assign = As(it->second.assign_stmt); - if (!assign) return false; - auto call = As(assign->value_); - if (!call || !call->op_) return false; - const auto& reg = OpRegistry::GetInstance(); - if (!reg.IsRegistered(call->op_->name_)) return false; - auto reuse_idx = reg.GetEntry(call->op_->name_).GetOutputReusesInputArg(); - if (!reuse_idx.has_value() || *reuse_idx >= call->args_.size()) return false; - auto in_var = AsVarLike(call->args_[*reuse_idx]); - if (!in_var) return false; - auto out_tile = GetTileTypeWithMemRef(var->GetType()); - auto in_tile = GetTileTypeWithMemRef(in_var->GetType()); - if (!out_tile || !in_tile) return false; - return MemRef::SameAllocation(GetDefinedMemRef(out_tile), GetDefinedMemRef(in_tile)); - } - - // For an IfStmt whose branches yield an in-place accumulator on one side and a - // fresh seed on the other (a different L0C buffer), retarget the seed onto the - // accumulator buffer. Scoped to Acc — the ISA case with no legal reconciling - // move. A declined retarget is a hard error (see CoalesceAccumulatorIfPhis). - void TryCoalesceAccIfPhi(const IfStmtPtr& if_stmt) { - if (!if_stmt->else_body_.has_value() || if_stmt->return_vars_.empty()) return; - auto then_yield = FindYieldStmt(if_stmt->then_body_); - auto else_yield = FindYieldStmt(if_stmt->else_body_.value()); - if (!then_yield || !else_yield) return; - - for (size_t i = 0; i < if_stmt->return_vars_.size(); ++i) { - if (i >= then_yield->value_.size() || i >= else_yield->value_.size()) continue; - auto then_var = AsVarLike(then_yield->value_[i]); - auto else_var = AsVarLike(else_yield->value_[i]); - if (!then_var || !else_var) continue; - - const bool then_acc = IsInplaceAccumulatorProducer(then_var); - const bool else_acc = IsInplaceAccumulatorProducer(else_var); - if (then_acc == else_acc) continue; // need exactly one in-place accumulator - - const VarPtr& acc_var = then_acc ? then_var : else_var; - const VarPtr& seed_var = then_acc ? else_var : then_var; - - auto acc_tile = GetTileTypeWithMemRef(acc_var->GetType()); - auto seed_tile = GetTileTypeWithMemRef(seed_var->GetType()); - if (!acc_tile || !seed_tile) continue; - if (acc_tile->GetMemorySpace() != MemorySpace::Acc) continue; // Acc-only (no legal move) - - auto acc_memref = GetDefinedMemRef(acc_tile); - if (MemRef::SameAllocation(acc_memref, GetDefinedMemRef(seed_tile))) continue; // already shared - - auto seed_def = defs_.find(seed_var); - if (seed_def == defs_.end() || seed_def->second.kind != VarDef::kAssign) continue; - // The seed must be a Call producer we can retype; a bare-Var / tuple rename - // cannot be retargeted. Leave the phi untouched; YieldFixup will reject - // the residual Acc mismatch because no legal copy exists. - auto seed_assign = As(seed_def->second.assign_stmt); - if (!seed_assign || !As(seed_assign->value_)) continue; - - // The `check_liveness=false` bypass below is only sound when branch - // exclusivity actually applies, which requires BOTH: - // (a) the seed producer is lexically *inside* this IfStmt's branch — a - // pre-if value yielded through the branch runs unconditionally and - // would clobber the accumulator the sibling in-place branch reads; and - // (b) the accumulator buffer is dead *within the branch* after the seed - // (exclusivity covers only cross-branch and post-if reads, not a - // same-branch tail read between the seed producer and the yield). - // When either fails, leave the phi untouched here. YieldFixup will fail - // loudly rather than emit an unsupported Acc->Acc move. - const auto& seed_anc = seed_def->second.ancestors; - const bool in_branch = std::any_of(seed_anc.begin(), seed_anc.end(), - [&](const StmtPtr& a) { return a.get() == if_stmt.get(); }); - if (!in_branch) continue; - if (!IsTargetDeadAtAssign(seed_def->second, acc_memref->base_.get(), /*stop_at=*/if_stmt.get())) { - continue; - } - - // Now safe: (a)+(b) plus exclusivity cover every read of acc_memref, so we - // bypass the global liveness (which would false-decline on the legitimate - // post-if phi consumer). A remaining decline is a genuine "cannot coalesce - // this Acc phi" — fail loud, since no legal Acc->Acc move exists. - const bool ok = RetargetAssign(seed_var, seed_def->second, acc_memref, acc_tile->GetMemorySpace(), - /*check_liveness=*/false); - INTERNAL_CHECK_SPAN(ok, seed_var->span_) - << "Internal error: cannot coalesce L0C accumulator across a peeled if-phi — seed producer '" - << seed_var->name_hint_ - << "' refused retarget onto the accumulator buffer, which would force an illegal " - "Acc->Acc tile.move."; - } - } - /// Current (possibly-rewritten) MemRef base of `var`. const Var* CurrentBase(const VarPtr& var) { auto it = rewrites_.find(var); @@ -521,15 +392,8 @@ class TopDownRetargeter { } /// Retype a Var defined by an AssignStmt. - /// - /// `check_liveness` gates the general dead-at-assign check (IsTargetDeadAtAssign). - /// It is true for the normal loop-carry retarget path. It is set false only by - /// CoalesceAccumulatorIfPhis: coalescing an IfStmt phi's two branch yields onto one - /// buffer is always safe (the phi is redefined by exactly one branch at runtime, so - /// the branches are mutually exclusive and the target's downstream liveness cannot be - /// violated by a branch-local producer). The op-legality checks below still apply. bool RetargetAssign(const VarPtr& var, const VarDef& def, const MemRefPtr& target, - std::optional target_memory, bool check_liveness = true) { + std::optional target_memory) { auto assign = As(def.assign_stmt); INTERNAL_CHECK_SPAN(assign, var->span_) << "Internal error: kAssign VarDef must carry an AssignStmt"; auto call = As(assign->value_); @@ -603,9 +467,8 @@ class TopDownRetargeter { if (CallReadsBase(*call, target->base_.get(), read_arg_count)) return false; } - // Unconstrained: check liveness, then plan retype. (Skipped for if-phi - // branch coalescing, where branch exclusivity is a stronger guarantee.) - if (check_liveness && !IsTargetDeadAtAssign(def, target->base_.get())) return false; + // Unconstrained: check liveness, then plan retype. + if (!IsTargetDeadAtAssign(def, target->base_.get())) return false; PlanRewrite(var, target, target_memory); return true; } @@ -765,15 +628,7 @@ class TopDownRetargeter { /// This check does NOT special-case IfStmt siblings: it never scans the other /// branch of an enclosing IfStmt. That is correct — branches are mutually /// exclusive — but it is a conservative side effect, not modelled exclusivity. - /// - /// `stop_at`, when non-null, bounds the walk to a single enclosing scope: the - /// walk halts (returns "dead") upon reaching that statement instead of - /// continuing into its parent body. `CoalesceAccumulatorIfPhis` passes the - /// enclosing `IfStmt` so the scan covers only the seed's *branch tail* (a - /// same-branch read between the seed producer and the yield) while ignoring - /// the mutually-exclusive sibling branch and the legitimate post-if phi - /// consumers — the reads it must *not* treat as conflicts. - bool IsTargetDeadAtAssign(const VarDef& def, const Var* target_base, const Stmt* stop_at = nullptr) { + bool IsTargetDeadAtAssign(const VarDef& def, const Var* target_base) { if (def.ancestors.empty()) return true; // `child_on_path` is the direct descendant of the current ancestor that @@ -796,10 +651,6 @@ class TopDownRetargeter { } } - // Branch-scoped boundary: stop at the caller-supplied enclosing statement - // (e.g. the accumulator if-phi) rather than walking into its parent body. - if (stop_at && anc.get() == stop_at) return true; - // Stop once we've scanned the body of the enclosing ForStmt: the // retyped value is consumed by that loop's yield, so anything outside // the loop cannot observe it. @@ -2749,11 +2600,16 @@ class YieldFixupMutator : public IRMutator { std::vector else_move_stmts; std::vector new_return_vars = if_stmt->return_vars_; + // Arms are read as Var-*like*: `As` never matches an `IterArg` + // (ir-kind-traits.md), so a phi that merely forwards an enclosing loop's carry + // used to be skipped entirely — leaving its return_var on the buffer Step 3's + // reuse had already retargeted away, and making the ForStmt half below report a + // divergent Acc carry that is not in fact divergent. for (size_t i = 0; i < new_return_vars.size(); ++i) { VarPtr then_var = - (then_yield && i < then_yield->value_.size()) ? As(then_yield->value_[i]) : nullptr; + (then_yield && i < then_yield->value_.size()) ? AsVarLike(then_yield->value_[i]) : nullptr; VarPtr else_var = - (else_yield && i < else_yield->value_.size()) ? As(else_yield->value_[i]) : nullptr; + (else_yield && i < else_yield->value_.size()) ? AsVarLike(else_yield->value_[i]) : nullptr; auto then_tile = then_var ? GetTileTypeWithMemRef(then_var->GetType()) : nullptr; auto else_tile = else_var ? GetTileTypeWithMemRef(else_var->GetType()) : nullptr; @@ -2821,11 +2677,20 @@ class YieldFixupMutator : public IRMutator { auto source_tile = GetTileTypeWithMemRef(source->GetType()); INTERNAL_CHECK_SPAN(source_tile, source->span_) << "Internal error: YieldFixup tile.move source must be a TileType with MemRef"; - INTERNAL_CHECK_SPAN( + // An Acc (L0C) accumulator that reaches a control-flow merge or a loop carry in a + // different buffer than the one the merge lands in cannot be reconciled: nothing + // reads L0C except the FIXPIPE drain, whose destinations are L1 and GM, so no + // Acc->Acc copy exists on any target. Report the two buffers rather than + // manufacturing a move the hardware cannot perform. + const auto& source_memref = GetDefinedMemRef(source_tile); + CHECK_SPAN( !(source_tile->GetMemorySpace() == MemorySpace::Acc && target_memory.value() == MemorySpace::Acc), source->span_) - << "Internal error: MemoryReuse cannot reconcile divergent L0C accumulator buffers with " - "tile.move; accumulator control-flow values must be coalesced before YieldFixup."; + << "Accumulator (L0C) tiles cannot be moved between buffers: there is no Acc-to-Acc copy " + "on this target. '" + << source->name_hint_ << "' is bound to " << MemRefLabel(source_memref) + << " but the value it flows into is bound to " << MemRefLabel(target_memref) + << "; both must name one buffer."; auto& op_reg = OpRegistry::GetInstance(); std::vector> kwargs = { {"target_memory", std::any(target_memory.value())}}; @@ -3110,18 +2975,10 @@ FunctionPtr TransformMaterializeSemanticAliases(const FunctionPtr& func) { // // PTOAS needs only the ForStmt YieldFixup half: addr-less codegen already // re-points a branch-local producer at the if-phi handle. DSA-RP emits - // explicit addresses, so it must first coalesce peeled accumulator if-phis, - // then materialize both IfStmt and ForStmt fixups, and finally repair bare-Var - // identity copies before lifetime analysis and placement. + // explicit addresses, so it must materialize both IfStmt and ForStmt fixups, + // then repair bare-Var identity copies before lifetime analysis and placement. const auto* ctx = PassContext::Current(); if (ctx != nullptr && ctx->GetMemoryPlanner() == MemoryPlanner::DsaRP) { - TopDownRetargeter acc_coalescer; - auto acc_rewrites = acc_coalescer.CoalesceAccumulatorIfPhis(new_body); - if (!acc_rewrites.empty()) { - RetypeApplier applier(std::move(acc_rewrites)); - new_body = applier.VisitStmt(new_body); - } - YieldFixupMutator yield_fixup(/*fixup_if_stmts=*/true); new_body = yield_fixup.VisitStmt(new_body); new_body = NormalizeIdentityCopyBuffersMutator().VisitStmt(new_body); @@ -3208,21 +3065,6 @@ FunctionPtr TransformMemoryReuse(const FunctionPtr& func) { new_body = align.VisitStmt(new_body); } - // Step 3.75: Coalesce peeled loop-carried accumulator if-phis so YieldFixupMutator - // does not reconcile them with an illegal Acc->Acc tile.move. See - // TopDownRetargeter::CoalesceAccumulatorIfPhis. Must run after all ForStmt-carry - // coalescing (Steps 0/3/3.5) so the accumulator branch is on its final buffer, and - // before YieldFixupMutator (Step 4) so it observes one buffer per phi. A no-op when - // no accumulator if-phi exists (e.g. non-pipelined kernels). - { - TopDownRetargeter acc_coalescer; - auto acc_rewrites = acc_coalescer.CoalesceAccumulatorIfPhis(new_body); - if (!acc_rewrites.empty()) { - RetypeApplier applier(std::move(acc_rewrites)); - new_body = applier.VisitStmt(new_body); - } - } - // Step 4: Fix ForStmt/IfStmt yield/return_var MemRef mismatches YieldFixupMutator yield_fixup; new_body = yield_fixup.VisitStmt(new_body); diff --git a/tests/ut/ir/transforms/test_auto_tile_matmul_acc_mn.py b/tests/ut/ir/transforms/test_auto_tile_matmul_acc_mn.py index 2991e2a0df..8b5e52a3c3 100644 --- a/tests/ut/ir/transforms/test_auto_tile_matmul_acc_mn.py +++ b/tests/ut/ir/transforms/test_auto_tile_matmul_acc_mn.py @@ -455,8 +455,12 @@ def test_canonical_split_k_boundary_codegen_uses_box_aligned_physical_width(): r"valid_col = %c16_index : !pto\.tile_buf loop for index in lhs_extracts) diff --git a/tests/ut/ir/transforms/test_memory_reuse.py b/tests/ut/ir/transforms/test_memory_reuse.py index 143c76ce7d..472a9f77cb 100644 --- a/tests/ut/ir/transforms/test_memory_reuse.py +++ b/tests/ut/ir/transforms/test_memory_reuse.py @@ -18,7 +18,7 @@ import pypto.language as pl import pytest -from pypto import DataType, InternalError, backend, ir, passes, testing +from pypto import DataType, backend, ir, passes, testing from pypto.backend import BackendType from pypto.ir.op import tile from pypto.ir.pass_manager import OptimizationStrategy, PassManager @@ -1987,13 +1987,17 @@ def main( def test_divergent_acc_phi_rejects_acc_to_acc_move(self): """YieldFixup must not manufacture an unsupported Acc-to-Acc copy. - The divergent phi cannot be safely coalesced because one seed is - produced before the branch. Reject it before codegen rather than emit a - type-correct ``tile.move`` that PTOAS cannot lower for distinct L0C - buffers. + Nothing reads L0C except the FIXPIPE drain, whose destinations are L1 and + GM, so there is no Acc-to-Acc copy on any target. A divergent accumulator + is reported with both buffer identities rather than reconciled with a + ``tile.move`` no hardware can perform. """ - with pytest.raises(InternalError, match="cannot reconcile divergent L0C accumulator buffers"): + with pytest.raises(ValueError, match="no Acc-to-Acc copy") as excinfo: _run_pipeline(_divergent_acc_phi_program()) + # The report must name both sides so the author can see which two buffers + # the accumulator ended up split across. + message = str(excinfo.value) + assert "mem_acc_5" in message and "mem_acc_6" in message, message class TestControlFlow: @@ -2819,28 +2823,20 @@ def visit_for_stmt(self, stmt): # type: ignore[override] valid_shape = tile_type.get_effective_tile_view().valid_shape assert [dim.value for dim in valid_shape if isinstance(dim, ir.ConstInt)] == [16, 16] - def test_pipelined_kloop_accumulator_coalesces_to_one_acc_buffer(self): + def test_pipelined_kloop_accumulator_lands_on_one_acc_buffer(self): """A stage-2 pipelined K-loop matmul (as AutoTileMatmulL0 emits) whose - L0C accumulator is large (176x176x4 = 121KB, fp32). After - LowerPipelineLoops peels it into the multi-if-block shape, MemoryReuse - must coalesce the whole accumulator chain (tile.create init + first-block - matmul seed + the per-block matmul_acc + the if phis + the loop yield) - onto ONE Acc allocation. - - Regression (fixed by TopDownRetargeter::CoalesceAccumulatorIfPhis): the - peeled epilogue if-phi has a live in-place ``matmul_acc`` branch (on the - accumulator buffer) and a dead ``if k==0`` fresh-``matmul`` seed branch on - a different buffer. YieldFixupMutator used to reconcile them by copying the - accumulator onto the seed buffer via an acc->acc ``tile.move`` -- a 2nd - co-live 121KB L0C buffer that overflows the 128KB L0C, and an Acc->Acc tmov - that ptoas rejects on every target. This reproduced the 512x512x192 bf16 - compile failure. The fix retargets the seed onto the accumulator buffer so - both branches share it and no move is emitted (mad_acc's shared-%dst - semantics). Runs the real ``lower_pipeline_loops`` so the peeled SSA shape - matches production (hand-authored post-peel IR coalesces fine, so the real - pass is required to trigger the gap). Runs with BASIC verification: the - coalescing makes the peeled IR round-trip-clean, so the check is on legal - IR, not just buffer count. + L0C accumulator is large (176x176x4 = 121KB, fp32) must land on ONE Acc + allocation after LowerPipelineLoops peels it into the multi-block shape. + + The first K step selects overwrite-vs-accumulate through ``init_cond`` + rather than an ``if k == 0`` phi over a fresh ``tile.matmul`` and an + in-place ``tile.matmul_acc``. A predicated accumulate is in place on both + paths (``set_output_reuses_input(0)``), so peeling replicates a chain that + never leaves the accumulator buffer: no phi, nothing to reconcile, and no + opportunity for a second co-live 121KB L0C buffer that would overflow the + 128KB L0C. Runs the real ``lower_pipeline_loops`` so the peeled SSA shape + matches production, with BASIC verification so the check is on legal IR + and not just a buffer count. """ @pl.program @@ -2868,19 +2864,16 @@ def kernel( sb: pl.Tile[[64, 176], pl.BF16, pl.Mem.Right] = pl.tile.extract( rhs_mat, ko, 0, shape=[64, 176], target_memory=pl.Mem.Right ) - if ko == 0: - c_first: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.tile.matmul(sa, sb) - c_phi: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.yield_(c_first) - else: - c_acc: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.tile.matmul_acc(c_iter, sa, sb) - c_phi: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.yield_(c_acc) - c: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.yield_(c_phi) + c_acc: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.tile.matmul_acc( + c_iter, sa, sb, init_cond=(ko == 0) + ) + c: pl.Tile[[176, 176], pl.FP32, pl.Mem.Acc] = pl.yield_(c_acc) result: pl.Tensor[[176, 176], pl.FP32] = pl.store(c, [0, 0], out) return result def assert_coalesced(after: ir.Program, planner: str) -> None: - # The whole accumulator chain must coalesce onto ONE Acc allocation. A - # phantom acc->acc tile.move (failed coalescing) leaves a 2nd Acc base. + # The whole accumulator chain must stay on ONE Acc allocation; a second + # Acc base would mean the predicated accumulate stopped being in place. acc_bases = {b for b in _collect_tile_memref_bases(after).values() if "acc" in b} assert len(acc_bases) == 1, ( f"{planner}: expected ONE Acc allocation (accumulator coalesced), " @@ -2893,8 +2886,8 @@ def assert_coalesced(after: ir.Program, planner: str) -> None: f"{ir.python_print(after)}" ) - # BASIC verification: the coalescing fix makes the peeled IR round-trip - # clean, so this exercises the legality check, not just the buffer count. + # BASIC verification: the peeled IR must be round-trip clean, so this + # exercises the legality check, not just the buffer count. with passes.PassContext([], passes.VerificationLevel.BASIC): legacy_after = passes.memory_reuse()( passes.materialize_semantic_aliases()( @@ -2904,8 +2897,8 @@ def assert_coalesced(after: ir.Program, planner: str) -> None: assert_coalesced(legacy_after, "PYPTO") # DSA-RP skips MemoryReuse, so MaterializeSemanticAliases itself must run - # the same accumulator coalescing -> yield fixup -> identity-copy - # normalization sequence before lifetime analysis. + # the same yield fixup -> identity-copy normalization sequence before + # lifetime analysis. with passes.PassContext( [], passes.VerificationLevel.BASIC, @@ -2916,16 +2909,16 @@ def assert_coalesced(after: ir.Program, planner: str) -> None: ) assert_coalesced(dsa_after, "DSA_RP") - def test_accumulator_if_phi_seed_retargets_to_accumulator_buffer(self): - """Structural before/after for CoalesceAccumulatorIfPhis on a minimal - accumulator if-phi (no loop/peel needed to reproduce). + def test_accumulator_if_phi_divergence_is_reported(self): + """A minimal accumulator if-phi whose two arms land on different L0C + buffers is reported, not silently retargeted onto one of them. ``then`` is a fresh ``matmul`` seed on its own Acc buffer; ``else`` is an in-place ``matmul_acc`` on the accumulator buffer (aliasing ``prev``). - Pre-fix, MemoryReuse reconciled them with a phantom Acc->Acc ``tile.move`` - onto a 2nd Acc buffer. The fix retargets the seed onto the accumulator - buffer, so both branches and the phi share ONE Acc allocation and no move - is emitted. Pinned structurally. + Merging them means writing the seed into ``prev``'s buffer, which is only + sound if nothing outside the ``if`` still needs ``prev`` — a property the + compiler cannot establish from the branch alone. The author decides, so the + report names both buffers instead. """ @pl.program @@ -2954,51 +2947,10 @@ def main( result: pl.Tensor[[16, 64], pl.FP32] = pl.store(phi, [0, 0], out) return result - # Both branches' matmul/matmul_acc AND the phi land on mem_acc_5; the seed's - # own buffer (mem_acc_6) is retargeted away and its alloc dropped; no tile.move. - @pl.program - class Expected: - @pl.function - def main( - self, - lhs: pl.Tensor[[16, 64], pl.BF16, pl.MemRef("mem_ddr_0", 0, 2048)], - rhs: pl.Tensor[[64, 64], pl.BF16, pl.MemRef("mem_ddr_1", 0, 8192)], - cond: pl.Scalar[pl.INDEX], - out: pl.Out[pl.Tensor[[16, 64], pl.FP32, pl.MemRef("mem_ddr_2", 0, 4096)]], - ) -> pl.Tensor[[16, 64], pl.FP32]: - mem_mat_3: pl.Ptr = pl.tile.alloc(pl.Mem.Mat, 2048) - mem_mat_4: pl.Ptr = pl.tile.alloc(pl.Mem.Mat, 8192) - mem_acc_5: pl.Ptr = pl.tile.alloc(pl.Mem.Acc, 4096) - sa: pl.Tile[[16, 64], pl.BF16, pl.MemRef(mem_mat_3, 0, 2048), pl.Mem.Mat] = pl.tile.load( - lhs, [0, 0], [16, 64], [16, 64], target_memory=pl.Mem.Mat - ) - sb: pl.Tile[[64, 64], pl.BF16, pl.MemRef(mem_mat_4, 0, 8192), pl.Mem.Mat] = pl.tile.load( - rhs, [0, 0], [64, 64], [64, 64], target_memory=pl.Mem.Mat - ) - prev: pl.Tile[[16, 64], pl.FP32, pl.MemRef(mem_acc_5, 0, 4096), pl.Mem.Acc] = pl.tile.matmul( - sa, sb - ) - if cond < 1: - seed: pl.Tile[[16, 64], pl.FP32, pl.MemRef(mem_acc_5, 0, 4096), pl.Mem.Acc] = ( - pl.tile.matmul(sa, sb) - ) - phi: pl.Tile[[16, 64], pl.FP32, pl.MemRef(mem_acc_5, 0, 4096), pl.Mem.Acc] = pl.yield_( - seed - ) - else: - acc: pl.Tile[[16, 64], pl.FP32, pl.MemRef(mem_acc_5, 0, 4096), pl.Mem.Acc] = ( - pl.tile.matmul_acc(prev, sa, sb) - ) - phi: pl.Tile[[16, 64], pl.FP32, pl.MemRef(mem_acc_5, 0, 4096), pl.Mem.Acc] = pl.yield_( - acc - ) - result: pl.Tensor[[16, 64], pl.FP32, pl.MemRef("mem_ddr_2", 0, 4096)] = pl.tile.store( - phi, [0, 0], out - ) - return result - - After = _run_pipeline(Before) - ir.assert_structural_equal(After, Expected) + with pytest.raises(ValueError, match="no Acc-to-Acc copy") as excinfo: + _run_pipeline(Before) + message = str(excinfo.value) + assert "mem_acc_5" in message and "mem_acc_6" in message, message def test_gating_vec_inplace_if_phi_is_not_acc_coalesced(self): """Gating: the accumulator coalescer is Acc-scoped. A structurally @@ -3051,45 +3003,15 @@ def main( f"expected the pre-existing Vec->Vec move (coalescer must skip non-Acc):\n{ir.python_print(After)}" ) - def test_pre_if_acc_seed_not_coalesced_onto_accumulator(self): - """Branch-locality guard: the accumulator coalescer must only retarget a - seed that is defined *inside* the non-accumulator branch. - - Here ``then`` yields ``pre`` — a *pre-if* Acc value (computed before the - ``if``, so it runs unconditionally) — and ``else`` accumulates in place - into ``prev``. Coalescing would retarget ``pre`` onto ``prev``'s buffer, - writing it before the ``if`` and clobbering the accumulator that the else - branch reads. Branch exclusivity does NOT hold for a pre-if producer, so - the coalescer must skip this phi (leaving it to YieldFixup) and keep - ``pre`` and ``prev`` on distinct buffers. - """ - # Branch-locality correctly prevents unsafe coalescing. Because Acc->Acc - # tile.move is unsupported, YieldFixup must then fail loudly instead of - # emitting invalid IR for this unlowerable control-flow shape. - with pytest.raises(InternalError, match="cannot reconcile divergent L0C accumulator buffers"): - _run_pipeline(_divergent_acc_phi_program()) - - def test_seed_branch_write_only_clobber_blocks_acc_coalesce(self): - """Safety gate for the accumulator-if-phi coalescer's branch-tail - liveness scan: a *write* to the accumulator buffer after the seed (with - no read) must block coalescing, not just a read. - - The seed branch computes the fresh ``seed`` (its own Acc buffer) and then - a *write-only* op ``_clob`` that lands on the accumulator buffer - (``mem_acc_5``) — a fresh matmul reading only Mat operands, so it never - *reads* ``mem_acc_5``. If the coalescer retargeted ``seed`` onto - ``mem_acc_5`` (as the read-only scan would allow), ``_clob`` — sequenced - after ``seed`` — would overwrite the yielded value before the phi is read. - The scan must therefore reject a later write of the target base, not only - a later read (``SubtreeReadsBase`` alone misses the write-only clobber - because the read collector skips the LHS definition of each stmt). + def test_divergent_acc_phi_reported_on_fully_lowered_input(self): + """The Acc divergence report also fires on already-allocated input. Authored in fully-lowered form (explicit allocs + MemRefs) and run through - ``memory_reuse`` alone: the clobbering alias is a specific buffer layout - that only surfaces after allocation, so it cannot be expressed through the - high-level ``init_mem_ref`` path (which hands every tile a distinct base). - The required safety decline leaves divergent Acc buffers. Because no - legal Acc->Acc move exists, YieldFixup must reject this unlowerable shape. + ``memory_reuse`` alone, so the phi's two arms are pinned to distinct L0C + bases by the input itself rather than by an earlier pass — the shape a + hand-managed split-K reaches when its branch arms disagree. ``_clob`` is a + write-only op on the accumulator buffer, kept because it makes the arms' + buffers unmergeable by any means, not merely inconvenient to merge. """ @pl.program @@ -3143,7 +3065,7 @@ def kernel( ) return result - with pytest.raises(InternalError, match="cannot reconcile divergent L0C accumulator buffers"): + with pytest.raises(ValueError, match="no Acc-to-Acc copy"): passes.memory_reuse()(Before) def test_retargeter_declines_when_target_still_live(self): @@ -3284,8 +3206,16 @@ def main( pl.yield_(tile_c) ) else: + # `acc_0` is the loop's IterArg, on mem_vec_2. The phi lives on + # mem_vec_3, and PYPTO-planner codegen emits no copy of its own + # for a branch arm (it relies on this pass having aliased them), + # so the else path must write mem_vec_3 here or the phi buffer is + # never written on that path. + acc_0_mv: pl.Tile[[64, 64], pl.FP32, pl.MemRef(mem_vec_3, 0, 16384), pl.Mem.Vec] = ( + pl.tile.move(acc_0, target_memory=pl.Mem.Vec) + ) if_result: pl.Tile[[64, 64], pl.FP32, pl.MemRef(mem_vec_3, 0, 16384), pl.Mem.Vec] = ( - pl.yield_(acc_0) + pl.yield_(acc_0_mv) ) _use: pl.Tensor[[64, 64], pl.FP32, pl.MemRef("mem_ddr_1", 0, 16384)] = pl.tile.store( acc_0, [0, 0], output