Internal Wiki

主题

融合成本模型 ​

本页中的每个 offset、值和地址都从 libtpu-0.0.40-cp314 wheel 中的 libtpu.so 按字节精确读取（BuildID md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，未 strip — 每个符号都是 demangle 后的 C++ 名称）。section map：.text/.rodata VMA == file offset。所有地址都是虚拟地址。其他 libtpu build 会不同。

摘要 ​

TPU 指令融合是优先级驱动的，而不是一次自底向上的 sweep。计算中的每个 producer 都由成本函数赋予一个 double 优先级，插入有序队列，并先融合最高优先级的 producer；每次融合后，受影响的 neighbour 会重新打分。本页记录队列可运行的两个可互换成本模型的数值打分 — 浮点优先级公式、它消费的按 op 成本权重，以及覆盖分数的 VMEM 硬门禁 — 还包括用于给 multi-output（sibling）fusion 排序的独立 profit 数。它不记录决定一次融合在结构上是否合法的模式谓词；这些内容 — ShouldFuseImpl 中约 30 个 FusionDecision 拒绝点、slice-like / output-fusion / duplicate-expensive 门禁、custom-call registry hook — 位于 Fusion Patterns。二者关系是：谓词级联是过滤器（合法/非法）；成本模型是在幸存者上的排序器（有多好），二者共同咨询一个共享硬门禁（FusionWouldExceedVmemCapacity）。

这里有两个成本模型，而且它们确实是不同函数，由 dispatcher 按 producer（以及按 edge）选择。current 模型 — 默认模型 — 将融合定价为节省的 HBM 字节（以 TensorCore cycles 表示）减去新增计算（粗粒度 opcode 权重阶梯），并按该融合会复制多少昂贵的 conv/reduce-window op 进行缩放。bundle-aware 模型定价实际节省的 VLIW bundle cycles（total_unfused − total_fused），它能捕获线性模型看不到的跨功能单元 packing。二者都馈入同一个优先级队列，并遵守同一个 VMEM 硬门禁。

熟悉 LLVM 的读者应把握一个类比和一个差异。类比：优先级循环是一个贪心 list scheduler，其“ready”集合是 producer 队列，而优先级函数就是这个成本模型 — 很像 LLVM MachineScheduler 的 SchedBoundary 选择最佳 SUnit。差异：分数中没有 register-pressure 项。Register/VMEM pressure 是一个二元准入门禁（FusionWouldExceedVmemCapacity），既在优先级预 pass 中检查，也在谓词级联中再次检查；能放下的候选仅按节省的 bytes/cycles 排序，放不下的候选会被直接拒绝（优先级 -1.0），绝不会被软惩罚。

重新实现的契约如下：

  • 优先级键是插入 std::map 的 3-tuple (double primary, double secondary, long tie)；队列取出最大键（最高优先级优先）。两个 double 在插入前都用 fatal CHECK 捕获 NaN。
  • Current-model 分数：priority = mem_reduce − compute × conv_rw_count。 mem_reduce 是以 TC cycles 计的 HBM 字节节省；compute 是 opcode 权重阶梯结果；conv_rw_count 是该融合复制的 Conv + ReduceWindow op 数量，因此复制昂贵 op 会按比例受罚。
  • Bundle-aware 分数：priority = total_unfused − total_fused，二者都是来自 ResourceVector 机制的 bundle cycles；这是唯一能看到 VLIW packing 的模型。
  • 三个优先级哨兵： -1.0 = do-not-fuse；FLT_MAX（current）/ 100.0 boost（bundle）= must-fuse；must-fuse 的 long tie-break 是 0x3ffffffffffffffe（current）/ 100（bundle）。
  • VMEM 是硬门禁，不是分数项。 任何 user 的融合若会超出 VMEM，就强制 producer 优先级为 -1.0。
  • 计算权重阶梯是粗粒度 opcode switch，映射到少数标量 tier（默认 1.0、4.0、10.0、42.0，加上 conv-flop 和 fusion-recurse escape），每个都乘以 Target::ChunksIn(shape)。
  • Multi-output（sibling）fusion 有自己的 profit 数 = 共享操作数只读取一次而非每个 sibling 读取一次所节省的字节数，在独立优先级队列中排序，并受每个 reduce 输出 4 MiB cap、max-operands cap 和 HBM-pressure cap 约束。
  • 打分代码与代际无关。 每代行为只通过 Target 常量、按代的 CycleTable latency 和 flag 默认值进入 — 公式和 switch 是一套共享实现。
