Merge pull request #8 from liuzidi/feature-vmi-shengtao

[liuzidi]

VMI VF Fusion — fusion_region 技术方案

面向 fa softmax 链(ColMax → tmuls → ColExpand-sub → tmuls → texp → ColSum → tcvt)的 US-005 mem2reg。
本文档只讲当前方案与技术点。

1. 整体思路

VMI tilelib 的 compute TileOp 展开后是单层 scf.for(内层 VL 循环)。把每个 compute TileOp 的 inline body 包进一个 pto.fusion_region 容器,然后在 unified VMI 层( lowering 之前)做两件事:

1. LoopFusion:把同 header、且数据依赖允许的相邻 scf.for 融成一个。
2. LoadStoreElision:把融合后同 iteration 内的 UB( tile buffer)vstore→vload 往返 forward 成纯 SSA 寄存器传递,消掉两端。

做完再复用老路 PTOFlattenFusionRegion 拆容器,继续 lower 到 VPTO。

                        ┌─────────────────────────────────────────────┐
                        │  inline: 每个 compute TileOp body            │
                        │  包进 pto.fusion_region + 空 pto.yield        │
                        └─────────────────────────────────────────────┘
                                          │
                                          ▼
                ┌─────────────────────────────────────────────────┐
                │  PTOPlanVmiFusionRegion                          │
                │  按 F1/F2/F3 + UB-disjoint 合并相邻 region 成组  │
                │  (splice body + 合并 yield)                     │
                └─────────────────────────────────────────────────┘
                                          │
                                          ▼
                ┌─────────────────────────────────────────────────┐
                │  PTOVmiLoopFusion                                │
                │  同 header scf.for,按依赖判 hoist/sink          │
                │  能融则融,卡住(reduce-final 依赖)则断开        │
                │  prelude/init-producer hoist 出 loop            │
                └─────────────────────────────────────────────────┘
                                          │
                                          ▼
                ┌─────────────────────────────────────────────────┐
                │  Canonicalizer + CSE                             │
                │  合并 mask(Pure trait 无 MEI,CSE 只合 mask)    │
                └─────────────────────────────────────────────────┘
                                          │
                                          ▼
                ┌─────────────────────────────────────────────────┐
                │  PTOVmiLoadStoreElision                          │
                │  扫顶层段 + fused for body                      │
                │  vstore→vload 配对 forward,消两端              │
                │  逃逸 UB(region 外 mte 读)的 store 保留        │
                └─────────────────────────────────────────────────┘
                                          │
                                          ▼
                ┌─────────────────────────────────────────────────┐
                │  PTOFlattenFusionRegion (复用老路)              │
                │  region 内联回父 block,删容器                  │
                └─────────────────────────────────────────────────┘
                                          │
                                          ▼
                   VMILowerUnifiedToLegacy → ... → VMIToVPTO

2. 容器:为什么用 fusion_region 而不是新建容器

复用老路 pto.fusion_region(SingleBlock + pto.yield)。它带 yield/outputs,语义和老路 a5 fusion lifecycle 一致;PTOFlattenFusionRegion 收尾 pass 也直接复用(只 move ops + 删 region,对内部 op 零假设,vmi/vpto 通用)。

fa 的 compute TileOp 是 DPS(in-place UB),几乎无值逃逸 region → region 是空 yield + 空 results,flatten 时 yield/result 那步空转。

3. Plan:圈定可融合范围(F1/F2/F3 + UB-disjoint)

PTOPlanVmiFusionRegion 把相邻 region 合成一个 group region:

• F1:可融合 VMI op 可共享一个 region。
• F2:ColReduce 结果要整个 region 跑完才完整,reduce loop 与 consumer elementwise loop 保持各自 for-row,但可共享 region(是否真融由 LoopFusion 判)。
• F3:mte_*/set_flag/wait_flag/mem_bar/pipe_barrier/vecscope/unknown-effect op 关闭当前 group。
• UB-disjoint:候选 region 的 UB 集合(compile-time 地址 + memref type)必须与 group 内已有 UB 相同或不相交;部分重叠拒绝。

合并 = 把候选 region body(除 yield)splice 进 group region 的 yield 之前,删候选 region。fa 7 个 region → 1 个 group region。

关键坑:FusionRegionOp 有 pto.yield terminator(老 FusionScopeOp 是 NoTerminator)。plan 的 boundary walk 要把 pto::YieldOp 加进跳过列表,否则 region 自己的 yield 被误判成 sync boundary,每个 region 都自成一组。

4. LoopFusion:能不能融 = 之间指令能不能 hoist/sink

判定准则(核心)

两个相邻 scf.for(同 header)能不能融成一个,取决于它们之间夹着的指令能不能合法地往上或往下提:

• 能 hoist(放到 fused for 之前):输入在 for 之前可得——SSA 输入不是任何 member 的 result,UB 读取不被任何 member 写(否则读的是循环产物)。
• 能 sink(放到 fused for 之后):输出不被任何 member 内部读——SSA result 不被 member body 用,UB 写不被任何 member 的 vload 读。
• 既不能 hoist 又不能 sink → 卡住 → 两个 for 不能融,断开成各自 run。

