Dot / Conv → MXU 降低

本页中的所有地址、符号和偏移都适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，未剥离，.text VA == 文件偏移）。其他版本会不同；请将每个 VA 都视为绑定到该版本。

摘要

每个 TPU 矩阵乘积都通过同一条下降路径到达脉动阵列：HLO kConvolution → 由 LLO matprep / matmul / matres op 组成的分块循环嵌套。这里没有独立的 dot 路径，上游 HLO pass 会把 kDot 重写为 kConvolution，因此同一个 lowering 同时服务两者（该重写和维度编号映射见 RaggedDot 和卷积几何降低）。本页记录这条下降路径的后半段：按窗口运行、用于选择脉动分块的 tile-cost comparator，用于从 19 种 MXU 发射策略中选择一种的 EmitFunctorEnum dispatch，以及通过把两个操作数行打包进一个 MXU latch 来将 latch 数量减半的 MxuLatchPacker 后置 pass。这三部分是重新实现者无法仅从 op 表面猜出的内容。

参考框架是教科书式 im2col/Winograd 卷积降低，而这里的分歧是彻底的。TPU MXU 是一个 128×128 脉动阵列，由一个锁存在其 gain 寄存器中的驻留操作数和一个通过 staging 寄存器推送的移动操作数供给；二进制中没有 im2col、没有 Winograd、没有 FFT。因此降低是一个分块问题（如何把 activation/kernel/output 立方体切成适配 VMEM 和阵列的 ≤128×128 窗口）、一个策略问题（哪个维度落到 lanes、sublanes 还是 K 归约）、以及一个打包问题（如何在小 tile 间摊薄 latch 成本）。这三个问题分别由 SpatialMajorConvolution::IterateThroughWindowConfigs（成本搜索）、MatrixMultiplyAccumulateFunctor::GetEmitFunctorEnumAndLoweringDecisions（19 路选择器）和发射后的 MxuLatchPacker / MxuDiagonalPacker pass 解决。MXU op 表面本身，即 VectorMatmulMubrOp、VectorLatchOp 的 ODS shape，以及四个寄存器枚举（MatmulMode/GainLatchMode/GainMatrixRegister/MatrixStagingRegister），由 tpu → LLO ODS 降低负责；本页使用这些构建工厂，而不是重新推导它们。

对重新实现而言，契约是：

tile-cost comparator：按窗口的 MXU-cycle 公式（precision-pass ×1/×2 乘数）、VMEM-footprint 预算、WindowConfig_CostModelType 枚举，以及选择获胜 WindowConfig 的 cycles → VMEM-fit → granule 平局裁决阶梯。
EmitFunctorEnum dispatch：19 个值的 convolution_util::EmitFunctor 枚举（ordinal 与 jump table 按字节精确一致）、GetEmitFunctorEnumAndLoweringDecisions 决策树（depthwise / batch-group / reduce-window / dense 分支），以及 enum → member-function pointer 派发。
MatrixMultiplyAccumulateFunctor 循环体：按 accumulator-window 的 MXU 序列（AccumFirstZero → Accumulate），以及 K>128 如何变成多 pass 累加。
MxuLatchPacker：贪心 same-mode 相邻成对打包（PackLatches）、操作数 pack/unpack 辅助函数、num_latches == num_latchpreps 不变量，以及兄弟 MxuDiagonalPacker 的 quadrant 共驻留。


降低入口	`ConvolutionEmitter::Create` — `0x130d86c0`（234 B；megacore dispatch → `SpatialMajorConvolution`）
降低驱动	`SpatialMajorConvolution::Emit` — `0x13178340`（3550 B）
tile-cost 搜索	`SpatialMajorConvolution::IterateThroughWindowConfigs` — `0x13167f20`（23877 B，939 BB）
成本冻结	`SpatialMajorConvolution::SetupBestConfig` — `0x13172580`（1770 B）
MXU codegen	`MatrixMultiplyAccumulateFunctor::operator()` — `0x1310cd80`（5398 B）
策略选择器	`GetEmitFunctorEnumAndLoweringDecisions` — `0x1310c720`（1312 B）→ 16-bit `(decision<<8)\|ord`
策略派发	`GetEmitFunctorFromEmitFunctorEnum` — `0x130e8de0`（`switch`，19 cases）
模式比较器	`ConvMatmulModes::operator<` — `0x130e12a0`（39 B；权重表 `0xae0f480`）
latch 打包器	`xla::jellyfish::PackLatches` — `0x10f726c0`（8559 B；`mxu_latch_packer.cc`）
发射的 LLO ops	`vmatprep.subr` / `vmatprep.mubr`（+ `.msk`），逐代 `vmatmul`，`vmatres`，融合的 `kVectorMatmulLmr`
MXU 几何	128×128 脉动阵列；`LaneCount`=128（`0x1d60f400`），`SublaneCount`=8（`0x1d60f300`）；v5+ 上每 TC 有 2 个 MXU

降低入口 — ConvolutionEmitter::Create

目的

ConvolutionEmitter 是一个单方法工厂：它解析 megacore core mask，并构造处理每种 conv shape 的唯一 lowering 类。这里没有策略 switch，策略选择被推迟到下游的成本搜索和按窗口选择器。

入口点

text

ConvolutionEmitter::Create (0x130d86c0, 234 B)        ── megacore dispatch
  ├─ megacore_util::GetMegacoreCoreMask                ── core_mask ∈ {1=secondary, 3=both}
  ├─ lowering_util::SecondaryCoreRegion (if mask==1)   ── swap LLO region to secondary core
  └─ make_unique<SpatialMajorConvolution>(...)         ── the only lowering class
        └─ FusedSpatialMajorConvolution                ── subclass when conv is the root of an output fusion
