HLO 预处理 Pass

下面的符号名、VA 和 build-id 适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会不同；请将每个 VA 都视为绑定到该版本。

摘要

HLO 预处理 pass 是一组有序的 xla::HloPassInterface pass，它们在 libtpu 收到的 XLA 程序离开 HloInstruction 世界并进入 MLIR 之前运行。它们把前端自由形式的图清理成编译器其余部分可以假定的规范形态：调度注解被合法化，线性代数和 RNG custom-call 被分解为原始算术，动态形状被静态填充，亚字节和高精度 dtype 被加上边界，sharding 被传播并且程序按设备做 SPMD 分区，layout 被分配，最后一个 pass 将 TPU frontend_attribute 字符串物化为带类型的 backend_config，使 MLIR 导入时看不到任何 TPU 私有内容。这里的每个 pass 都消费 HLO 并发出 HLO；移交给 MLIR 是下一阶段，由 overview.md 负责。

本页是这组预处理 pass 的完整有序枚举，也就是 compile-phases.md 委托到的表。它为每个 pass 固定了：它的 C++ 类、添加它的 builder 函数、它所在的 pipeline-builder 阶段（1–6，与 compile-phases.md 的编号一致）、它消费的 HLO 不变式、它保证输出满足的 HLO 不变式、它是 TPU 专用还是开源 XLA，以及在恢复到入口符号时给出的字节锚定入口符号。这五个 builder 函数位于 deepsea_compiler_hlo_passes.cc 的 (anonymous namespace) 中，并由 DeepseaCompilerBase::RunHloPasses（0x1093a420）到达；后者是公开 C-ABI 导出 TpuCompiler_RunHloPasses 的函数体，也是 separate-compilation 入口 xla::CompilePhase1HloOptimizations（0xf84ee00）的函数体。

这是枚举页面。它不重新推导顶层阶段顺序或每个阶段的入口符号；这条主干位于 compile-phases.md，这里只链接而不重复。它也不重新推导逐 pass 的转换算法（实际的图重写）；这些内容有各自的页面：algebraic-simplifier.md 讲 TpuAlgebraicSimplifier，sharding-propagation.md / auto-sharding-spmd.md 讲 sharding 流程。包含 372 个条目的 HloPassInterface RTTI 目录及 name() 字符串位于 hlo-pass-registry.md。

对于重新实现，预处理 pass 的契约是：

一种容器类型。 xla::HloPassPipeline 是唯一的 HLO pass 容器；不存在 TPU 私有的 pipeline 类。每个 TPU pass 都直接派生自 xla::HloPassInterface（因此 RunImpl(HloModule*, execution_threads) 返回 absl::StatusOr<bool>），并通过 AddPass<T>(...) 添加。
六个阶段，五个 builder。 顺序分解为六个 pipeline-builder 阶段。阶段 1–5 是 HLO→HLO；阶段 6 是最终的 pre-MLIR HLO 域。五个反编译出的 builder 函数添加这些 pass：PreOptimizationPipeline、AddAutoShardingAndRelatedPasses、AddTpuPartitioningPasses、HloOptimizeThroughLayoutAssignment、PostOptimizationPipeline。
预处理 pass 合法地引入 op。 Expander 会发出新 op；partitioner 会发出 collective。因此接受门禁（TpuHloSupportChecker）在阶段 4 运行，位于 expansion 和 partitioning 之后，而不是放在最前面；更早检查会拒绝实际上可以编译的程序。
定点循环。 多个 pass 被包在 xla::HloPassFix<P> 中并重复运行直到收敛；HloDCE 和 HloCSE 无论是否被包裹，都会在大多数阶段之间重新运行。
不变式检查器在每个 pass 后重新运行。 MaybeAddInvariantCheckers（0x10944600）通过单独的 AddInvariantChecker<T> API 在每个嵌套 pipeline 的头部添加 HloVerifier、LegalizeSchedulingAnnotations（作为检查器）和 HloCycleDetection；它们在每个 pass 后重新验证，而不是只验证一次。


