Internal Wiki

主题

融合模式 ​

本页中的所有地址均适用于来自 libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。所有结论都来自对未剥离、完整符号 ELF 的静态分析（nm -C 可解析每个方法）；原始源码是闭源的。

摘要 ​

xla::jellyfish::TpuInstructionFusion 是 TPU 后端的主融合 pass。它是 xla::InstructionFusion 的子类（与 XLA 的 CPU/GPU 后端使用同一个基类），但它运行得很晚 — 在 LayoutAssignment 以及 DotCanonicalizer/ConvolutionFolding 之后 — 因此等它看到模块时，每个 kDot 都已经变成了 kConvolution，每个卷积都有了确定布局，elementwise 链也已经形状规范化。因此，这个 pass 匹配的模式空间比上游 xla::InstructionFusion 窄得多，并且围绕 GPU 后端从来没有的一条轴线组织这些模式：融合后的 op 落到哪个 TPU 引擎上 — PE（MXU/脉动阵列）、ACT（MXU 后的向量 ALU），或者二者打包到同一个 VLIW bundle 中。

本页负责说明决策逻辑和被识别的形状：ShouldFuseImpl 谓词级联（lambda 集合 $_0–$_30，用于投票判断 producer/consumer 对是否可融合）、该 pass 重写为 kFusion 的 HLO 形状目录（conv+bias+activation、attention/softmax、layer-norm/RMSNorm 尾部、copy/packing 融合、collective 融合），以及无论候选看起来多有利都能否决它的合法性 gate。它不负责数值化的收益评分 — TpuPriorityFusionQueue::CalculateProducerPriority* 内的浮点优先级公式和按代际划分的 MXU 延迟表位于 fusion-cost-model.md。两页共享一个事实：TpuInstructionFusion::GetFusionQueue 返回 TpuPriorityFusionQueue，所以候选在发生任何重写前会同时经过谓词过滤（本页）和优先级排序（成本模型页）。

下面的结构是：先介绍谓词模型（ShouldFuse → ShouldFuseImpl 如何得到 FusionDecision），然后是硬性合法性 gate（无论如何都不融合的检查），再是按目标引擎分组的已识别形状目录，然后是把已识别 HLO 子图转换成 InputFusionOp/OutputFusionOp MLIR 载体的输入/输出 op-fuser 分发，最后是多输出和 collective 模式族。

对于重新实现，契约是：

  • 决策流水线： ShouldFuse(consumer, operand_index) → ShouldFuseImpl → 一个 FusionDecision（融合 / 带原因字符串的拒绝），先由硬性合法性谓词 gate，再由优先级队列的成本排名 gate。
  • 谓词集合： ShouldFuseImpl 内部的 $_0–$_30 $_* lambda 谓词（发出了 30 个不同的 lambda 符号；符号表中不存在 $_12）以及具名合法性方法（FusionFitsInVmem、ProducerCanBeLoopFused、CheckReduceBroadcastIntoReduceWindowFusionRequirements、MOF 的 TooMany*/TooMuch* 家族）。
  • 已识别形状： 该 pass 重写的 HLO match 表达式，按识别方法和目标 lowering 标注，并说明每个形状在哪些代际上可用。
  • op-fuser 分发： 已识别子图如何通过匿名命名空间的 OutputFusionOp::Create<MlirOp> 模板作为 OutputFusionOp/InputFusionOp 携带，以及为什么超越函数 activation 会绕过这条路径。
Pass 类xla::jellyfish::TpuInstructionFusion : public xla::InstructionFusion
RunImpl0x13080dc0
决策入口ShouldFuse 0x13089b20 → ShouldFuseImpl 0x13086660（约 1779 行反编译代码）
谓词 lambda在 ShouldFuseImpl 内定义的 $_0–$_30（nm -C 中 30 个不同的 lambda 符号；未发出 $_12）
融合队列GetFusionQueue 0x13083c40 → TpuPriorityFusionQueue（匿名命名空间）
VMEM gateFusionFitsInVmem 0x13084b40；预算 xla_jf_fusion_max_vmem_mib（默认 15；以 double 存储，字节模式 0x402E000000000000）
流水线阶段"Main fusion"（Phase 5），见 compile-phases.md
成本评分（链接）CalculateProducerPriorityWith{Current,BundleAware}CostModel — fusion-cost-model.md
源文件platforms/xla/service/jellyfish/tpu_instruction_fusion.cc（.rodata 中的字符串）

决策流水线 ​

目的 ​

InstructionFusion（基类）遍历每个 computation，并针对每条 (consumer, operand) 边询问子类：这个 operand 是否应该融合进这个 consumer？ 答案是一个 FusionDecision — 要么是 "yes"，要么是人类可读的原因字符串。TpuInstructionFusion 覆盖了三个 hook；实质内容在一个私有方法中。

入口点 ​

