Mosaic VectorLayout

本页中的所有地址、符号、field offsets、op-name 字符串和错误字符串均适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，构建 libtpu_lts_20260413_b_RC00）。field offsets 和 CHECK 行号都来自反编译二进制，按字节精确。其他版本会有所不同。

摘要

mlir::tpu::VectorLayout 是 Mosaic 后端的原子：一个 56 字节的值类型，说明一个逻辑 vector<…> SSA value 如何被打包进 TPU 的硬件向量寄存器。Mosaic kernel 内部每个 vector 类型的 operand 和 result 都携带一个这样的值，作为 in_layout / out_layout 数组 attributes 附加，而 Mosaic pipeline 的整个尾部就是消费它们的机制。本页记录该值类型本身（(sublane, lane) tiling 代数、offset/replication 模型、ImplicitDim rank-collapse 模型、逐 vreg packing 与 vreg-count 数学）、把带 layout 标注的 op 通过 49 项规则表转成 native-vreg ops 的 applyLayoutOp 分派，以及 relayout driver，即 disassemble → shift → re-tile → reassemble 引擎，它把 layout 变化物理物化为 lane/sublane shuffles。

范围边界：选择每个 value layout 的 inference pass（VectorLayoutInferer，生产者侧）位于 Mosaic Layout Inference；本页负责 struct 和 applier。二进制中两者被有意拆开：inference 从不插入 relayout，因此 applyLayoutOp 可以断言 producer-out 等于 consumer-in。

对于重新实现，layout-algebra 契约是：

VectorLayout 是 56 字节 POD：{offsets[2], tiling[2], bitwidth, implicit_dim}。 Offsets 是 optional<int64>（缺失 offset 表示沿该硬件轴复制）；tiling 是 (sublane_tile, lane_tile)；bitwidth 是 ≤ 32 的 2 的幂；implicit_dim 记录哪些逻辑 minor dims 被折叠。0x13249ba0 处的构造函数强制六个 invariants（layout.h:205-210）。
一个 vreg 容纳 packing = 32/bitwidth 个 tiles。 tilesPerVreg 和 tileArrayShape 从 layout 加 logical shape 精确推导该 value 占用多少 native vregs，以及每个 vreg 的具体 MLIR vreg type（getNativeVregOrVmaskType）。
applyLayoutOp 是基于 op-name 的 StringMap 分派。 49 条基础规则加一组 out-of-tree extension，并为任何携带 hasElementwiseMappableTraits 的 op 提供 elementwise_op_rule fallback。每条规则都把其 op 的逻辑 vectors 展开为逐 vreg native-shape ops。
layout 变化就是 relayout：disassemble → changeOffsets → changeTiling → changeImplicitDim → assemble。 Lane/sublane shifts（column/row shift）、re-tiling 和 implicit-dim insert/drop 是三个变化原语；relayout pass 保证 applier 只会在确实需要时看到 tpu.relayout op。


值类型	`mlir::tpu::VectorLayout`；ctor `mlir::tpu::VectorLayout::VectorLayout` @ `0x13249ba0`（源码 `…/mosaic/dialect/tpu/layout.h`）
Struct 大小	56 bytes；fields at `+0/+8/+16/+24` (offsets), `+32/+40` (tiling), `+48` (bitwidth, `i8`), `+52` (implicit_dim, `i32`)
Invariants	`layout.h:205-210` — 单 bit `bitwidth ≤ 32`、`tiling[i] > 0`、`offset ≥ 0`、sublane offset `< tiling[0]`
文本语法	`<bw>,{<o0>,<o1>},(<t0>,<t1>)[,<implicit>]`；print @ `0x14a94d80`，parse @ `0x14a95b40`，`printImplicitDim` @ `0x14a94b40`
Vreg-count 数学	`tilesPerVreg` @ `0x1325cec0`；`tileArrayShape` @ `0x14a94160`；`getNativeVregOrVmaskType` @ `0x14b766e0`（`vreg_util.cc:58`）
Applier 分派	`mlir::tpu::applyLayoutOp` @ `0x1325bca0`；`applyLayoutFunc` @ `0x1325cc80`；规则 StringMap `rules()::$_0` @ `0x1325b100`（49 项）
Elementwise fallback	`elementwise_op_rule` @ `0x1325c900`（任何不在表中但带 `hasElementwiseMappableTraits` 的 op）
Relayout driver	`mlir::tpu::relayout` @ `0x1325a480`（`apply_vector_layout.cc:9865`）；`disassemble` @ `0x132466a0`，`assemble` @ `0x132462c0`
变化原语	`changeOffsets` @ `0x1324bac0`，`changeTiling` @ `0x1324c880`，`changeImplicitDim` @ `0x13253b80`；row/col shift @ `0x13248c80`/`0x13249d40`
置信度	Confirmed（由字节锚定），除非某行或标注另有说明