```text

### 算法

```c
// ConvolutionEmitter::Create — 0x130d86c0
unique_ptr<ConvolutionEmitter> Create(conv, getter, region, ...):
    core_mask = megacore_util::GetMegacoreCoreMask(conv, target)
    if core_mask == kUseSecondaryCore:                 // == 1
        region = lowering_util::SecondaryCoreRegion(backend_cfg, target)
    CHECK(core_mask == 1 || core_mask == 3)            // else FATAL convolution_emitter.cc:781
    return make_unique<SpatialMajorConvolution>(
        conv, getter, region, /*core_mask=*/core_mask, ...)

NOTE — core_mask 是 2-bit core enable：1 = 仅 secondary core（region 被替换），3 = 两个 core（megacore split）。任何其他值都会 FATAL。Megacore splitting 会选择一个迭代维度，并在两个 TC core 间把它的 bound 减半（ChooseMegacoreDimAndIterationBounds）；如果没有维度足够大，conv 会以 single-core 运行。split dim 会反馈到成本搜索中，因为 megacore_active_dims 会影响选择器对“all dims static”的测试（见选择器树）。

注意事项

二进制中缺少 Winograd / Im2Col / FFT 字符串，是单策略设计的正面证据：存在的唯一 im2col 字符串属于不相关的 NVPTX intrinsic 表面。构建多算法 conv dispatcher 的重新实现者，相对于这个目标来说是在过度工程化：MXU 的自然适配是 direct spatial convolution，而所有 shape specialization 都通过成本搜索和 19 路 emit 选择器发生在同一个 SpatialMajorConvolution 类内部。

Tile-Cost Comparator

目的

IterateThroughWindowConfigs 枚举每个合法的 (kernel_window, output_window) tile tuple，用 MXU cycles 和 VMEM footprint 为每个候选打分，并把幸存者流式传给 visitor callback；SetupBestConfig 将获胜者冻结进 WindowConfig proto。这就是脉动 tile 选择，即把逻辑 matmul 切成 ≤128×128 阵列大小窗口的方式。成本模型错误不会产生错误答案，但会产生要么溢出 VMEM（编译失败）、要么浪费脉动吞吐的分块。

入口点

text

SpatialMajorConvolution::Emit (0x13178340)
  ├─ EmitZeroElementCases / EmitZeroByteCase           ── size-0 / all-padding short-circuits
  ├─ IterateThroughWindowConfigs (0x13167f20)          ── enumerate + score candidates
  │     ├─ CalculateWindowMxuCycles (0x1315fe60)        ── PRIMARY key: MXU cycles
  │     │     ├─ CalculateClassicWindowCost (0x131626c0)
  │     │     └─ CalculateDepthwiseWindowCycles (0x13161d80)
  │     ├─ GetConvPrecision (0x131916e0)               ── precision index → ×2.0 / ×1.0 multiplier
  │     └─ VmemToUseForPotentialWindows                ── SECONDARY key: VMEM footprint
  └─ SetupBestConfig (0x13172580)                      ── freeze winning WindowConfig
```text

### 算法 — 成本公式

主键是窗口的 MXU cycles。核心公式从 `CalculateWindowMxuCycles` 中恢复（`0x131605cf` 处的 `vmulsd`/`vdivsd` 链，紧跟 `0x131605a9` 处的 `GetConvPrecision` 调用之后）：

```c
// CalculateWindowMxuCycles — 0x1315fe60 (classic at 0x131626c0, depthwise at 0x13161d80)
long window_cycles(window, hlo, operand, target):
    conv_count = (int) matmul_step_count(window)        // # systolic steps this window issues
    precision  = GetConvPrecision(hlo, operand, target) // 0 or 1   @ 0x131916e0
    mult       = precision_mult_table[precision]        // @0xa2c6050 = { 2.0, 1.0 }
    cycles_d   = conv_count * mult / divisor + base_cycles   // base from MxuLatencyTable
    return (long) cycles_d

0xa2c6050 处的 precision_mult_table 是双元素 double 数组 {2.0, 1.0}。索引 0（×2.0）是低精度 / split-accumulate pass（fp32-emulated 和其他 two-pass dtype）；索引 1（×1.0）是原生 single-pass（bf16、fp8）。
绝对 per-format base cycles（bf16 ≈ 211，fp8 ≈ 204）来自逐代 MxuLatencyTable，见 MXU 延迟概览。这个 pass 在其上叠加 precision multiplier；divisor/base_cycles 操作数依赖 window shape，尚未完全解开（关于确切 divisor 项为 LOW confidence）。

GOTCHA — kF32 在 latency table 中携带与 kBf16 相同的 per-format base（211），但 fp32 吞吐大约减半。吞吐惩罚存在于 precision multiplier 中，而不是 base 中：fp32 走 GetConvPrecision == 0 → ×2.0。只读 per-format latency table 而不读 precision multiplier 的重新实现者，会把 fp32 建模为 bf16 速度，并错误分块每个 fp32 conv。

算法 — 预算过滤器

VMEM footprint 既是次级键，也是硬过滤器。对每个候选，它会求和三个驻留 tile：

// VmemToUseForPotentialWindows — inline @0x131684c6, uses xla::Product (0x20cf5200)
long vmem_bytes(window):
    bytes = Product(kernel_tile) + Product(activations_tile) + Product(output_accumulator)
    vmem_unit = Target::MemUnitFromBytes(bytes)         // 0x1d61bfe0; quantized to chunk granules
    return vmem_unit

budget = min( DefaultScopedVmemBytes(target,module) * LoweringVmemLimitScalingFactor(cfg,env),
              GetHloScopedVmemBytes() )
         - already_used_module_vmem