text
TpuInstructionFusion::RunImpl  (0x13080dc0)
  ├─ pre-passes (rewrite the graph before priority fusion):
  │    CreateFusionsAroundConvolutions   (0x1307c2c0)  ── wrap each bare Conv in a kCustom fusion
  │    BitcastConvOperands               (0x1307c960)  ── fold Bitcast into Conv window-config
  │    PrefuseReduceBroadcastReuse        (0x1307d9a0)  ── keep Reduce result in MXU latch
  │    MoveReduceBroadcastTogether        (0x1307f9c0)  ── reorder so Broadcast can fuse Reduce
  │    CreateLoopFusionAroundFusionLowerableHlo (0x13080780)
  │    DoExtendedAnalysisForCustomCallConsumerFusion (0x13082680)
  └─ base InstructionFusion::Run drives the queue:
       GetFusionQueue (0x13083c40) ── returns TpuPriorityFusionQueue
         └─ for each dequeued (producer, consumer):
              ShouldFuse              (0x13089b20)   ── public hook
                └─ ShouldFuseImpl     (0x13086660)   ── the predicate cascade
              ShouldFuseIntoMultiOutput (0x13089ce0) ── MOF variant
              ChooseKind              (0x13084a60)   ── kInput / kOutput / kLoop / kCustom
              Fuse                    (0x1308d820)   ── perform the rewrite

算法 ​

ShouldFuseImpl 是本页的核心。它的函数体（反编译后约 1779 行）是一长串 guard 谓词级联；第一个投票为 "no" 的谓词会返回一个携带原因字符串的 FusionDecision。谓词实现为 $_* lambda（引用了 $_0–$_30）以及对具名合法性方法的调用。下面的结构根据反编译中恢复到的决策字符串和调用点重建。

c
function ShouldFuseImpl(consumer, operand_index):          // 0x13086660
    producer = consumer->operand(operand_index)
    // 0. base-class structural gate, run as a lambda predicate $_0
    //    (FusionDecision(producer, consumer, AliasInfo, InPlaceFusionOptions))
    d = run_lambda_$_0(producer, consumer)                  // line ~650: __policy_func<FusionDecision(...)>
    if d.is_no(): return d                                  // "Not fusing"

    // 1. scalar-constant fast paths — always fuse (graph cosmetics / index folding)
    if producer is scalar constant feeding consumer:
        return Fuse("Did fuse: fusing scalar constant to make graph look nicer.")
    if producer is scalar constant used as dynamic-slice index:
        return Fuse("Did fuse: fusion scalar constant to dynamic slice index.")
    if producer is scalar constant used as DUS index:
        return Fuse("Did fuse: fusion scalar constant to fusible dynamic update slice index.")

    // 2. output-fusion enable gate
    if consumer is output-fusion candidate and !output_fusion_enabled:
        return NoFuse("No fusing; output fusion is disabled.")

    // 3. VMEM capacity — the dominant hard gate
    if !FusionFitsInVmem(producer, consumer):               // 0x13084b40, line ~679
        return NoFuse("No fusing: result is a fusion which will use too "
                      "much VMEM for its operands.")        // line ~688

    // 4. duplication cost — refuse to clone an expensive producer
    if NumProducerDuplicationsIfFused(producer, consumer) > 0    // 0x130896a0
       and producer is expensive:
        return NoFuse("No fusing: producer is duplicated and expensive.")

    // 5. RNG single-use rule
    if producer->opcode == kRng and producer has multiple users:
        return NoFuse("no fusing: rng is used by multiple users")

    // 6. elementwise guard (some elementwise producers are not output-fusable)
    if Should-not-fuse-elementwise(producer, consumer):
        return NoFuse("No fusing: Should not fuse with elementwise")

    // 7. slice-like keep-unfused flag
    if producer is slice-like and FLAGS_xla_tpu_keep_slice_like_instructions_unfused:
        return NoFuse("No fusing: slice-like instruction kept unfused due to "
                      "flag xla_tpu_keep_slice_like_instructions_unfused")

    // 8. convolution-input legality: only "trivial" inputs may fuse into a conv
    if consumer is convolution-like and producer is non-trivial:
        return NoFuse("Refusing to fuse a non-trivial inputs into a "
                      "convolution-like. Producer: ...")

    // 9. effective-scalar reachability checks (DUS / fusible-user routing)
    if producer produces effective scalar without fusible DUS user reachable:
        return NoFuse("... produces an effective scalar and does not have a "
                      "fusible DUS user reachable through consumer: ")

    // 10. bitcast / reduce-window window legality
    if !CheckReduceBroadcastIntoReduceWindowFusionRequirements(...):  // 0x1307ee60, line ~975
        return NoFuse("Cannot find window to lower for bitcast reduce fusion: ...")
    if dim-collapsing bitcast and not must_fuse mode:
        return NoFuse("Dim collapsing bitcast is fused only in must_fuse mode.")

    // 11. loop-fusion / gather-fusion routing (logged, not necessarily reject)
    can_loop  = ProducerCanBeLoopFused(producer, consumer, target_)  // 0x1307f800, line ~1202
    log("producer_can_be_loop_fused: ", can_loop)
    log("producer_can_be_gather_fused: ", ...)

    // 12. boundary predicate $_29 (structural, no log of its own):
    //     a producer that already forms the outer fusion boundary
    //     cannot be fused any further (string emitted by caller ShouldFuse)
    if run_lambda_$_29(producer, consumer, target_, reachability_map):  // 0x130899c0
        return NoFuse(/* "Producer forms the fusion boundary. ..." */)

    return Fuse("Fusing producer: ... into consumer: ...")