1. VectorLayout 值类型

VectorLayout 回答一个问题：给定一个逻辑 vector<…xT> value，其每个元素位于硬件 (sublane × lane) register grid 的哪里？ 它不描述 memref（那是 TiledLayoutAttr，见 Mosaic Layout Inference）；它描述的是在 tpu-dialect ops 之间流动的vector SSA value。

0x13249ba0 处的构造函数精确固定了字节布局。前 32 字节通过一次来自参数块的 vmovups ymm0 写入（两个 optional<int64> offsets），随后分别存储 tiling、bitwidth 和 implicit_dim：

// VectorLayout::VectorLayout(bitwidth, offsets[2], tiling[2], implicit_dim) @0x13249ba0
__asm { vmovups ymmword ptr [rdi], ymm0 }   // +0..+31: offsets_[0..1] {value, has}
*(_QWORD *)(_RDI + 32) = a3;                 // +32: tiling_[0]  (sublane tile)
*(_QWORD *)(_RDI + 40) = a4;                 // +40: tiling_[1]  (lane tile)
*(_BYTE  *)(_RDI + 48) = a2;                 // +48: bitwidth_   (i8)
*(_DWORD *)(_RDI + 52) = a5;                 // +52: implicit_dim (i32 enum)
```text

### 1.1 Field layout（按字节精确）

| off | field | type | 含义 |
|----|----|----|----|
| `+0` | `offsets_[0].value` | `int64` | tile 内第一个元素的 2nd-minor（**sublane**）offset |
| `+8` | `offsets_[0].has` | `bool` | `false` ⇒ value 在 sublanes 间**复制** |
| `+16` | `offsets_[1].value` | `int64` | tile 内第一个元素的 minor（**lane**）offset |
| `+24` | `offsets_[1].has` | `bool` | `false` ⇒ value 在 lanes 间**复制** |
| `+32` | `tiling_[0]` | `int64` | sublane tile size（例如 f32 为 8，bf16 为 16） |
| `+40` | `tiling_[1]` | `int64` | lane tile size（例如 128） |
| `+48` | `bitwidth_` | `int8` | element bit width ∈ {1,2,4,8,16,32} |
| `+52` | `implicit_dim` | `int32` | `ImplicitDim` enum（§1.4） |

置信度：Confirmed — 每个 offset 都直接从上面的构造函数 store 序列读取；`optional<int64>` 形状（8 字节 value + 1 字节 `has` flag，padding 到 16 字节槽）由 invariant checks 中的 `vmovups`/`value_or` 访问确认。

### 1.2 构造函数 invariants

构造函数在六个条件上 `CHECK` 失败，每个条件都带有其 `layout.h` 源码行号（从 `LogMessageFatal` 调用点逐字恢复）：

| line | condition | rationale |
|----|----|----|
| 205 | `llvm::has_single_bit<unsigned>(bitwidth_) && bitwidth_ <= 32` | bitwidth 必须是 2 的幂，且不宽于 register lane word |
| 206 | `tiling_[0] > 0` | sublane tile 为正 |
| 207 | `tiling_[1] > 0` | lane tile 为正 |
| 208 | `offsets_[0].value_or(0) >= 0` | sublane offset 非负 |
| 209 | `offsets_[1].value_or(0) >= 0` | lane offset 非负 |
| 210 | `offsets_[0].value_or(0) < tiling_[0]` | sublane offset 保持在一个 tile row 内 |

> **NOTE — lane offset *不*受 tiling 约束。** Invariant 210 只约束 *sublane* offset（`offsets_[0] < tiling_[0]`）。lane offset（`offsets_[1]`）有意不由构造函数约束：lane offset 可以跨越多个 lane tiles，tile-array 数学（§2）会把它折叠到 vregs 之中。这种不对称是真实的，二进制检查 208/209/210，但从不检查 `offsets_[1] < tiling_[1]`。

因此，一个 `VectorLayout` 的语义读法是：*该 value 的元素以 sublane tile `tiling_[0]` 和 lane tile `tiling_[1]` 打包进 vregs；第一个逻辑元素位于 tile 内的 `(offsets_[0], offsets_[1])`；缺失 offset 表示该 value 沿对应硬件轴 broadcast（复制）；`implicit_dim` 记录哪些逻辑 minor dims 被折叠，以适配严格 2-D-tiled 的 vreg 模型。*

### 1.3 文本语法

`VectorLayout::print` @ `0x14a94d80` 和 `VectorLayout::parse` @ `0x14a95b40` 是精确互逆。语法：

```text
<bitwidth>,{<o0>,<o1>},(<t0>,<t1>)[,<implicit>]