between-ops 按依赖连通分量分组

两个 for 之间的指令不是一条条独立判,而是按数据依赖(SSA def-use,或同 UB 的 store→load)划分连通分量。每个分量必须整体 hoist 或整体 sink(拆散会断内部数据流)。一个分量既不能 hoist 又不能 sink → 断开。

fa 的实例

for1(ColMax vmax)  ── 之间夹: vload→vmuls→vstore (scale ColMax) ──  for2(ColExpand-sub vsub)...
        │                          │                                        │
        │              vload 读 ColMax final UB(ColMax 写)→ 不能 hoist        │
        │              vstore 写 scaled ColMax,被 for2 读     → 不能 sink        │
        │              分量卡住 → 断开,for1 单独                              │
        ▼                                                                      ▼
   for1 单独               tmuls(scale ColMax) 原位            for2~6 融成 1 个

最终 fa 的 6 个 for → 2 个:

• for1 = ColMax(reduce carry)单独
• for2 = ColExpand-sub + tmuls + texp + ColSum + tcvt 融成一个(ColSum 的 vadd 是 reduce carry,输入是同轮 texp 结果,可融)

中间 ColSum 的 vstore(写 ColSum final)被 sink 到 for2 之后(tcvt 读别的 UB,不读 ColSum 的)。

UB 身份追溯

判 hoist/sink 要查"两个 op 是不是操作同一 UB"。UB 身份 = pto.castptr ← memref ← pto.pointer_cast(addr-const) 追到 (地址常量, memref type) 对。两个 castptr 链追溯到同一 addr 常量 + 同 type = 同 UB。

prelude / init-producer hoist

member for 之前的 loop-invariant op(pointer_cast/castptr/create_mask)和 init-arg 的定义 op(vbrc 等)不 clone 进 body,而是 moveBefore fused for——天然支配 body,且避免 body 内冗余物质化。后续 CSE 再合并跨 member 的同 UB 重复。

5. LoadStoreElision:消 UB 往返

扫描范围

elideInRegion 在两层跑线性扫描:

1. region 顶层段(for 之间的 straight-line,含 prologue/between/epilogue)——消 fused for 之外的 UB 往返,如 ColMax final vstore → tmuls vload。
2. 每个 fused for body——消同 iteration 内的 UB 往返。

顶层段里遇到 scf.for 自动 flush trackedStores(for body 可能读写 UB,之前的 store 可能 stale)。

配对 forward

线性扫描:遇到 vstore 记进 TrackedStore{base, offset, mask, value};遇到 vload 查匹配(base + offset + mask 三元组相等)→ forward(load.result.replaceAllUsesWith(store.value))+ erase 两端;vstore 覆盖同址则 erase 旧 store。

• canonical base:通过 pointer_cast→castptr 追溯 UB,两个 castptr 链到同 addr+type 判等。
• mask:vmi.vload 无 mask operand,从 consumer op 推断(consumer 都带同 mask 才约束;mask-free consumer 如 vcvt 不约束 mask——std::optional<Value>,空 = 无约束,任意 mask 的 store 都匹配)。
• 透明性:遇非 vload/vstore 的 op,若它读写 tracked UB 或有副作用(mte_*/sync/for 等)→ flush trackedStores;纯计算(vmuls/vexp 等)透传。

逃逸边界:哪些 store 不能消

一个 vstore 写的 UB 若被 region 外读(如 region 后的 mte_ub_gm 把 UB 写回 GM),则它是 region 的输出边界,不能消。fa 里 scaled ColMax 的 vstore 写的 32768 UB 要被 region 外 mte_ub_gm 读输出给 host → 保留;但 ColExpand-sub 读它的 vload 可 forward 成 scaled ColMax 的 vreg。

6. 关键技术点小结

技术点	做法
容器	复用老路 pto.fusion_region(带 yield),PTOFlattenFusionRegion 收尾
融合范围	Plan 的 F1/F2/F3 + UB-disjoint(compile-time addr+type)
能否融	between-ops 按依赖连通分量,整体判 hoist/sink,卡住就断
UB 身份	castptr←memref←pointer_cast(addr-const),(addr, type) 对
mask 配对	vload 无 mask,从 consumer 推断;mask-free consumer 不约束
reduce-final 依赖	tmuls 读 ColMax final + 写被后段读 → 卡住 → ColMax 单独 for
region escape	被 region 外 mte 读的 store 保留,只消其 load
prelude hoist	loop-invariant op moveBefore fused for,不 clone 进 body

7. fa 实测结果

• 7 region → 1 group region(plan)
• 6 scf.for → 2 个(ColMax 单独 + ColExpand~tcvt 融 1 个)(loop-fusion)
• vload/vstore 往返大幅消除(elision):ColMax→tmuls 的 UB 往返消成 SSA,scaled-ColMax store 保留(给 mte)、load forward 成 vreg
• 全链 EXIT=0,object 产出,全量 lit 1604 PASS