优先级 dispatcherTpuPriorityFusionQueue::CalculateProducerPriority(HloInstruction*) @ 0x1308fa20
Current 模型…::CalculateProducerPriorityWithCurrentCostModel @ 0x13096160
Bundle-aware 模型…::CalculateProducerPriorityWithBundleAwareCostModel @ 0x130954c0
入队 + NaN guard…::EnqueueToProducerPriorityQueue(tuple<double,double,long>, HloInstruction*) @ 0x1308fb20
计算项…::NormalizedComputationCost(HloInstruction*, long) @ 0x130989a0
内存项…::GetNormalizedMemoryCostReductionIfFusing @ 0x13099700
VMEM 硬门禁CostModel::FusionWouldExceedVmemCapacity(HloInstruction*, HloInstruction*) @ 0x130c4a80
Bundle cycles…::GetHloCycles @ 0x13097a00；CostModel::GetCyclesIfFused @ 0x130aba40
MOF profitTpuMultiOutputFusion::GetProfit(HloInstruction*, HloInstruction*) @ 0x110dd0a0
源码文件tpu_instruction_fusion.cc（priority）；cost_model/cost_model.cc（cycles, VMEM）
置信度CONFIRMED（字节锚定），除非某行另有说明

分数在 Pass 中的位置 ​

该 pass 是 TpuInstructionFusion，一个 xla::InstructionFusion 子类。基类驱动通用优先级循环；TPU 子类提供两件事：队列（GetFusionQueue @ 0x13083c40 构造 TpuPriorityFusionQueue）和合法性谓词（ShouldFuseImpl，记录于 Fusion Patterns）。本页的成本模型就是队列的优先级函数。

text
InstructionFusion::Run (base)
  ├─ GetFusionQueue(comp)                      ── build TpuPriorityFusionQueue
  │     └─ for each producer:
  │           CalculateProducerPriority(p)     ── THIS PAGE: the score
  │           EnqueueToProducerPriorityQueue   ── insert (double,double,long) → map
  └─ loop:
        DequeueNextInstructionAndOperandsToFuseInOrder()  ── pop LARGEST key
        ├─ ShouldFuseImpl(consumer, operand_idx)  ── FUSION-PATTERNS PAGE: legal?
        │     └─ shares the VMEM hard gate with the score (below)
        ├─ if legal: Fuse(), then
        │     OnFusingInstruction / InvalidateCachedCostModelState
        │     re-score affected neighbours (priorities change once an edge internalises)
        └─ else: drop / mark -1.0
```text

重新实现必须保留两个后果。首先，顺序是**全局贪心并增量重新打分**：节省最多 bytes/cycles 的融合先发生，而这次融合会改变其 neighbour 的优先级（internalised edge 不再节省其 HBM 流量），因此这些 neighbour 会在下一次 dequeue 前重新打分。其次，成本模型和合法性谓词是**彼此分离但共享且仅共享一个门禁的关注点** — `FusionWouldExceedVmemCapacity`。分数从不编码“这是否合法”；它只在合法者之间排序。如果把两者合并，要么会给非法融合排序，要么会拒绝合法但低价值的融合。

主融合循环没有固定指令数 cap；它一直运行到队列为空 — 每个 producer 要么已经融合，要么被标记为 `-1.0` / boundary。自然停止条件是“没有剩余候选既有正优先级又能放入 VMEM。”（Multi-output fusion 有显式预算 cap；见其章节。）

---

## 优先级键 — 有序 Map 中的 3-Tuple

`EnqueueToProducerPriorityQueue` (@ `0x1308fb20`) 插入到 `std::map<std::tuple<double,double,long>, HloInstruction*>`。mangled 签名无歧义：`…EnqueueToProducerPriorityQueueENSt3__u5tupleIJddlEE…` — `tuple<d,d,l>`。`std::map` 升序排序；循环 dequeue **最后一个**（最大）元素，因此最高优先级先融合。

```c
// EnqueueToProducerPriorityQueue(tuple<double,double,long> key, HloInstruction*)  @ 0x1308fb20 (decompiled)
// — both doubles are NaN-trapped before insertion:
__asm { vucomisd xmm0, xmm0 }                    // get<0>(key) == itself ?
if (parity)                                       // NaN
  LogMessageFatal(".../tpu_instruction_fusion.cc", 1678, "!std::isnan(std::get<0>(key))");
__asm { vucomisd xmm0, xmm0 }                    // get<1>(key) NaN-check (second double)
// ... then __tree __emplace_unique into the map keyed on the tuple.
tuple 槽类型含义
get<0> (primary)double优先级分数（下面的 current 或 bundle 模型公式）；排序键
get<1> (secondary)double优先级值的重新存储副本（稳定性 re-score）；用作相同 primary 的区分器
get<2> (tie)long确定性 tie-break：InputSizeAt 累加 / chunk count（current）；current-model must-fuse boost 用 0x3ffffffffffffffe，bundle-model must-fuse boost 用 100；确保跨运行的稳定顺序