```text

超出预算的候选会被拒绝，并记录 `"Giving up on potential kernel window ... because they exceed VMEM limits"`（或 `"... they'd run out of VMEM"`）。在 **scavenging mode** 中，重试计数器（`"compilation retry count: "`）会在重新编译时逐步放宽 scoped-VMEM cap。

### 算法 — 平局裁决阶梯

visitor callback 会收到 `(window6, cycles, MemUnit, granules, WindowConfig_CostModelType)`；消费者按以下顺序保留最佳值：

| 排名 | 键 | 规则 |
|---|---|---|
| 1 | MXU cycles | 最低者获胜 |
| 2 | VMEM fit（`MemUnit`） | cycles 相等时，浪费 scratchpad 最少的候选获胜 |
| 3 | `best_granules` | cycles+VMEM 相等时，tile granule 最匹配 `Target::ChunkGranules()` 的候选获胜（保留最大连续 DMA） |
| 4 | enumeration order | 最后手段：output-window 外层循环、kernel-window 内层循环 → 最先枚举到的相等候选获胜 |

`SetupBestConfig` 会记录冻结的选择：`"Chosen kernel window: <…> output window: <…> cycles <N> hlo <name> best_granules <G> max vmem: <M>"`。

> **NOTE —** 成本*组成部分*（cycles、MemUnit、granules、CostModelType）从 callback 签名（`IterateThroughWindowConfigs` 的 `std::function<void(...)>` visitor 类型）和 producer 恢复；conv consumer 是 `ConvolutionEmitter::ComputeWindowConfig`（`0x130dbf60`）以及 `SpatialMajorConvolution::ComputeWindowConfigInternal` 中的 `$_0` min-selection lambda，它们没有单独反汇编，因此 rank-2..4 顺序是根据传递的字段和 `best_granules` 日志字段推断的（MEDIUM）。Rank 1（cycles）为 HIGH。

### 数据表

| 调节项 / 表 | 地址 | 值 / 作用 |
|---|---|---|
| `precision_mult_table` | `0xa2c6050` | `{2.0, 1.0}` doubles；索引 = `GetConvPrecision` |
| `WindowConfig_CostModelType` | proto enum | `COST_MODEL_TYPE_INVALID`(0)，`COST_MODEL_TYPE_CLASSIC`(1)，`COST_MODEL_TYPE_ML_PGN_V1`(2) |
| `Target::LaneCount` | `0x1d60f400` | 128 |
| `Target::SublaneCount` | `0x1d60f300` | 8 |
| `Target::MemUnitFromBytes` | `0x1d61bfe0` | bytes → quantized `MemUnit` |
| `xla::Product(Span<long>)` | `0x20cf5200` | tile-byte product |

> **QUIRK —** callback 携带一个 `WindowConfig_CostModelType`，因此消费者知道应按 classic cycle/VMEM model 还是 `COST_MODEL_TYPE_ML_PGN_V1` learned-model 字段排序。在此构建中，learned 路径是数据表 fallback：没有随二进制发布 `LearnedCostModelClient`（见 [Learned Cost Model Client](../cost/learned-cost-model-client.md)），所以 `ML_PGN_V1` 会解析为同一组 classic 数字。重新实现可以只实现 `CLASSIC`，并且对此二进制保持行为精确。

---

## EmitFunctorEnum Dispatch

### 目的

分块被冻结后，`MatrixMultiplyAccumulateFunctor::operator()` 必须选择*如何*把 conv 维度映射到 MXU 的 lanes / sublanes / K-reduction 上。这个选择是 `convolution_util::EmitFunctor` 枚举中的 19 种策略之一。`GetEmitFunctorEnumAndLoweringDecisions` 是选择器（选择哪种策略）；`GetEmitFunctorFromEmitFunctorEnum` 是派发（enum → member-function pointer）。这是 dot/conv 版本的 instruction-selection table：固定 enum、switch，以及读取父对象中 shape/dtype 布尔值的决策树。

### 19 个值的枚举

`GetEmitFunctorFromEmitFunctorEnum`（`0x130e8de0`）是一个带有 `case 0..18` 的 `switch (ord)`；越界 default 会在 `matrix_multiply_accumulate_functor.cc` 第 `586` 行 FATAL。反编译确认 `case 18` 是最高 ordinal（19 种策略）。伴随的 `EmitFunctorToString`（`0x130e88a0`）是一个并行 `switch (ord)`，其 `case 0..18` 从内联 `.rodata` 常量构建调试名称；其 ordinal-3/4 字符串片段（`"...wSublane"`、`"...Lane"`）与下表按字节匹配，default 在 `matrix_multiply_accumulate_functor.cc` 第 `495` 行 FATAL。两者都是普通 switch，而不是按表索引派发；下表中的 member-function pointer 是直接从 dispatch 反编译中的 `case` 分支目标读取的。

| Ord | `EmitFunctor` 值 | member fn | MXU 策略 |
|----:|---|---|---|
| 0 | `kBatchGroupDepthwiseInputBatchInLanesOutputBatchInSublanes` | `0x130e8f40` | grouped depthwise：input batch 在 lanes，out 在 sublanes |
| 1 | `kBatchGroupDepthwiseInputBatchInSublanesOutputBatchInSublanes` | `0x130e9b80` | grouped depthwise：两个 batch 都在 sublanes |
| 2 | `kDepthwiseAllBatchInLanes` | `0x130ea960` | depthwise：所有 batch 在 lanes（每 channel 1 个 latch） |
| 3 | `kReduceWindowSublane` | `0x130eb860` | `kReduceWindow`-as-conv，sublane-major reduce |
| 4 | `kReduceWindowLane` | `0x130ebd80` | `kReduceWindow`-as-conv，lane-major reduce |
| 5 | `kDepthwiseInputBatchInLanes` | `0x130ec2a0` | depthwise：input batch 在 lanes |
| 6 | `kDepthwiseAllBatchInSublanesPacked` | `0x130ed3c0` | depthwise：所有 batch 在 sublanes，packed |
| 7 | `kDepthwiseInputBatchInSublanes` | `0x130ef2e0` | depthwise：input batch 在 sublanes |
| 8 | `kInputFeaturePackedInputBatchInLanes` | `0x130f01a0` | feature-packed K reuse，input batch lanes |
| 9 | `kInputBatchInLanes` | `0x130f0740` | classic：input batch 卷入 lanes |
| 10 | `kAllInputFeaturePackedInSublanesOutputBatchInSublanes` | `0x130f48a0` | 完整 K 在 sublanes，packed；out batch sublanes |
| 11 | `kAllInputFeatureInSublanesOutputBatchInSublanes` | `0x130f5d00` | 完整 K（≤128）在 sublanes，单个 matmul 覆盖 K |
| 12 | `kAllInputFeatureInSublanesOutputBatchInSublanesXposeReuse` | `0x130f7e20` | 同 11，复用转置 activations |
| 13 | `kOutputBatchInLanesKernelOutputFeatureInLanes` | `0x130fb360` | dual-lane：out batch + kernel out-feature 在 lanes |
| 14 | `kOutputBatchInLanesInputBatchInSublanes` | `0x130fee80` | transposed MAC：out batch lanes，in batch sublanes |
| 15 | `kOutputBatchInLanesKernelOutputFeatureInSublanes` | `0x131021a0` | hybrid：out batch lanes + out feature sublanes |
| 16 | `kAllBatchInSublanes` | `0x131055c0` | batch 打包进 sublanes（常见默认值） |
| 17 | `kInputBatchInSublanesOutputBatchInSublanesPacked` | `0x131064c0` | in+out batch 都在 sublanes，packed |
| 18 | `kOutputBatchInSublanes` | `0x13108b60` | output batch 在 sublanes（广义默认值） |

> **QUIRK —** enum 的 *ordinal* 顺序与其*分组*无关。Ordinals 0–7 是 depthwise / grouped / reduce-window 家族，8–18 是 dense dot/conv 核心，但选择器会以不同于 ordinal 的顺序到达它们（它可能先选 ord 18，再测试 ord 8）。重新实现应由**决策树**驱动，而不是 ordinal 序列；ordinal 只用于选择两个并行 switch 中的 `case` 分支（`GetEmitFunctorFromEmitFunctorEnum` → member fn，`EmitFunctorToString` → debug name）。

### 算法 — 选择器

`GetEmitFunctorEnumAndLoweringDecisions`（`0x1310c720`）返回 16-bit 值：低字节 = 上述 ordinal，高字节 = 第二个“lowering decision” bool，供 `UpdateLoweringStrategyWithWindowInfo` 使用。它从父 decision-state struct 和冻结 window 中读取 shape/dtype 布尔值。分支辅助函数已由二进制确认（`GetReduceWindowType`、`Target::LaneCount`/`SublaneCount` ×2、`ShouldPackInputFeature` ×2、`WindowCoversEntireEffectiveInputFeature`）：

```c
// GetEmitFunctorEnumAndLoweringDecisions — 0x1310c720
// returns (decision_byte << 8) | emit_ordinal
pair<EmitFunctor,bool> select(state, window):
    if state.is_batch_group_depthwise:                  // routes to grouped depthwise
        require !reduce_window && megacore_span.empty() && !output_batch_in_lanes
        if input_batch_in_lanes:  return (k0, decision=1) // BatchGroupDepthwise InLanes/InSub
        else:                     return (k1, decision=1) // BatchGroupDepthwise InSub/InSub
    else if state.is_depthwise:
        if input_batch_in_lanes:
            ord = kDepthwiseInputBatchInLanes (5)
            if output_batch_in_lanes:
                if reduce_window:
                    rw  = fusion_util::GetReduceWindowType(hlo)   // 0/1/2  @0x1454d4a0
                    ord = 4 - rw                          // 4=ReduceWindowLane, 3=…Sublane
                else:
                    ord = kDepthwiseAllBatchInLanes (2)
            return (ord, ...)
        else:                                             // input batch in sublanes
            if feature_packed && !out_lanes && kernel_of_lanes:
                return kDepthwiseAllBatchInSublanesPacked (6)
            else:
                return kDepthwiseInputBatchInSublanes (7)
    else:                                                 // ---- DENSE dot/conv core ----
        if input_batch_in_lanes:
            if output_batch_in_lanes:                     // both in lanes
                pack = ShouldPackInputFeature(out_feature, target)  // @0x13194220
                return pack ? kInputFeaturePackedInputBatchInLanes (8)
                            : kInputBatchInLanes (9)       // decision=0
            else:                                         // input lanes, output sublanes
                return xpose_reuse ? k13 : <14/15 family>
        else:                                             // input batch in sublanes
            if output_batch_in_lanes:                     // out lanes, in sublanes
                ord = kAllBatchInSublanes (16)            // decision=1
                if out_feature >= SublaneCount(8):
                    if WindowCoversEntireEffectiveInputFeature(window):  // @0x1315b2c0; K fits one window
                        ord = (spatial_extent < 2 || all_dims_static)
                                ? kAllInputFeatureInSublanesOutputBatchInSublanes (11)
                                : kAllInputFeatureInSublanesOutputBatchInSublanesXposeReuse (12)
                    else:                                 // K-tiled (K>128)
                        ord = ShouldPackInputFeature(...) ? k10 : k16   // decision=0
                return (ord, ...)
            else:                                         // both in sublanes
                ord = kOutputBatchInSublanes (18)         // decision=1, the broad default
                if out_batch >= LaneCount(128) && (out_batch % 128) && pack_eligible:
                    ord = kInputBatchInSublanesOutputBatchInSublanesPacked (17)
                return (ord, ...)

