TPU 编译器

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。

摘要

TPU 编译器是 libtpu.so 中把序列化 XLA 程序转换为 TensorCore 执行的打包 VLIW bundle 的那一半。它不是单体 lowering。它是一个五阶段、多方言的下降过程，从跨 PjRt 边界传入的可移植 StableHLO bytecode 开始，最终变成按代际区分的 bundle byte stream，中途穿过六个连续 IR 层级和四类优化机制。整个下降过程可以从一个类进入 — xla::jellyfish::DeepseaCompilerBase — 其 RunHloPasses (0x109152a0) 和 RunBackend 家族驱动 HLO 层工作；也可以从 phase registry xla::TpuCompiler::RegisterAllPhases (0xf849ec0) 进入，该 registry 将五个 xla::CompilePhase* 入口点接入 PjRt partial-program 流水线。

熟悉 LLVM 的读者可以保持一个分层类比。流水线前半段是 XLA proper — 基于 HloInstruction 的 HLO 优化器，包含 layout assignment、fusion、memory-space allocator（MSA，XLA 中类似寄存器分配但面向 HBM/VMEM 层级的机制）以及 latency-hiding scheduler — 结构上与上游 XLA 为 GPU 提供的机制相同，只是以 TPU cost model 和 TPU layout model 重新定向。后半段是 MLIR — 程序从 HloInstruction 跨入 MLIR 世界（StableHLO/CHLO/MHLO 在摄入时已经处于这里；tpu 方言是 TPU 专用目标方言），并经 dialect-conversion legalizer 下降到 jellyfish::Llo* 低层 IR，也就是 TPU 的 LLVM MachineInstr 等价物。从 LLO 开始，调度流水线打包 bundles，并由按代际区分的 encoder 序列化字节。还有一条侧门 — Mosaic kernel compiler — 允许手写 kernel 通过 tpu_custom_call HLO 直接进入 tpu 方言，完全绕过 HLO 优化器。

这是 Part V 的导览页。它固定 IR 层栈，命名五个编译阶段和运行在其中的四个 pass 家族，说明它们在二进制中的位置，并映射 Part V 的其余页面，使读者可以跳到负责每个阶段细节的页面。它不会复述任何阶段的算法 — HLO pass registry、layout assignment、MSA、fusion、MLIR legalizers、tpu→LLO ODS lowering 和 Mosaic 都有各自页面，并在下文交叉引用。

对于重新实现，导览契约是：

IR 层栈及其 crossing 点。 六个层级 — StableHLO/CHLO/VHLO bytecode → XLA HLO → MHLO/tpu-dialect MLIR（via TLP）→ LLO → scheduled Bundles → packed bytes — 以及两个属于格式 crossing 而非普通 pass 的边界：摄入时的 StableHLO↔HLO legalization，以及 Phase 2a 内的 HLO→MLIR crossing。
五阶段主干及每阶段消费/产出。 CompilePhase0StablehloToHlo 到 CompilePhase3Linking，由 RegisterAllPhases 注册，每个都是 PjRtPartialProgramProto-in / PjRtPartialProgramProto-out 阶段。
四个 pass 家族及其位置。 HLO optimization passes（Phase 1）、约束 codegen 的分析（layout assignment、MSA）、lowering legalizers（Phase 2a/2b），以及 schedule+pack 后端（Part VIII，LLO 之后）。
Mosaic 侧通道。 为什么存在第二个 IR 生产者，以及它的输出如何在 tpu 方言处重新加入主下降过程。