NaN guard 是 tpu_instruction_fusion.cc:1678 处的硬 CHECK（以及针对第二个 double 的第二个 vucomisd）。任何可能产生 NaN 的公式 — 例如 bytes-per-cycle 转换中的 0/0 — 都会中止编译，而不是破坏 map 顺序。因此重新实现必须保证分数有限，或同样断言。

NOTE — 队列选择最大键。 因为 std::map 是升序，“最高优先级”= 最后一个元素。打分为 -1.0 的 producer 排在底部，等效于“do not fuse”。打分为 FLT_MAX 的 producer 排在顶部 — must-fuse boost。

Dispatcher — Current vs Bundle-Aware ​

CalculateProducerPriority (@ 0x1308fa20) 是基类循环调用的单一入口。它按 producer 和按可融合 edge 选择模型：

c
// CalculateProducerPriority(HloInstruction* producer)  @ 0x1308fa20 (decompiled, condensed)
if (producer == nullptr) {                                  // queue init / no producer
  if (ShouldAlwaysUseBundleAwareCostModel(flags))           // @0x130d3cc0
    return WithBundleAwareCostModel(producer);
  return WithCurrentCostModel(producer);
}
if (!producer->IsOutputFusion()) {                          // opcode[+12] != kFusion(0x81)
  // not an output-fusion ROOT → priority -1.0, UNLESS a force-priority bit is set:
  if (!(producer[+0xd] & 0x8) &&
      !any_user(u : u.opcode==0x81 /*kFusion*/ && (u[+0xd] & 0x8)))
    return -1.0;                                            // do-not-fuse sentinel
}
// per-edge: if ANY qualifying user wants the bundle model, use it:
return ShouldUseBundleAwareCostModel(user, flags)           // @0x130d3c00
           ? WithBundleAwareCostModel(producer)
           : WithCurrentCostModel(producer);
```text

第二个冗余 dispatch 位于 current 模型**内部**（`0x13096177`）：`if (flags[+0xc] == 1) tail-jump to bundle-aware` — 即 flag offset `+0xc` 处的 `xla_tpu_use_bundle_aware_cost_model_for_fusions` 字节会强制使用 bundle 模型，即使是通过“current”路径进入。因此模型选择是分层的：全局 flag、按 producer 的 output-fusion 门禁、按 user 的 edge 查询。**默认是 current 模型。**

| 选择器 | 函数 / 字段 | 效果 |
|---|---|---|
| global force | `ShouldAlwaysUseBundleAwareCostModel` @ `0x130d3cc0` | 所有 producer 使用 bundle 模型 |
| flag byte | `flags[+0xc]` (`xla_tpu_use_bundle_aware_cost_model_for_fusions`) | 在 current 入口内部强制 bundle 模型 |
| per-edge | `ShouldUseBundleAwareCostModel(user)` @ `0x130d3c00` | 任何符合条件的 user 将 edge 选入 bundle 模型 |
| output-fusion gate | `producer->IsOutputFusion()` + force-bit `[+0xd]&0x8` | 没有该 bit 的非 output-fusion producer → `-1.0` |

---

## Current 成本模型 — `mem_reduce − compute × conv_rw_count`

`CalculateProducerPriorityWithCurrentCostModel` (@ `0x13096160`) 是默认排序器。它计算一个线性估计：节省的 HBM 字节减去复制的计算。反编译控制流：

```c
// CalculateProducerPriorityWithCurrentCostModel(producer)  @ 0x13096160 (decompiled, condensed)
users = GetFusibleUsers(producer);                                   // 0x13097220
// (1) VMEM HARD GATE — pre-emptive. If ANY user's fusion would exceed VMEM → reject.
for (user : users)
  if (CostModel::FusionWouldExceedVmemCapacity(producer, user))      // 0x130c4a80
    return {priority = -1.0, long = -1};

conv_rw_count = GetConvAndRWCountsInNestedFusion(producer);          // 0x1454d240  → [rbp-0x88]

// (2) MUST-FUSE shortcuts → FLT_MAX:
if (producer.opcode == 0x1a /*collective*/ &&
    GetTpuCompEnv(producer)[+0x1206] != 0)                          // byte 4614
  return {priority = FLT_MAX, long = 0x3ffffffffffffffe};
if (IsNonFusionCollective(producer) && UserDirectedFuseInfo.IsMustFuse)
  return {priority = FLT_MAX};   // logs "Giving collective producer with must fuse attribute highest priority: "