注意 — 承载大部分逻辑的是 lambda，而不是具名方法。 ShouldFuseImpl 定义了 $_0–$_30（nm -C 中保留了 30 个不同的 lambda 符号；索引 $_12 未发出）。其中若干是没有自身日志字符串的微小结构谓词（例如 0x130899c0 处的 $_29，它接收 (producer, consumer, consumer, Target&, HloReachabilityMap*)，是融合边界测试）。具名方法（FusionFitsInVmem、ProducerCanBeLoopFused、CheckReduceBroadcastIntoReduceWindowFusionRequirements）是重量级 gate；lambda 是先运行的廉价结构过滤器。[置信度：HIGH — 调用点和字符串都在反编译结果中；每个 $_* 的精确逐 lambda 函数体没有单独展开。]

陷阱 — HardSwish 反融合字符串不在 TPU fusion pass 中。 字符串 " HardSwish pattern was found, so fusion failed." 解析到 tensorflow::grappler::Remapper::Optimize（0x105aa960），这是编译进同一个 .so 的 CPU TensorFlow 图重写器，而不是 TpuInstructionFusion。TPU 侧的功能性结论仍然成立 — HardSwish 没有直接 ACT ALU opcode，因此不可 output-fusable（见下面的 activation 讨论）— 但重新实现者不能根据这条日志行判断 TPU 行为。[置信度：CONFIRMED by symbol resolution。]

函数映射 ​

函数地址作用
RunImpl0x13080dc0Pass 入口；运行 pre-pass 后再运行基类 Run
ShouldFuse0x13089b20公共 hook；发出 "Not fusing MOF" / boundary 字符串
ShouldFuseImpl0x13086660谓词级联（约 1779 行）
ShouldFuseIntoMultiOutput0x13089ce0决策的 MOF 变体
FusionFitsInVmem0x13084b40VMEM 容量 gate
NonBroadcastOperandsSize0x13084d80ShouldFuseImpl 内调用的 operand-size 度量（约第 895 行）
ProducerCanBeLoopFused0x1307f800kLoop-fusion 合法性
CheckReduceBroadcastIntoReduceWindowFusionRequirements0x1307ee60RWB window 合法性
RwbPreliminaryCandidateCheck0x13085700RWB 预过滤器
NumProducerDuplicationsIfFused0x130896a0重复成本计数
$_29 lambda0x130899c0融合边界结构谓词
ChooseKind0x13084a60选择 kInput/kOutput/kLoop/kCustom
GetFusionQueue0x13083c40返回 TpuPriorityFusionQueue

相关开关 ​

开关类型默认值它 gate 的内容
xla_jf_conv_input_fusionBOOLtrue启用输入侧（conv 上方）的 elementwise 链融合
xla_jf_conv_output_fusionBOOLtrue启用输出侧（conv 下方）的链；关闭时触发 "output fusion is disabled" 拒绝
xla_jf_conv_reshape_fusionBOOLtrue允许 ReshapeFuser 将 reshape 吸收到 conv fusion 中
xla_tpu_keep_slice_like_instructions_unfusedBOOLfalse为 true 时，slice-like producer 被否决（谓词 7）
xla_jf_fusion_max_vmem_mibDOUBLE15FusionFitsInVmem gate 强制执行的 VMEM 预算（MiB）。默认 15.0 已在 AbslFlagDefaultGenForxla_jf_fusion_max_vmem_mib::Gen 中确认（0x402E000000000000）
xla_jf_enable_final_priority_fusionBOOLtrue驱动 TpuPriorityFusionQueue 的优先级遍历

硬性合法性 Gate ​

目的 ​

有些检查是无条件的：即使某个候选被优先级队列排在第一，只要失败其中之一就会被丢弃。它们分为 (a) ShouldFuseImpl 内部的结构/合法性否决，以及 (b) TpuMultiOutputFusion 内部的多输出融合（MOF）成本谓词。每个否决都会发出一个独立的 .rodata 字符串，这是重新实现者判断一次融合为什么没有形成的真值来源。

拒绝谓词 ​