Pipeline 驱动器	`xla::jellyfish::DeepseaCompilerBase::RunHloPasses` @ `0x1093a420`
Sep-comp 入口	`xla::CompilePhase1HloOptimizations` @ `0xf84ee00`
Builder — 阶段 1	`(anon)::PreOptimizationPipeline(Target const&, unique_ptr<HloModule>, long, CompilationStats*)`
Builder — 阶段 2	`xla::jellyfish::AddAutoShardingAndRelatedPasses(...)` @ `0x10939c40`
Builder — 阶段 3	`xla::jellyfish::AddTpuPartitioningPasses(HloModule, AliasInfo const, PartitioningPipelineConfig)` @ `0x1278a440`
Builder — 阶段 4	`(anon)::HloOptimizeThroughLayoutAssignment(Target const&, RunHloPassesConfig const&, long, HloModule, ComputationLayout, CompilationStats*)` — 函数体在 `…::$_0::operator()` @ `0x1094ada0`（std::function trampoline `InvokeObject<…$_0>` @ `0x1094ad80`）
Builder — 阶段 5/6	`(anon)::PostOptimizationPipeline(Target const&, AliasInfo const, HloModule, long, CompilationStats*, bool, bool)` @ `0x1093fd40`
Pipeline 开启器	`(anon)::CreateHloPipeline(...)` @ `0x1093efe0`；嵌套变体 `CreateNestedHloPipeline` @ `0x1093a180`、`CreateNestedHloPipelineFix` @ `0x10953e20`
不变式检查器	`(anon)::MaybeAddInvariantCheckers(...)` @ `0x10944600`
Gather/scatter 辅助函数	`(anon)::AddGatherScatterExpanderPasses(Target const&, HloModule*, HloPassPipeline&)` @ `0x1095a040`
接受门禁	`xla::TpuHloSupportChecker::RunImpl` @ `0x11071480`（vtable `_ZTVN3xla20TpuHloSupportCheckerE`）
最后一个 HLO 域 pass	`xla::jellyfish::ConvertFrontendAttributesToBackendConfig`
源单元	`platforms/xla/service/jellyfish/deepsea_compiler_hlo_passes.cc`（字符串锚定）
置信度	CONFIRMED（字节锚定），除非某行或标注另有说明

Pipeline 级架构

RunHloPasses（0x1093a420）是单体驱动器。它构建并运行一系列 xla::HloPassPipeline 容器；更细的编号阶段映射到五个 builder 函数，这些函数已通过符号表中的精确签名确认：

text

RunHloPasses (0x1093a420)  ==  CompilePhase1HloOptimizations (0xf84ee00)
  ├─ PreOptimizationPipeline ........................ PHASE 1  input scrub
  │
  ├─ CreateHloPipeline (0x1093efe0)  opens the main pipeline
  │    ├─ AddAutoShardingAndRelatedPasses (0x10939c40) PHASE 2  sharding
  │    ├─ AddTpuPartitioningPasses (0x1278a440) ...... PHASE 3  SPMD prep
  │    ├─ HloOptimizeThroughLayoutAssignment (0x1094ada0) PHASE 4 pre-layout
  │    │    └─ LayoutAssignment  ........... (see layout-assignment.md)
  │    └─ PostOptimizationPipeline (0x1093fd40) ...... PHASE 5  post-layout
  │                                                    PHASE 6  pre-MLIR HLO
  └─ MLIR conversion  ..................... (see overview.md)