关键辅助函数：

fusion_util::GetReduceWindowType(hlo)（0x1454d4a0）→ 0/1/2；ord = 4 - type 映射到 ReduceWindowLane(4) / ReduceWindowSublane(3)。
convolution_util::ShouldPackInputFeature(out_feature, target)（0x13194220）→ 当且仅当 input-feature count 小于 Target::SublaneCount()（==8）时为 true，即 sub-8 K 归约会被 feature-packed。dense both-in-lanes 分支中的 xor $0x9,al 映射 {pack?1:0} → {8:9}。
WindowCoversEntireEffectiveInputFeature(window)（0x1315b2c0）→ 当且仅当一个 MXU window 吸收完整 K 维度（无 K-tiling）时为 true。

GOTCHA — 高 decision byte 不是策略。它是第二个 bool（大多数路径为 1，在 packed-K 和 dual-lane 分支为 0），传给 UpdateLoweringStrategyWithWindowInfo(.., decision)，后者为所选 window 修补 ConvolutionLoweringStrategy，最可能是 transpose-reuse / sublane-tile toggle。这个 bit 已恢复；人类可读字段名尚未恢复（LOW）。重新实现者必须把它传到 strategy-patch 步骤，即使其精确含义在这里未解决。
NOTE — 实用解码：普通 bf16 batched dot，K≤128，两个 batch 都在 sublanes，K 折叠到一个 window → ord 11（或带 transpose-reuse 时为 12）。同一 dot 在 K>128（K-tiled）时 → ord 18（kOutputBatchInSublanes），即规范默认值，并跨 pass 累加 K（下一节）。