否决来源方法决策字符串（.rodata）
VMEM 容量（整个 fusion）FusionFitsInVmem 0x13084b40"No fusing: result is a fusion which will use too much VMEM for its operands."
Nested-dot VMEMShouldFuseImpl + xla_tpu_nested_dot_fusion_vmem_fraction"Nested dot fusion would exceed vmem capacity"
Custom-call VMEMDoExtendedAnalysisForCustomCallConsumerFusion 0x13082680"Custom Fusion would exceed vmem capacity"
Producer 被重复且昂贵ShouldFuseImpl + NumProducerDuplicationsIfFused"No fusing: producer is duplicated and expensive."
RNG 有多个用户ShouldFuseImpl"no fusing: rng is used by multiple users"
Elementwise 不可 output-fusableShouldFuseImpl"No fusing: Should not fuse with elementwise"
Slice-like 因开关保持未融合ShouldFuseImpl"No fusing: slice-like instruction kept unfused due to flag ..."
非平凡输入进入 convShouldFuseImpl"Refusing to fuse a non-trivial inputs into a convolution-like. ..."
找不到 bitcast-reduce windowCheckReduceBroadcastIntoReduceWindowFusionRequirements"Cannot find window to lower for bitcast reduce fusion: ..."
Dim-collapsing bitcast 位于 must_fuse 之外ShouldFuseImpl"Dim collapsing bitcast is fused only in must_fuse mode."
Effective scalar 没有可融合 DUS 用户ShouldFuseImpl"... produces an effective scalar and does not have a fusible DUS user ..."
Output fusion 全局关闭ShouldFuseImpl"No fusing; output fusion is disabled."
MOF 产生环xla::MultiOutputFusion::Perform 0x14bdb5a0"multi-output fusion creates a cycle"
结果 operand 过多TpuMultiOutputFusion::TooManyResultOperands 0x110ddec0（无字符串 — 静默成本预检查）
Reduce-output MOF 过多TpuMultiOutputFusion::TooMuchReduceOutputMultiOutput 0x110e1060"TooMuchReduceOutputMultiOutput: "
融合后 HBM 压力过高TpuMultiOutputFusion::IsHBMPressureHighIfFused 0x110de640（无字符串）
已有 must_fuse（CCF）TpuUserGuidedFusionVerifier::VerifyMustFuseCalls" already has must-fuse attribute, skipping ..."

陷阱 — VMEM 是最容易让移植意外失败的 gate。 TPU 没有 GPU SM 那样的通用寄存器文件；fusion 会把整个融合区域物化到 VMEM scratch 中，而且operand window 的并集（不只是输出）必须放得下。FusionFitsInVmem（0x13084b40）直接求和 operand-window 字节数 — Target::TileBytes × tile count，加上每个 DUS operand 的 fusion_util::MinFusedOperandBytes、GetUnalignedDUSMinimumVmemOperandBytes，以及 reduce 输出的 ReduceEmitter::EvaluateReduceOutput — 并在该总量（按 DefaultScopedVmemBytes 和 xla_jf_fusion_max_vmem_mib 预算缩放）溢出时拒绝；当组合后的 operand 数超过 0x100（256）时也会硬拒绝。只检查输出大小的朴素移植会接受溢出 VMEM 并导致误编译的 fusion。回收开关（xla_tpu_scavenge_vmem_for_fusions）允许队列在其他 fusion 释放 VMEM 后重试被拒绝的 fusion — 因此在一次 pass 内 VMEM 拒绝并不一定是最终结果。该重试的 cycle-budget 机制归 fusion-cost-model.md 所有。

注意事项 ​

MOF 谓词（TooManyResultOperands、TooMuchReduceOutput、IsHBMPressureHighIfFused）gate 的是fan-out 方向（一个 producer → 多个 consumer，或多个 producer → 一个 tuple root），并且在查询 GetProfit（0x110dd0a0）之前由 TpuMultiOutputFusion::LegalToFuse（0x110ddc20）检查。fan-out 限制和 operand-count 上限由 xla_tpu_multi_output_fusion_limit 与 xla_tpu_multioutput_fusion_max_operands 设置。

已识别形状 — PE 锚定（conv / matmul + epilogue） ​

目的 ​

主导性的 TPU fusion 形状是一个卷积（或已被 DotCanonicalizer 重写为卷积的 matmul）加上融合在其下方的 elementwise epilogue，和/或融合在其上方的 elementwise 链。本节中的每个形状都通过 conv emitter（FusedSpatialMajorConvolution::EmitOneChunk）lower，并在同一个 per-chunk 循环中把 epilogue 发射到 ACT 引擎。

形状目录 ​