{/} 包住两个 offsets；复制的（缺失的）offset 打印为 *（0x2A）。(/) 包住 tiling。
<implicit> 是可选的，由 printImplicitDim @ 0x14a94b40 发射：NONE → 省略，MINOR → -1，SECOND_MINOR → -2，MINOR_AND_SECOND_MINOR → -2,-1。parse 接受相同 tokens。

规范示例：

text

32,{0,0},(8,128)            f32, offset (0,0), native tiling 8x128, 1 tile/vreg
16,{0,0},(16,128)           bf16, packed sublane tile 16, native tiling
32,{*,0},(8,128)            f32 replicated across sublanes (sublane offset absent)
16,{0,0},(16,128),-1        bf16 with the minor logical dim implicit
```text

置信度：Confirmed — print/parse 符号存在于引用的 VAs；`printImplicitDim` @ `0x14a94b40` 确认存在。

### 1.4 ImplicitDim 模型

`VectorLayout::ImplicitDim`（`+52` 处的 `int32`）是一个 4 值 enum。它让 Mosaic 能在一个*始终*至少为 2-D-tiled 的模型中表示低 rank 逻辑 vector（1-D、scalar），而无需为每条规则特判：

| value | name | printed | use |
|----|----|----|----|
| 0 | `NONE` | (omitted) | 最后两个逻辑 dims 直接映射到 `(sublane, lane)` |
| 1 | `MINOR` | `-1` | minor（lane）逻辑 dim 是 implicit（size 1）；1-D / lane-broadcast value |
| 2 | `SECOND_MINOR` | `-2` | second-minor（sublane）逻辑 dim 是 implicit |
| 3 | `MINOR_AND_SECOND_MINOR` | `-2,-1` | 两者都是 implicit；提升到 vreg 的 scalar |

`implicitShape(shape)` @ `0x14a94080`（以及 `insertImplicit<>` helpers，即 `insertImplicit<long>` @ `0x132958a0`、`insertImplicit<bool>` @ `0x13295a40`）会在 implicit 位置重新插入 `popcount(implicit_dim_bits)` 个 size-1 dims，因此无论 value 的逻辑 rank 如何，所有 tile 数学都运行在 `≥2-D` 的 "implicit shape" 上。置信度：HIGH — enum 值从 `printImplicitDim` 的输出 tokens 恢复；`implicitShape`/`insertImplicit` 机制由符号锚定，但没有逐个反编译每个 helper 的展开。

---

## 2. Vreg 计数与 Native-Type 数学

一个 `VectorLayout` 加一个 logical shape 会决定：(a) 一个 vreg 能容纳多少 sub-tiles，(b) 整个 value 占用多少 vregs，(c) 每个 vreg 的具体 MLIR type。这三个函数是 apply rules 用来把逻辑 vector 转成 native vregs 的具体 `xla::Array<Value>` 的工具。

### 2.1 tilesPerVreg

`VectorLayout::tilesPerVreg(array<long,2> target_shape)` @ `0x1325cec0`：

```text
packing      = 32 / bitwidth_                          // 1(f32) / 2(bf16) / 4(i8) / 8(i4)
tilesPerVreg = packing * sublane * lane / (tiling_[0] * tiling_[1])   // remainder MUST be 0

