Mosaic 布局推断
本页中的所有地址、符号、op-name 字符串和错误字符串均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so（build-id89edbbe81c5b328a958fe628a9f2207d，buildlibtpu_lts_20260413_b_RC00）。其他版本会有所不同。
摘要
Mosaic 的布局阶段是在应用器（apply-vector-layout）之前运行的一个三遍求解器。本页覆盖这三遍：16 阶段流水线中的 InferMemRefLayout（阶段 4）、TilingPropagation（阶段 9）和 InferVectorLayout（阶段 10）。它们共同计算物理 memref 分块，以及应用 pass 会消费的逐值 in_layout/out_layout VectorLayout 属性。VectorLayout 值类型本身，包括 (sublane, lane) 分块对、偏移、ImplicitDim、tilesPerVreg/tileArrayShape，位于 Mosaic VectorLayout；本页只记录这些布局是如何被选择的。
这个模型类似带半格的数据流分析，但有一个决定一切的设计选择：InferVectorLayout 从不插入 relayout。 LLVM pass 在求解寄存器类别时可能会拆分一个值或插入拷贝，而这个 pass 会为每个值精确选择一个 VectorLayout，以操作数驱动的贪心方式进行，并用布局格上的 VectorLayout::join meet 来调和多个操作数需求。当调和确实不可能时（例如 transpose 喂给 matmul），它只会发出一个与消费者 in_layout 不匹配的生产者 out_layout，并让独立的 RelayoutInsertion pass（阶段 11）弥合差距。这就是为什么 apply pass 可以断言 in == out，并把问题归因于缺少 relayout-insertion 运行。
本页按流水线顺序组织为三个 pass：先是 InferMemRefLayout（它产生其他部分读取的 memref TiledLayoutAttr），然后是 TilingPropagation（它把该分块穿到 load/store/DMA 位置），再到 InferVectorLayout（op-TypeID 分发和逐 op 规则）；随后是它们共享的 join/generalizes 格，以及一个端到端的 matmul 示例。
对于重新实现，契约是：
- memref 分块公式。
getTilingFactor，即 sublane-tile 选择器：packing = 32/bw、(32/bw)*sublanes的 “large 2nd-minor” 乘数表、可整除性/二次幂回退，以及inferLayout在上一层应用的leading_tile_rowsarg-attr 覆盖（完全绕过getTilingFactor）。这是后续 load/store 向量规则要匹配的分块的生产者。 - 传播不动点。
propagateTiling在 26 项 op-name→ruleStringMap上的 worklist、看穿erase_layout的消费者规则、memref_slice的 tile-stride 折叠，以及延迟的EraseLayoutOp删除。 - 向量 op-TypeID 分发。
inferBlock的约 35 分支 TypeID switch（applyStringMap的生产者侧对应物）、bitwidth-cast 预分发，以及 elementwise/extension 回退。 - 逐 op 规则。 共享的
bitwidth → native-tiling骨架，以及达到重新实现粒度的 matmul / elementwise / load-store / broadcast / rotate 规则主体。 - 格。
VectorLayout::join（选择更具体兼容布局的 meet）和generalizes（≤ 关系），以及为什么 join 失败是 relayout-insertion 信号而不是错误。
| Pass 1（阶段 4） | mlir::tpu::InferMemRefLayoutPass::runOnOperation @ 0x132c1820（create: 0x132c0f00） |
| Pass 2（阶段 9） | mlir::tpu::TilingPropagationPass::runOnOperation @ 0x132e0dc0（create: 0x132e0900） |
| Pass 3（阶段 10） | mlir::tpu::InferVectorLayoutPass::runOnOperation @ 0x132c3600（create: 0x132c2c20） |
| 向量分发核心 | (anon)::VectorLayoutInferer::inferBlock @ 0x132c3dc0（反编译约 1400 行；约 35 分支 TypeID switch） |
| Memref 分块核心 | inferLayout @ 0x132bef00 → getTilingFactor @ 0x132bed80 |
| 传播不动点 | propagateTiling @ 0x132e10a0；规则表 rules() @ 0x132e15e0（26 项） |
| 格 | mlir::tpu::VectorLayout::join @ 0x14a957c0；generalizes（≤）；源码 layout.h:320 |
| 输入/输出 | 读取 kernel-arg/alloca memrefs + leading_tile_rows arg-attr；写入 memref TiledLayoutAttr 和逐 op 的 in_layout/out_layout ArrayAttrs |
| 关键设计规则 | 推断为每个值选择一个布局；它不插入 tpu.relayout，这由 RelayoutInsertion（阶段 11）完成 |
| 置信度 | HIGH（由符号/字符串锚定），除非某行或标注另有说明 |
这三个 Pass 位于何处
Mosaic 布局阶段由四个 pass 组成；本页覆盖其中三个，apply-vector-layout 是第四个。数据流严格为生产者→消费者：
stage 4 infer-memref-layout InferMemRefLayoutPass ── memref args/allocas get a TiledLayoutAttr
│ (sublane×lane physical tiling + packing tile; insert erase_layout / reinterpret_cast)
▼
stage 9 tiling-propagation TilingPropagationPass ── push the tiling through erase_layout / slice /
│ casts to every load/store/DMA; drop the EraseLayoutOps once consumed
▼
stage 10 infer-vector-layout InferVectorLayoutPass ── every vector op gets in_layout / out_layout
│ (reads the memref tiling at load/store sites; join-reconciles multi-operand demands)
▼
stage 11 relayout-insertion RelayoutInsertionPass ── where producer.out_layout != consumer.in_layout,
│ insert tpu.relayout (so the applier never reconciles) [not this page]
▼
stage 12 apply-vector-layout ApplyVectorLayoutPass ── materialize native vregs + shuffles [not this page]
```text
重新实现者必须内化的架构事实是：阶段 4/9/10 都是*分析*，它们添加注解，从不重排数据。它们唯一会修改的 IR 是 memref 类型（阶段 4 重新给 args/allocas 定型并插入 `erase_layout`/`reinterpret_cast`）、`erase_layout` 周围的操作数连线（阶段 9 重新穿线然后擦除它），以及 `in_layout`/`out_layout` 属性（阶段 10）。所有实际的 vreg 物化和 relayout shuffle 都发生在之后的阶段 11–12。
---
## Pass 1 — InferMemRefLayout（阶段 4）
### 目的
为每个内核参数和 `alloca` memref 分配一个 `mlir::tpu::TiledLayoutAttr`：物理 sublane×lane 分块加上 tile stride，包括 sub-32-bit 的 packing tile。它精确地在向量 pass 之前运行，这样阶段 10 的 load/store 向量规则就有一个具体的 memref 分块，可用来与其 vreg 布局匹配；它们不会发明分块，而是读取这个分块。
### 入口点
```text
InferMemRefLayoutPass::runOnOperation (0x132c1820) ── requires hardware_generation (struct +114)
└─ inferFunc (0x132c0560) ── per-func; single-block required
├─ inferMemref (0x132bfd60) [per arg] ── semaphore→contiguous; else inferLayout
│ └─ inferLayout (0x132bef00) ── the core tiling math
│ └─ getTilingFactor (0x132bed80) ── the sublane-tile chooser
└─ inferOp (0x132c01a0) [per body op] ── memref.alloca / tpu.alloca_semaphore results算法
inferFunc 将函数的 memref 参数和所有 alloca 结果重新定型为 tiled 类型，然后通过一个 erase_layout（或 reinterpret_cast）为每个值建立别名，使函数体其余部分在传播重新穿线之前仍看到 untiled 视图：
function inferFunc(func, gen, target_shape, flags): // 0x132c0560
require func.body.hasOneBlock() // "Functions should only have a single block"
core_type = TPUDialect::GetCoreTypeAttr(func)
is_arg_in_smem = (core_type == 0x100000001) // default MemorySpace for inferMemref
for arg in func.arguments where isa<MemRefType>(arg):
lead_rows = consume_arg_attr(arg, "leading_tile_rows") // i32 override; removed after read
tiled = inferMemref(arg.type, gen, target_shape, flags,
is_arg=1, lead_rows, default_space=is_arg_in_smem)
set_arg_type(arg, tiled)
if canReinterpretToUntiledMemref(tiled): // 2-D, 32-bit, contiguous, ...
alias = insert tpu.reinterpret_cast(arg) // fold tile strides into linear strides
else:
alias = insert tpu.erase_layout(arg) // typed-but-untiled body alias
replaceAllUsesExcept(arg, alias, /*except=*/alias)
rebuild FunctionType with the new arg types
for op in func.body: inferOp(op, gen, target_shape, flags, default_space) // 0x132c01a0
```text
`inferOp` 处理 memref **结果**：对于 `memref.alloca` 和 `tpu.alloca_semaphore`，它会在结果 memref 上调用 `inferMemref` 并用一个 `EraseLayoutOp` 包裹它，然后递归进入嵌套 region。所有其他 op 原样返回 success。
`inferMemref` 解析 memory space，然后委托分块决策：
```c
function inferMemref(memref, gen, target_shape, flags, is_arg, lead_rows, space): // 0x132bfd60
if memref has a non-tpu MemorySpace attr: return memref // already placed; leave it
if elem is semaphore / dma-semaphore:
return TiledLayoutAttr::getContiguous(...) in space=4 // semaphore_mem, untiled
layout = inferLayout(memref, ctx, target_shape, flags, is_arg, lead_rows) // 0x132bef00
checkTiles(layout) // 0x132bfac0 — validates the tiling
return MemRefType(elem, shape, layout, resolved_MemorySpaceAttr)inferLayout 是核心分块数学。两种模式分别是一维（沿 minor 轴一个 tile）和二维及以上（由 getTilingFactor 选择 sublane tile，lane tile = target_lane），并为 sub-32-bit 元素类型追加 packing tile：
function inferLayout(memref, ctx, target_shape, flags, is_arg, lead_rows): // 0x132bef00
require target_shape.size() == 1 or == 2 // infer_memref_layout.cc:112
bw = bitwidth(elem)
if layout already isa<TiledLayoutAttr>:
if lead_rows requested and mismatches existing sublane tiling:
error "Trying to infer memref layout with sublane tiling <X>, but the memref"
" already has sublane tiling <Y>"
return existing
if layout isa<AffineMapAttr>: require identity else "Non-identity affine layout"
elif layout isa<StridedLayoutAttr>: ok
else: error "Unrecognized layout annotation"
require isIntOrFloat(elem) // "Invalid element type for memref"
if shape empty: return TiledLayoutAttr::get(ctx, 0,0,0,0) // scalar
lane = target_shape[-1] // e.g. 128
if rank == 1:
tile = ((32/bw) << (gen < 4)) * lane // {2,...} packing descriptor 0x200000000
if bw < 32: append packing tile (32/bw, 1)
return single-tile layout along the minor axis
else: // rank >= 2
if lead_rows != 0: // leading_tile_rows arg-attr overrides
sublane = lead_rows
else:
sublane = getTilingFactor(shape[-2], gen, target_shape[0], flags, bw, is_arg, /*1d=*/0)
L = TiledLayoutAttr::getContiguous(ctx, {sublane, lane}, shape) // computes tile strides
if bw < 32: append packing tile (32/bw, 1)
return L
// non-power-of-2 or bw > 32 → "Unsupported bitwidth: <bw>"
```text
`getTilingFactor` 选择 sublane tile。它的第三个参数是**目标 sublane 数**（`target_shape[0]`，通常为 8）：`inferLayout` 在这里传入 `*a3`，而不是 `leading_tile_rows`。packing factor `32/bw` 是下限；逐 bitwidth 的 “large 2nd-minor” 乘数（packing factor 乘以目标 sublane，受三个 `TpuTilingFlags` 字节、`is_arg` 和 `gen` 门控）可以提高它；可整除性/大小回退会让它对实际维度保持合法。`leading_tile_rows` arg-attr 在上一层 `inferLayout` 中处理，存在时会直接设置 sublane tile，此时不会为该 arg 调用 `getTilingFactor`：
```c
function getTilingFactor(dim, gen, sublanes, flags, bw, is_arg, is_1d): // 0x132bed80
require isPowerOf2_32(bw) // infer_memref_layout.cc:53
require 2 <= bw <= 32 // :54 / :55
packing = 32 / bw
base = max(packing, sublanes)
factor = base
if not is_1d: // "large 2nd-minor" multiplier
switch tzcnt(bw): // bw2→1, bw4→2, bw8→3, bw16→4
case bw==2 : candidate = 16 * sublanes // unconditional
case bw==4 : candidate = 8 * sublanes if flags[2]
case bw==8 : candidate = 4 * sublanes if flags[1]
case bw==16: candidate = 2 * sublanes if (flags[0] or (not is_arg and gen >= 6))
if candidate set: factor = candidate
if dim % factor != 0: factor = base // divisibility fallback
if dim < factor: // walk powers of two up
f = packing << (gen < 4)
while f < min(dim, base): f *= 2
factor = f
return factor每个分支中的乘数都是 (32/bw) * sublanes，也就是扩大 sublane tile，使该 tile 容纳一个完整 vreg 数量的打包子元素。对于常见的 sublanes == 8：bw=2 → 128，bw=4 → 64，bw=8 → 32，bw=16 → 16；对于 f32（bw=32），switch 不会进入，tile 保持为 base = max(1, 8) = 8。
GOTCHA — memref 分块依赖 generation 和 arg，而不仅依赖 bitwidth。
getTilingFactor会读取gen、三个TpuTilingFlags字节、is_arg标志和目标 sublane 数。只从32/bw推导 sublane tile 的重新实现，对 f32（factor 8）会给出正确答案，但对 bf16/int8/int4 会出错；在这些类型上，当 flags 和可整除性允许时，“large 2nd-minor” 分支会扩大 tile（sublanes=8的 bf16 →2*8 = 16）。门控随 bitwidth 不同：bw=2 分支无条件，bw=4 需要flags[2]，bw=8 需要flags[1]，bw=16 在设置了flags[0]或 memref 是gen >= 6上的非 arg 时采用宽 tile（即flags[0] or (not is_arg and gen >= 6)）。这些 flags 恢复为flags[0..2]；其人类可读名称是从用途推断的（MEDIUM）。QUIRK —
erase_layout与reinterpret_cast由一个未反编译的谓词决定。 当canReinterpretToUntiledMemref成立时，inferFunc插入tpu.reinterpret_cast（把 tile stride 折叠为线性 stride 的 untiled 视图），否则插入tpu.erase_layout（typed-but-untiled 别名）。调用点已恢复；其函数体（2-D/32-bit/contiguous 资格测试）未反编译。可把选择视为：可 reinterpret 的 contiguous 32-bit 2-D →reinterpret_cast；其他所有情况 →erase_layout（精确谓词 LOW）。
函数表
| 函数 | 地址 | 作用 |
|---|---|---|
InferMemRefLayoutPass::runOnOperation | 0x132c1820 | pass 入口；需要 hardware_generation（+114）；调用 inferFunc |
createInferMemRefLayoutPass | 0x132c0f00 | 工厂（gen、target span、TpuTilingFlags）；pass 结构 0x328 B |
inferFunc | 0x132c0560 | 逐函数 arg 重新定型 + 插入 erase_layout/reinterpret_cast |
inferOp | 0x132c01a0 | 为 alloca/alloca_semaphore 做 memref 结果推断 |
inferMemref | 0x132bfd60 | semaphore→contiguous；否则 inferLayout + checkTiles |
inferLayout | 0x132bef00 | 1-D/2-D/packing 分块数学 |
getTilingFactor | 0x132bed80 | sublane-tile 选择器（上面的公式） |
checkTiles | 0x132bfac0 | 验证已解析的分块 |
Pass 2 — TilingPropagation（阶段 9）
目的
InferMemRefLayout 把每个 tiled memref 放在 erase_layout/reinterpret_cast 别名后面，因此函数体仍引用一个 untiled 视图。TilingPropagation 是将物理分块向前推送到每个 load / store / DMA / memref-shape op 的不动点：把它们的 memref 操作数改写为直接引用 tiled memref（看穿 erase_layout），然后在每个 EraseLayoutOp 不再有剩余 use 后删除它。此 pass 之后，每个 memory op 都引用 tiled memref，这正是阶段 10 load/store 向量规则要读取的内容。
入口点
TilingPropagationPass::runOnOperation (0x132e0dc0) ── builds a PropagationContext from options
└─ propagateTiling(ctx, entry_block) (0x132e10a0) ── the worklist fixpoint
├─ rules() (0x132e15e0) ── 26-entry op-name → rule StringMap
├─ propagate_layout_to_consumer_rule (0x132e40a0) ── 20 consumer ops: see-through erase_layout
└─ tpu_memref_slice_rule (0x132e1b20) ── slice: re-thread + fold tile strides
```text
`PropagationContext` 携带 `{target_sublane, target_lane}`（结构 `+0x150`）、收集到的 `EraseLayoutOps` 向量（`+16`），以及 `sparse_core` `cl::opt` bool（结构 `+472`）。`createTilingPropagationPass({sublane,lane}, sparse_core)`（`0x132e0900`）构造一个嵌套在 `func.func` 上、大小为 `0x228` 字节的 pass。
### 算法
```c
function propagateTiling(ctx, block): // 0x132e10a0
for op in block, IN ORDER:
if isa<EraseLayoutOp>(op):
ctx.erase_list.push(op) // defer removal to the end
for region in op.regions: for b in region:
propagateTiling(ctx, b) // depth-first into nested regions
rule = rules().lookup(op.name) // xxh3_64 + StringMapImpl::FindKey
// miss → StringMap default-bucket entry
if rule(ctx, op) == failure: return failure
for erase_op in ctx.erase_list: // cleanup
if erase_op.use_empty() or all_uses_in_sc_tpu(erase_op): // "sc_t"/"up" literal test
erase_op.erase()
else:
return error "Failed to propagate the layout to all operations"
return success“all uses in sc_tpu” 谓词是字面字符串测试 (*str ^ 0x745F6373) | (*(str+4) ^ 0x7570)，也就是 op-name 前缀等于 "sc_t" 然后 "up"（sc_tpu）。SparseCore op 合法地保留 untiled 别名，因此它们不会阻止擦除。
中央规则会看穿 erase_layout：
function propagate_layout_to_consumer_rule(ctx, op): // 0x132e40a0
for operand in op.operands:
def = operand.getDefiningOp()
if isa<EraseLayoutOp>(def):
op.setOperand(operand_index, def.getOperand(0)) // use the TILED memref directly
return success // always succeeds
```text
硬规则是 `tpu_memref_slice_rule`。它必须重新给 slice 分块，并且当它折叠一个 4 级 tiled 32-bit memref 的 2nd-minor 维时，用发出的算术把 tile stride 折叠进 slice 的 base index：
```c
function tpu_memref_slice_rule(ctx, slice): // 0x132e1b20
src = slice.source
if not isa<EraseLayoutOp>(src.getDefiningOp()): return success // no-op
tiled_src = src.getDefiningOp().getOperand(0)
verifyOffsetAndSizeTileAlignment(slice)
if collapses 2nd-minor of a 4-level-tiled 32-bit memref:
// fold the tile strides into the physical base offset:
new_idx = (idx / tile_size) * tile_stride + (idx % tile_size)
// emitted via arith.constant + divui + remui + muli + addi (createOrFold)
require isGuaranteedDivisible(off, stride, 128)
elif contiguous & reinterpretable:
require isGuaranteedDivisible(off, 8, 128) // "Slice offset for 32-bit 1D
// memrefs must be a multiple of 8"
emit ReinterpretCastOp over an untiled TiledLayoutAttr
new_slice = rebuild MemRefSliceOp on tiled_src
result = wrap new_slice in a fresh EraseLayoutOp; ctx.erase_list.push(it)
return success另外五个 memref-shape 规则（squeeze/reshape/bitcast/reinterpret_cast/memref.cast，0x132e2a60..0x132e3ee0）类似地把 tiled memref 重新穿过 cast，并重新计算结果 tile stride；其逐规则 stride 数学经过抽样，但没有逐个反编译（HIGH）。
26 项规则表
rules()（0x132e15e0）是一个 op-name→rule StringMap，是 apply pass 的 49 项 rewrite 表在传播类型上的同级物。与其列出 26 行几乎相同的内容，不如说明其形状：6 个 memref-shape op 带有专门规则；另外 20 个 “consumer” op 共享 propagate_layout_to_consumer_rule。
| 分组 | Op-name | 规则 fn | 地址 |
|---|---|---|---|
| memref slice | tpu.memref_slice | tpu_memref_slice_rule | 0x132e1b20 |
| memref squeeze | tpu.memref_squeeze | tpu_memref_squeeze_rule | 0x132e2a60 |
| memref reshape | tpu.memref_reshape | tpu_memref_reshape_rule | 0x132e33a0 |
| memref bitcast | tpu.memref_bitcast | tpu_memref_bitcast_rule | 0x132e37a0 |
| reinterpret cast | tpu.reinterpret_cast | tpu_reinterpret_cast_rule | 0x132e3e60 |
| memref cast | memref.cast | memref_cast_rule | 0x132e3ee0 |
| consumers（×20） | tpu.{load,store,strided_load,strided_store,vector_load,vector_store,vector_load_idx,vector_store_idx,enqueue_dma,enqueue_indirect_dma,wait_dma2,wait_indirect_dma,log_buffer,fetch_and_add_sync,sem_signal,sem_wait,sem_read}, memref.{store,load,reinterpret_cast} | propagate_layout_to_consumer_rule | 0x132e40a0 |
四个 op-name 字符串 tpu.vector_store、tpu.enqueue_indirect_dma、tpu.wait_indirect_dma、tpu.fetch_and_add_sync 是从 0x869EF8B / 0x879F7AE / 0x879F798 / 0x8730A53 处的 .rodata 解析出的。
NOTE — 分发机制与 apply pass 共享。 查找过程是
xxh3_64bits+StringMapImpl::FindKey，miss 时落到 default-bucket fallback，与 apply pass 的applyLayoutOpStringMap完全相同。重新实现者可以为两个生产者侧 map 使用同一个StringMap实现。
函数表
| 函数 | 地址 | 作用 |
|---|---|---|
TilingPropagationPass::runOnOperation | 0x132e0dc0 | pass 入口；构造 PropagationContext |
createTilingPropagationPass | 0x132e0900 | 工厂（{sublane,lane}、sparse_core）；结构 0x228 B |
propagateTiling | 0x132e10a0 | worklist 不动点 + 延迟的 EraseLayoutOp 删除 |
rules() | 0x132e15e0 | 26 项 op-name→rule StringMap |
propagate_layout_to_consumer_rule | 0x132e40a0 | 为 20 个 consumer op 看穿 erase_layout |
tpu_memref_slice_rule | 0x132e1b20 | slice 重新分块 + tile-stride 折叠 |
tpu_memref_{squeeze,reshape,bitcast}_rule | 0x132e2a60/0x132e33a0/0x132e37a0 | 重新穿过 cast + 重新计算 stride |
tpu_reinterpret_cast_rule / memref_cast_rule | 0x132e3e60 / 0x132e3ee0 | 重新穿过 cast + 重新计算 stride |
Pass 3 — InferVectorLayout（阶段 10）
目的
为每个 op 标注 in_layout ArrayAttr（每个操作数一个 VectorLayout，标量为 kNoLayout）和 out_layout ArrayAttr（每个结果一个）。这正是 apply pass 读取的精确属性的生产者。它实现为一个 op-TypeID switch，是 apply pass 的 op-name StringMap 在生产者侧的对应物，但由于 inferer 是单个翻译单元，它使用直接的 TypeIDResolver<> 指针相等。
入口点
InferVectorLayoutPass::runOnOperation (0x132c3600) ── requires hardware_generation (pass +114)
│ constructs a VectorLayoutInferer on the stack:
│ +0 hardware_generation +24 target_sublane (8) +32 target_lane (128)
│ +TpuTilingFlags (+214/+164/+568), large-2nd-minor bool (+1056), reduce bool (+1138)
│ requires single-block func ── "Only one block functions supported"
└─ VectorLayoutInferer::inferBlock(block, fn::return_handler) (0x132c3dc0) ── per-op driver
```text
### 算法
`inferBlock` 是逐 op driver。对终结符之前的每个 op，它会运行 guard，收集操作数布局，判断标量或向量，运行 bitwidth-cast 预分发，然后进入 TypeID switch，最后走 fallback：
```c
function inferBlock(block, terminator_handler): // 0x132c3dc0
for op in block until terminator:
// 1. pre-attached guard
if (op has in_layout or out_layout) and not isa<AssumeLayoutOp>(op):
error "layout attributes already attached"
if isa<AssumeLayoutOp>(op):
require op has both layouts // "expect layout attributes in tpu::AssumeLayoutOp"
continue // pass through
// 2. collect operand layouts (skipped for broadcast / extract_strided_slice)
if not isa<vector::BroadcastOp, ExtractStridedSliceOp>(op):
operand_layouts = getLayoutFromOperands(op) // 0x132c59a0; kNoLayout for scalars
else:
set inferer+44 lane-replicate flag if an operand is replicated along an axis
// 3. scalar test — no VectorType among operands OR results
if no operand and no result is a VectorType:
setInLayout(op, kNoLayout per operand) // 373/374 CHECK path
if op.getNumResults() == 0: continue
// 4. bitwidth-cast pre-dispatch (extsi/extf/sitofp/uitofp/trunci/truncf/fptosi/fptoui/extui)
if isConversionOp(op):
bin = operand_bitwidth(op) // Float8EXMY counts as 8; i1 special
bout = result_bitwidth(op)
if bout > bin: return inferExt(op) // 0x132c5be0 widening
elif bout < bin: return inferTrunc(op) // 0x132c6600 narrowing
else: return inferElementwise(op) // extui i1 (bout == bin)
// 5. select / cmp guards
if isa<arith.select, arith.cmpi, arith.cmpf>(op):
return inferElementwise(op) // "Only one side of arith/cmp is a vector?"
// 6. op-TypeID switch (≈35 cases) → infer(<Op>) [table below]
switch TypeID(op): ...
// 7. fallback
else if hasElementwiseMappableTraits(op): inferElementwise(op)
else if extensions::canInferVectorLayout(op): // 0x13246280
extensions::inferVectorLayout(op) // 0x132462a0 out-of-tree ops
else: error "Not implemented: Unsupported operation: <op> in infer-vector-layout pass"
// 8. post-checks
re-assert 373/374 CHECKs; clear inferer+44 lane-replicate flag每个 op 会重新断言两个不变量（infer_vector_layout.cc:373/374）：每个有结果的 op 都带有 out_layout；每个有操作数的 op 都带有 in_layout。setInLayout（0x14b75c60）、setOutLayout（0x14b75e40）和 setLayout（0x14b75fa0/0x14b75fe0/0x14b76020）是写入 ArrayAttrs 的自由函数；kNoLayout 是一个模块全局哨兵 VectorLayout，标记标量/非向量操作数槽位。
QUIRK —
broadcast和extract_strided_slice会有意跳过操作数布局收集。 对这两个 op，inferBlock不会调用getLayoutFromOperands；相反，当某个操作数沿某个轴被 replicated/singleton 时，它会设置 inferer 的+44lane-replicate 标志。该标志会让后续getLayout读取强制为源采用 lane-replicated 布局。对这些 op 运行统一操作数收集路径的重新实现会读取错误的源布局。QUIRK — extension 布局由 packing 方向选择，而不是由 op 名称选择。
i1的arith.extui有bout == bin，会路由到inferElementwise；而arith.extsi/tpu.extf变宽（inferExt），arith.trunci/tpu.truncf变窄（inferTrunc）。分发器比较操作数和结果 bitwidth，并忽略 op 名称是 ext 还是 trunc，因此 sign-extension 和 fp-cast 布局遵循 sub-32-bit packing 的变化，而不是遵循拼写。
Op-TypeID 分发表
这个 switch 有约 35 个分支。作为 apply pass 的 49 个 rewrite 规则在生产者侧的对应物，它把每个 op TypeID 映射到一个 infer(<Op>) 方法。多个 op 共享一个方法（已注明），四个 op 内联处理：
| Op | infer 方法 | 地址 | 备注 |
|---|---|---|---|
arith.constant | infer(arith::ConstantOp) | 0x132c78c0 | |
cf.assert | inline | — | setInLayout = kNoLayout |
memref.load | infer(memref::LoadOp) | 0x132c7de0 | 不同于 tpu.load |
tpu.load | infer(LoadOp) | 0x132cb6c0 | |
tpu.store | infer(StoreOp) | 0x132cbb20 | |
tpu.strided_load / tpu.strided_store | infer(Strided{Load,Store}Op) | 0x132cbd80 / 0x132cc240 | |
tpu.matmul（+ push_rhs/acc_lhs/pop） | infer(Matmul*Op) | 0x132cc740 / 0x132cccc0 / 0x132cce00 / 0x132ccf40 | |
tpu.rotate / tpu.dynamic_rotate | infer({Rotate,DynamicRotate}Op) | 0x132ca600 / 0x132ca840 | |
tpu.concatenate | infer(ConcatenateOp) | 0x132cacc0 | |
tpu.erase_layout | inline | — | setLayout = kNoLayout（in+out） |
tpu.iota | infer(IotaOp) | 0x132cd0a0 | |
tpu.gather | infer(GatherOp) | 0x132cd380 | 不同于 dynamic |
tpu.dynamic_gather | infer(DynamicGatherOp) | 0x132cd3e0 | |
tpu.reduce_index | infer(ReduceIndexOp) | 0x132cd6e0 | |
tpu.bitcast | infer(BitcastOp) | 0x132cdb40 | |
tpu.trace | infer(TraceOp) | 0x132ce020 | |
tpu.prng_random_bits | infer(PRNGRandomBitsOp) | 0x132ce160 | |
tpu.region | infer(RegionOp) | 0x132ce300 | 携带 region |
scf.if / scf.for / scf.while | infer(scf::{If,For,While}Op) | 0x132c8140 / 0x132c8c40 / 0x132c9760 | 携带 region |
vector.broadcast | infer(vector::BroadcastOp) | 0x132ce520 | |
vector.extract | infer(vector::ExtractOp) | 0x132cece0 | |
vector.multi_reduction | infer(MultiDimReductionOp) | 0x132cf660 | |
vector.shape_cast / tpu.reshape | inferReshape | 0x132d09a0 | 共享 |
vector.extract_strided_slice | infer(ExtractStridedSliceOp) | 0x132d2d80 | 跳过操作数收集 |
tpu.vector_load / tpu.vector_store | infer(Vector{Load,Store}Op) | 0x132cf2c0 / 0x132d2340 | → inferLoadStoreVectorLayout |
tpu.transpose | infer(TransposeOp) | 0x132d26c0 |
只有两个向量内存 op，即 tpu.vector_load 和 tpu.vector_store，会经过共享的 inferLoadStoreVectorLayout（0x132d48c0）；标量 memref.load / tpu.load / tpu.store 以及 strided_* op 各有自己的 infer(...) 主体，并不会调用它。
函数表
| 函数 | 地址 | 作用 |
|---|---|---|
InferVectorLayoutPass::runOnOperation | 0x132c3600 | pass 入口；构造 VectorLayoutInferer |
createInferVectorLayoutPass | 0x132c2c20 | 工厂（gen、{sublane,lane}、TpuTilingFlags、bool） |
VectorLayoutInferer::inferBlock | 0x132c3dc0 | 逐 op TypeID 分发（反编译约 1400 行） |
getLayoutFromOperands | 0x132c59a0 | 收集每个操作数的生产者 out_layout |
getLayout | 0x132d3260 | 读取某个值在结果索引处的 out_layout |
inferExt / inferTrunc | 0x132c5be0 / 0x132c6600 | 变宽 / 变窄 cast 布局 |
inferElementwise | 0x132c70e0 | layout-preserving elementwise 规则（使用 join） |
inferLoadStoreVectorLayout | 0x132d48c0 | 读取 memref 分块，匹配 vreg 布局 |
verifyMemoryTiling | 0x132d3580 | 强制合法的 memory-op 分块 |
infer(MatmulOp) | 0x132cc740 | MXU-packed 操作数/acc/result 布局 |
infer(vector::BroadcastOp) | 0x132ce520 | replicated-axis 选择 |
infer(RotateOp) / infer(IotaOp) / infer(TransposeOp) / inferReshape | 0x132ca600 / 0x132cd0a0 / 0x132d26c0 / 0x132d09a0 | 逐 op 规则 |
其余约 25 个 infer(...) 主体 | 0x132c78c0..0x132d2d80 | 地址 + 签名已恢复；主体已抽样 |
setInLayout / setOutLayout / setLayout | 0x14b75c60 / 0x14b75e40 / 0x14b75fa0 | 写入 in_layout/out_layout ArrayAttrs |
extensions::{can,}inferVectorLayout | 0x13246280 / 0x132462a0 | out-of-tree op 回退 |
逐 Op 规则骨架
每个逐 op 规则都共享一个骨架。重新实现者只要内化这个骨架，就能重构约 25 个只抽样、未完全反编译主体的规则：
- Bitwidth。 计算每个向量操作数/结果的元素 bitwidth
bw。Float8EXMYType按 8 计；i1特判。 - Native sublane tiling =
target_sublane * 32 / bw，即每个 sublane 容纳32/bw个子元素的 packed tile：f32 为 8，bf16 为 16，int8/fp8 为 32，int4 为 64。 - Native lane tiling =
target_lane（128）。 - Offset = 新产生值的
{0,0}。只有当 op 在某个轴上 broadcast 或 reduce 时，该轴才会变为 replicated（offset 缺失）。 implicit_dim按 rank 决定：rank-2+ →NONE(0)；1-D →MINOR(1) 或SECOND_MINOR(2)，取决于该向量位于 lane 轴还是 sublane 轴；标量提升 →MINOR_AND_SECOND_MINOR(3)。
以下是四个达到重新实现粒度的规则和次要规则。
infer(MatmulOp) — 0x132cc740
matmul 规则把操作数强制为 MXU latch 期望的 packed native tiling，并把 accumulator/result 强制为 f32 (8,128) 分块。它 CHECK "Expected 32-bit acc in tpu::MatmulOp" 和 "Expected 32-bit result in tpu::MatmulOp"，然后发出三个 in_layout 和一个 out_layout：
function infer(MatmulOp op): // 0x132cc740
bw_lhs = bitwidth(op.lhs); bw_rhs = bitwidth(op.rhs)
require bitwidth(op.acc) == 32 // "Expected 32-bit acc in tpu::MatmulOp"
require bitwidth(op.res) == 32 // "Expected 32-bit result in tpu::MatmulOp"
in_layout[lhs] = { bw_lhs, {0,0}, (32*target_sublane/bw_lhs, target_lane), NONE }
in_layout[rhs] = { bw_rhs, {0,0}, (32*target_sublane/bw_rhs, target_lane), NONE }
in_layout[acc] = { 32, {0,0}, (target_sublane, target_lane), NONE }
out_layout = { 32, {0,0}, (target_sublane, target_lane), NONE }
// 32*sublane/bw computed as 32 * inferer[+24] / bw
```text
对于 bf16 操作数，这会得到 packed 的 `16,{0,0},(16,128)` 和 f32 结果 `32,{0,0},(8,128)`，正是 [apply pass](mosaic-vectorlayout.md) 在其 bf16 matmul worked example 中消费的属性。解释操作数*为什么*必须 packed 的 MXU tile-cost 推理位于 dot/conv cost 页面。
### inferElementwise — `0x132c70e0`
Layout-preserving：所有非 `i1` 的向量/标量操作数加上结果必须共享一个 bitwidth；向量操作数布局会通过 `VectorLayout::join` 折叠为一个公共布局，该布局同时成为每个向量操作数的 `in_layout` 和 `out_layout`：
```c
function inferElementwise(op): // 0x132c70e0
require op.getNumResults() == 1 // "only one result supported"
require op.getNumOperands() > 0 // "elementwise ops with no operands unsupported"
bw = -1
for v in operands + [result] where v is non-i1 vector/scalar:
if bw == -1: bw = bitwidth(v)
else require bitwidth(v) == bw // "Mismatched bitwidth in elementwise for non-i1 ..."
require isa<VectorType, scalar>(v) // "expected only vector and scalar operands"
L = none
for vop in vector operands:
l = getLayout(vop) // "missing vector layout"
L = (L is none) ? l : VectorLayout::join(L, l, shape)
if join failed: L = native_layout(bw) // re-derive a fresh native layout
if no vector operand but result is vector:
L = { bw, {0,0}, (32*target_sublane/bw, target_lane) } // synthesize native
set every vector operand's in_layout = L; out_layout = L; scalar operands → kNoLayoutjoin 失败在这里不是错误，它会为该操作数回退到重新推导的 fresh native layout，最终的布局不匹配稍后由 RelayoutInsertion 弥合。
inferLoadStoreVectorLayout — 0x132d48c0
load/store 规则不会发明分块，而是读取 memref 的 TiledLayoutAttr（由 InferMemRefLayout 设置、由 TilingPropagation 传播），并让 vreg 布局与它匹配：
function inferLoadStoreVectorLayout(memref, vec, indices, ...): // 0x132d48c0
require rank(memref) == rank(vec) // "memref and vector rank mismatch"
require rank(vec) > 0 // "rank 0 vectors unsupported"
bw = bitwidth(elem) // "Unsupported bitwidth"
if rank == 1:
require 1D tiling // "Expected 1D tiling in 1D loads/stores"
// else "Unsupported tiling for 1D load/store"
layout = { offset {0, base_offset mod (lane*packing)},
tiling (1, lane*packing), implicit MINOR(2) }
else: // rank >= 2
require 2D tiling // "Expected 2D tiling in 2D+ loads/stores"
if memref 2nd-minor dim <= one sublane-tile OR vec minor dim == 1:
layout.offset = {0,0} // sublane-replicated fast path
else:
layout.offset = { base_idx mod tile (sublane), base_idx mod tile (lane) }
if bw==32 and vec minor dim==1 and canReinterpretToUntiledMemref:
layout = sublane=1 broadcast layout
verifyMemoryTiling(...) // 0x132d3580 (CHECKs below)
if canReinterpretToUntiledMemref(memref):
fold tile strides into offset using leading_tile_rows
```text
`verifyMemoryTiling`（`0x132d3580`）强制："Loads of types wider than 32-bit unsupported"、"Only three-level tiling supported for 1D memory ops narrower than 32-bit"、"Invalid first-level tile in 1D memory op"。
### infer(vector::BroadcastOp) — `0x132ce520`
标量源会产生一个 result native layout（rank-1 → `MINOR`(2)，否则 `NONE`），并令 `in_layout = kNoLayout`。向量源会通过 `getLayout` 读取源布局，用两个方向的 `VectorLayout::generalizes` 测试等价性，然后在输出上把 broadcast 轴标为 replicated，只在源和目标 dim 大小匹配的位置保留 offset：
```c
function infer(vector::BroadcastOp op): // 0x132ce520
require rank(result) > 0 // "rank 0 vectors unsupported"
if scalar source:
out = native_layout(bw, implicit = rank1 ? MINOR(2) : NONE)
in_layout = kNoLayout
else:
src = getLayout(op.source) // "missing vector layout"
// "unsupported broadcast source type"
require generalizes(src, dst) or generalizes(dst, src) // equivalence
for axis i: keep offset[i] iff src_dim[i] == dst_dim[i] // else replicated (absent)
out = src with broadcast axes replicated次要规则（rotate、iota、transpose、multi_reduction、reshape）
这些规则遵循同一个 bitwidth → native-tiling 骨架；其精确 offset/implicit-dim 选择经过抽样，但不是每个都完全反编译（HIGH）：
| Op | 地址 | 行为 |
|---|---|---|
infer(RotateOp) | 0x132ca600 | 仅 32-bit（"not implemented: Rotate with non-32-bit data"）；native (sublane,lane)，in == out |
infer(IotaOp) | 0x132cd0a0 | native tiling；iota dim 在该轴上设为 replicated（offset 缺失） |
infer(MultiDimReductionOp) | 0x132cf660 | native tiling；reduced dim 在该轴上设为 replicated |
infer(TransposeOp) | 0x132d26c0 | 交换 (sublane,lane) 角色 → 产生与操作数不同的 out_layout（因此由 RelayoutInsertion 弥合） |
inferReshape | 0x132d09a0 | 重新计算 implicit-dim；可能强制 {0,0} re-layout |
NOTE — 携带 region 的 op（
scf.for/scf.while/scf.if/tpu.region）未被完全追踪。 它们的入口点是0x132c8140/0x132c8c40/0x132c9760/0x132ce300。统一 loop body 中 iter-arg 和 yield 布局的跨迭代不动点未反编译（这些主体 LOW）。重新实现必须添加一个 loop-body 布局统一步骤，这是上面的逐 op 表没有捕获的。
格：join 和 generalizes
因为 InferVectorLayout 从不插入 relayout，它会用 VectorLayout 上的半格（源码 layout.h）来调和多个操作数布局。VectorLayout 结构本身位于 Mosaic VectorLayout；这里仅说明 meet 和 ≤ 关系。
function VectorLayout::join(a, b, shape): // 0x14a957c0
require shape.size() >= layout_rank(implicit_dim) // layout.h:320
if a more_general_than b: // a replicated where b concrete,
return b // same bw/tiling/implicit, singleton dim
if b more_general_than a:
return a
require bitwidth(a) == bitwidth(b) // 16-byte tiling pair XOR + ptest
require implicit_dim(a) == implicit_dim(b)
out = a
for each dim: // merge offsets
if one offset replicated, other concrete: take the concrete
elif both concrete and equal: take it
else (concrete and different): return invalid() // no join
return out
```text
`generalizes(this, other, shape)` 是 ≤ 关系（同样有 `layout.h:320` CHECK）：当且仅当 `this` 至少与 `other` 一样泛化时为真，也就是 `this` 的每个 concrete 轴都被 `other` 匹配，而 `this` 的每个 replicated 轴都接受 `other` 的 concrete 值。`inferElementwise` 使用 `join`；`infer(BroadcastOp)` 在把 broadcast 轴标为 replicated 之前，用两个方向的 `generalizes` 测试等价性。
`getLayout(Value)`（`0x132d3260`）读取一个值的布局：它找到定义 op，并在结果索引处读取其 `out_layout` ArrayAttr（CHECK `infer_vector_layout.cc:2152/2154/2157` — `"op"`/`"op_result"`/`"out_attrs.size() > result_index"`）。当 inferer 的 `+44` lane-replicate 标志被设置时，它会强制 lane offset 为 replicated，使 broadcast/extract 源被读取为 lane-broadcast。
> **GOTCHA — `join` 失败是功能，不是错误。** 当两个操作数确实需要不同布局时，`join` 返回一个*无效*布局（`_RDI[56] = 0` 的 “no value” 标志）。`inferElementwise` 把这视为重新推导 fresh native layout 的信号，产生的生产者/消费者不匹配随后由 `RelayoutInsertion`（阶段 11）弥合。把 join 失败当作硬错误的重新实现，会拒绝生产编译器可以正常编译的 kernel。
---
## 示例：一个 bf16 Matmul Kernel
该 kernel（`{sublane=8, lane=128}`）：
```mlir
%a : memref<512x256xbf16, #tpu.memory_space<vmem>>
%b : memref<256x128xbf16, #tpu.memory_space<vmem>>
%o : memref<512x128xf32, #tpu.memory_space<vmem>>
%va = tpu.vector_load %a : vector<512x256xbf16>
%vb = tpu.vector_load %b : vector<256x128xbf16>
%acc = tpu.matmul %va, %vb : vector<512x128xf32>
tpu.vector_store %acc, %o阶段 4 — infer-memref-layout。 每个 arg memref 都获得一个 TiledLayoutAttr。没有 leading_tile_rows arg-attr 时，对于 bf16（bw=16），getTilingFactor(512, gen, sublanes=8, …, is_arg=1, 16)：packing=2，并且（假设 flags[0] 已设置，因为这些是 kernel args）bw==16 的 “large 2nd-minor” 分支给出 2*sublanes = 16（可整除 512），再加上 packing tile (2,1)：
%a → tiles [(16,128),(2,1)], contiguous strides (512x256 bf16)
%b → tiles [(16,128),(2,1)] (256x128 bf16)
%o → f32 (bw=32): sublane tile 8 → tiles [(8,128)] (512x128 f32)
```text
每个 arg 的 tiled type 都会通过一个 `tpu.erase_layout` 为函数体建立别名。
**阶段 9 — tiling-propagation。** 两个 `tpu.vector_load` 和一个 `tpu.vector_store` 匹配 `propagate_layout_to_consumer_rule`；它们的 memref 操作数被重新穿线到 tiled memref（看穿 `erase_layout`）。三个 `EraseLayoutOp` 现在没有 use，因此被擦除。
### 阶段 10 — infer-vector-layout。
```text
tpu.vector_load %a → inferLoadStoreVectorLayout reads %a's tiling (16,128)+(2,1)
out_layout = 16,{0,0},(16,128) (bf16 native, packed sublane 16)
in_layout = kNoLayout (memref operand)
tpu.vector_load %b → out_layout = 16,{0,0},(16,128)
tpu.matmul %va,%vb → infer(MatmulOp):
in_layout[lhs] = 16,{0,0},(16,128)
in_layout[rhs] = 16,{0,0},(16,128)
in_layout[acc] = 32,{0,0},(8,128)
out_layout = 32,{0,0},(8,128)
── lhs/rhs MATCH the loads' out_layouts ⇒ no relayout
tpu.vector_store → in_layout[%acc] = 32,{0,0},(8,128) (matches matmul out) ⇒ no relayout阶段 11 — relayout-insertion。 每个 producer.out_layout == consumer.in_layout ⇒ 不会插入 tpu.relayout；apply-pass 的 in == out 不变量成立。
阶段 12 — apply-vector-layout 消费的正是这些 attrs：bf16 → vector<8x128x2xbf16> vregs；f32 → vector<8x128xf32>；matmul → MXU latch/matpush/matres。
NOTE — 改一个 op 就会出现 relayout。 如果一个
tpu.transpose喂给 matmul lhs，infer(TransposeOp)会发出一个无法与 matmul 所需 lhsin_layoutjoin的out_layout；阶段 11 随后会在两者之间插入tpu.relayout。这正是求解器和应用器分离为两个 pass 的原因：求解器记录分歧，relayout-insertion pass 解决它。
未追踪的内容
- matmul / elementwise / load-store / broadcast / rotate / iota 之外的约 25 个逐 op
infer(...)主体：地址 + 签名已恢复，骨架已在代表性子集上验证；reduce_index/dynamic_gather/prng_random_bits/concatenate/reshape/transpose/extract_strided_slice的精确 offset/implicit-dim 选择经过抽样，未完全反编译（HIGH）。 - 除
tpu_memref_slice_rule之外的 5 个专门 memref-propagation 规则主体：已验证会重新穿线 tiled memref 并重新计算 tile stride；逐规则 stride 数学未逐个反编译（HIGH）。 TpuTilingFlags字节语义：门控getTilingFactor的 “large 2nd-minor” 乘数的 3 个 flag 字节已恢复为flags[0..2]；其字段名从用途推断（MEDIUM）。canReinterpretToUntiledMemref/canReinterpretToUntiledContiguousMemref，即选择erase_layout还是reinterpret_cast的谓词：调用点已恢复，函数体未反编译（精确谓词 LOW）。scf.for/while/if+tpu.regionregion-layout 推断（iter-arg/yield 统一不动点）：入口点已恢复，跨迭代不动点未反编译（这些主体 LOW）。- SparseCore 镜像
mosaic_sc::InferVectorLayoutPass（0x132ed380+ 逐 op infer 方法0x132ed7a0..0x132ee1e0）：确认存在且与 TensorCore inferer 平行，主体未在此反编译。
交叉引用
- Mosaic 概览 — 16 阶段流水线；这里的三个 pass 是其中的阶段 4、9、10。
- Mosaic VectorLayout —
(sublane, lane)VectorLayout结构、ImplicitDim、tilesPerVreg/tileArrayShape、relayoutdriver，以及消费本页产生属性的 49 项applyLayoutOp表。 - tpu MLIR 方言 — 这些 pass 会推断其布局的 ops（
tpu.matmul、tpu.vector_load/store、tpu.erase_layout、tpu.memref_slice、…）。 - MHLO → XTile → tpu Lowering — 通用 HLO 路径，以及不存在 MHLO→
tpulegalizer 的证明（Mosaic 是唯一的tpu生产者）。 - Layout Assignment — host-side HLO layout assignment，与这个 in-kernel Mosaic 分块求解器不同。
- Binary:
extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so（build-id89edbbe81c5b328a958fe628a9f2207d） - 索引项： Part V — Compiler: Lowering & Optimization Passes / MLIR lowering chain — 返回索引