模式HLO match 表达式识别方法目标 lowering
Conv+BiasAdd(Conv(act,kernel), Broadcast(bias))TpuInstructionFusion output-fusionMXU 结果 + vadd.f32/vadd.bf16 epilogue，同一 chunk
Conv+Bias+ReLUMaximum(Add(Conv(a,w),Broadcast(b)),Broadcast(0))output-fusion → TpuLoopFusionEnhancer同上 + vmax.f32(_,0) / 按代际的 Relux
Conv+Bias+SigmoidLogistic(Add(Conv(a,w),Broadcast(b)))output-fusion（elementwise 尾部）同上 + ShiftedSigmoid ACT op（按代际）
Conv+Bias+TanhTanh(Add(Conv(a,w),Broadcast(b)))output-fusion（elementwise 尾部）同上 + vtanh ACT op
Conv+Bias+GELUMul(Add(Conv,b), Mul(0.5, Add(1, Erf(...))))output-fusion + ElementwiseOutputFuser同上 + verf + add/mul 链（ACT）
Conv+Activation（无 bias）<act>(Conv(a,w)), <act>∈{Relu,Tanh,Sigmoid,Exp,Erf}output-fusion（无 bias 变体）与 Conv+Bias+act 相同，但没有 bias add
MatMul+Bias+ReLU（dense）post-DotCanonicalizer：与 Conv+Bias+ReLU 相同TpuInstructionFusion相同的 conv lowering
MatMul→LayerNorm 尾部Subtract(x,ReduceMean(x)); Mul(...); Rsqrt(ReduceMean(Square(...)))TpuInstructionFusion + TpuLoopFusionEnhancerPE matmul + ACT（vrsqrt,vmul,vadd）bundle
MatMul→RMSNorm 尾部Mul(x, Rsqrt(ReduceMean(Square(x)) + eps))TpuInstructionFusionACT：vrsqrt.f32 + vmul.f32
AttentionMatMul（Q·Kᵀ）Dot(Q,K)（softmax 在单独 fusion 中）TpuInstructionFusion（dot-as-conv）标准 MXU lowering
Attention + SoftmaxDot 后接 Exp/Sum/Div 链由 xla_tpu_enable_multi_level_nested_dot_fusion gate启用时为一个嵌套 kFusion，否则为两个
DotDot（A·B·C）Dot(Dot(A,B),C) → 两个 Conv由 xla_tpu_dot_dot_fusion gate一个嵌套 super-fusion；需要 B+C 的 VMEM
RWB（Reduce-Window-Broadcast）Broadcast(ReduceWindow(input,window)) 匹配为 convRwbPreliminaryCandidateCheck + CheckReduceBroadcastIntoReduceWindowFusionRequirements，由 xla_tpu_rwb_fusion gate作为带 broadcast output-fusion 的 Conv lower

怪癖 — 每个 matmul fusion 都是 convolution fusion。 因为 DotCanonicalizer 在主 fusion pass 之前运行，TpuInstructionFusion 从不会看到 kDot。dense 层（MatMul+Bias+ReLU）和 conv 层（Conv+Bias+ReLU）在这个 pass 运行时已经是同一个 fusion 形状，并共享一条 lowering 路径。在 fuser 内按 kDot 与 kConvolution 分支的重新实现会发现 kDot 分支是死的 — canonicalizer 已经把它折叠了。只匹配 kConvolution。

形状重构 pre-pass ​

四个 RunImpl pre-pass 会在任何 priority fusion 运行之前重塑图，使更多候选变得可匹配：

Pre-pass地址它重写的内容
CreateFusionsAroundConvolutions0x1307c2c0将每个裸 Conv 包装在一个 kCustom fusion 中，使其 operand window 显式化
BitcastConvOperands0x1307c960将 Bitcast(a)/Bitcast(w) 折叠进 conv 的 window-config（避免物化 reshape）
PrefuseReduceBroadcastReuse0x1307d9a0将 Reduce 结果保留在 MXU latch 中，使下游 Broadcast 可摊销它
MoveReduceBroadcastTogether0x1307f9c0重新排序，使 Reduce 及其远处的 Broadcast 相邻且可融合

已识别形状 — 多输出和 Collective ​

多输出（TpuMultiOutputFusion） ​

TpuMultiOutputFusion : public xla::MultiOutputFusion（继承的 RunImpl 位于 0x14bdaa80）融合 fan-out 方向：一个 producer 供给多个 consumer，或多个 producer 绑定到一个 tuple root。它暴露两个驱动器 — DoProducerConsumerMultiOutputFusion（0x110e43e0）和 DoAdvancedMultiOutputFusion（0x110e3300）。

模式HLO 形状驱动器Lowering
MultiOutput conv一个 Conv → {ReLU, ReLU_grad}DoProducerConsumerMultiOutputFusion一个带 tuple root 的 kFusion；复用 PE+ACT bundle
MultiOutput reduce两个共享 operand 的 Reduce opDoAdvancedMultiOutputFusion一个 kFusion，PE 侧 + ACT 侧 reduce 位于同一 iteration

MOF 合法性链是 ShapesCompatibleForFusion（0x110dcca0）→ IsFusible（0x110dce20）→ LegalToFuse（0x110ddc20，包含 cycle check）→ GetProfit（0x110dd0a0）。TooMany*/TooMuch*/IsHBMPressureHighIfFused 谓词（见硬性 gate）在这条链内部运行。

Collective 和 copy/packing 形状 ​