// (3) THE FORMULA:
compute    = NormalizedComputationCost(producer, 0);                 // 0x130989a0 → [rbp-0x40]
mem_reduce = GetNormalizedMemoryCostReductionIfFusing(producer,users);// 0x13099700 → [rbp-0x50]
multiplier = conv_rw_count;                                          // [rbp-0x88]
priority   = mem_reduce - compute * multiplier;                      // vmulsd ; vsubsd

// (4) PRED single-bit packing correction:
if (producer.shape.element_type == PRED)
  priority *= (double) ElementPackingFactor(topology, PRED, ShouldPackPREDAsSingleBit(...));

// (5) NEGATIVE-PRIORITY FALLBACK:
if (priority < 0.0)
  priority = ActualCostReduction(producer, users);                  // 0x1309a060

两个 .text 位置按字节精确固定算术：vmulsd xmm0, xmm0, [rbp-0x88] (0x13096b32)，然后 vsubsd xmm0, xmm1, xmm0 (0x13096b3f) — mem_reduce(var_50) − compute(var_40)×multiplier(var_88)；0x13096b97 处的 PRED 乘法 vmulsd xmm1, xmm1, xmm0；随后调用 ActualCostReduction fallback。

项的解释。 两项都以 cycle 为单位，因此可直接比较：

text
priority =  (HBM bytes saved, converted to HBM-transfer CYCLES)
          − (added MXU/vector compute, in compute CYCLES)
              × (number of Conv/ReduceWindow ops the fusion would DUPLICATE)
```text

`conv_rw_count` 乘数是关键的 TPU 特定塑形：复制一个*昂贵* op（lowering 到 MXU 工作的 convolution）会按融合产生多少副本成比例惩罚，而融合纯 elementwise（其中 `compute` 是便宜的 `1.0` tier，且没有昂贵 op 被复制）几乎不增加成本。这里有意**没有 register-pressure 项** — 那是第 (1) 步中的 `FusionWouldExceedVmemCapacity` 硬门禁。

**PRED 修正**（第 4 步）处理布尔张量：当 topology 将 `PRED` 打包为单 bit（8 个 boolean/byte）时，字节 footprint — 因而内存节省项 — 会按 packing factor（`≈8`）缩放，因此大的 boolean producer 不会被低估。`ShouldPackPREDAsSingleBit` @ `0x1d6b0080`，`ElementPackingFactor` @ `0x1d6b03e0`。

**负优先级 fallback**（第 5 步）使用 `ActualCostReduction` (@ `0x1309a060`) 重新定价，它会在计入 nested-fusion overhead 和 tuple-result 惩罚后重新计算真实 reduction。这会拯救仍然净节省内存但被线性估计误拒的计算密集融合。

### 内存项 — 以 TC cycles 表示的 HBM 字节节省

`GetNormalizedMemoryCostReductionIfFusing` (@ `0x13099700`) 返回融合消除的 HBM 流量，以 TensorCore cycles 表示：

```text
reduction =  GetNormalizedMemoryCost(producer)            // producer's write to HBM (saved)
           + Σ_users  InputSizeAt(producer, user)         // each user's read of producer (saved)
           − cost_of_fused_result_writes_and_reads        // edges that survive the fusion

按 tensor 的 byte→cycle 转换（来自 GetNormalizedMemoryCost @ 0x1309a620）：

c
bytes        = OutputSize(inst);                           // 0x13097f20 — granule-aligned Σ subshapes
bytes_per_cy = HbmFullChipBytesPerSecond()                 // 0x1d6172a0
             / LogicalDevicesPerChip(core_type)            // 0x1d615b00
             / (TensorCoreFrequencyInMegaHertz() * 1e6);   // 0x1d615b60 × (qword_A2E0208 = 1.0e6)
hbm_cycles   = bytes * GranuleBytes() / bytes_per_cy;      // 0x1d617f80
```text

因此内存项字面上就是“融合此 producer 消除了多少 TensorCore cycles 的 HBM 流量”。输出很大且被多个 user 读取的 producer 会产生很大的正 reduction。`OutputSize` 遍历所有输出 subshape（处理 `kTuple`）；`InputSizeAt` (@ `0x1309a180`，非缓存版 `0x1309ad40`) 是 consumer 从 producer 读取的按 edge 字节数 — 也就是融合 internalise 的量 — 缓存在 `(producer, consumer)` 上。

### 计算项 — `NormalizedComputationCost`

`NormalizedComputationCost` (@ `0x130989a0`) 是**计算惩罚**：每个 op 一个 `double`，表示粗粒度“开销类别”而非 cycle-accurate latency（LLVM 类比是 `TTI::getInstructionCost` 返回 tier，而不是计数）。它是一个 `switch (opcode)`（反编译为 `switch (*((_BYTE*)inst + 12))`），映射到少数标量权重，每个权重乘以 `Target::ChunksIn(shape)` (@ `0x1d619900`，覆盖该 shape 的 MXU chunk-granule tile 数)：

```text
cost = Σ_operands ChunksIn(operand.shape) × W_operand  +  ChunksIn(root.shape) × W_root

