编译流程导览
本页所有地址均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so（wheel 版本0.0.40；运行时报告的0.103无法在二进制中静态验证，因此 build-id 是明确锚点：89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，未剥离符号 — 完整 C++ 符号；.textVA == file offset）。其他 wheel 版本会不同；应把每个 VA 都视为绑定到该版本。
摘要
本页跟踪一个程序 — 单个 bf16 matmul — 从框架跨过 PjRt 边界交给编译器的 StableHLO module，一直到 TPU TensorCore 发射的打包 VLIW bundle 字节。它是编译器的定向地图：一条叙事主线，命名每个 IR 层级、该层使用的 IR/dialect、该层真正重要的少数 pass，以及拥有完整细节的深入页面。先读这里；再跳入每个阶段链接到的页面。
这条下降路径不是教科书式 MLIR 级联 HLO → MHLO → tpu → LLO。在这个 shared object 内有两棵不同的 lowering tree，且只有其中一棵是 TPU device path。device path（我们的 matmul 采用的路径）是 HLO → [≈97 HLO pre-passes] → HLO → [layout / fusion / scheduling / memory] → HLO → [jellyfish per-op Emitters via LloRegionBuilder] → LLO → [bundle packer] → per-gen bundle bytes。通用 tensor IR 由数百个 C++ *Emitter class 直接 lower 到 LLO — 它从不会变成 tpu-dialect MLIR。mlir::tpu（Mosaic）dialect 和 XTile dialect 都随二进制一起发布，但位于 device path 之外：tpu 只会作为从 Pallas/Mosaic kCustomCall 导入的序列化 MLIR module 出现，而 XTile 是 XLA CPU/GPU tiled-fusion codegen dialect，matmul lowering 永远不会到达它。对这个编译器最常见的地图正是在这一点上出错；下面的逐阶段小节会标出位置。
参考框架是 XLA-on-LLVM。如果你熟悉 XLA 的 HLO optimization pipeline、layout assignment、LatencyHidingScheduler 和 MemorySpaceAssignment（MSA），你已经理解七个阶段中的四个 — 它们是开源 pass，以 TPU Target 运行，并与约 137 个 TPU-private xla::jellyfish::* pass 交织。真正 TPU-specific 的下降是最后两个阶段：HLO→LLO emitter wall 和按代 bundle packer。LLO 是 TPU 的低层 VLIW IR — 扁平的 LloInstruction，每个 machine op 一个，是 wire 之上的最后一层 IR。硅相关选择（哪种 MXU data-format、哪种 native vreg shape、每个 bundle 多少字节）很晚才进入，几乎都通过一个 xla::jellyfish::Target object 穿针引线，每个阶段都会查询它。
对重新实现而言，定向契约如下：
- 七个 IR 层级以及驱动每层的单一函数入口：
CompilePhase0StablehloToHlo（0xf84de60）→RunHloPasses（0x1093a420）→LayoutAssignment（0x169bf440）→LatencyHidingScheduler（0x136321a0）→RunMemorySpaceAssignment（0x12fc3080）→ jellyfish*Emitter/LloRegionBuilder→BundlePacker（0x13b206a0）。 - 双树拆分：device path（HLO → LLO direct emit）vs. 已打包但不在路径上的 XTile（CPU/GPU）和 Mosaic-
tpu（Pallas-import）树。 Target作为硅开关：data-format、native vreg shape 和 bundle width 都是从同一个Target参数读取的按代值，而不是分支到某个全局变量。
| 前门输入 | 跨 PjRt 边界的 StableHLO MLIR bytecode（混合 CHLO/VHLO） |
| Wire 输出 | 按代 VLIW bundle bytes — JF 41 B / PF 51 B / VF·GL 64 B |
| 单体驱动器 | DeepseaCompilerBase::RunHloPasses @ 0x1093a420 |
| Separate-compile 阶段 | phase0_stablehlo_to_hlo · phase1_hlo_opts · phase2a_tlp_lowering · phase2b_deduped_lowering · phase3_linking |
| 硅开关 | xla::jellyfish::Target const& — 传入每个阶段 |
| 我们跟踪的 matmul | 一个 bf16 × bf16 → f32 kDot，收缩一个维度 |
阶段地图一览
我们的 matmul 会经过七个 IR 层级、每层使用的 dialect、该处的关键 pass，以及拥有完整叙述的页面。Tree 标记该层是否位于 TPU device path（D）上，或位于某棵已打包但不在路径上的树（X = XLA CPU/GPU XTile，M = Mosaic-tpu Pallas import）。
| # | IR 层级 | Dialect / 形式 | 关键 pass（入口 VA） | Tree | 归属深入页面 |
|---|---|---|---|---|---|
| 0 | 前门 bytecode | StableHLO + CHLO/VHLO | CompilePhase0StablehloToHlo（0xf84de60） | D | HLO 摄入 |
| 1 | HLO pre-passes | XLA HLO | RunHloPasses（0x1093a420）；TpuHloSupportChecker（0x11071480） | D | HLO Pre-Passes · Compile Phases |
| 2 | Layout / sharding | XLA HLO + layouts | LayoutAssignment（0x169bf440）；TpuLayoutAssignment（0x110ace00） | D | Layout Assignment · Sharding Propagation |
| 3 | Fusion + scheduling | XLA HLO（scheduled） | LatencyHidingScheduler（0x136321a0）；TpuInstructionFusion | D | LHS Core · Fusion Patterns |
| 4 | Memory-space assignment | XLA HLO + memspace | RunMemorySpaceAssignment（0x12fc3080） | D | MSA Overview |
| 5 | HLO → LLO emit | LLO（VLIW IR） | jellyfish *Emitter via LloRegionBuilder；MatrixMultiplyAccumulateFunctor::operator()（0x1310cd80） | D | Dot/Conv → MXU Lowering · LLO Opcode Enum |
| 6 | Bundle packing | vector<Bundle> | BundlePacker::runOnMachineFunction（0x13b206a0） | D | LLO Bundle Packing |
| 7 | Bundle bytes | 按代 wire word | Encoder<gen>::EncodeBundleInternal | D | Bundle Model |
| — | Mosaic tpu import | mlir::tpu MLIR | createLowerToLLOPass（0x11203ba0）；通过 GetMlirModuleOpFromCustomCall（0x13e327a0）导入 | M | tpu → LLO ODS · Mosaic Overview |
| — | XLA XTile codegen | xla::xtile MLIR | StablehloLowerToXtilePass（0x15060560） | X | MHLO/XTile/tpu Lowering |
GOTCHA — 第 0–7 行是被跟踪 matmul 实际走过的路径。最后两行是已打包但不在路径上：普通
kDot永远不会访问tpudialect 或 XTile。tpudialect 只会出现在框架作为序列化 MLIR 交入的 Pallas/Mosaic kernel 中；XTile 是 XLA CPU/GPU backend 的 tiled-fusion codegen，只依赖 LLVM/CPU dialect set。把其中任意一个误认为 “TPU lowering of MHLO”，是阅读这个编译器时最常见的错误。
Stage 0 — 前门：StableHLO bytecode → HLO
我们的 matmul 不是以 XLA HLO 进入的。它以可移植 MLIR bytecode 进入 — 一个混有 CHLO 和 VHLO op 的 StableHLO module，由框架桥（JAX、TF/XLA bridge 或 PyTorch/XLA）序列化并跨过 PjRt 边界。dot 是一个 stablehlo.dot_general，带有 contracting-dimension numbers 和 bf16 operand element type。
编译器的第一步是一次格式跨越，不是优化。xla::CompilePhase0StablehloToHlo（0xf84de60）将 bytecode 解析为内存中的 mlir::ModuleOp，运行一个有序的 MLIR pipeline 来 legalize CHLO→StableHLO→MHLO（createChloLegalize* / createStablehloLegalizeToHlo pass），遍历 MHLO module 并发射 xla::HloProto，随后再把该 proto 解析回 XLA 其余部分所使用的 xla::HloModule / HloInstruction graph。此阶段之后，我们的 stablehlo.dot_general 变成 opcode 为 kDot 的 HloInstruction。从这里向下直到 Stage 5 的所有内容都在该 HLO graph 上操作 — MHLO/StableHLO/VHLO 只存在于此边界，以及作为 JAX-bridge wire format 存在。
深入页面拥有两侧的 bytecode 格式、CHLO/VHLO version-skew 处理，以及 proto round-trip。
→ HLO Ingestion · 下降路径的框架：How to Read This Book
Stage 1 — HLO pre-passes：≈97 行清理
HLO graph 现在运行长长的 pre-lowering pipeline，该 pipeline 由 xla::jellyfish::DeepseaCompilerBase::RunHloPasses（0x1093a420）构造并分派 — 这个函数位于公开 C ABI export 背后，也位于 separate-compilation 流程中的 xla::CompilePhase1HloOptimizations（0xf84ee00）背后。pipeline 容器是普通的 xla::HloPassPipeline；不存在 TPU-private pipeline class。大约 97 个 pass 运行在六个概念阶段中，每个 pass 都直接派生自 xla::HloPassInterface，并由 MaybeAddInvariantCheckers（0x10944600）在每个 pass 之后重新验证结构（HloVerifier、scheduling-annotation 合法性、cycle detection）。
这些 pass 中的大多数会忽略我们的 matmul — 它们展开 custom-call（TpuCholeskyExpander、TpuQrExpander、RNG、gather/scatter）、分解 dynamic shape（DynamicPadder），并执行规范化。两个 pass 会直接触及它。xla::jellyfish::XPrecisionRewriter 会在 dot 携带 x6 / x9 / x128 高精度 annotation 时触发，把它拆成由较低精度 dot 组成的 2、3 或 8 步 accumulation chain；普通 bf16 dot 原样通过。规范接收 gate xla::TpuHloSupportChecker::RunImpl（0x11071480，pass name tpu_hlo_support_checker）会遍历每个 computation，并通过 ShapeUtil::ValidateShapeWithOptionalLayout 验证每个 result Shape — 它从不修改 module，只拒绝 TPU backend 无法 lower 的 HLO。我们的 bf16 dot 通过该 gate。在约 97 个 pass 中，约 33 个是明确的 TPU-private（xla::jellyfish::*）；其余是以 TPU Target 运行的开源 XLA。
→ HLO Pre-Passes · Compile Phases · HLO Pass Registry · Algebraic Simplifier
Stage 2 — Layout 与 sharding：数据落在哪里
Sharding 先运行（在 pipeline 两半之间）：xla::ShardingPropagation 处理手动 annotation，或 xla::TpuAutoSharding + xla::sdy::ShardyXLA 处理 auto / Shardy 流，随后 xla::jellyfish::TpuSpmdPartitioner 执行每个 partition 的 rewrite。我们的单设备 matmul 是 Replicated，所以 sharding 对它是 no-op，但多芯片 matmul 会在这里获得自己的 collective。
然后 layout assignment 为每个数组决定物理 minor-to-major order。驱动器是开源的 xla::LayoutAssignment::RunImpl（0x169bf440），并由 xla::jellyfish::TpuLayoutAssignment::RunImpl（0x110ace00）针对 TPU 专门化，后者增加 TPU-specific constraint — 包括 gather/scatter 的 indices-layout constraint，以及 MXU 偏好的 tiling rule。对我们的 matmul 而言，这是一个关键决策：为 LHS、RHS 和 result 选择的 layout 决定 systolic array 是否可以不经 transpose 直接 stream operand，而 tiling assignment（TpuTilingAssignment、WindowConfigAssignment）固定 MXU emitter 稍后会循环遍历的 tile shape。此阶段是硅相关选择第一次进入的位置 — native tile geometry — 从 Target 读取。
→ Layout Assignment · Sharding Propagation · Auto-Sharding / SPMD
Stage 3 — Fusion 与 scheduling：工作的顺序
layout 固定后，fusion 会把 producer/consumer op 分组进 fusion computation。TPU 的主 fusion pass 是 xla::jellyfish::TpuInstructionFusion（pass name tpu_fusion），它在 “Pre main fusion” / “Main fusion” / “Post main fusion” pipeline 阶段中驱动 priority-fusion queue。如果我们的 matmul 后面跟着 elementwise bias-add 或 activation，fusion 会将它们折叠成一个 fusion node，使 emitter 能把结果保留在 vreg 中，而不是往返内存。
随后 scheduling 将 graph 线性化。规范 pass 是 xla::LatencyHidingScheduler::RunImpl（0x136321a0，name latency-hiding-scheduler）— 这是其他地方也使用的同一个 XLA LHS，在这里以 TPU cost model 运行，以便把我们的 matmul 的 MXU 工作与供给和排出它的 DMA 重叠。对于 greedy list scheduling 重叠不足的情况，还存在一个带 ILP 风味的变体和一个 per-layer scheduler（LatencyHidingLayerScheduler）；scheduler 阶段运行 async_scheduling HloPassPipeline（携带 LHS 本身的经字节验证的 pipeline name），位于 base memory-scheduling pass 之后。这里产生的 schedule 是很久之后的 bundle packer 消费 IR 顺序的来源 — 二者是在不同粒度上回答同一个“何时运行什么”的问题（HLO instruction order vs. VLIW slot fill）。
NOTE — macro scheduler（LHS，本阶段）和 micro scheduler（Stage 6 的 per-region greedy bundle packer）彼此独立。LHS 重排 HLO instruction，以便在整个 computation 范围内隐藏 latency；packer 在一个 region 内按 IR 顺序局部填充 VLIW slot。二者都不是 modulo scheduling — inner-loop software pipelining 是第三条独立路径。
→ LatencyHidingScheduler Core · LHS ILP Variant · Fusion Patterns · Fusion Cost Model
Stage 4 — Memory-space assignment：VMEM vs. HBM
已经 scheduled 的 HLO graph 仍然没有说明每个 buffer 在哪里。Memory-space assignment（MSA）把每个 value 放进 HBM、VMEM 或 CMEM，并插入 copy/prefetch，在需要之前把数据移入快速片上 scratch。TPU 驱动器是 xla::jellyfish::RunMemorySpaceAssignment（0x12fc3080），它配置并运行开源引擎 xla::memory_space_assignment::MemorySpaceAssignment::Run（0x1dc2e200）及其 best-fit repacker。对我们的 matmul 而言，MSA 决定 LHS 和 RHS tile 被预取到 VMEM，使 MXU 不会因 HBM latency 停顿，并根据从 Target 读取的按代预算确定 VMEM working set 大小。每个版本的默认预算和 HBM reservation policy 位于各自页面。
→ MSA Overview · MSA Per-Version Defaults · MSA Reservation / HBM Policy · VMEM Allocator
Stage 5 — HLO → LLO：emitter wall（以及 MHLO 不会去哪里）
这是货真价实的 TPU-specific 下降，也是最容易被误描述的阶段。scheduled、laid-out、memory-assigned 的 HLO graph 由数百个 C++ *Emitter class 直接 lower 到 LLO — TPU 的低层 VLIW IR；每个 class 拥有一个 HLO op family，并通过 xla::jellyfish::LloRegionBuilder 在 region 中构造 LLO。通用程序不存在 MHLO→tpu-dialect conversion pass；tpu MLIR dialect 只由 Pallas/Mosaic frontend 产生，并被单独导入（见下方 callout）。
我们的 bf16 matmul 通过与 convolution 共用的一条下降路径到达 systolic array：上游 pass 将 kDot rewrite 为 kConvolution，因此二者共享同一个 lowering。per-window tile-cost comparator 选择 systolic tiling，emission-strategy dispatch（GetEmitFunctorFromEmitFunctorEnum 0x130e8de0，一个 19-case switch）选择某个 MXU strategy，而 codegen 主体 xla::jellyfish::MatrixMultiplyAccumulateFunctor::operator()（0x1310cd80，接受 LloRegionBuilder）生成由 LLO matmul op 组成的 tiled loop nest。matmul atom 是每个 native chunk 严格的 3 或 4 条 instruction 序列 — llo.vmatprep → llo.vmatmul →（可选 pack）→ llo.vmatres（kVectorMatprep* / kVectorMatmul / kVectorMatres opcode band）。data-format 选择（bf16×bf16→f32 vs. 某个 f8 或 int8 变体）以及 MXU register-bank assignment 在这里根据 operand element type 和 Target 选出 — 这是硅相关行为第二次进入的位置，现在粒度是 instruction。融合到 matmul 上的 elementwise op 通过 scalar/vector LLO atom table lower；LLO-level allocation（LloAllocation，interval-tree live-range allocator）和 region analysis 作为此阶段的一部分运行。
NOTE — device-path HLO matmul emitter（
MatrixMultiplyAccumulateFunctor::operator()0x1310cd80，构建到LloRegionBuilder中）与 Mosaic-path atom emittermlir::tpu::(anonymous namespace)::EmitMatmuls（0x11241460）/EmitLatches（0x112403c0）都会产生相同的kVectorMatprep*/kVectorMatmul/kVectorMatresband，但它们是从不同输入到达的不同入口。EmitMatmuls/EmitLatches接受mlir::ConversionPatternRewriter，并且只从createLowerToLLOPass（0x11203ba0）内部的mlir::tpu::MatmulOp/MatmulAccLhsOpconversion pattern 调用 — 也就是下面不在路径上的 Mosaictpu→LLO import，而不是我们的bf16HLO matmul。两条下降路径在 LLO level 汇合到同一个 LLO atom；它们不共享函数。GOTCHA —
tpu和 XTile dialect 不在这条路径上。 通用 MHLO 永远不会 lower 到tpuMLIR。mlir::tpu（Mosaic）dialect 只会作为序列化 MLIR module 进入，该 module 由xla::jellyfish::mlir_utils::GetMlirModuleOpFromCustomCall（0x13e327a0）从 Pallas/MosaickCustomCall中提取；该 module 随后被 canonicalize、推断 vector layout（createInferVectorLayoutPass0x132c2c20、createApplyVectorLayoutPass0x1325cda0），再由mlir::tpu::createLowerToLLOPass（0x11203ba0）lower 到 LLO — 在 LLO level 汇入我们 matmul 的流。另一路，xla::xtiledialect（StablehloLowerToXtilePass0x15060560）是 XLA CPU/GPU tiled-fusion codegen，只依赖 LLVM/CPU dialect set；它被打包进libtpu.so，但 jellyfish TPU lowering 从不调用它。
→ Dot/Conv → MXU Lowering · LLO Opcode Enum · Slot — MXU · Mosaic-import path: tpu → LLO ODS · Mosaic Overview · XTile（off-path）：MHLO/XTile/tpu Lowering
Stage 6 — Bundle packing：LLO ops → VLIW slots
扁平的 LloInstruction 列表 — 我们的 matmul 现在是由 vmatprep/vmatmul/vmatres 加上其 prefetch DMA 以及任何融合 elementwise op 组成的 loop nest — 会被 BundlePacker::runOnMachineFunction（0x13b206a0，LLVM-backend tpu-bundle-packer pass）打包进定宽 VLIW bundle。算法是forward greedy list scheduling：按 IR 顺序遍历 LLO op，根据 read-after-write dependency 计算每个 op 最早合法的 bundle，询问按代 SlotTracker 找到第一个剩余容量可容纳该 op 的 BundleRequirement 的 bundle，若没有合适 bundle 则追加空 bundle，然后提交。每一代都有不同的 slot grid — MXU slot 数量、vector source port、immediate slot、predicate field — 因此相同的 LLO stream 会在不同硅上打包成不同数量的 bundle。空 slot 会被盖上 kNeverExecute = 31 predicate，使缺失 slot 成为定义明确的 no-op。输出是由 typed sub-field 构成的 vector<Bundle>。
→ LLO Bundle Packing · Bundle Modulo Scheduling · Scheduler Resource Model · Bundle-Aware Cost
Stage 7 — Bundle bytes：按代 wire word
每个 typed Bundle 都由按代 encoder 序列化为 raw bytes。这是硅发挥作用的最后位置，也是影响最大的位置：bundle width 本身是按代常量，由以 (TpuVersion, TpuSequencerType) 为键的 codec-metadata table inline 返回。
| 代（代号） | Bundle width | Codec-metadata anchor |
|---|---|---|
| Jellyfish（≤ v3） | 41 B / 328 bit | JellyfishCodecMetadata::BundleSizeBytes @ 0x1ecf7460 → 41 |
| Pufferfish（v4） | 51 B / 408 bit | EncoderPfTensorCore::BundleSizeBytes @ 0x1d227740 → 0x33 |
| Viperfish（TPU v5 / v5e / v5p） | 64 B / 512 bit | ViperfishCodecMetadata::BundleSizeBytes @ 0x1ee71320 → 64 |
| Ghostlite（TPU v6 lite / v6e） | 64 B / 512 bit | GhostliteCodecMetadata::BundleSizeBytes @ 0x1eeb7640 → 0x40 |
encoder 本身不仅在宽度上不同，在机制上也分叉。Jellyfish 是 direct-pack encoder：EncoderJf::EncodeBundleInternal 用 shl/and/or 算术构造一个 53 字节 scratch struct，并剥掉前 12 字节（output byte N == struct byte 0x0C+N）。Pufferfish 以及每个 V5+ generation 则会 memset 一个零 buffer，并用共享 bit-packing primitive BitCopy(dst, dst_bit, src, src_bit, nbits)（0x1fa0a900）写入每个字段，因此字段的绝对 bundle bit 就是其 BitCopy 参数。同一个 64-byte 宽度由 Viperfish、Ghostlite 和 6acc60406（三个 V5+ generation）共享，但 slot bit layout 不同（Ghostlite 将 opcode 从 7 bit 扩到 8 bit，并把 scalar/sequencer region 后移 +3 bit）。我们的 matmul 的 vmatmul op 会落在 VectorExtended/MXU slot 中，其精确 bit 位置属于这些按代 map 之一。
此阶段之后，程序就是一串准备供 TensorCore fetch 和 issue 的 bundle bytes — 下降完成。
→ Bundle Model Overview · Jellyfish 41-B Bundle · Pufferfish 51-B Bundle · Viperfish 64-B Bundle · Ghostlite Bundle · MC-Emitter · 代号词汇表：Codename Cheat-Sheet
从这里阅读本书其余部分
上面的每个阶段都是一篇深入页面的一段概述；本节是通往完整算法的索引。
| 如果你想了解… | 从这里开始 |
|---|---|
| 精确的 ≈97-pass 顺序和逐 pass rewrite | HLO Pre-Passes |
阶段分类（phaseN_*、monolithic vs. separate compile） | Compile Phases |
| layout 如何选择，以及 MXU tiling constraint | Layout Assignment |
| scheduler 定价所用的 latency cost model | LatencyHidingScheduler Core · Cost Overview |
| buffer 位于何处，以及 prefetch 如何插入 | MSA Overview |
| 完整 matmul → MXU emission 以及 19 种 strategy | Dot/Conv → MXU Lowering |
| LLO opcode space | LLO Opcode Enum |
Mosaic/Pallas tpu-dialect import path | tpu → LLO ODS · Mosaic Overview |
| bit-exact bundle wire format | Bundle Model Overview 以及按代 bundle 页面 |
| “Jellyfish / Pufferfish / Viperfish / Ghostlite” 映射到什么 | Codename Cheat-Sheet |
交叉引用
- HLO Ingestion — Stage 0；本页开头的 StableHLO→HLO 格式跨越
- HLO Pre-Passes — Stage 1；完整的 ≈97 行 pre-lowering pipeline
- Compile Phases — 每个阶段背后的 monolithic vs. separate-compilation phase taxonomy
- Layout Assignment — Stage 2；物理 layout + MXU tiling constraint
- LatencyHidingScheduler Core — Stage 3；macro scheduler
- Fusion Patterns — Stage 3；
TpuInstructionFusion和 fusion bucket - MSA Overview — Stage 4；HBM/VMEM/CMEM placement
- Dot/Conv → MXU Lowering — Stage 5；完整的 matmul emitter wall
- LLO Opcode Enum — 我们的 matmul 在 Stage 5 变成的 LLO IR
- tpu → LLO ODS — 不在路径上的 Mosaic
tpu-dialect import，它在 LLO 重新汇入 - MHLO/XTile/tpu Lowering — 不在路径上的 XTile（CPU/GPU）树以及双树拆分
- LLO Bundle Packing — Stage 6；greedy slot-fill list scheduling
- Bundle Model Overview — Stage 7；按代 wire-format model
- Codename Cheat-Sheet — 本页用于硅相关行为的 generation name
- How to Read This Book — 证据约定和全书使用的 confidence label
- Glossary — 本导览依赖的 IR 和 unit 术语（LLO、MXU、VMEM、bundle）定义
- Subsystem Map — 这里跟踪的 compiler stage 位于九个 subsystem 中的位置