CHECKs（逐字恢复）：0 != bitwidth（"bitwidth cannot be 0"，layout.h:245）以及 layout.h:250 处的整除 guard。对于 native tiling（t0 = sublane，t1 = lane），公式简化为 tilesPerVreg == packing：一个 vreg 容纳 packing 个 sublane×lane tiles，沿 sub-32-bit packing 轴堆叠。因此 bf16 native → 2 tiles/vreg，int8 → 4，i4 → 8，f32 → 1。置信度：Confirmed — 0 != bitwidth / layout.h:245/250 CHECK 字符串在反编译中按字节精确。

2.2 tileArrayShape

VectorLayout::tileArrayShape(bool, bool, shape, {sublane,lane}) @ 0x14a94160 返回每个 dim 上的 vreg 数。遍历 implicit shape 时，对于末尾两个 dims：

text

n_2nd_minor_tiles = ceil_div( offsets_[0].value_or(0) + shape[-2], tiling_[0] )
n_minor_tiles     = ceil_div( offsets_[1].value_or(0) + shape[-1], tiling_[1] * tilesPerVreg )
```text

复制轴会强制其 tile-count 为 1（反编译中的 `*(…)=1` 分支）。Leading dims 原样穿过；随后 `implicit_dim` 剥离 implicit axes（case 1 丢弃 `dim[-1]`；case 2 把 `dim[-1]` 折入 `dim[-2]` 的位置再丢弃；case 3 丢弃两者）。`CHECK`s：`layout.cc:410`（`src_shape.size() >= layout_rank()`）、`layout.cc:423`（`src_shape.size() >= 2`）。置信度：HIGH — ceil-div 公式和 replicated→1 分支来自反编译 walker；精确的 implicit-dim folding cases 从结构读取，但未逐 case 穷尽追踪。

### 2.3 getNativeVregOrVmaskType

`getNativeVregOrVmaskType(elemTy, layout_bitwidth, {sublane,lane})` @ `0x14b766e0`（`vreg_util.cc:58`）产生具体 MLIR vreg type：

| bitwidth | native vreg type |
|----|----|
| 32 | `vector<sublane × lane × T>`（2-D） |
| < 32 | `vector<sublane × lane × (32/bw) × T>`（3-D，尾随 dim 是 packing axis） |
| 1 (`i1`) | 一个 `vmask`（bitwidth-1 special path） |

`CHECK vreg_util.cc:58`：`bitwidth == layout_bitwidth`（反编译中按字节精确）。apply pass 后，每个 `vector` value 都正好是这些 native shapes 之一，因此 LowerToLLO 会把它 1:1 映射到 LLO `VregType`（[tpu → LLO ODS](tpu-to-llo-ods.md)）。置信度：Confirmed — `bitwidth == layout_bitwidth` / `vreg_util.cc:58` CHECK 按字节精确；type-shape 分支来自 bitwidth 分派。

---

## 3. applyLayoutOp — 逐 Op 分派

`mlir::tpu::applyLayoutOp(ApplyVectorLayoutContext&, Operation&)` @ `0x1325bca0` 是由 `applyLayoutFunc` @ `0x1325cc80` 对每个 op 调用的 driver（后者要求单 region、单 block 的 `FuncOp`，有两个独立检查，按字节精确："Expected FuncOp to have a single region" 和 "Expected FuncOp to have a single block"，并遍历每个 op）。对于每个 op，`applyLayoutOp`：

1. **读取附加 layout。** `getOutLayouts(op)` 随后 `getInLayouts(op)` 读取 inference pass 附加的逐 op `out_layout` / `in_layout` array-of-`VectorLayoutAttr`。反编译 guard `if (v79 != 1) return 0;`，也就是说没有 out-layouts 时会提前返回。

2. **强制 in==out invariant。** 对每个 vector operand，它 `CHECK` *operand-is-vector ⇔ in-layout has_value*，并且**producer 的 out-layout 等于该 op 的 in-layout**。不匹配时发出（按字节精确）：

   > `Invariant violation: Input layout does not match output layout - did you forget to run relayout-insertion?`

   这是架构接缝：inference 为每个 value 选择一个 layout，*单独的* relayout-insertion pass 桥接分歧，因此 applier 永远不需要调和。逐 operand in==out 循环会在 op *不是* `tpu.assume_layout` 时启用（这是唯一允许 consumer in-layout 与 producer out-layout 不一致的 op）。

   对 offset-in-first-tile guard 还有一个**单独**豁免（`"Not implemented: Input offsets outside of the first tile"`，按字节精确）。该 guard 会对九个可合法消费 offset 不在第一个 tile 内的 operand 的 ops **跳过**：`tpu.truncf`、`tpu.relayout`、`tpu.reshape`、`tpu.concatenate`、`vector.shape_cast`、`vector.extract_strided_slice`、`vector.broadcast`、`arith.trunci`、`arith.extsi`（反编译中的精确 `TypeIDResolver` 集合）。这是二进制中两个不同的豁免，不是同一个共享列表。

3. **按 op-name 分派。** 在惰性构建的 `rules()` StringMap（`xxh3_64bits` + `StringMapImpl::FindKey`）中查找 op-name。命中时，调用规则 fn-ptr（vtable `+24`）。未命中时：如果 op 有 `hasElementwiseMappableTraits` → `elementwise_op_rule` @ `0x1325c900`；否则发出 `"Not implemented: Unsupported operation: <op> in apply-vector-layout pass"`。

> **NOTE — apply rule 实际做什么。** 每条规则都会通过 `disassemble`/`Each<Value>` 把其 op 的逻辑 vectors 展开为逐 vreg native-shape ops（使用 §2 的数学产生 native vregs 的 `xla::Array<Value>`），并发射 LowerToLLO 随后会 1:1 消费的具体 sublane/lane shuffles，如 `tpu.rotate`、`tpu.relayout`、`broadcast_in_sublanes`。规则表就是 "`tpu`-dialect op → native-vreg op" lowering。逐规则 shuffle *body*（例如 `tpu_matmul_rule` 到底如何 tile K，或 `vector_transpose_rule` 如何发射 helper）属于逐 op 反编译，本页没有完成，置信度 LOW。

### 3.1 49 项规则表

由 `rules()::$_0::operator()` @ `0x1325b100` 构建（49 项基础条目），随后与 `extensions::rules()` @ `0x13246180` 合并（out-of-tree ops）。Op-name → rule fn（按注册顺序）。所有 49 个 op-name 字符串都在 StringMap 的 string pool 中按字节精确确认：

| op-name | rule fn | @VA |
|----|----|----|
| `arith.constant` | `arith_constant_rule` | `0x132620a0` |
| `arith.extsi` | `arith_extsi_rule` | `0x13262bc0` |
| `arith.extui` | `arith_extui_rule` | `0x13262ec0` |
| `arith.trunci` | `arith_trunci_rule` | `0x13263780` |
| `func.return` | `func_return_rule` | `0x13263960` |
| `scf.for` | `scf_for_rule` | `0x13263a60` |
| `scf.while` | `scf_while_rule` | `0x13265940` |
| `scf.condition` | `scf_condition_rule` | `0x13266f40` |
| `scf.if` | `scf_if_rule` | `0x13267380` |
| `scf.yield` / `tpu.yield` | `yield_rule`（shared） | `0x13268900` |
| `tpu.rotate` | `tpu_rotate_rule` | `0x13268d40` |
| `tpu.dynamic_rotate` | `tpu_dynamic_rotate_rule` | `0x13269c00` |
| `tpu.concatenate` | `tpu_concatenate_rule` | `0x1326a900` |
| `tpu.pack_elementwise` | `tpu_pack_elementwise_rule` | `0x1326c2a0` |
| `tpu.unpack_elementwise` | `tpu_unpack_elementwise_rule` | `0x1326c420` |
| `tpu.iota` | `tpu_iota_rule` | `0x1326c520` |
| `tpu.gather` | `tpu_gather_rule` | `0x1326d440` |
| `tpu.dynamic_gather` | `tpu_dynamic_gather_rule` | `0x1326df60` |
| `tpu.reduce_index` | `tpu_reduce_index_rule` | `0x1326e860` |
| `tpu.load` | `tpu_load_rule` | `0x13270500` |
| `tpu.store` | `tpu_store_rule` | `0x13270ca0` |
| `tpu.strided_load` | `tpu_strided_load_rule` | `0x13271440` |
| `tpu.strided_store` | `tpu_strided_store_rule` | `0x13271680` |
| `tpu.vector_store` | `tpu_vector_store_rule` | `0x132718e0` |
| `tpu.matmul` | `tpu_matmul_rule` | `0x132727a0` |
| `tpu.region` | `tpu_region_rule` | `0x13274480` |
| `tpu.bitcast` | `tpu_bitcast_rule` | `0x132750a0` |
| `tpu.trace` | `tpu_trace_rule` | `0x132755c0` |
| `tpu.assume_layout` | `tpu_assume_layout_rule` | `0x13275840` |
| `tpu.prng_random_bits` | `tpu_prng_random_bits_rule` | `0x13275ec0` |
| `tpu.relayout` | `tpu_relayout_rule` | `0x1325aea0` |
| `tpu.reshape` / `vector.shape_cast` | `reshape_rule`（shared） | `0x13276760` |
| `tpu.fptosi` / `tpu.fptoui` | `tpu_fptoi_rule`（shared） | `0x13278960` |
| `tpu.sitofp` / `tpu.uitofp` | `tpu_itofp_rule`（shared） | `0x13278e80` |
| `tpu.extf` | `tpu_extf_rule` | `0x132792c0` |
| `tpu.truncf` | `tpu_truncf_rule` | `0x13279680` |
| `vector.broadcast` | `vector_broadcast_rule` | `0x13279900` |
| `vector.extract` | `vector_extract_rule` | `0x1327bd00` |
| `tpu.vector_load` | `tpu_vector_load_rule` | `0x1327d020` |
| `vector.multi_reduction` | `vector_multi_reduction_rule` | `0x1327ddc0` |
| `vector.extract_strided_slice` | `vector_extract_strided_slice_rule` | `0x1327f400` |
| `tpu.transpose` | `vector_transpose_rule` | `0x1327f960` |
| `tpu.matmul_push_rhs` | `tpu_matmul_push_rhs_rule` | `0x13281e60` |
| `tpu.matmul_acc_lhs` | `tpu_matmul_acc_lhs_rule` | `0x13282020` |
| `tpu.matmul_pop` | `tpu_matmul_pop_rule` | `0x132821e0` |

（按 op-name 计数并把共享 fn-ptrs 算入 = 49 项。"shared" = 一个 rule fn 注册在多个 op-name 下。）不在表中但携带 `hasElementwiseMappableTraits` 的 ops，如 `arith.addf`、`arith.mulf`、`arith.select`、`math.exp` 等，会路由到 `elementwise_op_rule`；其他所有 op 都是 "Unsupported operation"。

置信度：Confirmed，op-name 集合（所有 49 个字符串在 pool 中按字节精确）和分派（`applyLayoutOp` + invariant 字符串按字节精确）均已确认。逐规则 VA 从 StringMap fn-ptr 表恢复。

> **GOTCHA — 这些不是 inference rules。** `applyLayoutOp` 表是 layout 的*消费者*。一个单独的 `PropagationContext` 类型 StringMap（`rules()` @ `0x132e15e0`）驱动生产者侧的 *memref-tiling propagation*，而逐 op 的 `VectorLayoutInferer::infer(...)` 分派负责*选择* layouts；两者都记录在 [Mosaic Layout Inference](mosaic-layout-inference.md)。memref 侧 ops（`tpu.memref_slice`、`memref.cast`、`tpu.reinterpret_cast` 等）只出现在那个 propagation map 中，从不出现在这里。

---

## 4. relayout Driver — 物化 Layout 变化

当两个 values 确实需要不同 layouts 时，relayout-insertion pass 会在它们之间发射一个 `tpu.relayout` op；它的 apply rule（`tpu_relayout_rule` @ `0x1325aea0`）调用中央变化引擎。任何必须在自身 lowering 内调和 layout mismatch 的规则也会直接调用该引擎。

`mlir::tpu::relayout(ApplyVectorLayoutContext&, OpBuilder&, value, src_layout, dst_layout)` @ `0x1325a480`（`apply_vector_layout.cc:9865`）：

1. **Replication-compatibility check。** 如果一个逻辑 dim 是非 singleton，且在 `dst` 中复制但在 `src` 中未复制，则非法（不能从具体数据凭空制造复制数据）：

   > `Invalid relayout: Non-singleton logical dimension is replicated in destination but not in source for <v> : <src> -> <dst>`

   （反编译中按字节精确）。

2. **Disassemble。** `disassemble(builder, src_layout, value, {sublane,lane})` @ `0x132466a0` 把 value 分解成 native vregs 的 `xla::Array<Value>`（使用 §2 的 `tilesPerVreg`/`tileArrayShape`）。

3. **Element-type split + 三个变化原语。** Masks（`i1`）走 `relayoutMasks` @ `0x13258d40`；其他所有类型走 `relayoutVregs` @ `0x13257a80`。两者都会依次应用三个变化原语：

   | primitive | @VA | what it emits |
   |----|----|----|
   | `changeOffsets` | `0x1324bac0` | lane/sublane shifts — `doRowShiftRelayout` @ `0x13248c80`（sublane），`doColumnShiftRelayout` @ `0x13249d40`（lane） |
   | `changeTiling` | `0x1324c880` | 把 vregs re-tile 到目标 tiling |
   | `changeImplicitDim` | `0x13253b80` | 插入 / 丢弃 implicit dims |

4. **Post-change assert。** 变化原语之后 layout 必须已经等于 `dst`；`apply_vector_layout.cc:9865` 处的 `CHECK src == dst` 强制完整调和。

5. **Reassemble。** `assemble(builder, vecTy, dst_layout, vregArray, {sublane,lane})` @ `0x132462c0` 重建 relaid-out value。

因此 relayout 是：*explode to native vregs → lane/sublane shuffle（column/row shift）+ re-tile + implicit-dim adjust → reassemble。* 变化原语是 layout 代数的最底层，即实际的 gather/roll/broadcast register shuffles，表达为 `changeOffsets`/`changeTiling`/`changeImplicitDim` 三元组。

置信度：Confirmed — `relayout`、`disassemble`/`assemble`、`relayoutVregs`/`relayoutMasks` 以及全部三个变化原语（加上 `doRowShiftRelayout`）都以真实符号存在于引用的 VAs；`"Invalid relayout"` 字符串按字节精确。每个变化原语内部的 shuffle *codegen* 由符号锚定，但本页未反编译 body，置信度 HIGH。

> **NOTE — relayout op set 的意图。** 三个变化原语映射到三类 register moves：`changeOffsets` → **roll/shift**（重新对齐第一个元素位置的 lane 和 sublane rotates）；`changeTiling` → **gather/re-pack**（tile shape 变化时在 vregs 之间移动 sub-tiles）；`changeImplicitDim` → **broadcast/squeeze**（插入或折叠 size-1 轴，在 replicated axis 上就是 broadcast）。如果一个 layout 只在 `{offsets, tiling, implicit_dim}` 的某一项上不同于其 consumer，只会触发对应原语；当 `src == dst` 于该方面时，其余原语是 no-op。

---

## 5. Worked Example — 一个 bf16 matmul kernel

一个 Mosaic kernel `main`，`{sublane = 8, lane = 128}`，lowering `tpu.matmul`：

```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