这些形状由运行在 "Pre main fusion" 阶段（B2）的专用 pass 在 TpuInstructionFusion 之前识别，或者是它的 gated 子模式。

模式HLO 形状识别 pass / gate
Async collectiveAllGatherStart → … → AllGatherDoneAsyncCollectiveFusion::RunImpl 0x109b4ec0
AllReduce+ScatterAllReduce(x) → Slice(x, my_shard)TpuAllReduceScatterFusion::RunImpl 0x127acd40 → 内部 FusionOp::kAllReduceScatter（由 AsyncPincerFusionEmitter::EmitAllReduceScatterFusion 发出）
Mosaic kernelCustomCall(target="tpu_custom_call") + neighboursMosaicFusion::RunImpl 0x10f12500（由 HloPassFix<MosaicFusion> 驱动）
Megacore conv+AR跨 core 的 Conv 与 AllReduce 配对MegacoreFusion::RunImpl 0x110d8f00
Copy（data-format）layout-only Copy ≥ xla_tpu_copy_fusion_threshold 字节TpuInstructionFusion + xla_tpu_enable_copy_fusion
Copy-permute-minor交换第 2-minor 维的 Copy+ xla_tpu_enable_copy_permute_minor_fusion
Pad/Unpad copyratio < xla_tpu_copy_fusion_pad_unpad_ratio 的 Copy(Pad(x))TpuInstructionFusion
Bf16-packed matmul两个共享 operand 的 bf16 opFusedSpatialMajorConvolution::EmitPackedBf16Chunk 0x130e3120
Int8（x8）-packed matmul两个共享 operand 的 int8 matmulconv lowering-strategy ls_.generate_*_x8_packed_* 标志（例如 generate_x8_packed_vmatmuls）
DS_CC_DUSDynamicSlice → CustomCall → DynamicUpdateSliceTpuInstructionFusion-AdvancedDS_CC_DUS（日志字符串）

注意 — collective 和 Mosaic 形状在主 pass 之前被识别。 AsyncCollectiveFusion、TpuAllReduceScatterFusion 和 MosaicFusion 运行在 "Pre main fusion" 流水线（Phase B2）中，因此当 TpuInstructionFusion 在 "Main fusion"（B3）中运行时，这些已经是它会当作不透明节点处理的 kFusion/tpu_custom_call。逐 pass 的位置归 compile-phases.md 所有；本页只负责 match 形状。

Op-Fuser 分发 — 将形状携带到 MLIR ​

目的 ​

一旦 TpuInstructionFusion 判定一个子图是一个 kFusion，MLIR lowering 就必须在 InputFusionOp（conv 上方的 elementwise 链）或 OutputFusionOp（conv 下方的链）region 内逐 op 重建该子图。一小组 FusionOpFuser 子类负责逐 op 翻译。这里是已识别形状从 HLO 模式变成具体 MLIR 载体的地方。

识别 gate 是 OutputFusionOp::Create，不是 switch ​

最重要的重新实现事实：ElementwiseOutputFuser::CanFuse（0x10f36aa0）不包含硬编码 opcode switch。它调用匿名命名空间的 OutputFusionOp::Create(op)，并返回 create 是否成功。

c
function ElementwiseOutputFuser::CanFuse(op):              // 0x10f36aa0
    OutputFusionOp::Create(&tmp, op)   // anon-namespace factory; sets tmp.ok flag
    if tmp.ok == 1:
        if tmp.has_cleanup_fn: run_cleanup(); return 1     // op IS recognized
    return 0                                               // op not in the carrier set

OutputFusionOp::Create<MlirOp> 正好为下面 16 个 MLIR arith/math op 实例化模板（RTTI 确认 InputFusionOp 侧也是同一个集合 — 该载体是对称的）：

text
arith::AddFOp  arith::AddIOp  arith::SubFOp  arith::SubIOp
arith::MulFOp  arith::MulIOp  arith::DivFOp  arith::DivSIOp
arith::MaximumFOp  arith::MaxSIOp  arith::MinimumFOp  arith::MinSIOp
arith::NegFOp  math::AbsFOp  math::AbsIOp  math::ExpOp

陷阱 — 超越函数 activation 会绕过 elementwise 载体。 Tanh、Sqrt、Rsqrt、Erf、Logistic、Sin、Cos、Atan2 不在上面的 16-op OutputFusionOp 集合内，即使 Conv+Tanh 和 Conv+Sigmoid 被列为已识别形状。它们会被直接 lower 为专用 LLO op（llo.vtanh、llo.vrsqrt、llo.verf，……），而不是通过 elementwise OutputFusionOp 路径。只根据 16-op 载体驱动 epilogue 的重新实现会静默丢掉每个超越函数 activation。拆分方式是：载体持有代数 op（add/mul/min/max/abs/exp/neg），LLO dialect 将超越函数 op 作为一等指令持有。

Fuser 映射 ​