MatrixMultiplyAccumulateFunctor 循环

目的

这是内部 code generator：给定冻结 window 和选定的 EmitFunctor，它按 accumulator tile 发出实际的 LLO matprep / matmul / matres 流，将 K>128 处理为 multi-pass accumulation，并运行 output-fusion epilogue。functor 每个 chunk 构造一次；operator() 驱动它。

算法

// MatrixMultiplyAccumulateFunctor::operator() — 0x1310cd80
void operator()(builder):
    ret      = GetEmitFunctorEnumAndLoweringDecisions()   // 0x1310c720
    ord      = ret & 0xff;  decision = (ret >> 8) & 0xff   // split at 0x1310cf91
    functor  = GetEmitFunctorFromEmitFunctorEnum(ord)      // 0x130e8de0 → member fn ptr
    name     = EmitFunctorToString(ord)                    // 0x130e88a0 (debug)
    UpdateLoweringStrategyWithWindowInfo(.., decision)     // 0x13167e80 (patch strategy)

    ResetRegistries(builder)                               // clear latch-group + rotation maps
    for acc_window in 0 .. AccWindowCount():
        ctx = InterLatchGroupContext(acc_window)
        (this->*functor)(acc_window, builder)              // issues loads + matprep/matmul/matres
        PossiblyDoOutputFusion(acc_window, builder)        // bias/activation epilogue
    MaterializeLatchGroups(tile_emitter_map_)              // flush deferred latch-group emitters
```text

### K>128 multi-pass accumulation

128×128 阵列每个 pass 最多归约 K（input-feature）维度中的 128 个元素。对 K>128，lowering 会分块 K，并把部分 `matres` 结果累加进 VMEM accumulator：

```c
// per output tile, K tiled into n passes:
AccumFirstZero(off0, ..., builder)    // 0x13124e80 — first K-tile: matmul, matres → acc (no add)
for k in 1 .. n:
    Accumulate(..., builder)          // 0x13123f20 — matres → tmp; VaddF32/VaddS32(acc, tmp)
PossiblyDoOutputFusion(acc, builder)  // 0x... — bias add / activation after K fully reduced

Accumulate 对 fp accumulation 使用 LloRegionBuilder::VaddF32，对 int 使用 VaddS32。在 schedule 允许处，后续 MxuLmrTransform / MxuResultAddJoin pass 会合并 matres+Vadd 链。

text

LLO shape per output tile (K = 256 → 2 passes, bf16):
  ; --- K-tile 0 (AccumFirstZero) ---
  vlatch        rhs_K0   (GainLatchMode _PACKED_BF16)   ; stationary weights → MXU latch
  vmatprep.mubr lhs_K0   (MatpushTarget MSRA)           ; moving activations → staging
  vmatmul                (DoneWithGains NORMAL, fmt kBf16)
  vmatres       acc                                     ; first matres → accumulator (no add)
  ; --- K-tile 1 (Accumulate) ---
  vlatch        rhs_K1   (_PACKED_BF16)
  vmatprep.mubr lhs_K1   (MSRB)                          ; alternate staging register
  vmatmul                (NORMAL, kBf16)
  vmatres       tmp
  VaddF32       acc, tmp                                ; K-accumulate
  ; --- epilogue ---
  DoOutputFusion(acc)                                   ; bias / activation if fused