inference 之后（见 Mosaic Layout Inference），这些 values 携带如下 VectorLayout：

text

%va  : 16,{0,0},(16,128)     bf16 native: packed sublane tile 16 (= 8 sublanes x 2 packing), lane 128, tilesPerVreg = 2
%vb  : 16,{0,0},(16,128)
%acc : 32,{0,0},(8,128)      f32 native: 8x128, tilesPerVreg = 1
```text

通过 `tileArrayShape`（§2.2）得到 vreg 数：

```text
%va : ceil((0+512)/16) x ceil((0+256)/(128*2)) = 32 x 1 = 32 vregs
%vb : ceil((0+256)/16) x ceil((0+128)/(128*2)) = 16 x 1 = 16 vregs
%acc: ceil((0+512)/8)  x ceil((0+128)/(128*1)) = 64 x 1 = 64 vregs

随后 applyLayoutOp 遍历 block：

tpu.vector_load → tpu_vector_load_rule @ 0x1327d020：展开为 32 / 16 个逐 vreg native loads，每个都是 vector<8x128x2xbf16>（来自 bw = 16 时的 getNativeVregOrVmaskType，§2.3）。
tpu.matmul → tpu_matmul_rule @ 0x132727a0：在 native vregs 上物化 matmul（MXU latch/matpush/matres 序列由下游 LowerToLLO + MMA functor 产生）。
tpu.vector_store → tpu_vector_store_rule @ 0x132718e0：64 个 native vector<8x128xf32> stores。