Fuser地址（GetFusionOp / CanFuse,Fuse）处理的 HLO op
BinaryOpFuser0x10f1be60Add/Sub/Mul/Div/Max/Min（F 和 I 变体）
UnaryOpFuser0x10f2bf00math::ExpOp、AbsFOp、AbsIOp、arith::NegFOp
TernaryOpFuser0x10f2b7a0Select / Clamp（FMA 风格 ternary）
CompareFuser0x10f1e760arith::CmpFOp、arith::CmpIOp（所有谓词）
ConvertFuser0x10f1f100ExtF/TruncF/SIToFP/UIToFP/FPToSI/FPToUI（约 13 种 case）
BroadcastFuser0x10f1c5a0kBroadcast → vector::BroadcastOp / llo.vbcast_sublane_chunk
ReduceFuser0x10f20120kReduce → vector::MultiDimReductionOp → llo.vmax.{x,s}lane.*
ReshapeFuser0x10f22680kReshape → vector::ShapeCastOp / tensor::CollapseShapeOp
ConvertOutputFuser0x10f2c480 / 0x10f2c4a0输出侧 kConvert（f32 → bf16 / f8 downcast）
ElementwiseOutputFuser0x10f36aa0 / 0x10f36e20上述 16-op 载体集合（通过 OutputFusionOp::Create）
ReduceOutputFuser0x10f37920 / 0x10f37940由 fusion root 消费的 kReduce
CustomCallOutputFuser0x10f2dac0 / 0x10f2dae0target 位于 xla_tpu_nested_dot_fusion_supported_custom_ops 中的 kCustomCall

为什么有些 activation 永远不会融合 ​

独立于载体拆分，三类 activation 在任何代际上都永远不可 output-fusable，必须 lower 为读取 VMEM 中 matmul 输出的独立 HLO computation：

  • HardSwish — 没有直接 ACT ALU opcode（"HardSwish pattern was found" 日志字符串属于 CPU TF Remapper，而不是 TPU fusion pass — 见算法中的陷阱）。
  • LeakyReLU / ELU / SELU / Mish / Swish — 通过 select 或 multi-op 多项式合成；成本模型会拒绝多项式展开。
  • Softmax 当 xla_tpu_enable_multi_level_nested_dot_fusion=false 时 — 保持为单独 fusion。

某个给定代际拥有哪些超越函数 opcode（因此哪些 activation 可直接融合，哪些会回退到多项式）归 activation-inventory 分析所有；对本页来说，相关事实是已识别形状（例如 Conv+Bias+GELU）只有在目标代际拥有 Erf ACT opcode 时才可匹配。

ReduceEmitter::EmitReduction — ReduceFuser 背后的轴分发器 ​

ReduceFuser / ReduceOutputFuser（Fuser 映射中的行）识别 kReduce 并将它交给 xla::jellyfish::ReduceEmitter，其顶层入口是 EmitReduction（0x13e16240，位于 platforms/xla/service/jellyfish/lowering/reduce_emitter.cc）。这个函数决定一个逻辑 reduce 会变成哪种物理 reduction primitive — 它是一个很薄的分发器，而不是循环嵌套：它进行验证，计算两个轴 bit，然后尾调用五个专用 emitter 中恰好一个。由 0x13e16240–0x13e16720（下一个符号 EmitPrologue；1248 字节）的窗口反汇编恢复。

c
function ReduceEmitter::EmitReduction(this, reduce_map, window, dims,
                                      output_span /*Span<LloValue* const>*/, builder):
    // 1. Copy the window's working LloValue* array into a heap scratch buffer
    n   = window->[0x358]                          // element count
    src = window->[0x350]                          // LloValue** base
    if n >> 0x3d != 0: throw length_error          // count*8 overflow guard
    buf = operator new(n*8); memcpy(buf, src, n*8) // freed at function exit

    // 2. CHECK every output operand lives in VMEM
    for v in output_span:                          // unrolled x4 + scalar tail
        ms = (v->[0xb] >> 2) & 0x1f                 // LloValue.memory_space bitfield
        CHECK_EQ(ms, MemorySpace::kVmem /*==3*/)    // reduce_emitter.cc:1400
                                                    // "output_span->memory_space() == MemorySpace::kVmem"

    // 3. Fast path: a "keep sublanes" flag on the emitter
    if this->[0x610] != 0:
        return EmitKeepSublanesReduction(...)       // 0x13e11f40 (tail call)

    // 4. Otherwise compute the two axis bits from the window's reduced-dimension list
    kind = this->[0x2e8]                            // shape dimension-storage kind
    CHECK(kind == 3 || kind == 5)                   // else FATAL shape.h, reduce_emitter.cc:843
    ndim = (*dim_ptr) >> 1                          // tagged length; loss-of-tag → /2
    CHECK(ndim >= 1)                                // "input_rank >= 1", reduce_emitter.cc:1407

    last      = ndim - 1
    lane_bit  = bit_test(this->minor_mask  @ [0x590], last)   // reduces minormost (lane) axis?
    sub_bit   = (ndim == 1) ? 1 : bit_test(same_mask, ndim-2) // reduces second-minor (sublane) axis?

    // window shape sign-probes pick the operand-shape pointer to forward (0x268/0x270, 0x168/0x170)

    // 5. Four-way dispatch on (lane_bit, sub_bit)
    if  sub_bit && lane_bit: r = EmitLaneAndSublaneReduction(...)  // 0x13e122c0
    elif lane_bit:           r = EmitLaneReduction(...)            // 0x13e12d80
    elif sub_bit:            r = EmitSublaneReduction(...)         // 0x13e15260
    else:                    r = EmitMajorReduction(...)           // 0x13e15b00  (leading/major-dim)

    free(buf)
    return r