编译器类	`xla::jellyfish::DeepseaCompilerBase`（Deepsea = internal TPU backend name）
HLO pass driver	public entry `xla::jellyfish::DeepseaCompilerBase::RunHloPasses(unique_ptr<HloModule>, StreamExecutor*, CompileOptions const&)` @ `0x109152a0` — `xla::Compiler` virtual override，它委托给实际构建并运行 `HloPassPipeline` 的内部 pipeline-running overload `RunHloPasses(Target const&, RunHloPassesConfig const&, …)` @ `0x1093a420`（C 导出 `TpuCompiler_RunHloPasses` @ `0xeabcd80` 是 PjRt 侧 wrapper）。见编译阶段。
Phase registry	`xla::TpuCompiler::RegisterAllPhases` @ `0xf849ec0`
Phase entry points	`CompilePhase0StablehloToHlo` `0xf84de60` · `CompilePhase1HloOptimizations` `0xf84ee00` · `CompilePhase2aTlpLowering` `0xf850840` · `CompilePhase2bDedupedLowering` `0xf852180` · `CompilePhase3Linking` `0xf852f40`
Phase I/O 类型	`absl::Span<const xla::PjRtPartialProgramProto>` in / out（PjRt partial-program protocol）
IR 层级	StableHLO/CHLO/VHLO → XLA HLO → TLP/MHLO + `tpu` dialect → LLO (`jellyfish::Llo*`) → `Bundle` → packed bytes
存在的 MLIR 方言	`stablehlo`, `chlo`, `vhlo`, `mhlo`, `tpu`（RTTI 中已验证：`mlir::{mhlo,tpu,chlo,vhlo,stablehlo}::*`）
tpu→LLO lowering	`mlir::tpu::createLowerToLLOPass` @ `0x11203ba0` (`LowerToLLOPass`); `mlir::tpu::LowerPassBase` family
LLO → bundles	`PackBundles` @ `0x10a30a20` (back end; see Part VIII)
Mosaic 侧门	`tpu_custom_call` HLO → `tpu` dialect via `mlir::tpu::CanonicalizeMosaicPass` / `MosaicSerdePass`; LLO emission driven by `jellyfish::CustomCallEmitter::Emit` @ `0x111ef740` (`MosaicEmitter::EmitWindow` @ `0xfaadcc0` is its per-window helper)
置信度	CONFIRMED（byte-anchored），除非某行或 callout 另有说明

IR 层栈

TPU 程序不是通过一个 IR lowering 的，而是穿过六层栈下降，其中两个位于上游 XLA 的 HloInstruction 世界，四个位于 MLIR。两个世界之间的 crossing 是本页最重要的结构事实：前半段是 XLA，后半段是 MLIR，程序是在边界处重构的，而不只是被遍历。

text

┌──────────────────────────────────────────────────────────────────────────┐
│  LEVEL 0 — PORTABLE WIRE IR (MLIR bytecode, crosses the PjRt boundary)     │
│    StableHLO · CHLO (high-level "client" HLO) · VHLO (versioned)           │
│    mlir::{stablehlo,chlo,vhlo}::*   — what JAX/the bridge actually ships    │
└──────────────────────────────────────────────────────────────────────────┘
        │  Phase 0 — CompilePhase0StablehloToHlo  (format crossing, not a pass)
        │  ConvertStablehloToHlo / StablehloLegalizeToHlo / ChloLegalizeToHlo
        ▼
┌──────────────────────────────────────────────────────────────────────────┐
│  LEVEL 1 — XLA HLO   (HloModule / HloInstruction, the HloOpcode enum)      │
│    the optimizer's home: HLO opt passes, layout assignment, fusion, MSA    │
│    DeepseaCompilerBase::RunHloPasses  0x109152a0                           │
└──────────────────────────────────────────────────────────────────────────┘
        │  Phase 2a — CompilePhase2aTlpLowering  (HLO → MLIR crossing)
        │  TlpLowering → TLPFunction (TPU-Level Program); enters MLIR
        ▼
┌──────────────────────────────────────────────────────────────────────────┐
│  LEVEL 2 — MLIR HIGH/MID   (MHLO + TLP) → tpu target dialect               │
│    mlir::mhlo::* → mlir::tpu::*   ;  Mosaic kernels join HERE via          │
│    tpu_custom_call (CanonicalizeMosaicPass / MosaicSerdePass)             │
└──────────────────────────────────────────────────────────────────────────┘
        │  mlir::tpu::LowerPassBase family → createLowerToLLOPass 0x11203ba0
        │  (DialectConversion legalizers; ApplyVectorLayoutPass; LowerToMloPass DMA)
        ▼