```text

> **NOTE —** 相邻 matprep 会在 `MSRA → MSRB` 之间交替（double-buffered moving-operand staging register），使 MXU 能在当前 K-tile drain 时 stage 下一个 K-tile；*数值* GMR index 和 A/B 交替是这个 allocator 的 per-tile assignment，而不是上游 `GetGainLatchModeAndScalingFactor` selector。latch/matmul/accum 三元组由 `ReorderLatchesMatmulsAndAccums(lhs_mode, rhs_mode, lhs_ty, rhs_ty, bool)`（`0x1311f880`）排序，它会按 latch 调用 `GetPackedGainLatchMode`，并调用 `GroupLatchGroupsForX4`（`0x1311c8a0`）把 4 个 latch 成组为一个 x4-packed issue。slot 编码见 [Matprep、IAR 和 Latch 子槽](../isa/slot-matprep-iar-latch.md)。

### 函数映射

| 函数 | 地址 | 作用 |
|---|---|---|
| `operator()` | `0x1310cd80` | per-chunk MXU sequence driver |
| `GetEmitFunctorEnumAndLoweringDecisions` | `0x1310c720` | 19-way strategy selector |
| `GetEmitFunctorFromEmitFunctorEnum` | `0x130e8de0` | enum → member fn ptr |
| `EmitFunctorToString` | `0x130e88a0` | debug name（parallel jump table） |
| `AccumFirstZero` | `0x13124e80` | 第一个 K-tile，accumulator seeded zero |
| `Accumulate` | `0x13123f20` | 后续 K-tile，`VaddF32`/`VaddS32` |
| `ReorderLatchesMatmulsAndAccums` | `0x1311f880` | (latch, matmul, accum) ordering + gain staging |
| `CreateMatprepOrLatch` | `0x1311bf00` | 发射 latch / matprep（携带 `GainLatchMode`、`PrimitiveType`） |
| `CreateMatprepOrMatmul` | `0x1311bb40` | 发射 matprep / matmul（携带 `MatmulMode`、`MatmulDataFormat`） |
| `GroupLatchGroupsForX4` | `0x1311c8a0` | 把 4 个 latch 组合成一个 x4 packed issue |

---

## ConvMatmulModes — 模式对比较器

### 目的

发射前，`GetMatmulModes` 会构建 `{lhs_mode, rhs_mode}` 对的候选列表（两个操作数的 stationary/moving feed role），然后排序，使最便宜的可行对被优先消费。比较器是单键加权偏好。这是一个小但容易漏掉的部分：*mode* 选择（哪个操作数是 stationary、transposed、packed）与上面的*策略*选择以及下面的 *dtype format* 选择正交。

### 算法

`ConvMatmulModes` 是一个 2-byte struct `{uint8 lhs_mode; uint8 rhs_mode;}`，每个 byte 都是 `xla::jellyfish::MatmulMode` ordinal（0..15）。比较器按字节精确如下（在 `0x130e12a0` 反编译）：

```c
// ConvMatmulModes::operator< — 0x130e12a0 (39 B)
// weight table W[16] @ 0xae0f480
bool operator<(const ConvMatmulModes& a, const ConvMatmulModes& b) {
    return (unsigned) (W[a.lhs_mode] + W[a.rhs_mode])
         < (unsigned) (W[b.lhs_mode] + W[b.rhs_mode]);   // setb — unsigned-less on the int sums
}

反编译字面读取为 dword_AE0F480[*a1] + dword_AE0F480[a1[1]] < dword_AE0F480[*a2] + dword_AE0F480[a2[1]]。16-entry 权重表（0xae0f480，越低 = 越便宜 / 越偏好）：

MatmulMode ordinal	0	1	2	3	4	5	6	7	8	9	10	11	12	13	14	15
weight	5	4	3	2	1	40	40	30	30	20	10	10	40	40	40	40

因此偏好层级是：{0..4}（权重 5..1，便宜的 stationary-feed 组，ord 4 最便宜）、{10,11}（=10，中等）、{7,8,9}（30/30/20），以及 {5,6,12,13,14,15}（=40，最昂贵的 transposed / multi-quadrant feeds）。

GetMatmulModes()（0x130df600）将每个操作数的 LHS mode list 与 RHS mode list 做笛卡尔积（每个 list 来自 0x130dfbe0 处的 GetMatmulModes(long operand)，它会调用 GetConvPrecision），跳过退化 {2,2} 对，然后按 operator< 做 std::stable_sort（sorter 0x130dfd80）。反编译确认了 stable_sort + operator< 调用。因为排序是稳定的，当两个 pair 总和相等时，次级平局规则是“先生成者获胜”，生成顺序是 LHS-mode-outer × RHS-mode-inner。

QUIRK — 比较器是一个单一加法键，而不是字典序 (lhs, rhs)。两个很不同的 pair {4,10} 和 {1,2} 分别求和为 11 和 3，该模型确实把 mode pair 当作一个成本数，因此如果总和更低，那么 {stationary=cheap, moving=mid} 可能输给一个 stationary 更贵但 moving 更便宜的 pair。symbolizeMatmulMode（0x13e4e660）中 MatmulMode ordinal 的名称（stationary/moving/transposed/packed roles）无法干净提取；权重按字节精确，名称未解决（LOW）。成本侧 modifier 视图见 Matmul Mode Modifiers。

MxuLatchPacker — PackLatches

目的

conv emitter 会为每个操作数行生成一个 latch，但 bf16/int8 操作数会把两个元素行打包进单个 MXU latch，从而将 latch 指令数减半。PackLatches 是执行这种打包的 post-emission pass，会对相邻 same-mode latch 贪心处理。它在 MxuSequenceCollector 和 MxuAssigner（把 sequence 绑定到物理 MXU bank）之后、LMR fusion 之前运行。

入口点

text

LLO post-emission MXU stack (after MatrixMultiplyAccumulateFunctor):
  MxuSequenceCollector  ── gather (matprep,matmul,matres) runs into MxuSequence objects
  MxuAssigner           ── bind sequences to physical MXU bank; SetLatchIndices; MSR bounce
  MxuSequenceTrimmer    ── drop unused matres/matprep
  PackLatches (0x10f726c0) ──────── greedy same-mode adjacent latch packing   [THIS UNIT]
  CombineSequencesToFitMxuSize / PackSingleSequence (mxu_diagonal_packer.cc) ── quadrant co-residence
  MxuLmrTransform       ── fuse (latch,matmul,matres) → LMR; partial-quadrant rewrite
  MxuAccumulation / MxuResultAddJoin / MxuLatencyBalancing