子 emitter地址当 reduce 触及这些轴时触发…
EmitKeepSublanesReduction0x13e11f40设置了 this->[0x610] keep-sublanes 标志（必须保留 sublane 布局的 segmented/windowed reduce）
EmitLaneAndSublaneReduction0x13e122c0同时触及 lane（minormost）和 sublane（second-minor）轴 — 完整的 2D in-tile reduction
EmitLaneReduction0x13e12d80只触及 lane 轴（llo.v*.xlane.* cross-lane tree）
EmitSublaneReduction0x13e15260只触及 sublane 轴（llo.v*.slane.*）
EmitMajorReduction0x13e15b00不触及任何 in-tile 轴 — 对 leading/major 维的 reduction，跨 tile 累加

注意 — lane/sublane 拆分在这里决定，而不是在 ReduceFuser 中。 ReduceFuser 只知道它有一个 kReduce；是 EmitReduction 通过以 minormost（ndim-1）和 second-minor（ndim-2）位置对 this->[0x590] 处的 reduced-dimension mask 进行 bit-test，将 reduced HLO 维度映射到物理 xlane/slane 轴。四路 fan-out 解释了为什么 ReduceFuser 行列出 llo.vmax.{x,s}lane.* — 这两个后缀正是 lane-only 和 sublane-only emitter；LaneAndSublane 与 Major case 会组合或顺序执行它们。[置信度：HIGH — 每个地址都是 0x13e16240 窗口中解析出的直接 call/jmp 目标；CHECK 字符串（output_span->memory_space() == MemorySpace::kVmem、input_rank >= 1）从 FATAL 站点的 .rodata 读取。]

怪癖 — VMEM 驻留是硬 CHECK，而不是 fallback。 在任何轴逻辑之前，EmitReduction 会对每个输出 operand 断言 CHECK_EQ(v->memory_space(), kVmem)（reduce_emitter.cc:1400）。reduce emitter 没有 spill 路径：输出落在 HBM 或非 VMEM 空间中的 reduce 是编译器 invariant 违规，会 abort，而不是慢路径。这是上面 FusionFitsInVmem gate 在 lowering 侧的镜像 — fusion 必须已经保证 reduce 输出适合放在 VMEM 中，否则这个 CHECK 会触发。[置信度：HIGH — 对 byte-0xb 的 (>>2)&0x1f memory_space 字段与 0x3 比较，并在第 0x578=1400 行经 MakeCheckOpString<MemorySpace,MemorySpace> FATAL。]

相关组件 ​

组件关系
TpuPriorityFusionQueue对本页谓词允许的候选排序；数值公式见 fusion-cost-model.md
DotCanonicalizer / ConvolutionFolding在 fusion 前运行；把每个 kDot 变成本 pass 匹配的 kConvolution — 见 dot-conv-mxu-lowering.md
LayoutAssignment在 fusion 前运行；fusion 合法性依赖已确定的 tile layout — 见 layout-assignment.md
MXU LMR transform（mxu_lmr_transform.cc；ReplaceMatmulsWithMatmulLmrs，由 xla_tpu_use_interleaving_lmr_transform gate）Post-lowering；将 (matprep+matmul+matres)→LMR 折叠，释放让 PE+ACT bundle-packing 在物理上可行的 ACT VLIW slot
TpuLoopFusionEnhancer在主 pass 后运行；扩展已有 kLoop fusion 边界，以吸收更多 elementwise leaf

交叉引用 ​

  • compile-phases.md — 顶层阶段顺序；"Main fusion" 是 Phase 5，位于 layout assignment 之后
  • fusion-cost-model.md — 数值收益部分：CalculateProducerPriority* 系数、MXU 延迟表、VMEM scavenging cycle budget
  • dot-conv-mxu-lowering.md — DotCanonicalizer/ConvolutionFolding/SpatialMajorConvolution；为什么这个 pass 只会看到 convolution
  • layout-assignment.md — 在 fusion 前完成并约束 fusion 合法性的 layout fixing
  • overview.md — 编译器流水线概览，以及 fusion pass 家族的位置
  • ../cost/overview.md — EstimateFusionCost/GetHloCycles 查询以获得 VMEM 和 cycle 估计的成本分析子系统