因为每个 producer-out layout 都等于 matmul 要求的 operand in-layout（lhs/rhs 为 16,{0,0},(16,128)，acc/result 为 32,{0,0},(8,128)），relayout-insertion 不会添加任何 tpu.relayout op，applyLayoutOp 中的 in==out invariant（§3，步骤 2）成立。若是 tpu.transpose 给 matmul lhs 供给数据，inference 会发出与 matmul 的 lhs in-layout 不匹配的 out-layout，relayout-insertion 会在它们之间拼接一个 tpu.relayout，而 tpu_relayout_rule 会驱动 §4 的变化原语来物理重新对齐数据。

置信度：HIGH — layout 字符串、vreg 数和规则分派直接来自按字节精确的 struct/dispatch 数学；示例数字是推导而来，不是运行时 trace 读取。

6. SparseCore 镜像

反编译还包含一个并行的 mlir::mosaic_sc::VectorLayout 值类型，它有自己的 parse @ 0x132fc7c0 / print（以及 0x132fa0c0/0x132f9fa0 处的 VectorLayoutAttr parse/print），由 LowerToMlo 之前的 SparseCore layout solver 使用。它镜像 TensorCore 的 mlir::tpu::VectorLayout（相同的 (sublane, lane) 代数），但位于不同的符号 namespace。上文记录的 TensorCore 路径是通用 Pallas/Mosaic kernels 会走到的路径；SparseCore 镜像在此为完整性注明，本页未反编译。置信度：HIGH — 已确认存在 mosaic_sc::VectorLayout 符号；其 body 未追踪。

交叉引用

Mosaic Overview — import/serde 接缝、CustomCallEmitter::Emit，以及运行 infer-vector-layout → apply-vector-layout 的 16 阶段 RunMLIRPasses pipeline。
Mosaic Layout Inference — 生产者侧：逐 op 的 VectorLayoutInferer 规则，它们选择本页 applier 消费的 in_/out_layout attrs、InferMemRefLayout memref tiling、tiling-propagation fixpoint，以及 join/generalizes lattice。
The tpu MLIR Dialect: Ops and the Op-Model Contract — 这些规则操作其 layouts 的 tpu-dialect op definitions。
MHLO → XTile → tpu Lowering — 为什么普通 HLO 永远不会变成 tpu ops（apply table 只在 Mosaic 分支中抵达）。
tpu → LLO ODS — 下一段下降：apply-vector-layout 产生的 native vregs 会 1:1 映射到 LLO VregType。
Binary: extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）— back to index

Mosaic VectorLayout ​

摘要 ​

1. VectorLayout 值类型 ​

2.2 tileArrayShape ​

6. SparseCore 镜像 ​

交叉引用 ​

Mosaic VectorLayout

摘要

1. VectorLayout 值类型

2.2 tileArrayShape

6. SparseCore 镜像

交叉引用