```text

### 算法

```c
// xla::jellyfish::PackLatches(Span<unique_ptr<MxuSequence>>, Target, CycleTable,
//                             PackLatchesOptions) — 0x10f726c0 (8559 B), mxu_latch_packer.cc
void PackLatches(sequences, target, cycle_table, opts):
  for each MxuSequence run, walk latch instructions in pairs (cur, nxt):
    m_cur = cur.latch_mode()                 // LloInstruction::latch_mode  @0x1d4e7500
    m_nxt = nxt.latch_mode()
    if m_cur != m_nxt:                        // "matching pair failed because latch modes don't match"
        emit cur UNPACKED; advance by 1
        continue
    // packable pair — also check MSR target compat ("merging matpushes")
    check cur.matrix_staging_register() vs nxt.matrix_staging_register()  // @0x1d4e7b80
    // pack two operand vregs into one packed vreg:
    if bf16: packed = VpackCBf16(lo, hi)      // @0x1d5669a0
    else    : packed = VmpackCLow(lo, hi)     // @0x1d55d520  (int8/byte)
    CreatePackedVlatchOrVlatchprep(builder, opcode, packed,
                                   packed_GainLatchMode, mask)   // @0x10f74c20
    // at the consuming matmul, split back:
    hi = VunpackUpperCB8ToBf16(packed)        // @0x1d5696e0
    lo = VunpackLowerCB8ToBf16(packed)        // @0x1d569600
    align via VslaneRotateAZ(v, SublaneCount) // @0x1d54ee00
    advance by 2
  // odd trailing latch → "skipping last latch" (left unpacked)
  CHECK(num_latches == num_latchpreps)        // matrix_register.h pairing invariant
  LloMutator::SubstituteWithRegionAndDestroy(latch, packed_region)

guard helper：MaskAndVregForLatchOrLatchprep（0x10f74aa0）为部分有效 latch 构建 predicate mask + source vreg；VectorS8FeedingF32Latch（0x10f74840）检测必须避免二次打包的 int8-operand-into-fp32-latch 情况。

函数映射

函数	地址	作用
`PackLatches`	`0x10f726c0`	贪心 same-mode 相邻 latch 打包器
`CreatePackedVlatchOrVlatchprep`	`0x10f74c20`	发射单个 packed latch
`MaskAndVregForLatchOrLatchprep`	`0x10f74aa0`	predicate mask + source vreg
`VectorS8FeedingF32Latch`	`0x10f74840`	int8→fp32-latch guard
`LloRegionBuilder::VpackCBf16`	`0x1d5669a0`	pack bf16 pair → one vreg
`LloRegionBuilder::VmpackCLow`	`0x1d55d520`	pack int8/byte → one vreg
`VunpackUpper/LowerCB8ToBf16`	`0x1d5696e0` / `0x1d569600`	consumer 处的 result split
`LloInstruction::latch_mode`	`0x1d4e7500`	packing equivalence key
`LloInstruction::matrix_staging_register`	`0x1d4e7b80`	MSR compat check

GOTCHA — PackLatches 是邻接贪心，不是最优：只有当 (cur, nxt) 的 latch_mode 完全匹配时才打包，且从不跨 mode boundary 重排。一个 [bf16, bf16, int8, bf16] run 会打包第一对，未打包发射 int8，并让尾部 bf16 保持未打包（"skipping last latch"），即使重排可以配对两个 bf16 端点。执行全局 min-cost pairing 的重新实现会偏离此二进制的输出。num_latches == num_latchpreps CHECK 是硬不变量：每个 packed latch 都必须保持 latch/latchprep 数量平衡。

MxuDiagonalPacker — 兄弟 quadrant 打包器

CombineSequencesToFitMxuSize（0x10f6eb80）和 PackSingleSequence（0x10f6fd20）不是 latch packer 的一部分；它们位于 mxu_diagonal_packer.cc 中，把短 MXU sequence 打包进 128×128 阵列未使用的 quadrant：

CombineSequencesToFitMxuSize：block-diagonal packing（策略 "PackToDiagonal" / "ShardToDiagonal"）：把操作数 A 锁存进 upper-left quadrant，把操作数 B 锁存进 lower-right quadrant，使一组 matmul steps 计算两个独立乘积。兼容性 CHECK："Incompatible matmul format"、"Incompatible number of matmuls"、"Incompatible number of latches"、"Incompatible size for matres sequences"。
PackSingleSequence：single-quadrant packing（策略 "PackToSingleQuadrant"）：把 K 和 N 适配 64×64（或 64×128）quadrant 的一个 sequence 打包进一个角。要求 matmul count 可被 2 整除（"Rejecting, number of matmuls must be divisible by 2"）且 format uniform。

它们共同为本会低效占用阵列的窄 K/N tile 将 MXU 吞吐翻倍。

Per-dtype Format — GetMatmulDataFormat

目的

与 strategy 和 mode 正交，per-matmul MatmulDataFormat 为每种 dtype 固定实际的 MXU op 变体。GetMatmulDataFormat（0x1307be40）是一个两阶段决策：strategy-packing-bit 优先，然后是 PrimitiveType jump table。

算法

// convolution_util::GetMatmulDataFormat(PrimitiveType prim,
//   const ConvolutionLoweringStrategy& s, const HloInstruction&, const Target&) — 0x1307be40
int GetMatmulDataFormat(prim, s, hlo, target):
    // STAGE 1 — packing-bit precedence (int8/bf16 packed paths short-circuit)
    if any_bf16_pack_flag(s):  return 1 or 2          // bf16 packed-format codes
    if any_x8_pack_flag(s):    /* continue, format 2 */
    // STAGE 2 — dense dtype dispatch via jt[prim-2] @ 0xae0d6f4 (22 entries, cmp 0x15)
    switch (prim - 2):
        int8     -> 6
        int4     -> 6 (signed) / 8 (unsigned)
        F8E5M2   -> 5 / 7
        F8E4M3Fn -> 3 (native) / 9 (converted)
        fp8 fnuz -> 10
        F32      -> 4
        default  -> FATAL (convolution_util.h line 297)