8. softmax 融合前后对比(伪代码)

fa softmax 链 ColMax → tmuls(scale) → ColExpand-sub → tmuls → texp → ColSum → tcvt,所有 for header 都是 c0→c128 step1。

融合前(Inline 后,6 个独立 for + 中间顶层 op)

// ── prologue: pointer_cast/mask/vbrc(-inf) ──

for1 = scf.for c0..c128 iter_args(acc = -inf) {     // ColMax (reduce)
    row = vload(input_x[i])
    acc = vmax(acc, row)
    yield acc
}
vstore max = for1.result  -> UB[32768]               // ColMax final 存 UB (顶层 epilogue)

// tmuls(scale ColMax) —— 顶层,不在任何 for 里
scaled = vload UB[32768]                              // 读 ColMax final
scaled = vmuls(scaled, scale)
vstore scaled -> UB[32768]                            // 覆盖 32768 = scaled ColMax

// ── prologue: pointer_cast/mask/vbrc(0)/broadcast vload ──

for2 = scf.for c0..c128 {                             // ColExpand-sub (elementwise)
    x = vload(input_x[i])
    sub = vsub(x, broadcast_scaled)                   // 读 32768 的 broadcast
    vstore sub -> input_x[i]
}
for3 = scf.for c0..c128 { tmuls }                     // tmuls(scale x)
for4 = scf.for c0..c128 { texp }                      // texp
for5 = scf.for c0..c128 iter_args(acc=0) {           // ColSum (reduce)
    e = vload(input_x[i])
    acc = vadd(acc, e)
    yield acc
}
vstore sum = for5.result -> UB[33024]                  // ColSum final (顶层 epilogue)

for6 = scf.for c0..c128 { tcvt }                      // tcvt f32->bf16

融合后(LoopFusion + Elision 后,2 个 for)

// ── hoisted prologue: pointer_cast/mask/vbrc ──  (loop-invariant,移出 loop)

for1 = scf.for c0..c128 iter_args(acc = -inf) {      // ColMax 仍单独
    row = vload(input_x[i])
    acc = vmax(acc, row)
    yield acc
}
// tmuls 原位(不能 hoist:读 ColMax final;不能 sink:被 for2 读 → 卡住,断开)
max = for1.result                                     // ColMax final 直接 SSA(UB 往返已消)
scaled = vmuls(max, scale)                            // 不再 vstore→vload 中转
vstore scaled -> UB[32768]                            // 保留:region 外 mte_ub_gm 要读输出

// ── hoisted prologue ──

for2 = scf.for c0..c128 iter_args(acc = 0) {         // for2~6 融成 1 个
    x  = vload(input_x[i])
    s  = vsub(x, scaled)          // ColExpand-sub, scaled 直接 SSA(UB 往返消)
    s  = vmuls(s, scale)          // tmuls,        同轮 SSA 传递
    e  = vexp(s)                  // texp
    acc = vadd(acc, e)            // ColSum carry
    bf = vcvt(e)                  // tcvt
    vstore bf -> x_exp[i]
    yield acc
}
vstore sum = for2.result -> UB[33024]                 // ColSum final sink 到 for2 后(顶层)

要点

• for1(ColMax)单独:tmuls(scale ColMax)读 ColMax final(要整个 for1 跑完才有)→ 不能 hoist;其 vstore 被 for2 读 → 不能 sink → 卡住 → 断开。reduce 与消费其 final 的 loop 不能融成同一 iteration。
• for2~6 融成一个:它们之间是逐行 SSA 依赖(行 i 算完给下一行 i),可同 iteration 执行;ColSum 的 vadd 是 reduce carry(输入同轮 texp 结果),跟着进 fused for。
• ColSum 的 vstore sink:写 33024,tcvt 读别的 UB → 不影响 tcvt → sink 到 for2 后。
• scaled ColMax 的 vstore 保留:32768 要被 region 外 mte_ub_gm 读输出给 host(scaled gmax 是 fa 的输出之一)→ region escape,不能消;但其 load 被 forward 成 SSA。

9. 文件清单

• lib/PTO/Transforms/PTOPlanVmiFusionRegion.cpp — F1/F2/F3 + UB-disjoint 合并
• lib/PTO/Transforms/PTOVmiLoopFusion.cpp — 单层融合,依赖连通分量判 hoist/sink
• lib/PTO/Transforms/PTOVmiLoadStoreElision.cpp — 顶层段 + for body UB 往返消除
• lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp — inline 时 wrap fusion_region
• tools/ptoas/ptoas.cpp appendVMISemanticPipeline — 挂载顺序
• ptodsl/ptodsl/vmi_tilelib.py + _tile_template_tracing.py — VMI candidate(emit_col_reduce/col_expand/tcvt)
• test/lit/vpto/vmi_fusion_region_softmax.pto — 回归测试
• test/samples/FlashAttention/fa_dn_softmax_128x64.pto — fa 样本