```text

两个结构性事实塑造了整个集合：

- **`xla::HloPassFix<P>` 将一个 pass 包装在定点循环中。** 二进制中的 vtable 证据显示它为 `HloPassPipeline`、`HloDCE`、`AllReduceReassociate`、`ReduceWindowRewriter`、`ReduceScatterReassociate`、`WhileLoopConstantSinking`、`AllReduceReduceScatterReorder`、`jellyfish::TpuReduceWindowRewriter` 和 `jellyfish::MosaicFusion` 实例化。定点机制为 CONFIRMED。据称，其未收敛时崩溃的行为由 `--xla_tpu_crash_if_hlo_pass_fix_did_not_converge` 控制；该 flag 字符串在抽样的字符串表中**未**定位到。[置信度：崩溃 flag 为 LOW；`HloPassFix` 为 CONFIRMED。]

- **不变式检查器会在每个嵌套 pipeline 头部重新添加。** `MaybeAddInvariantCheckers`（`0x10944600`）使用 `AddInvariantChecker<T>`（不是 `AddPass<T>`），因此每个检查器都会在每个 pass 之后重新运行：

| 检查器类 | 目的 |
|---|---|
| `xla::HloVerifier` + `xla::jellyfish::TpuVerifierMetadata` | 完整 HLO 结构 / 形状验证（shape-size 通过封装在 `std::function<int64_t(Shape const&)>` 中的 `HloVerifierOpts::shape_size` 成员函数指针） |
| `xla::LegalizeSchedulingAnnotations`（作为检查器） | 调度注解图的良构性 |
| `xla::HloCycleDetection` | 拒绝包含循环依赖的 HLO 图 |

> **注意 — 反编译交叉检查。** 所有五个 builder 签名都已在 `*_functions.json` 中逐字确认：`PreOptimizationPipeline`（lambda `$_0..$_8`，谓词 `$_7` 已确认是 `HloDomainIsolator` 参数）、`PostOptimizationPipeline`（lambda `$_0..$_12`）、`HloOptimizeThroughLayoutAssignment`、`AddAutoShardingAndRelatedPasses`、`AddTpuPartitioningPasses`（接受 `PartitioningPipelineConfig`），以及 `AddGatherScatterExpanderPasses`、`MaybeAddInvariantCheckers` 和 `CreateHloPipeline`。四个 TPU vtable（`TpuHloSupportChecker`、`jellyfish::TpuCallInliner`、`jellyfish::TpuInt2AutoUpDownCaster`、`jellyfish::TpuRngBitGeneratorTupleDecomposer`）存在于 `*_names.json` 中。[置信度：CONFIRMED。]

---

## 有序预处理 Pass 表

下面的主表是通过遍历每个 builder 的 `AddPass<T>` 实例化而恢复出的发现顺序枚举。**阶段**列匹配 [compile-phases.md](compile-phases.md) 上的 pipeline-builder 编号：阶段 1 = `PreOptimizationPipeline`，阶段 2 = sharding，阶段 3 = SPMD 准备，阶段 4 = `HloOptimizeThroughLayoutAssignment`（直到 layout assignment），阶段 5–6 = `PostOptimizationPipeline`（post-layout 精化，然后是最终的 pre-MLIR HLO 域）。编号是*发现*顺序，不是源码连续顺序；循环/定点 pass 只出现一次，但会重新运行。

**来源**列：`T` = TPU 专用（`xla::jellyfish::`、`xla::tpu::` 或 `xla::Tpu*` 前缀）；`O` = 开源 XLA pass（有些会用 TPU `Target` 调用）。

### 阶段 1 — PreOptimizationPipeline（输入清理）

| # | Pass 类 | HLO 输入不变式 | HLO 输出不变式 | 来源 | 备注 |
|---|---|---|---|---|---|
| 1 | `xla::LegalizeSchedulingAnnotations` | sched. annotation 可能是文本形式 | annotation 被规范化为 backend-config 形式 | O | 由 `non_mitigatable_gap_checking` config 控制 |
| 2 | `xla::CheckNoDataDependencyInSchedulingAnnotations` | instruction 上的 sched. annotation | 不变（verifier）；依赖环时报错 | O | 不变式检查器 |
| 3 | `xla::HloDomainIsolator`（谓词 `$_7`） | sharding domain 可能是隐式的 | 显式 `kDomain` op 包裹 sharding region | O | 仅当 `EnableDomainPasses()` |
| 4 | `xla::DynamicIndexSplitter` | dynamic-index op，多维 index | dynamic-index op，拆分后的标量 index | O | |
| 5 | `xla::BatchNormExpander(true,true,true)` | `BatchNorm{Training,Inference,Grad}` | 仅有原始算术 op | O | 重写为 log/sqrt/divide |
| 6 | `xla::TpuCholeskyExpander` | `Cholesky` custom-call | tile 递归 dot + triangular-solve | T | `xla::CholeskyExpander` 的子类 |
| 7 | `xla::TpuQrExpander` | `QrDecomposition` | Givens rotation HLO 图 | T | `xla::QrExpander` 的子类 |
| 8 | `xla::LuDecompositionExpander` | `LuDecomposition` | scatter/gather 分解 | O | |
| 9 | `xla::TpuEighExpander` | `Eigh` | Jacobi rotation HLO 图 | T | `xla::EighExpander` 的子类 |
| 10 | `xla::FftExpander(Target const&)` | `Fft` op | Cooley–Tukey radix-2，TPU tile size | O* | 通过 `Target` 参数感知 TPU |
| 11 | `xla::TpuTriangularSolveExpander` | `TriangularSolve` | 递归 blocked solve | T | `xla::TriangularSolveExpander` 的子类 |
| 12 | `xla::jellyfish::TpuRngBitGeneratorExpander` | `RngBitGenerator` | Philox / ThreeFry HLO 序列 | T | 启用 impure-RNG 时 |
| 13 | `xla::RngBitGeneratorExpander` | `RngBitGenerator` | 开源默认算法 | O | 未选择 12 时的替代路径 |
| 14 | `xla::jellyfish::TpuRngBitGeneratorTupleDecomposer` | tuple 输出的 `RngBitGenerator` | 每个输出对应独立 `GetTupleElement` | T | |
| 15 | `xla::HloDCE` | 任意 | dead instruction/computation 被移除 | O | 在大多数阶段之间重新运行 |
| 16 | `xla::jellyfish::TpuInt2AutoUpDownCaster` | 带 int2 的混合精度算术 | int2 op 被 `Convert` ↔ int8 加边界 | T | MXU wire 最小为 int8 |
| 17 | `xla::jellyfish::TpuCallInliner(MustFuseInlineMode)` | 标记为 `must_fuse` 的 call site | callee 被内联到 caller | T | 由 `xla_tpu_impure_inline_must_fuse_early` 控制 |
| 18 | `xla::jellyfish::UserGuidedFusionIdAssigner` | `frontend_attribute: fusion_id` 字符串 | fusion_id 成为整数 backend-config | T | 供后续 fusion pass 使用 |
| 19 | `xla::ConditionalCanonicalizer` | 任意嵌套的 `Conditional` | 规范形式（单一 branch root） | O | |
| 20 | `xla::DynamicDimensionSimplifier` | dynamic-dim op | 冗余 dynamic-dim op 被折叠 | O | |
| 21 | `xla::DynamicPadder(DynamicPadderOptions)` | 动态形状 | 静态形状 + padding mask | O | 运行动态形状 lowering（2× `AddPass`） |
| 22 | `xla::jellyfish::PreX64RewriterOptimizations(Target const&)` | 任意 | sub-op 被预重写以适配 x64 host codegen | T | |
| 23 | `xla::ScatterExpander(ScatterExpander::Mode)` | `Scatter` op | 遍历 index 的 `While` 循环 | O | mode = slow-path vs simplified |
| 24 | `xla::MapInliner` | `Map` op | callee 被内联为 elementwise 序列 | O | |
| 25 | `xla::HloDomainRemover("sharding", ApplyDomainSharding)` | sharding `kDomain` bracket | bracket 被移除，sharding 保留为 attribute | O | 总是在阶段 2 之前 |
| 26 | `xla::jellyfish::TpuHloPrecisionTracer` | instruction precision | 缺失处填充 `precision_config` | T | |
| 27 | `xla::BitcastDtypesExpander` | 仅 dtype 的 `Bitcast` | reinterpret-cast HLO 图 | O | |
| 28 | `xla::jellyfish::XPrecisionRewriter(kX128Precision)` | 带 `x128` precision 的 matmul | 8 步 accumulation chain | T | 高精度 dot 分解 |
| 29 | `xla::jellyfish::XPrecisionRewriter()` | 带 x6 / x9 precision 的 matmul | 2 步 / 3 步 accumulation chain | T | 运行两次（累积；已确认 2× `AddPass`） |
| 30 | `xla::ComparisonExpander({{S64,S32}})` | s64 comparison | hi/lo halves 上的 s32 comparison | O | 配置为 S64→S32 |

`O*` = 开源类，但通过 `Target` 参数进行 TPU 专门化。

### 阶段 2 — AddAutoShardingAndRelatedPasses（sharding）

| # | Pass 类 | HLO 输入不变式 | HLO 输出不变式 | 来源 | 备注 |
|---|---|---|---|---|---|
| 31 | `xla::FlattenCallGraph` | 嵌套 call | 单层 call graph | O | 在每个 sharding pass 之前重新运行 |
| 32 | `xla::jellyfish::TpuCallInliner` | 小的可复用 computation | 为 sharding 内联到 caller | T | |
| 33 | `xla::HloDCE` | 任意 | dead instruction/computation 被移除 | O | |
| 34 | `xla::ShardingPropagation` | 部分 sharding annotation | 每个 op 都有 sharding（或 `Replicated`） | O | manual-sharding 流程（4× `AddPass`） |
| 35 | `xla::TpuAutoSharding(AutoShardingOption, Target*, AliasInfo*)` | 未注解 HLO | 每个 op 上都有 sharding | T | auto-sharding 流程；`auto_sharding` flag 族 |
| 36 | `xla::sdy::ShardyXLA(PropagationOptions, ...)` | Shardy 格式的 sharding op | HLO sharding annotation | O | 检测到 JAX/Shardy 前端 |
| 37 | `xla::jellyfish::TpuRngBitGeneratorTupleDecomposer(nullptr, bool)` | sharding 后的 tuple RNG 输出 | 去 tuple 化 | T | sharding 后重新运行 |
| 38 | `xla::TupleSimplifier` | 冗余 tuple/GTE chain | 已简化 | O | |

### 阶段 3 — AddTpuPartitioningPasses（SPMD 准备）

| # | Pass 类 | HLO 输入不变式 | HLO 输出不变式 | 来源 | 备注 |
|---|---|---|---|---|---|
| 39 | `xla::spmd::SpmdPrepare` | 带 sharding 的 HLO | 规范化以便 SPMD partitioning | O | |
| 40 | `xla::ConvOperandSwapper` | operand layout 互换的 conv | 规范 operand 顺序 | O | |
| 41 | `xla::jellyfish::TpuSpmdConcatRewriter(Target const&)` | sharded `Concatenate` | replicated concatenate + slice | T | |
| 42 | `xla::HloConstantSplitter(bool)` | 跨 sharding 共享的 constant | 按 sharding 复制 constant | O | |
| 43 | `xla::jellyfish::TpuPartitionAssignment(Target const&, long)` | sharded HLO | 附加 partition/device-id metadata | T | `name()` → `"tpu-partition-assignment"` |
| 44 | `xla::jellyfish::ConvolutionFolding(Target const&, bool)` | 带可折叠 bias/activation 的 conv | 融合的 conv-bias-activation | T | 运行两次（pre-sharding，in-layout） |
| 45 | `xla::jellyfish::TpuSpmdPartitioner(...)` | sharded HLO | 带 collective 的 per-partition HLO | T | `SpmdPartitioner` 的 TPU 子类 |
| 46 | `xla::RecognizeReduceWindow` | `ReduceWindow` 形状的图 | 显式 `ReduceWindow` op | O | |
| 47 | `xla::CollectivePermuteCSE` | 重复的 `CollectivePermute` | 已去重 | O | |
| 48 | `xla::WholeGraphManualPass` | 整图上的 `manual` sharding | passthrough sharding annotation | O | |

### 阶段 4 — HloOptimizeThroughLayoutAssignment（直到 layout）

| # | Pass 类 | HLO 输入不变式 | HLO 输出不变式 | 来源 | 备注 |
|---|---|---|---|---|---|
| 49 | `xla::ZeroSizedHloElimination` | operand/output 大小为零的 op | 被消除 / empty constant | O | |
| 50 | `xla::TpuHloSupportChecker` | 任意 | 不变（validator）；不支持的 Shape 报错 | T | **规范的 TPU HLO 接受测试** |
| 51 | `xla::ConvertMemoryPlacementToInternalAnnotations` | `frontend_attribute: memory_space` | backend-config memory_space 整数 | O | |
| 52 | `xla::HloModuleDCE` | module-level dead computation | 被移除 | O | |
| 53 | `xla::ConvolutionTypeCanonicalizer` | 混合类型 conv | 带显式 convert 的统一类型 conv | O | |
| 54 | `xla::ConvolutionPrecisionNormalizer` | 只有 operand precision 的 conv | 规范化的 precision_config | O | |
| 55 | `xla::BroadcastCanonicalizer` | broadcast，非规范 dim 顺序 | 规范 broadcast | O | |
| 56 | `xla::TransposeFolding` | 可吸收到 dot/conv 中的 transpose | 带置换 operand 的 dot/conv | O | |
| 57 | `xla::ConvertOperandFolding` | 可吸收到 dot/conv 中的 `Convert` | 带混合精度 operand 的 dot/conv | O | |
| 58 | `xla::HloCSE(bool)` | 重复的 pure op | 共享的单一 op | O | 在大多数阶段之间重新运行 |
| 59 | `xla::HloPassFix<xla::jellyfish::TpuReduceWindowRewriter>` | 带非平凡 window 的 `ReduceWindow` | 重复重写直到规范 | T | 定点 |
| 60 | `xla::jellyfish::TpuAlgebraicSimplifier(Target, AlgSimpOptions)` | 任意 | 已做代数简化，感知 TPU | T | `xla::AlgebraicSimplifier` 的超集 — 见 [algebraic-simplifier.md](algebraic-simplifier.md) |
| 61 | `xla::GatherOptimizer(Target const&)` | `Gather` op | TPU 友好的 Gather（拆分/分解） | O* | |
| 62 | `xla::AllReduceSimplifier` | `AllReduce`，退化 replica group | 已简化 | O | |
| 63 | `xla::jellyfish::TpuAllGatherSimplifier(Target const&)` | `AllGather`，退化 replica group | 已简化 | T | |
| 64 | `xla::AllToAllDecomposer(bool, int)` | 带 split-dim 的 `AllToAll` | per-partition slice + AllToAll | O | |
| 65 | `xla::jellyfish::RaggedAllToAllExpander(long)` | `RaggedAllToAll` custom-call | dense AllToAll + scatter | T | |
| 66 | `xla::SortSimplifier` | 冗余 Sort operand | 已简化 Sort | O | |
| 67 | `xla::jellyfish::TpuReduceRewriter(bool)` | 多输出 `Reduce` | per-output reduce | T | |
| 68 | `xla::jellyfish::TpuDegenerateDimensionRewriter` | 带 size-1 batch dim 的 op | size-1 dim 通过 reshape 消除 | T | |
| 69 | `xla::jellyfish::TpuBroadcastRewriter` | broadcast，不利的目标 dim | reshape + broadcast 到 TPU 偏好的 dim | T | |
| 70 | `xla::jellyfish::TpuReduceRewriter` | 同 67，无 flag 变体 | per-output reduce | T | |
| 71 | `xla::ReduceWindowResizer` | `ReduceWindow`，非 pow-2 window | 通过 padding 调整 window 大小 | O | `name()` → `"reduce-window-resizer"` |
| 72 | `xla::WhileLoopConstantSinking(bool)` | 携带 constant 的 While | constant 下沉到 body 中 | O | |
| 73 | `xla::WhileLoopSimplifier(bool)` | trip 为 constant 的 While | 展开 / 简化 | O | |
| 74 | `xla::WhileLoopConcatCodeMotion(long)` | While，invariant 的 `Concatenate` | concatenate 被提升出来 | O | |
| 75 | `xla::HloConstantFolding` | 可折叠的 constant op | 折叠为 literal | O | |
| 76 | `xla::jellyfish::TpuConditionalSimplifier(Target const&)` | 带 TPU pattern 的 `Conditional` | 已简化或重写 | T | |
| 77 | `xla::DeadDynamicUpdateSliceElimination` | 目标 dead 的 DUS chain | dead DUS 被丢弃 | O | |
| 78 | `xla::conditional_opt::ConditionalCodeMotion(...)` | cond. branch 中相同的 code | code 被提升到之前 / 下沉到之后 | O | |
| 79 | `xla::jellyfish::SortMerger` | 相邻且 key 兼容的 `Sort` | 合并为单个 Sort | T | |
| 80 | `xla::ScanExpander` | `Scan` HLO | partial reduce 的 While 循环 | O | |
| 81 | `xla::StableSortExpander` | 请求 stable 的 `Sort` | augmented-key Sort + post-strip | O | |
| 82 | `xla::InfeedTokenPropagation` | 没有显式 token edge 的 `Infeed` | 为顺序性串接 token edge | O | |
| 83 | `xla::jellyfish::InfeedDecomposer` | `Infeed` op | DMA + token sequence（TPU host-transfer） | T | |
| 84 | `xla::jellyfish::OutfeedDecomposer` | `Outfeed` op | DMA + token sequence | T | |
| 85 | `xla::megascale::compiler::TpuAllReduceMerger(Target, mapper)` | per-slice AllReduce | cross-slice AllReduce 被合并 | T | MegaScale 路径 |
| 86 | `xla::megascale::compiler::CrossSliceLegalizer(Target const&)` | cross-slice op | 针对 MegaScale topology 合法化 | T | MegaScale 路径 |
| 87 | `xla::TpuGatherScatterFlattener(Target, long)` | 高 rank gather/scatter | rank-flattened gather/scatter | T | 在 `AddGatherScatterExpanderPasses` 中 |
| 88 | `xla::TpuGatherExpander(Target const&)` | 可扩展的 `Gather` | slice 的 `While` 循环 | T | |
| 89 | `xla::TpuScatterExpander(Target const&)` | 可扩展的 `Scatter` | DUS 的 `While` 循环 | T | |

*Layout assignment（`Phase2PreLayoutAssignment` / `TpuLayoutAssignment`）在阶段 4 的尾部运行，由 [layout-assignment.md](layout-assignment.md) 负责。*

### 阶段 5–6 — PostOptimizationPipeline（post-layout 精化，然后 pre-MLIR）

| # | Pass 类 | 阶段 | HLO 输入不变式 | HLO 输出不变式 | 来源 | 备注 |
|---|---|---|---|---|---|---|
| 90 | `xla::HloDomainIsolator`（谓词 `$_7`） | 5 | 如果 `EnableDomainPasses()`，重新添加 domain bracket | `kDomain` bracket 被重新添加 | O | step 25 的逆操作 |
| 91 | `xla::HloCSE(bool=false)` | 5 | 重复的 pure op | 共享的单一 op | O | post-layout CSE |
| 92 | `xla::jellyfish::WrapFusionOutputForDebug` | 5 | fused HLO | fusion output 用 `kCopy` 包裹以便 debug capture | T | 条件性 |
| 93 | `xla::jellyfish::AlwaysCrash` | 5 | 任意 | **有意失败**（测试 pass） | T | 由 `xla_tpu_always_crash` 控制 |
| 94 | `xla::AddOriginalValue` | 5 | HLO post-layout | 每条 instruction 都标记 provenance metadata | O | `name()` → `"add-original-value"` |
| 95 | `xla::jellyfish::AddRandomHostOffloading(double)` | 6 | 任意 | 随机 instruction 用 host-offload custom-call 包裹 | T | debug pass，flag 控制 |
| 96 | `xla::jellyfish::ConvertFrontendAttributesToBackendConfig` | 6 | `frontend_attribute` 字符串 | 解析为带类型的 backend_config protobuf | T | **在 HLO 域中最后运行** |
| 97 | `xla::HloHostDeviceTypeCallWrapper(Options)` | 6 | host_compute call site | 用 type-call marker 包裹以便 MLIR import | O | 最终 pre-MLIR pass |

> **注意事项 — 接受门禁（#50）在阶段 4，而不是阶段 1。** 想要提早拒绝不支持 op 的重新实现，可能会倾向于把 `TpuHloSupportChecker` 前置。libtpu 并不这样做：它在 expander（阶段 1）和 SPMD partitioner（阶段 3）*之后*运行，因为这些 pass 会合法地引入新 op（expander output、partitioned collective），而这些新 op 本身也必须通过检查。在 expansion 之前检查会拒绝可编译程序。该 checker 从不修改；它遍历每个 `HloComputation`，并用 `ShapeUtil::ValidateShapeWithOptionalLayout` 验证每个 result `Shape`，在遇到第一个不支持的 shape 时返回 error `Status`。入口 `RunImpl` @ `0x11071480`。[置信度：RunImpl VA + vtable 为 CONFIRMED；`name()` 返回 `"tpu-hlo-support-checker"` 为 HIGH（该字面量未在抽样字符串表中出现）。]

---

## TPU 专用与开源拆分

在 97 个 pipeline 提到的 pass 中，33 个带有显式 TPU 前缀（`xla::jellyfish::`、`xla::tpu::` 或 `xla::Tpu*`）；其余是开源 XLA pass，其中一些使用 TPU `Target` 参数化。TPU 专用 pass 按角色聚集如下：

- **Custom-call 分解（expander）：** `TpuCholeskyExpander`、`TpuQrExpander`、`TpuEighExpander`、`TpuTriangularSolveExpander`、`TpuGatherExpander`、`TpuScatterExpander`、`jellyfish::TpuRngBitGeneratorExpander`（+ `TupleDecomposer`）、`jellyfish::RaggedAllToAllExpander`、`jellyfish::Infeed/OutfeedDecomposer`。
- **形状 / 规范化 rewriter：** `jellyfish::TpuBroadcastRewriter`、`jellyfish::TpuDegenerateDimensionRewriter`、`jellyfish::TpuReduceRewriter`（带 flag + 无 flag）、`HloPassFix<jellyfish::TpuReduceWindowRewriter>`、`xla::TpuGatherScatterFlattener`、`xla::TpuGatherSplit`、`jellyfish::TpuConvolutionTypeCanonicalizer`。
- **代数 / simplifier：** `jellyfish::TpuAlgebraicSimplifier`（开源 simplifier 的超集 — 见 [algebraic-simplifier.md](algebraic-simplifier.md)）、`PostFusionTpuSubgraphSimplifier`、`TpuTrivialFusionRemover`、`TpuTrivialInstructionUnfuser`。
- **Call-graph / 接受：** `jellyfish::TpuCallInliner`（感知 must-fuse；`name()` → `"tpu-call-inliner"` / `-must-fuse` / `-inner-must-fuse` / `-non-must-fuse`）、`TpuHloSupportChecker`。
- **Sharding / partitioning：** `TpuAutoSharding`（包装 `xla::AutoSharding`）、`jellyfish::TpuSpmdPartitioner`、`TpuPartitionAssignment`、`TpuSpmdConcatRewriter`。
- **Precision / dtype：** `jellyfish::TpuInt2AutoUpDownCaster`、`XPrecisionRewriter`（x6/x9/x128 dot precision）、`TpuHloPrecisionTracer`。
- **Custom-call / fusion 准备：** `jellyfish::MosaicFusion`（Pallas/Mosaic kernel）、`UserGuidedFusionIdAssigner`、`WrapFusionOutputForDebug`、`TpuCustomCallScopedVmemAdjuster`（已确认在 `PostOptimizationPipeline::$_12` 中 `AddPass`）。

二进制的 RTTI 中还存在另一组 TPU pass（仅 vtable），但在五个反编译出的 builder 中**没有**观察到它们；它们由 `RunHloPasses` 更深处的配置 flag 门控分支添加。这些包括 SparseCore offload 族（`SparseCoreComputeOffloader`、`OffloadGatherToSparseCore` 等）、async-collective 族（`TpuAsyncCollectiveCreator`、`AsyncCollectiveMerger`、`AsyncOpScheduler`）、fusion 族（`TpuInstructionFusion`、`TpuMultiOutputFusion`）、`FlashAttention`、`Gmm`（group-matmul）以及 int4 路径 `TpuInt4Rewriter`。[置信度：这些类存在为 CONFIRMED（vtable RTTI）；它们的 pipeline 位置和门控为 LOW。]

---

## HLO 边界处消费的 Custom-Call Target

预处理 pass 会消费 / 产生带有下列 target 字符串的 HLO `kCustomCall` op（从字符串段恢复）。它们构成 JAX/PyTorch 前端与 TPU 编译器之间的表面：

| Target 字符串 | 消费/产生者 | 角色 |
|---|---|---|
| `Sharding` | `ShardingPropagation` | sharding 边界标记 |
| `SPMDFullToShardShape` / `SPMDShardToFullShape` | SPMD lowering | sharding lowering 辅助项 |
| `mhlo.sharding` / `_XlaSharding`（attrs） | `sdy::ShardyXLA` | Shardy/MHLO import sharding marker — 作为 MLIR attribute 携带，而不是作为 `Sharding-mhlo` custom-call target |
| `RngBitGenerator` | `TpuRngBitGenerator*` | RNG，由 RNG expander 族分解 |
| `TopK` | — | 保持 opaque，稍后 lowering |
| `tpu_custom_call` | `MosaicFusion` / Mosaic emit | 已注册的 Pallas/Mosaic kernel custom-call target（`CustomCallRegistration::RegisterCompilationProperties("tpu_custom_call", …)`）；通用 TPU custom-call wrapper |
| `MoveToHost` / `MoveToDevice` | host-offload legalizer | host-offload marker |
| `Pin` | `jellyfish::PinPrecoloring` | precoloring marker |
| `inspect_sharding` | `jellyfish::RemoveInspectShardingCustomCall` | 仅 debug，已移除 |

`__cudnn$convForward` 和类似 GPU target **不会**被任何 TPU pass 消费；它们会被 `TpuHloSupportChecker`（#50）拒绝。

---

## 未恢复的内容

- **仅 vtable 的 TPU pass 的精确顺序。** SparseCore、async-collective、fusion、FlashAttention 和 int4 族出现在 RTTI 中，但没有在五个反编译出的 builder 中观察到；它们位于 `RunHloPasses` 更深处的 flag 门控分支，本次遍历没有覆盖。[置信度：位置为 LOW。]
- **逐 `TpuVersion` 的差异。** 反编译反映的是*一个*会根据 `Target` / `TpuCompilationEnvironment` flag 分支的 pipeline；例如 TPU v3 与 v6e 分别跳过哪些 pass 尚未隔离。[置信度：逐代差异为 LOW。]
- **Sharding 流程选择器逻辑。** manual / auto / Shardy 中哪个运行，由 flag 加 builder 内部的 frontend-attribute 检测决定；精确的 dispatch 分支尚未隔离（`auto_sharding` flag 族已确认存在于字符串表中，但精确选择器为 LOW）。见 [auto-sharding-spmd.md](auto-sharding-spmd.md)。
- **逐 pass 重写算法。** 本页恢复的是*名称 + 顺序 + I/O 不变式*，不是内部图重写。simplifier 算法有自己的页面（[algebraic-simplifier.md](algebraic-simplifier.md)）；其余仍待完成。
- **`--xla_tpu_crash_if_hlo_pass_fix_did_not_converge` flag 字符串**未在抽样字符串表中找到。[置信度：LOW。]

---

## 交叉引用

- [compile-phases.md](compile-phases.md) — 顶层有序阶段列表以及每个阶段的入口符号（本页是它委托到的表）。
- [overview.md](overview.md) — 编译器导览以及 HLO → … → LLO IR 层栈；负责阶段 6 之后的 MLIR 移交。
- [hlo-ingestion.md](hlo-ingestion.md) — StableHLO/MHLO 前端程序如何变成这些预处理 pass 消费的 HLO module。
- [hlo-pass-registry.md](hlo-pass-registry.md) — 包含 372 个条目的 `HloPassInterface` 类目录和 `name()` 字符串。
- [algebraic-simplifier.md](algebraic-simplifier.md) — `TpuAlgebraicSimplifier` 重写算法（pass #60）。
- [sharding-propagation.md](sharding-propagation.md) / [auto-sharding-spmd.md](auto-sharding-spmd.md) — 阶段 2/3 的 sharding 和 SPMD 流程（pass #34–48）。
- [optimization-barrier.md](optimization-barrier.md) — 整个预处理 pass 集中的 `OptimizationBarrier` 处理。
- [layout-assignment.md](layout-assignment.md) — 阶段 4 尾部的 `TpuLayoutAssignment`。
- [返回索引](../index.md)

HLO 预处理 Pass ​

摘要 ​

Pipeline 级架构 ​

HLO 预处理 Pass

摘要

Pipeline 级架构