vmulsd 位置使用的 .rodata 权重常量（CONFIRMED byte-exact）：

weight constvalueop 类别（本页范围）
qword_A2DF2301.0base / 默认便宜 elementwise
qword_A2DE8304.0mid class（reduce / cross-lane / logistic）
qword_A2DF49810.0divide
qword_A2DF1A042.0erf / matmul-conv heavy class

对于 convolution，成本是基于 flop_count 的估计（在 conv block 调用 HloCostAnalysis::flop_count），而不是扁平 tier；对于 fusion producer，成本是其函数体按 op 成本的递归和，按 unique_id 缓存。完整 opcode→weight switch — 包括 0.0 layout-op tier、2.0 first-parameter tier 和 dot-is-CHECK-fatal escape — 在专门的 NormalizedComputationCost 成本页枚举；本页记录优先级公式消费的四 tier 阶梯，并把完整 106 项 jump-table decode 委托给该页。

NOTE — dot 绝不能到达这里。 NormalizedComputationCost 在 dot opcode 上 CHECK-fatal：到 fusion 运行时，每个 dot 都已经被重写为 convolution（见 Dot/Conv MXU Lowering）。计算项感知 conv，而不感知 dot。

Bundle-Aware 成本模型 — total_unfused − total_fused ​

CalculateProducerPriorityWithBundleAwareCostModel (@ 0x130954c0) 使用 ResourceVector bundle-cost 机制，为融合实际节省的 VLIW bundle cycles 定价。因为该机制按被占用功能单元 lane 的最大值（而非总和）为 bundle 定价，它能捕获跨 engine overlap — 例如 producer 的 matprep 与 consumer 的 vector ALU 打包进同一个 bundle — 这是线性 current 模型看不到的。

c
// CalculateProducerPriorityWithBundleAwareCostModel(producer)  @ 0x130954c0 (decompiled, condensed)
users = GetFusibleUsers(producer);
if (fusion_util::IsMustFuseProducer(producer))               // 0x14559400
  return {priority = 100.0, long = 100};   // set directly to the 100.0 constant (qword_A2DF5C0)

// (A) UNFUSED: producer's own bundle cycles × #users, plus each user's own cycles.
total_unfused = GetHloCycles(producer) * producer.user_count;   // 0x13097a00 ; vmulsd
for (user : users)
  total_unfused += GetHloCycles(user);                          //          → [rbp-0x38]

// (B) FUSED: the cost of each merged (producer,user) bundle.
total_fused = 0;
for (user : users) {
  ResourceVector rv = {};
  total_fused += CostModel::GetCyclesIfFused(producer, user, opts, &rv);  // 0x130aba40 → [rbp-0x48]
}                                                               // cached on (prod_id, user_id)