┌──────────────────────────────────────────────────────────────────────────┐
│  LEVEL 3 — LLO   (jellyfish::LloModule / LloInstruction / LloRegion)       │
│    the "MachineInstr" of the TPU: opcodes, allocations, dependency graph   │
│    LSRAv2 live-range alloc + MSA placement resolve here                    │
└──────────────────────────────────────────────────────────────────────────┘
        │  Part VIII — schedule + pack  (PackBundles 0x10a30a20)
        ▼
┌──────────────────────────────────────────────────────────────────────────┐
│  LEVEL 4 — SCHEDULED BUNDLES   (vector<jellyfish::Bundle>)                 │
│    MXU/MRB assignment → bundle packing → modulo scheduling for loops       │
└──────────────────────────────────────────────────────────────────────────┘
        │  Phase 3 — CompilePhase3Linking ; per-gen Encoder*::EncodeBundle
        ▼
┌──────────────────────────────────────────────────────────────────────────┐
│  LEVEL 5 — PACKED PROGRAM   (TpuProgram, serialized byte stream)           │
└──────────────────────────────────────────────────────────────────────────┘
```text

### 两个世界和 crossing

Level 0 和 1 分别是 MLIR-bytecode 和 XLA-HLO；二者都早于 TPU target，属于上游 XLA。StableHLO/CHLO/VHLO 是 framework bridge（JAX）实际序列化的*可移植* wire IR — 它们是 MLIR 方言，但在此时它们是 payload，还未受 TPU 控制。Level 1 是经典 XLA HLO：`HloModule`/`HloInstruction` 图，XLA 的整个优化器都针对它编写，约有 100 个 `HloOpcode`。Phase 0 是把可移植 bytecode 导入 HLO 的格式 crossing；它是*等价表示之间的 legalization*，不是优化，因此它是独立阶段，而不是 Phase 1 内的 pass。

Level 2 到 5 是 TPU 专用的。Phase 2a（`CompilePhase2aTlpLowering`，`0xf850840`）是通过 TPU-Level Program（TLP）表示从 `HloInstruction` 跨入 MLIR 的 crossing。`tpu` 方言（覆盖 MXU matmul、DMA enqueue/wait、semaphores、vector relayout、pack/unpack、PRNG 和 trace）是 TPU 的机器方言。**Level-2 `tpu` 层恰好由一个生产者到达 — 导入的 Mosaic `tpu_custom_call` kernels — 而不是通过 lowering 通用 MHLO 到达**（二进制没有 MHLO→`tpu` pass）；通用 HLO 由 jellyfish `*Emitter` 家族直接发射到 LLO，完全不会进入 `tpu` 方言。见 [MHLO → XTile → tpu](mhlo-xtile-tpu-lowering.md) 和 [Mosaic 概览](mosaic-overview.md) 上的双树图。`tpu` 方言通过 `mlir::tpu::LowerPassBase` 家族 / `createLowerToLLOPass` (`0x11203ba0`) lower 到 LLO，也就是 `jellyfish::Llo*` 家族 — `LloModule`、`LloInstruction`、`LloRegion`、`LloOpcode`、`LloAllocation`、`LloDependencyGraph` — 它是真正的“机器指令”层：pre-bundle、post-MLIR，live-range allocation（LSRAv2）和 memory-space placement 在这里 resolve。[调度后端](../sched/overview.md)随后将 LLO 打包成 `Bundle`s，Phase 3（`CompilePhase3Linking`，`0xf852f40`）完成最终处理并序列化 `TpuProgram`。

> **NOTE — “TLP” 和 “MHLO” 是同一层，而不是两层。** TLP（TPU-Level Program）命名 Phase 2a import 机制 — 它仅有的直接锚点是 phase 符号 `CompilePhase2aTlpLowering` (`0xf850840`) 和 LLVM-backend predicate `llvm::TPU::isTLPFunction` (`0x13b67a60`)；MHLO 是程序到达时所表达的 MLIR 方言。它们不是层叠层级 — Level 2 是“MLIR mid-level”，TLP 命名 import 机制，而 MHLO/`tpu` 命名方言。重新实现者应把 Phase 2a 视为 “HLO → MLIR”，而不要寻找具有自身 op set 的独立 TLP IR；二进制中无法隔离出 `tlp.` op-model vtables 或独立的 `TlpLowering`/`TLPFunction` 类型（TLP 是否有独立 op enum 为 LOW confidence）。
>
> **QUIRK — 方言 op 表看起来很大，是因为 bytecode 方言完整随附。** RTTI 确认 StableHLO、MHLO（包括 `mhlo.async_bundle`）、CHLO（高层 “client” HLO）和 VHLO（versioned 方言，带 `_v1`/`_v2` 后缀 op 名）都作为完整 op 表存在 — 每个都贡献数千个 `mlir::{stablehlo,mhlo,chlo,vhlo}::*` 符号。只有一小部分在 TPU 路径上可达 — 特别是 VHLO 纯粹为*wire-format forward compatibility* 存在（较旧 runtime 忽略较新字段），大多数 VHLO/CHLO op 在 Phase 0 中 legalization 掉。为每个方言 op 构建 handler 的重新实现是在为 import surface 构建，而不是 codegen surface；codegen surface 是 HLO opcodes 加 `tpu` 方言 ops。（这里没有字节锚定精确的逐方言 op-enum 大小 — 对 `<dialect>.` op-name 前缀的字符串扫描会被截断字符串以及 type/attribute 助记符污染，因此不声称精确 op 计数。）

---

## 五个编译阶段

下降过程组织为由 `xla::TpuCompiler::RegisterAllPhases` (`0xf849ec0`) 注册的五个阶段。每个都是 `xla` 命名空间中的自由函数，接收并返回 `absl::Span<const PjRtPartialProgramProto>` — 即 PjRt *partial program* 协议，这使运行时可以在任何阶段后停止（用于 AOT、缓存或分阶段编译），并从序列化中间体恢复。这个 phase boundary 就是为什么每个阶段都是独立符号，而非内联块：每个阶段都可独立调用，且其输出可序列化。

| 阶段 | 入口点 | 地址 | Crossing | 负责内容 |
|---|---|---|---|---|
| **0 — Stablehlo→Hlo** | `CompilePhase0StablehloToHlo` | `0xf84de60` | wire bytecode → XLA HLO | StableHLO/CHLO/VHLO legalization，HLO import |
| **1 — Hlo optimizations** | `CompilePhase1HloOptimizations` | `0xf84ee00` | HLO → HLO | 整个 HLO pass pipeline（opt, layout, fusion, MSA, schedule） |
| **2a — Tlp lowering** | `CompilePhase2aTlpLowering` | `0xf850840` | HLO → MLIR (`tpu`) → LLO | HLO→MLIR crossing，dialect legalization，tpu→LLO |
| **2b — Deduped lowering** | `CompilePhase2bDedupedLowering` | `0xf852180` | LLO → LLO | lowered output 的跨 region deduplication |
| **3 — Linking** | `CompilePhase3Linking` | `0xf852f40` | LLO/bundles → packed | bundle finalization、encode、`TpuProgram` serialization |

单个最重的阶段是 Phase 1：这是 `DeepseaCompilerBase::RunHloPasses` (`0x109152a0`) 运行完整 HLO pass pipeline 的地方，也是下一节描述的四个 optimization-pass 家族所在的位置。Phase 2a 和 2b 是 lowering 下降 — 2a 完成进入 MLIR 并下降到 LLO 的主体 crossing，2b 做 deduplicate：`CompilePhase2bDedupedLowering` 名称及其 LLO-in/LLO-out 位置表明它会折叠相同的 lowered regions，使后端只打包每个唯一 region 一次（dedup 算法本身未被逆向 — 见 [编译阶段](compile-phases.md)；精确 dedup key 为 MEDIUM）。

> **GOTCHA — Phase 1 不是“先优化再 lower”；layout assignment 会改变 lowering 依赖的程序形状。** 在 Phase 1 内，layout assignment 和 MSA 是*会重写 HLO 的分析*（它们插入 copies、选择 tiled layouts 并分配 memory spaces），fusion 和 latency-hiding scheduler 针对 post-layout 程序运行。把 Phase 1 当作纯优化器、并把所有 layout/memory 决策推迟到 MLIR 侧的重新实现者，会产生 tile shapes 和 memory spaces 均未分配的 `tpu`-dialect 程序，`tpu`→LLO legalizer 无法 lower 它。HLO 半段*提交* layout 和 memory space；MLIR 半段*实现*它们。

---

## 四个 Pass 家族

在五个阶段内部，实际变换工作聚集为四个家族。每个家族都有自己的细节页面；本节放置它们并说明它们决定什么。

### 家族 1 — HLO optimization passes（Phase 1）

经典 XLA pass pipeline：algebraic simplification、sharding propagation 和 SPMD partitioning、dynamic-shape handling、optimization barriers、custom-call lowering，以及大量 canonicalizations。这些在 `RunHloPasses` 下的 `HloInstruction` 上运行，并由 HLO pass registry 注册/排序。它们是*target-retargeted upstream XLA*：XLA 为每个 backend 提供的同一批 pass 类，通过 `xla_jf_*` flag 家族（Jellyfish；主导家族，数百个 flags）、`xla_tpu_*`（泛 TPU codegen）以及按代际区分的前缀（`xla_sc_*` SparseCore、`xla_vf_*` Viperfish、`xla_gf_*` 6acc60406、`xla_pf_*` Pufferfish）驱动 TPU 专用启用。细节：[HLO Pass Registry](hlo-pass-registry.md)、[HLO 预 pass](hlo-pre-passes.md)、[代数简化器](algebraic-simplifier.md)、[分片传播](sharding-propagation.md)、[自动分片 / SPMD](auto-sharding-spmd.md)、[动态形状支持](dynamic-shape-support.md)、[优化屏障](optimization-barrier.md)、[Custom-Call Lowering](custom-call-lowering.md)。

### 家族 2 — codegen-gating 分析：layout assignment + MSA（Phase 1）

两个分析会提交 lowering 所依赖的物理决策。**Layout assignment**（`LayoutAssignment`）选择每个 tensor 的 tiled memory layout — 它是 TPU 中选择 SIMD 机器数据布局的类似物，只是 tiling 会与 MXU 的 128×128 systolic geometry 交互。**MSA**（`MemorySpaceAssignment`）是 XLA 的 memory-space assignment：它决定哪些 buffers 位于快速 VMEM 而非 HBM，以及何时跨层级 prefetch/evict，由 `xla_msa_*` flag 家族控制 — 它是 TPU 中最接近寄存器分配的机制，但面向内存层级而非寄存器。细节：[Layout Assignment](layout-assignment.md)、[MSA 概览](msa-overview.md)、[MSA AllocateSegment](msa-allocate-segment.md)、[MSA Per-Version Defaults](msa-per-version-defaults.md)、[MSA Reservation & HBM Policy](msa-reservation-hbm-policy.md)。

### 家族 3 — fusion（Phase 1）

Fusion 将 elementwise 和 reduction ops 聚类成单个 kernel，以减少内存流量，并由 TPU cost model 定价。`HloFusion` 和 `FusionPipeline` 驱动它；profitability 决策读取 Part VII cost model。细节：[融合模式](fusion-patterns.md)、[融合成本模型](fusion-cost-model.md)，以及 fusion 所馈送的 dot/conv → MXU lowering：[Dot / Conv → MXU Lowering](dot-conv-mxu-lowering.md)、[RaggedDot → Convolution](raggeddot-convolution.md)。

### 家族 4 — MLIR lowering legalizers + schedule/pack 后端（Phases 2a/2b/3）

MLIR 侧 lowering 是以 `tpu` → LLO（`createLowerToLLOPass`，`0x11203ba0`）为中心的一串 DialectConversion legalizers，并包含 vector-layout realization pass（`ApplyVectorLayoutPass`）和 DMA bridge（`createLowerToMloPass`）。**此二进制中没有 MHLO → `tpu` legalizer**（0 个 `*ToTpuPass`/`MhloToTpu`/`StablehloToTpu` 符号）：`tpu` 方言从 Mosaic `tpu_custom_call` 中*导入*并 lower 到 LLO，而通用（非 Mosaic）HLO 由 jellyfish `*Emitter` 家族直接发射到 LLO — 它永远不会变成 `tpu` ops。“MHLO → XTile → `tpu`” 对 Level 2 的解读其实是两棵互不连接的树；见 [MHLO → XTile → tpu](mhlo-xtile-tpu-lowering.md) 中的双树图。LLO 之后，schedule-and-pack 后端（`PackBundles`，`0x10a30a20`）分配 MXU sequences、打包 VLIW bundles，并对 loops 做 modulo-schedules — 这是 Part VIII 的主题。MLIR 侧细节：[MHLO → XTile → tpu](mhlo-xtile-tpu-lowering.md)、[tpu MLIR 方言](tpu-dialect-and-ops.md)、[tpu → LLO lowering](tpu-to-llo-ods.md)、[DialectConversion legalizer](dialect-conversion-legalizer.md)、[ConversionPatternRewriter](conversion-pattern-rewriter.md)、[LowerToMlo DMA bridge](lower-to-mlo-dma-bridge.md)、[MLIR op-model contract](mlir-op-model-contract.md)。调度/打包后端位于 [Part VIII — 调度](../sched/overview.md)；其定价位于 [Part VII — 成本模型](../cost/overview.md)。

---

## Mosaic 侧通道

并非每个 `tpu`-dialect 程序都来自 HLO 优化器。**Mosaic** 是一个独立 kernel compiler，允许手写 kernel（Pallas/JAX-side）直接表达为 `tpu` 方言，并作为 `tpu_custom_call` ops 嵌入 HLO 图。下降过程中，`mlir::tpu::CanonicalizeMosaicPass` 会规范化嵌入的 MLIR，`mlir::tpu::MosaicSerdePass` 处理其序列化，因此 kernel body 在 `tpu` 方言（Level 2）重新加入主下降过程。导入嵌入模块、验证其 ABI、运行 Mosaic pass pipeline 并把结果缝入 parent region 的 LLO-level driver 是 `jellyfish::CustomCallEmitter::Emit` (`0x111ef740`)；较小的 `jellyfish::MosaicEmitter`（`MosaicEmitter::EmitWindow` @ `0xfaadcc0`，接收 `LloRegionBuilder`）是其内部使用的 per-window emit helper，而不是 driver。见 [Mosaic 概览](mosaic-overview.md)。

对重新实现者而言，结构要点是 Mosaic 是*馈入同一目标方言的第二个 IR 生产者*，不是并行后端。Mosaic kernel 跳过整个 HLO 优化器（家族 1–3） — 它到达时已经是 `tpu` ops — 但它在 `tpu` 方言处与优化路径汇合，并从那里共享相同的 tpu→LLO→bundle 下降。Mosaic 携带自己的 layout 机制（`mlir::tpu::VectorLayout`、`TiledLayoutAttr`），因为其 kernels 声明自己的 tiling，而不是继承 layout assignment 的选择。`MosaicFusion` 允许周围 HLO 围绕 custom call 融合，`IsMosaicVerificationEnabled` (`0x14514d40`) gate 一个额外 verifier。细节：[Mosaic 概览](mosaic-overview.md)、[Mosaic layout inference](mosaic-layout-inference.md)、[Mosaic VectorLayout](mosaic-vectorlayout.md)。

> **NOTE — Mosaic 边界是真实 IR seam，也是 SparseCore 进入的位置。** SparseCore 路径（`mosaic_sc::MosaicSCDialect`、`LowerToSparseCoreLlvmPass`）复用同一 custom-call 入口和 dialect-conversion 机制来 target SparseCore engines，而非 TensorCore MXU。其 type lowering 是普通 `LLVMTypeConverter` 加上一个 SC address-space conversion lambda — 没有独立的 `SCTypeConverter` 符号，该名称只是这个增强 converter 的文档性句柄。只构建 TensorCore 下降的重新实现者会看到无法 lower 的 `tpu_custom_call`s；custom-call registry 会按 target 将每个 call 路由到其 emitter。见 [LowerToSparseCoreLlvm](lower-to-sparsecore-llvm.md) 和 [SCTypeConverter](sc-type-converter.md)，以及 [SparseCore 部分](../sparsecore/overview.md)。

---

## Part V 地图

Part V 的其余页面负责本概览只定位的细节。按 IR 层级分组：

| 层级 / 家族 | 页面 |
|---|---|
| **Front-end / phases** | [HLO 摄入](hlo-ingestion.md), [编译阶段 0–3](compile-phases.md), [编译缓存](compilation-cache.md) |
| **HLO opt passes** | [HLO Pass Registry](hlo-pass-registry.md), [HLO 预 pass](hlo-pre-passes.md), [代数简化器](algebraic-simplifier.md), [分片传播](sharding-propagation.md), [自动分片 / SPMD](auto-sharding-spmd.md), [动态形状支持](dynamic-shape-support.md), [优化屏障](optimization-barrier.md), [Custom-Call Lowering](custom-call-lowering.md) |
| **Layout / memory** | [Layout Assignment](layout-assignment.md), [MSA 概览](msa-overview.md), [MSA AllocateSegment](msa-allocate-segment.md), [MSA Per-Version Defaults](msa-per-version-defaults.md), [MSA Reservation & HBM Policy](msa-reservation-hbm-policy.md) |
| **Fusion / MXU** | [融合模式](fusion-patterns.md), [融合成本模型](fusion-cost-model.md), [Dot / Conv → MXU lowering](dot-conv-mxu-lowering.md), [RaggedDot → Convolution](raggeddot-convolution.md), [Loop tiling & unrolling](loop-tiling-unrolling.md) |
| **MLIR lowering** | [MHLO → XTile → tpu](mhlo-xtile-tpu-lowering.md), [tpu MLIR 方言](tpu-dialect-and-ops.md), [tpu → LLO lowering](tpu-to-llo-ods.md), [DialectConversion legalizer](dialect-conversion-legalizer.md), [ConversionPatternRewriter](conversion-pattern-rewriter.md), [LowerToMlo DMA bridge](lower-to-mlo-dma-bridge.md), [MLIR op-model contract](mlir-op-model-contract.md), [LlvmTpu intrinsic catalog](llvmtpu-intrinsic-catalog.md) |
| **Mosaic / SparseCore lowering** | [Mosaic 概览](mosaic-overview.md), [Mosaic layout inference](mosaic-layout-inference.md), [Mosaic VectorLayout](mosaic-vectorlayout.md), [LowerToSparseCoreLlvm](lower-to-sparsecore-llvm.md), [SCTypeConverter](sc-type-converter.md) |
| **输出** | [TpuProgram 序列化](tpu-program-serialization.md) |
| **相邻部分** | [Part VIII — 调度](../sched/overview.md), [Part VII — 成本模型](../cost/overview.md), [Part — SparseCore](../sparsecore/overview.md) |

---

## 置信度摘要

| 声明 | 证据 |
|---|---|
| 名为 `CompilePhase0..3` 的五个编译阶段（含 2a/2b 拆分） | symbols `CompilePhase0StablehloToHlo` `0xf84de60`, `…1HloOptimizations` `0xf84ee00`, `…2aTlpLowering` `0xf850840`, `…2bDedupedLowering` `0xf852180`, `…3Linking` `0xf852f40` |
| 阶段由一个 registry 注册，并接收/返回 `PjRtPartialProgramProto` spans | `TpuCompiler::RegisterAllPhases` `0xf849ec0`; phase signatures `absl::Span<const PjRtPartialProgramProto>` |
| HLO pass pipeline 由 `DeepseaCompilerBase::RunHloPasses` 驱动 | `RunHloPasses` `0x109152a0`; C wrapper `TpuCompiler_RunHloPasses` `0xeabcd80`; `RunBackendOnModuleInternal`, `LowerHloModuleImpl`（见于 `XLA_Timer*` guard variables） |
| IR 栈：StableHLO/CHLO/VHLO → HLO → MHLO/`tpu` → LLO → bundles | RTTI `mlir::{stablehlo,chlo,vhlo,mhlo,tpu}::*`; `jellyfish::Llo{Module,Instruction,Region,Opcode}` |
| Phase 0 是 legalization crossing | `ConvertStablehloToHlo`, `StablehloLegalizeToHlo`, `ChloLegalizeToHlo`, `HloLegalizeToStablehlo`, `VhloToVersion` |
| `tpu` 方言通过 `LowerToLLOPass`/`createLowerToLLOPass` lower 到 LLO | `mlir::tpu::createLowerToLLOPass` `0x11203ba0`; `mlir::tpu::LowerPassBase`, `ApplyVectorLayoutPass`, `createLowerToMloPass` `0x1322adc0` |
| Layout assignment、MSA、fusion 是 HLO 侧 codegen 家族 | `LayoutAssignment`, `MemorySpaceAssignment`, `HloFusion`, `FusionPipeline` symbols present |
| LLO 由后端打包为 bundles | `PackBundles` `0x10a30a20`; `jellyfish::BundlePacker`, `GlobalBundlePacker` |
| Mosaic 通过 `tpu_custom_call` 进入 `tpu` 方言 | `tpu_custom_call` strings, `CanonicalizeMosaicPass`, `CustomCallEmitter::Emit` (driver), `MosaicEmitter::EmitWindow` (helper), `MosaicSerdePass`, `MosaicFusion` |
| TLP 是 HLO→MLIR import vehicle，与 MHLO 同层（无单独 op enum） | phase symbol `CompilePhase2aTlpLowering` `0xf850840`, `llvm::TPU::isTLPFunction` `0x13b67a60`; no isolable `tlp.` op-model vtables or standalone `TlpLowering`/`TLPFunction` type |
| Phase 2b key 是跨 region 的 lowered-output deduplication | name `CompilePhase2bDedupedLowering` + LLO-in/LLO-out position; dedup key not traced |

---

## 交叉引用

- [HLO 摄入](hlo-ingestion.md) — Phase 0：StableHLO/CHLO/VHLO bytecode 如何被 legalized 并导入 XLA HLO。
- [编译阶段 0–3](compile-phases.md) — 本页只定位的逐阶段细节：每个 phase 的主体、`RegisterAllPhases` wiring，以及 Phase 2b dedup 机制。
- [编译缓存](compilation-cache.md) — `PjRtPartialProgramProto` 边界如何复用于缓存和 AOT。
- [HLO Pass Registry](hlo-pass-registry.md) — 家族 1：有序 HLO pass pipeline 及其由 `xla_*` flags gate 的启用。
- [Layout Assignment](layout-assignment.md) — 家族 2：MLIR lowering 所实现的 tiled-layout commitment。
- [MSA 概览](msa-overview.md) — 家族 2：跨 VMEM/HBM 层级的 memory-space assignment。
- [融合模式](fusion-patterns.md) — 家族 3：kernel clustering 及其 cost-model profitability gate。
- [MHLO → XTile → tpu](mhlo-xtile-tpu-lowering.md) — 家族 4：Phase 2a 内的 HLO→MLIR mid-level lowering。
- [tpu MLIR 方言](tpu-dialect-and-ops.md) — TPU target dialect，优化器路径和 Mosaic 路径的汇合点。
- [tpu → LLO lowering](tpu-to-llo-ods.md) — 家族 4：`createLowerToLLOPass` 下降到 LLO machine IR。
- [Mosaic 概览](mosaic-overview.md) — 通过 `tpu_custom_call` 直接馈入 `tpu` 方言的侧通道 kernel compiler。
- [TpuProgram 序列化](tpu-program-serialization.md) — Phase 3：打包 bundle stream 如何序列化为 `TpuProgram`。
- [Part VIII — 调度](../sched/overview.md) — 本页交接到的 LLO→bundle schedule-and-pack 后端。
- [Part VII — 成本模型](../cost/overview.md) — 为 fusion、layout、MSA 和 scheduling 决策定价的 cycle/latency model。
- [SparseCore 概览](../sparsecore/overview.md) — custom-call/Mosaic seam 路由到的第二个 target。
- **Binary:** `extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so` (build-id `89edbbe81c5b328a958fe628a9f2207d`)
- **Index entry:** Part V — Compiler: Lowering & Optimization Passes / Front-end and pipeline — [返回索引](../index.md)

TPU 编译器 ​

摘要 ​

IR 层栈 ​

TPU 编译器

摘要

IR 层栈