```text

| dtype | `MatmulDataFormat` | packed-VLATCH path | latch GainLatchMode group |
|---|---|---|---|
| bf16 | 1（`kBf16`） | `VpackCBf16` pairs | `_PACKED_BF16` |
| fp32 | 4（`kF32`） | none（½ throughput，×2.0 precision） | `_NO_XPOSE_F32`/`_HI`/`_LOW` |
| int8（x8） | 6 | `VmpackCLow` ×8 | `_NO_XPOSE_S8`/`_U8` |
| int4（x4） | 6/8 | x4 packed VLATCH | `_NO_XPOSE_S4`/`_U4`/`_NIBBLE` |
| fp8 E4M3Fn | 3/9 | `_PACKED_E4M3FN` | `_F8E4M3*_TO_BF16` |
| fp8 E5M2 | 5/7 | `_PACKED_E5M2` | `_F8E5M2_TO_BF16` |

> **NOTE —** jump-table *targets*（`0x1307bec0`…`0x1307bf1f`）按字节精确，但此构建的 `PrimitiveType` ordinal 是 libtpu 内部枚举（`prim - 2` 索引），不是上游 XLA 值。上面的 dtype 标签是根据结果 format code 与 `MatmulDataFormat` 枚举（1=bf16，4=F32，5=F8E5M2，6=int8，等等）的匹配推断的，因此 int4/fp8 行为 MEDIUM。`GetConvPrecision`（`0x131916e0`）会将 `primitive_util::SignificandWidth` 与 `PrecisionConfig` 比较，并返回供 ×2.0/×1.0 cycle multiplier 使用的 0/1 索引。
>
> **QUIRK —** int8 和 int4 **没有**独立的 `MatmulDataFormat` *enum members*，它们复用 format code 6/8，并通过 `ConvolutionLoweringStrategy` packing booleans 选择 x8/x4 *opcode variant*。因此 dtype 路径被拆分到两套机制中：fp 变体走 `MatmulDataFormat` enum，integer 变体走 strategy bool-vector。把所有 dtype 都建模为 enum members 的重新实现者，将没有地方放 int4 unsigned-vs-signed 区分，因为它存在于 strategy flags 中。

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `DotCanonicalizer`（HLO pass） | 在上游把 `kDot` → `kConvolution`；见 [RaggedDot 和卷积几何降低](raggeddot-convolution.md) |
| `FusedSpatialMajorConvolution` | 当 conv 是 output-fusion root 时使用的子类；按 chunk 发射 bias/activation epilogue |
| `MxuAssigner` | 在 `PackLatches` 前把 `MxuSequence` 绑定到物理 MXU bank 和 latch index |
| `MxuLmrTransform` | 在 packing 后把 `(latch, matmul, matres)` 三元组融合为 `kVectorMatmulLmr` |
| `RotatedPincerEmitter` family | flag-gated 的 collective matmul alternative outer wrapper；针对 per-shard MXU emission 调用标准 MMA functor（不是独立 MXU path） |

---

## 交叉引用

- [TPU 编译器](overview.md) — 此 lowering 所处的五阶段 dialect descent
- [tpu → LLO ODS 降低](tpu-to-llo-ods.md) — MXU build factories 以及这些 op 携带的四个 register enum（`MatmulMode`/`GainLatchMode`/GMR/MSR）；gain/staging 表
- [RaggedDot 和卷积几何降低](raggeddot-convolution.md) — 本页上游的 `kDot` → `kConvolution` 重写和维度编号映射
- [Fusion Patterns](fusion-patterns.md) — 在 conv 之上形成 bias/activation chain 的 output-fusion
- [Mosaic VectorLayout](mosaic-vectorlayout.md) — alternative Mosaic/Pallas matmul layout path
- [MXU Slot](../isa/slot-mxu.md) — 此 lowering 发射的 matmul op 的 bundle-slot bit encoding
- [Matprep、IAR 和 Latch 子槽](../isa/slot-matprep-iar-latch.md) — latch/matprep op 以及 MSR/GMR 字段的 slot encoding
- [MXU 延迟概览](../cost/mxu-latency-overview.md) — tile-cost 公式在其上叠加 ×1/×2 precision multiplier 的 per-format base-cycle table
- [Matmul Mode Modifiers](../cost/matmul-mode-modifiers.md) — `ConvMatmulModes` 比较器排序所依据的 `MatmulMode` 轴的 cost-side view
- [Learned Cost Model Client](../cost/learned-cost-model-client.md) — 为什么此构建中的 `ML_PGN_V1` cost-model type fallback 到 classic data tables

Dot / Conv → MXU 降低 ​

摘要 ​

降低入口 — ConvolutionEmitter::Create ​

目的 ​

入口点 ​

注意事项 ​

Tile-Cost Comparator ​

目的 ​

入口点 ​

算法 — 预算过滤器 ​

MatrixMultiplyAccumulateFunctor 循环 ​

目的 ​

算法 ​

MxuLatchPacker — PackLatches ​

目的 ​

入口点 ​

函数映射 ​

MxuDiagonalPacker — 兄弟 quadrant 打包器 ​

Per-dtype Format — GetMatmulDataFormat ​

目的 ​

算法 ​

Dot / Conv → MXU 降低

摘要

降低入口 — ConvolutionEmitter::Create

目的

入口点

注意事项

Tile-Cost Comparator

目的

入口点

算法 — 预算过滤器

MatrixMultiplyAccumulateFunctor 循环

目的

算法

MxuLatchPacker — PackLatches

目的

入口点

函数映射

MxuDiagonalPacker — 兄弟 quadrant 打包器

Per-dtype Format — GetMatmulDataFormat

目的

算法