// (C) PRIORITY = cycles saved.
priority = total_unfused - total_fused;                         // 0x13095d38 vsubsd
```text

`0x13095d38` 处的 `vsubsd xmm0, xmm0, [rbp+var_48]` 和 VLOG 字符串精确确认该公式：二进制记录 `" total_unfused_cycles: "`、`" total_fused_cycles: "`、`" new_model_priority: "` 和 `" producer.user_count: "`（"JFF" = JellyFish Fusion 日志族）。如果某次融合能让 producer 工作与 consumer 在更少 bundle 中重叠，就会使 `total_fused < total_unfused`，从而得到正优先级。

`GetHloCycles` (@ `0x13097a00`) = `CostModel::GetCycles` (@ `0x130aade0`) = `ResourceVector::MaxResourceCycles` bundle cost，按 `unique_id` 缓存。`GetCyclesIfFused` (@ `0x130aba40`) 对合并后的 producer+consumer bundle 执行相同的 `max` reduction 定价。二者在 [Bundle-Aware Cost](../cost/bundle-aware-cost.md) 中完整记录；fused-merge driver（`ScaleAndSumOutputFusionResourceVectors`，slot-9/11 `max`-combine）位于 [NormalizedComputationCost](../cost/normalized-computation-cost.md)。本页只负责消费它们的*优先级减法*。

must-fuse boost 将优先级**直接设置为**常量 `qword_A2DF5C0 = 100.0`（`IsMustFuseProducer` 为 true 时 `vmovsd xmm0, cs:qword_A2DF5C0` — 不是对计算分数的乘法），并将 `long` tie-break 设置为 `100`（`_R14 = 100`）；user-pinned fusion 会排在任何成本派生分数之上。该模型中没有 `INT64_MAX` 哨兵。

### 模型比较

| 项 | current model | bundle-aware model |
|---|---|---|
| 收益 | 节省的 HBM 字节（→ HBM cycles） | unfused − fused bundle cycles |
| 计算惩罚 | `NormComputeCost × conv_rw_count`（显式） | 折入 `total_fused` |
| VLIW bundle packing | **未**建模 | 已建模（`GetCyclesIfFused`） |
| VMEM / register pressure | 硬门禁（相同） | 硬门禁（相同） |
| must-fuse 分数 | `FLT_MAX` | `100.0`（直接设置） |
| tie-break `long` | input-size 累加 | `100`（搭配 `100.0` must-fuse boost） |
| 默认？ | **是** | 仅在 flag / per-edge 选中时 |

---

## 优先级哨兵

三个 `.rodata` double 和两个 `long` 承载带外排序信号。全部 CONFIRMED byte-exact（double 通过重新解释存储 bit pattern 验证）：

| sentinel | `.rodata` | bit pattern | 含义 |
|---|---|---|---|
| `-1.0` | `qword_A2DE728` | `0xbff0000000000000` | do-not-fuse — producer 不是 output-fusion root，或某个 user 的融合会超出 VMEM |
| `FLT_MAX` | `qword_A2E0530` | `0x47efffffe0000000` (`3.4028e+38`) | must-fuse，current 模型（带 must-fuse 属性的 collective）— 排在顶部 |
| `100.0` | `qword_A2DF5C0` | `0x4059000000000000` | must-fuse 优先级，bundle 模型（`IsMustFuseProducer`）— 直接设置，不是乘数 |
| must-fuse `long` (current) | — | `0x3ffffffffffffffe` | 与 `FLT_MAX` 配对的 tie-break |
| must-fuse `long` (bundle) | — | `100` (`0x64`) | 与 `100.0` boost 配对的 tie-break |
| µs conversion | `qword_A2E0208` | `1.0e6` | memory/tie-break 数学中的 cycles↔microsecond 因子 |

---

## VMEM 硬门禁（与谓词级联共享）

`CostModel::FusionWouldExceedVmemCapacity` (@ `0x130c4a80`) 是成本模型和合法性谓词共同咨询的唯一门禁 — 也是 VMEM pressure 从不作为分数项出现的原因。在**分数**中，它是 current 模型中的预 pass（上面的第 1 步）：任何 user 的融合若会超出 VMEM，就会强制 producer 为 `-1.0`。在**谓词**中（见 [Fusion Patterns](fusion-patterns.md)），相同检查会以 `"Nested dot fusion would exceed vmem capacity"`（`.rodata 0x84b1b50`）或 `"Custom Fusion would exceed vmem capacity"`（`0x84b1b7d`）拒绝，operand-fit 变体则发出 `"No fusing: result is a fusion which will use too much VMEM for its operands."`（`0xa0281ee`）。

设计意图：VMEM 是正确性/可行性约束（融合区域必须物理适配片上空间），不是质量权衡。将其编码为二元门禁而不是软惩罚，意味着一次融合要么能放下并按其价值排序，要么放不下并被排除 — 没有“略微超预算但价值很高”的中间地带。把 VMEM 建模为惩罚的重新实现会接纳无法 emit 的融合。

---

## Multi-Output（Sibling）Fusion — `GetProfit` + 预算门禁

`TpuMultiOutputFusion`（ctor @ `0x110dcbe0`，继承 `xla::MultiOutputFusion`）是一个独立 pass，它将一个 producer 与多个 consumer，或共享操作数的两个 sibling，合并为一个 tuple-rooted `kFusion`。它**不**使用上面的优先级队列公式；它有自己的 profit 数。

`TpuMultiOutputFusion::GetProfit(producer, consumer)` (@ `0x110dd0a0`) 返回**节省的字节数**，即共享操作数只读取一次而非每个 sibling 读取一次：

```text
profit = Σ Target::ShapeSize(shared_operand)      // bytes read once instead of once-per-sibling

通过 Target::ShapeSize (@ 0x1d61a8a0) 和 ShapeUtil::TrueNumDimensions (@ 0x20cdde20)；imul 形成 ShapeSize × count，并与从 TpuCompilationEnvironment 字节 offset 0x13c0 读取的阈值比较（GetTpuCompEnv(...)[+0x13c0]，反编译为 *(qword*)(TpuCompEnv + 5056)）。must-fuse 属性（UserDirectedFuseInfo::IsMustFuse）会不考虑 profit 强制该 pair。profit 馈入 MultiOutputFusion::AddToWorkList(p, c, profit) (@ 0x14bdf6e0) → 基类 priority_queue<ToBeFused>，按 profit 排序。VLOG："Fusing instr1=" / " instr2=" / ", the profit is ="。

合法性和预算门禁会在提交 merge 前拒绝（除非注明，否则 CONFIRMED）：

guardfunction规则
cycle checkMultiOutputFusion::Perform @ 0x14bdb5a0若 merge 产生 cycle 则拒绝 → "multi-output fusion creates a cycle" (0x86cc9cf)
max operandsTooManyResultOperands @ 0x110ddec0TotalBufferCount(p) > this[+0xc8] — cap 是 TpuMultiOutputFusion 对象上的 long 字段（ctor long 参数，来自 xla_tpu_multioutput_fusion_max_operands），不是 Target 字段 → 拒绝
reduce-output capTooMuchReduceOutputMultiOutput @ 0x110e1060每输出 reduce 字节 cap 0x400000 = 4 MiB（两个位置反编译为 > 0x400000）；CHECK "num_outputs >= instructions.size()" (0xa163eb7)
reduce-output cap (pair)TooMuchReduceOutput @ 0x110de000pairwise reduce 检查，但不是 4 MiB cap：将合并后的 reduce-output 字节（通过 ReduceEmitter::EvaluateReduceOutput、Target::TileBytes、MinFusedOperandBytes）与 0.8 × DefaultScopedVmemBytes（因子 qword_A2E0738 = 0.8）比较；当组合操作数计数 > 0x100 时短路为“not too much”
HBM pressureIsHBMPressureHighIfFused @ 0x110de640Σ ShapeSizeRecursive(outputs) × count > UserAllocationSharedMemoryLimitBytes(...) (@ 0x1d616680) → 拒绝
structural legalityLegalToFuse @ 0x110ddc20；ShapesCompatibleForFusion @ 0x110dcca0；IsFusible @ 0x110dce20shape 为共享迭代空间 tile-align；无 in-place 冲突；op 类别可融合

4 MiB reduce cap 防止合并那些单个输出已经足够大、会让合并 tuple 撑爆片上预算的 reduction；max-operands cap 约束单个 multi-output kFusion 的 fan-in；HBM-pressure cap 根据按 Target 的 shared-memory limit 约束总输出 footprint。该 pass 由 xla_tpu_multi_output_fusion_limit（flag 字符串 0x84f98d5）限制 — 这是每个 pass 可形成多少 multi-output fusion 的 cap — 并由 xla_jf_enable_multi_output_fusion / _advanced_… / _producer_consumer_… 控制。

NOTE — MOF profit 阈值（TpuCompEnv[+0x13c0]）和 max-operands cap（TpuMultiOutputFusion 对象字段 +0xc8，ctor long 参数）是运行时/flag 派生的，没有 constant-fold 到 .text；它们的绝对值来自编译环境 / flag（xla_tpu_multioutput_fusion_max_operands），不是来自 Target 常量。对这两个具体整数的置信度：LOW。消费它们的规则和load offset 是 CONFIRMED。

打分代码的代际不变性 ​

二进制中正好只有一条 TpuInstructionFusion 打分路径 — 一个 CalculateProducerPriority、一个 current 模型、一个 bundle 模型、一个 NormalizedComputationCost switch、一个 TpuMultiOutputFusion::GetProfit。不存在按 codename（viperfish / ghostlite / gfc / pufferfish）覆盖成本函数的实现。按代行为只通过数据进入分数：

  1. Target 常量和运行时字段 — ChunksIn（MXU 几何：每代 chunk count 不同 → NormalizedComputationCost 不同）、HbmFullChipBytesPerSecond、TensorCoreFrequencyInMegaHertz、GranuleBytes、ShapeSize、UserAllocationSharedMemoryLimitBytes，以及 MOF cap（TpuCompEnv[+0x13c0] profit 阈值、TpuMultiOutputFusion 的 +0xc8 max-operands 字段）。公式相同，数字按代 / 按 flag 不同。
  2. 按代 CycleTable latency — GetHloCycles / GetCyclesIfFused 为每个 op 存入的内容（bf16 matmul 在不同代之间相差约 26×），因此 matmul-bound fusion 的 bundle-aware 优先级会随芯片缩放，而减法本身相同。
  3. Flag 默认值 — xla_tpu_use_bundle_aware_cost_model_for_fusions、nested-dot / fp8 / packing-fusion enable。这些控制哪个模型运行以及哪些谓词分支准入，但打分代码是共享的。

这与 bundle-cost 发现相呼应（见 Bundle-Aware Cost）：代际不变代码作用于按代数据。

未字节钉住的内容 ​

  • 优先级 tuple 的 secondary double（current 模型）：从相同 priority slot 重新存储到多个 rbp offset；它是稳定性 re-score 还是重新存储的 ActualCostReduction 尚未字节确认。long tie 是 input-size/chunk 累加。(PARTIAL)
  • 绝对运行时整数 TpuCompEnv[+0x13c0]（MOF profit 阈值）和 TpuMultiOutputFusion 的 +0xc8 max-operands 字段，以及将内存项转换为绝对字节的 HbmFullChipBytesPerSecond / TC-frequency 值 — 它们是运行时/按代的，没有 constant-fold。（数字为 LOW；消费它们的公式/load offset 为 CONFIRMED。）
  • NormalizedComputationCost 内部完整 opcode→weight 分配 — 优先级公式使用的四个权重常量已确认；完整 106 项 jump-table decode（包括 0.0/2.0 tier 和 dot-fatal）由 NormalizedComputationCost 负责。
  • GetCyclesIfFused 的完整函数体 — 已确认它在 merged op 上返回 MaxResourceCycles bundle cost；其双 ResourceVector merge 记录在成本页面中。

函数映射 ​

symboladdress作用
TpuPriorityFusionQueue::CalculateProducerPriority0x1308fa20dispatcher：current vs bundle，output-fusion -1.0 门禁
…::CalculateProducerPriorityWithCurrentCostModel0x13096160mem_reduce − compute × conv_rw_count
…::CalculateProducerPriorityWithBundleAwareCostModel0x130954c0total_unfused − total_fused
…::EnqueueToProducerPriorityQueue0x1308fb20插入 tuple<double,double,long>；tpu_instruction_fusion.cc:1678 处 NaN CHECK
…::DequeueNextInstructionAndOperandsToFuseInOrder0x13090120pop 最大键；demote infrequent-conditional producer
…::OnFusingInstruction / InvalidateCachedCostModelState0x13090900 / 0x1309c120融合后重新给 neighbour 打分
…::NormalizedComputationCost0x130989a0计算项 — opcode 权重阶梯 × ChunksIn
…::GetNormalizedMemoryCostReductionIfFusing0x13099700内存项 — 节省的 HBM 字节 → TC cycles
…::GetNormalizedMemoryCost / OutputSize / InputSizeAt0x1309a620 / 0x13097f20 / 0x1309a180按 tensor byte→cycle、输出字节、按 edge 读取字节
…::ActualCostReduction0x1309a060负优先级 fallback 重新定价
…::GetHloCycles0x13097a00（→ CostModel::GetCycles 0x130aade0）按 op bundle cycles
CostModel::GetCyclesIfFused0x130aba40merged-bundle cycles
CostModel::FusionWouldExceedVmemCapacity0x130c4a80共享 VMEM 硬门禁
fusion_util::GetConvAndRWCountsInNestedFusion0x1454d240conv_rw_count 乘数
Target::ChunksIn / ShapeSize / HbmFullChipBytesPerSecond0x1d619900 / 0x1d61a8a0 / 0x1d6172a0按代成本常量
TransferSizeUtil::ShouldPackPREDAsSingleBit / ElementPackingFactor0x1d6b0080 / 0x1d6b03e0PRED single-bit packing 修正
TpuMultiOutputFusion::GetProfit0x110dd0a0sibling-fusion 字节节省 profit
…::TooManyResultOperands / TooMuchReduceOutputMultiOutput / IsHBMPressureHighIfFused0x110ddec0 / 0x110e1060 / 0x110de640MOF 预算门禁
MultiOutputFusion::Perform / AddToWorkList0x14bdb5a0 / 0x14bdf6e0MOF cycle check；按 profit 排序的 work list

交叉引用 ​

  • Fusion Patterns — 谓词半边：ShouldFuseImpl 的有序拒绝级联、slice-like / output-fusion / duplicate-expensive 门禁、custom-call registry hook。本页给幸存者排序；该页过滤。
  • Compiler Overview — fusion pass 在 lowering pipeline 中的位置。
  • Compile Phases — 将 fusion 放在 dot→conv 重写之后的 pass 顺序。
  • Dot/Conv MXU Lowering — 为什么 NormalizedComputationCost 在 dot 上 CHECK-fatal（dot 在 fusion 前被重写为 convolution）。
  • Cost Model Overview — 本页 bundle-aware 模型消费的 bundle/ResourceVector 成本机制。
  • TPU HLO Cost Analysis — conv-flop 计算权重背后的按 op cycle/flop 分析。
  • NormalizedComputationCost — 穷尽的 opcode→weight switch 和 GetCyclesIfFused fused-merge driver。
  • Bundle-Aware Cost — GetHloCycles / MaxResourceCycles，即线性 current 模型看不到的 bundle-packing 模型。
  • back to index