HLO 摄取

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

摘要

每个 TPU 程序都以可移植 MLIR 字节码进入编译器，也就是一个 StableHLO 模块（混有 CHLO 和 VHLO op），由框架桥接层（JAX、TF/XLA bridge 或 PyTorch/XLA）跨 PjRt 边界序列化而来。它不是以 XLA HLO 进入。编译器的第一个动作 xla::CompilePhase0StablehloToHlo (0xf84de60) 因此不是优化，而是一次格式跨越：它把字节码解析成内存中的 mlir::ModuleOp，运行有序的 MLIR pass 流水线以合法化 CHLO→StableHLO→MHLO，然后遍历 MHLO 模块发出 xla::HloProto，最后再把该 proto 解析回 XLA 其余部分所基于的 xla::HloModule/HloInstruction 图。本页负责说明这次跨越以及两侧的线格式。

熟悉 LLVM 的读者可以采用这个类比：Phase 0 是前端解析器加 bitcode reader，而不是一个 transform pass。这里同时存在三种不同表示，每种都有自己的序列化。输入侧：StableHLO/CHLO/VHLO MLIR（稳定、带版本的线 IR，相当于带前向兼容保证的 LLVM bitcode）。中间：xla.HloModuleProto（HLO 图的扁平、按 id 索引 protobuf 形式，是 XLA 自己的序列化，区别于 MLIR 字节码）。输出侧：实时的 HloModule 对象图。Phase 0 先把第一种表示转换为第二种（通过 MLIR PassManager 和 mlir::ConvertMlirHloToHlo, 0x16a64920），再把第二种反序列化为第三种（通过 xla::HloModule::CreateFromProto, 0x1e5dbe60）。只有在这之后，HLO optimizer（Phase 1，compile-phases.md）才开始。

本页记录三件事，并链接其余内容。(1) StableHLO→HLO 转换：xla::MlirToXlaComputation (0xf907d40) 构建并运行的 MLIR pass 流水线、逐 op 的 StablehloToHloOpConverter pattern，以及 CHLO/VHLO 处理。(2) 前端交入的 HLO proto schema：从二进制的 protodesc_cold 描述符池重建出的 HloModuleProto/HloComputationProto/HloInstructionProto id 图。(3) 编译入口：PjRt 的 phase-compile 和 TF/XLA bridge 的 CompileComputationToHlo 如何到达 Phase 0。摄取之后运行的枚举式 HLO pass 流水线见 compile-phases.md 和 hlo-pre-passes.md；IR 层栈总览见 overview.md。

对重新实现而言，摄取契约是：

• 三种序列化，两次转换。 StableHLO/CHLO/VHLO 字节码 →（legalize + ConvertMlirHloToHlo）→ HloModuleProto →（CreateFromProto）→ HloModule。把摄取当成一步的重新实现者会漏掉 proto 是一个真实中间体，运行时可以 dump 和缓存它。
• opcode 是字符串，不是 enum。 在 HLO-proto 线上，HloInstructionProto.opcode 是 string opcode = 2。不存在 xla.HloOpcode proto enum。这正是该格式具备前向兼容性的原因。
• HloInstructionProto 是一个约 83 字段的 union。 每个 op 专属属性都是单个宽消息上的可选字段；opcode 字符串选择哪个字段子集有意义。
• 图是扁平的按 id 索引 DAG。 没有嵌套 instruction 对象：数据边是 operand_ids int64 引用，调用边是 called_computation_ids 引用，root_id 命名每个 computation 的输出。
• CHLO 和 StableHLO↔MHLO legalizer 在 Phase 0 内运行，在 HLO 存在之前。 把 CHLO handler 构建进 HLO optimizer 的重新实现者处在错误层级；Phase 1 看到模块时 CHLO 已经消失。

Phase 0 入口	xla::CompilePhase0StablehloToHlo @ 0xf84de60
Phase 0 I/O	(CompileOptions, absl::Span<const PjRtPartialProgramProto>, const PjRtTopologyDescription&) → StatusOr<vector<PjRtPartialProgramProto>>
字节码解析	xla::ParseMlirModuleString(string_view, mlir::MLIRContext&) @ 0xf908580
转换驱动器	xla::MlirToXlaComputation @ 0xf907d40; xla::ConvertStablehloToHlo(mlir::ModuleOp) @ 0x16a3d200
MHLO → HloProto 发出	mlir::ConvertMlirHloToHlo(ModuleOp, HloProto*, …) @ 0x16a64920
逐 op converter	mlir::stablehlo::(anon)::StablehloToHloOpConverter<Op> — 121 个 matchAndRewrite 特化
CHLO legalizer	createChloLegalizeToStablehloPass, mlir::mhlo::createChloLegalizeToHighLevelMhloPass（opts getDefaultChloToHighLevelMhloOptions @ 0x16ad78e0）
HLO proto 解析	xla::HloModule::CreateFromProto(HloModuleProto const&, HloModuleConfig const&, …) @ 0x1e5dbe60; CreateModuleConfigFromProto @ 0x1e5e0480
Phase registry	xla::TpuCompiler::RegisterAllPhases @ 0xf849ec0
HLO proto schema	protodesc_cold（VA 0x0be8af30）：hlo.proto @ 0xc189a60, xla_data.proto @ 0xc1b7e20, xla.proto @ 0xc021470
置信度	除非某行或标注另有说明，否则为 CONFIRMED（byte-anchored）

---

编译入口

目的

Phase 0 是由 xla::TpuCompiler::RegisterAllPhases (0xf849ec0) 注册的五阶段分离编译流水线的开端。它可从两个前端表面到达，而两者交入的都是 StableHLO MLIR，从不是原始 HLO。重新实现者必须理解模块来自哪个表面，因为该表面决定字节码中已经包含哪些内容（sharding dialect、dim-args、layout-mode 属性）。

入口点

PjRt phase-compile  (PJRT_Api extension type 9 — see ../pjrt/ext-compile-phasecompile.md)
  │  serialized StableHLO module + CompileOptions  →  PjRtPartialProgramProto[]
  ▼
xla::CompilePhase0StablehloToHlo                                        0xf84de60
  ├─ xla::ParseMlirModuleString(view, MLIRContext&)                     0xf908580   parse bytecode → ModuleOp
  ├─ {GetArgLayoutModes / GetOutputLayoutModes}                                     read kArg/kOutLayoutModesAttr
  ├─ {GetArgMemoryKinds / GetOutputMemoryKinds}                                     read kArg/kOutMemorySpacesAttr
  ├─ xla::MlirToXlaComputation(ModuleOp, XlaComputation&, …, ChloOpts) 0xf907d40   the conversion (below)
  └─ → HloModuleProto (inside XlaComputation)  →  PjRtPartialProgramProto out

TF/XLA bridge alternative entry:
tensorflow::tpu::CompileComputationToHlo(                              0xf7cdba0
    TpuTopology, variant<MlirToHloArgs, FunctionToHloArgs>,
    CompileOnlyClient*, …) → XlaCompilationResult
  └─ MlirToHloArgs path also funnels through the StableHLO→HLO conversion
```text

这个签名有两个重要事实。第一，Phase 0 接收 `Span<const PjRtPartialProgramProto>` 并返回同类 vector，也就是*partial-program* 协议；它允许运行时持久化摄取后的模块，并稍后从 Phase 1 继续（这是编译缓存的基础；见 [compile-phases.md](compile-phases.md)）。第二，`PjRtTopologyDescription&` 参数表示拓扑（芯片代际、mesh shape）在摄取*期间*可用，这与 layout-mode 和 memory-space 属性解析（见下文）可能具备拓扑感知有关。

> **NOTE — bridge 路径和 PjRt 路径会汇合。** `tensorflow::tpu::CompileComputationToHlo` (`0xf7cdba0`) 接收一个 `variant<MlirToHloArgs, FunctionToHloArgs>`。`MlirToHloArgs` 分支携带 StableHLO 模块，并通过与 PjRt phase-compile 路径相同的 MLIR→HLO 转换机制；`FunctionToHloArgs` 分支是旧的 TF GraphDef→XLA 路径。只面向现代 JAX/PjRt 的重新实现只需要 StableHLO 分支。[Confidence: CONFIRMED 两个分支都存在；FunctionToHloArgs 内部未追踪 — LOW on that legacy path.]

### XLACallModule wrapper

跨边界携带的序列化模块在源自 TensorFlow 上下文中的 JAX 原生序列化时，是 **`XLACallModule`** op 的 payload。该 op 的文档字符串（位于 `0x1898480`）记录版本契约：最低支持的 `version` 是 2；从 v2 起 op 携带 StableHLO 文本*或*字节码；v3 添加 `platforms` 属性；v4 添加 StableHLO 兼容性保证；v5 允许 `stablehlo.custom_call`。这种版本化解释了为什么 VHLO（versioned HLO）op 会出现在导入表面：它们是旧运行时摄取新前端生成模块的机制。

---

## StableHLO → HLO 转换

### 目的

这是 Phase 0 的核心：把包含 StableHLO/CHLO/VHLO 的 `mlir::ModuleOp` 转成 `xla::HloProto`。其实现是一次常规 MLIR `PassManager` 运行，随后由一个 MHLO 遍历 emitter 发出 proto。存在两个相关驱动器：`xla::MlirToXlaComputation` (`0xf907d40`，Phase-0 路径，生成 `XlaComputation`) 和 `xla::(anonymous namespace)::ConvertStablehloToHloProtoInternal` (`0x16a3d400`，生成裸 `HloProto`)。它们共享流水线的*形状*（CHLO recompose → SymbolDCE → CHLO legalize → normalize → run → emit），但**不共享完全相同的 pass 集合**：`MlirToXlaComputation` 添加 `StablehloComplexMathExpander` 并使用默认 verifier；`ConvertStablehloToHloProtoInternal` 则（有条件地）添加 `StablehloTargetIndependentOptimization` + `StablehloSanitizeDiscardableAttributes`，并显式调用 `enableVerifier(false)`。两者最终都会把规范化后的模块遍历进 proto：`MlirToXlaComputation` 通过 `ConvertStablehloToHloWithOptions` → `ConvertStablehloToHloInternal` (`0x16a3d220`) → `ConvertStablehloToHloProtoInternal`，因此尽管前置 pass 链不同，proto emitter 是共享的。[Confidence: CONFIRMED both pipelines from the decompiled bodies.]

### 入口点

```text
xla::MlirToXlaComputation(ModuleOp, XlaComputation&, bool, bool,        0xf907d40
                          ExecutableBuildOptions*,
                          mhlo::ChloLegalizeToHighLevelMhloPassOptions const&)
  └─ public wrappers:
       xla::ConvertStablehloToHlo(ModuleOp)                             0x16a3d200
       xla::ConvertStablehloToHloWithOptions(ModuleOp, bool, bool)      0x16a3d3a0
       xla::ConvertStablehloToHloProto(ModuleOp, HloProto*)             0x16a3d3c0
       xla::ConvertStablehloWithManyArgsToHloProto(…)                   0x16a3d7c0

算法

转换会构建一个 mlir::PassManager，添加一串有序的 MLIR pass（大多嵌套在 func.func 下），运行它，然后发出 proto。下面的 pass 链从 MlirToXlaComputation (0xf907d40) 反编译体中的调用目标恢复；ConvertStablehloToHloProtoInternal (0x16a3d400) 变体的差异如上方“目的”所述。

function MlirToXlaComputation(module, out_computation, chlo_opts):   // 0xf907d40
    pm = PassManager(module.getContext(), "any")                     // mlir::PassManager (verifier left at default)

    // --- 0. Shardy fallback (only when GSPMD attrs/ops coexist with Shardy) ---
    if module has GSPMD attrs but Shardy is enabled:
        ExportShardyForGSPMD(module)                                 // disable Shardy, fall back to GSPMD propagation

    // --- 1. CHLO recompose, then SymbolDCE, then CHLO legalize ---
    pm.nest("func.func").addPass(stablehlo_ext::createChloRecomposeOpsPass())  // rebuild fused CHLO ops
    pm.addPass(createSymbolDCEPass())                                // drop unreferenced symbols (module-level)
    pm.nest("func.func").addPass(
        mhlo::createChloLegalizeToHighLevelMhloPass(chlo_opts))      // CHLO → high-level MHLO (top_k, erf, ragged…)
    pm.nest("func.func").addPass(
        stablehlo::createChloLegalizeToStablehloPass())              // remaining CHLO → StableHLO primitives

    // --- 2. StableHLO normalization ---
    pm.nest("func.func").addPass(
        stablehlo::createStablehloComplexMathExpanderPass())         // expand complex arithmetic
    pm.nest("func.func").addPass(
        stablehlo_ext::createSinkConstantsToControlFlowPass())       // push consts into while/case regions

    status = pm.run(module)                                          // BaseScopedDiagnosticHandler captures errors
    if !status.ok(): return status                                  // module now lives in MHLO + builtin dialects

    // --- 3. StableHLO → HloProto via the shared emitter (wraps ConvertMlirHloToHlo, 0x16a64920) ---
    hlo_proto = ConvertStablehloToHloWithOptions(module, …)          // → ConvertStablehloToHloProtoInternal → ConvertMlirHloToHlo
    out_computation = XlaComputation(hlo_proto.hlo_module())          // wrap proto in XlaComputation
    return out_computation
```text

有两个结构性说明。合法化是*自顶向下分阶段*完成的：CHLO（最高层 dialect，例如 `chlo.top_k`、`chlo.erf`、`chlo.ragged_dot`）先被 recomposed 并 lowering，一部分进入高层 MHLO op（它们有直接 HLO 等价物），一部分进入 StableHLO primitive；然后 StableHLO 层被规范化；最后整体被遍历进 proto。`ConvertMlirHloToHlo` 遍历（经由 `ConvertStablehloToHloProtoInternal` 到达）是真正发生 MHLO-op → `HloInstructionProto` 映射的地方，也就是程序离开 MLIR、成为 XLA HLO proto 的边界。

> **GOTCHA — 两个驱动器的 verifier 策略不同，并且 proto 路径并不是“每个 pass 后都开启”。** proto-emitting 驱动器 `ConvertStablehloToHloProtoInternal` (`0x16a3d400`) 显式调用 `pm.enableVerifier(false)`，它*不会*在 pass 之间重新验证。`MlirToXlaComputation` (`0xf907d40`) 构造 `PassManager` 时没有显式 `enableVerifier` 调用（继承 MLIR 默认值）。两个驱动器都会构造 `mlir::BaseScopedDiagnosticHandler`，它负责把 `pm.run` 期间抛出的 MLIR diagnostic 转换成 `absl::Status`（通过 `ConsumeStatus`）。重新实现者不应假设摄取路径启用了逐 pass 验证；诊断处理器，而非 verifier，才是把格式错误模块呈现为干净错误的机制。

### 逐 op converter 表

StableHLO→MHLO op 映射由模板化 pattern `mlir::stablehlo::(anonymous namespace)::StablehloToHloOpConverter<Op>` 实现，每个 StableHLO op 一个特化。二进制中存在 121 个不同的 `matchAndRewrite` 特化。这里不 dump 全部 121 个，而是用表描述转换*轴线*，也就是 converter 对每类 op 必须完成什么。

| Op 类别 | 代表性 op（已验证特化） | 转换动作 |
|---|---|---|
| 逐元素一元 | `AbsOp`, `CeilOp`, `CbrtOp`, `ClzOp`, `ConvertOp`, `CosineOp` | 1:1 到 MHLO/HLO op；复制结果类型 |
| 逐元素二元 | `AddOp`, `AndOp`, `DivOp`, `CompareOp`, `ComplexOp` | 1:1；`CompareOp` 携带 `comparison_direction`/`comparison_type` |
| Shape / 数据移动 | `BroadcastOp`, `BroadcastInDimOp`, `ConcatenateOp`, `DynamicBroadcastInDimOp` | 映射 `broadcast_dimensions` → `dimensions`；动态形式携带额外 shape operand |
| Reduction / windowed | `BatchNormTrainingOp`, `BatchNormInferenceOp`, `BatchNormGradOp` | 携带 `epsilon`, `feature_index`；expander 后续在 HLO pre-pass 中运行 |
| Matmul / conv | `DotOp`, `DotGeneralOp`, `ConvolutionOp`, `DynamicConvOp`, `CholeskyOp` | 映射 dimension-number attr → `dot_dimension_numbers` / `convolution_dimension_numbers`；携带 `precision_config` |
| 控制流 | `CaseOp`, `AsyncStartOp`, `AsyncDoneOp` | 映射带 region 的 op → `called_computation_ids` 边 |
| Collective | `AllGatherOp`, `AllReduceOp`, `AllToAllOp`, `CollectiveBroadcastOp`, `CollectivePermuteOp`, `CrossReplicaSumOp` | 携带 `channel_id`、replica grouping；`use_global_device_ids` |
| Gather / dynamic | `DynamicGatherOp`, `CreateTokenOp`, `AfterAllOp` | gather dim-number；token 类型 op 生成 `TOKEN` shape |
| Custom | `CustomCallOp`, `CompositeOp` | 保留 `custom_call_target`, `backend_config`, `api_version`；`CompositeOp` lowering 为 decomposition call |

> **QUIRK — 这 121 个 converter 特化只是约 182 个 StableHLO op 的*子集*，因为许多 StableHLO op 与其 MHLO 对应物相同，不需要 rewriter。** 为每个 StableHLO op 都构建 converter 的重新实现者会写出冗余的 identity rewriter；只构建这 121 个并假设其余 op 原样通过，更接近 libtpu 的实际结构。*需要* converter 的 op 是那些在 StableHLO 和 MHLO ODS 定义之间属性布局、region 结构或类型语义不同的 op。[Confidence: CONFIRMED 121 specializations; the exact StableHLO/MHLO divergence per op was not individually audited — HIGH on the category mapping.]

### 相关函数图

| 函数 | VA | 角色 |
|---|---|---|
| `xla::CompilePhase0StablehloToHlo` | `0xf84de60` | phase 入口；解析 + 转换 + 重新打包为 partial program |
| `xla::ParseMlirModuleString` | `0xf908580` | StableHLO 文本/字节码 → `mlir::ModuleOp` |
| `xla::MlirToXlaComputation` | `0xf907d40` | 转换驱动器（PassManager + emit） |
| `xla::ConvertStablehloToHlo` | `0x16a3d200` | 薄 wrapper，默认选项 |
| `xla::ConvertStablehloToHloWithOptions` | `0x16a3d3a0` | 暴露两个 bool flag 的 wrapper；tail-call `ConvertStablehloToHloInternal` |
| `xla::(anon)::ConvertStablehloToHloInternal` | `0x16a3d220` | 包装 `ConvertStablehloToHloProtoInternal`，返回 `XlaComputation` |
| `xla::(anon)::ConvertStablehloToHloProtoInternal` | `0x16a3d400` | 真正的 pass-pipeline + `ConvertMlirHloToHlo` emit（verifier disabled） |
| `xla::ConvertStablehloWithManyArgsToHloProto` | `0x16a3d7c0` | 多参数 bundle 变体 |
| `mlir::ConvertMlirHloToHlo` | `0x16a64920` | MHLO module 遍历 → `HloProto` |
| `mlir::mhlo::getDefaultChloToHighLevelMhloOptions` | `0x16ad78e0` | 默认 CHLO-legalization 选项 |
| `mlir::mhlo::StablehloLegalizeToHloPass::runOnOperation` | `0x16ae0320` | StableHLO→HLO pass（独立） |
| `mlir::mhlo::ChloLegalizeToHloPass::runOnOperation` | `0x16adbd00` | CHLO→HLO pass（独立） |

> **NOTE — `StablehloLegalizeToHloPass` (`0x16ae0320`) 和内联 converter 流水线共存。** 独立的 `mlir::mhlo::StablehloLegalizeToHloPass` 与 `ChloLegalizeToHloPass` 是已注册 pass（存在完整 `…PassBase` vtable：`getName`、`getArgument`、`clonePass`、`getDependentDialects`）。`MlirToXlaComputation` 驱动器并不按名称调用它们；它组装自己的 `createChlo…` / `createStablehlo…` pass 链。两条路径生成相同合法化结果。独立 pass 供反向和 round-trip 路径使用（`HloLegalizeToStablehloPass`, `0x16adcea0`，在 HLO 流水线*末尾*运行，以便为 MLIR descent 重新发出 StableHLO — 见 [compile-phases.md](compile-phases.md)）。重新实现者应把内联链视为摄取的权威路径。

---

## HLO Proto Schema（线契约）

### 目的

`ConvertMlirHloToHlo` 会发出 `xla.HloModuleProto`。这是 HLO 程序的**稳定序列化**：`HloModule::ToProto()` 生成它，`HloModule::CreateFromProto` 解析它，`xla_dump_hlo_as_proto` 写出它，并且前端程序在 Phase-0/Phase-1 边界实际就是以这种格式表示的。schema 从二进制中嵌入的 `protodesc_cold` 描述符池逐字段重建（section VA `0x0be8af30`，大小 `0x334180`）；三个 FileDescriptorProto 记录是 `hlo.proto` (`0xc189a60`)、`xla_data.proto` (`0xc1b7e20`) 和 `xla.proto` (`0xc021470`)。

### 图主干

```text
HloModuleProto
  ├─ string device_type = 21                 // "tpu"
  ├─ repeated HloComputationProto computations = 3
  │     ├─ string name = 1
  │     ├─ repeated HloInstructionProto instructions = 2   // FLAT list
  │     ├─ int64 id = 5
  │     └─ int64 root_id = 6                  // names the output instruction
  ├─ int64 entry_computation_id = 6
  ├─ ProgramShapeProto host_program_shape = 4 // entry signature
  ├─ HloScheduleProto schedule = 7            // per-computation id ordering
  ├─ HloInputOutputAliasProto input_output_alias = 8
  ├─ repeated bytes payloads = 22             // interned backend-config side-channel
  ├─ bool is_dynamic = 11                     // module has dynamic shapes
  ├─ OpSharding spmd_output_sharding = 12 / spmd_parameters_shardings = 14
  ├─ StackFrameIndexProto stack_frame_index = 17   // interned source provenance
  └─ FrontendAttributes frontend_attributes = 19

程序图是一个带 id 边的扁平 instruction 列表：没有嵌套 instruction 对象。每条数据边都是 int64 operand_ids（字段 36）引用，指向同级 instruction 列表；每条调用边都是 int64 called_computation_ids（字段 38）引用，指向模块的 computation 列表；id（字段 35）在 computation 内唯一，root_id（字段 6）命名输出。这种 id-graph 表示让 proto 无需指针修补就能完成序列化；它是按索引表示的 DAG，而不是树。

通用 instruction 记录

HloInstructionProto 是一个单一消息，约 83 个声明字段，字段号直到 99（从 protodesc_cold 中的描述符解析）。每个 op 专属属性都是自己的可选字段；opcode 字符串选择哪个子集有意义。下表描述该 union 的轴线（完整字段列表太宽，不适合 dump；这些是重新实现者必须复现的维度）。

字段组	代表性字段（编号）	由哪些 opcode 读取
身份 / 边	name(1), opcode(2), shape(3), id(35), operand_ids(36), control_predecessor_ids(37), called_computation_ids(38)	all
叶 payload	literal(8), parameter_number(9), delta(66), distribution(23), rng_algorithm(70)	constant, parameter, iota, rng, rng-bit-generator
Shape op	dimensions(14), slice_dimensions(17), dynamic_slice_sizes(20), padding_config(21), is_reverse(94)	reshape, transpose, slice, pad, reverse, …
Matmul / conv	dot_dimension_numbers(30), ragged_dot_dimension_numbers(90), convolution_dimension_numbers(16), window(15), feature_group_count(50), precision_config(51), conv_kind(97)	dot, ragged-dot, convolution
Collective	channel_id(26), replica_groups(49), oneof {collective_device_list(87), iota_collective_device_list(92), mesh_axes_replica_group_list(93)}, use_global_device_ids(71), source_target_pairs(52)	all-reduce, all-gather, all-to-all, collective-permute, …
Custom-call	custom_call_target(28), backend_config(43), backend_config_payload(99), custom_call_api_version(77), output_operand_aliasing(74)	custom-call (incl. tpu_custom_call)
精度控制	result_accuracy(91), is_associative(96), exponent_bits(18), mantissa_bits(19)	transcendentals, reduce-precision
Sharding	sharding(40), domain_entry_sharding(54), domain_exit_sharding(55)	any sharded op, domain
溯源	metadata(7), original_value(88), frontend_attributes(68)	all

QUIRK — HloInstructionProto.opcode 是字符串（string opcode = 2），不是 proto enum。 对整个 protodesc_cold 描述符池（约 770 个嵌入 .proto 文件）的穷尽扫描发现任何地方都没有 xla.HloOpcode 描述符，子串 HloOpcode 在池中一次也未出现。C++ HloOpcode enum 通过 HloOpcodeString ↔ StringToHloOpcode 对序列化为小写文本助记符："add"、"dot"、"convolution"、"fusion"、"all-reduce"、"dynamic-update-slice"、"custom-call"。这是最重要的序列化细节：它就是该格式能跨 XLA 版本前向/后向兼容的原因，新 opcode 是一个新字符串，无需在前端和后端之间协调 enum 编号。在线上定义数字 opcode enum 的重新实现，会悄悄偏离每一个真实 dump 出的模块。[Confidence: CONFIRMED — definitive negative result from the descriptor pool.]

NOTE — backend_config 有两种编码，新编码会 intern。 旧的 bytes backend_config = 43 仍然存在，但字段 99 backend_config_payload（xla.Payload）是新路径：Payload 是 bytes value = 1 或 int64 id = 2 的 oneof，其中 int64 id 索引 HloModuleProto.payloads（字段 22，repeated bytes）。这是一个 interning side-channel，使重复 backend config 在每个模块内只存一次。对 TPU 而言，ConvertFrontendAttributesToBackendConfig（最后一个 HLO pass，见 compile-phases.md）会在 MLIR descent 之前填充这些字段。

Proto 中的 dynamic shape 与 sharding

Dynamic shape 是结构性地编码在 ShapeProto 中，不是单独消息：is_dynamic_dimension(6) 是与 dimensions(3) 平行的 repeated bool（dimension 值是最大边界；bool 标记它运行时可变），HloModuleProto.is_dynamic(11) 是模块级标志，LayoutProto.dynamic_shape_metadata_prefix_bytes(15) 为运行时大小 metadata 前缀预留空间。DynamicPadder pre-pass 会消费这些字段并发出静态 shape 加 mask。

Sharding 有三个共存层，全部存在：经典 tile-based OpSharding（tile_assignment_dimensions，显式 tile_assignment_devices 或紧凑 iota_reshape_dims+iota_transpose_perm）；可从 OpSharding._named_sharding(14) 到达的 Shardy bridge NamedShardingProto（mesh-relative AxisRef shardings，由 ShardyXLA 在 use_shardy_partitioner=true 时消费）；以及模块级 spmd_output_sharding/spmd_parameters_shardings。Sharding 以 kCustomCall marker（"Sharding"、"SPMDFullToShardShape"、"SPMDShardToFullShape"）和 bracket uniform-sharding region 的 domain op 形式流入。

---

HLO Proto 解析（Proto → HloModule）

目的

一旦 ConvertMlirHloToHlo 生成了 HloModuleProto，实时 HloModule 对象图就由 xla::HloModule::CreateFromProto 重建。这是 ToProto() 的对称逆操作，也是按 id 索引 DAG 变成指针链接 HloInstruction 图的时刻。从这里开始，编译器其余部分操作的是 HloModule，而不是 proto。

入口点

xla::HloModule::CreateFromProto(HloModuleProto const&,                 0x1e5dbe60
                                HloModuleConfig const&, bool,
                                unique_ptr<CompilationEnvironments>,
                                bool, BufferAssignmentProto*)
  ├─ overload (HloModuleProto const&, HloModuleConfig const&,          0x1e5dbe20
  │            BufferAssignmentProto*, bool)
  └─ xla::HloModule::CreateFromProtoWithConfig(                        0x1e5e07e0
         HloModuleProtoWithConfig const&, …)

xla::HloModule::CreateModuleConfigFromProto(                           0x1e5e0480
    HloModuleProto const&, DebugOptions const&, ExecutionOptions const*)
  └─ builds the HloModuleConfig (entry layout, replica/partition counts,
     SPMD flags, MXU precision) that CreateFromProto consumes
```text

### 算法

```c
function CreateFromProto(proto, config):                               // 0x1e5dbe60
    module = HloModule(proto.name(), config)
    // 1. Rebuild every computation, resolving the id-graph:
    for comp_proto in proto.computations():                            // flat list
        builder = HloComputation::Builder(comp_proto.name())
        id_to_instr = {}
        for instr_proto in comp_proto.instructions():                  // in id order
            instr = HloInstruction::CreateFromProto(instr_proto, id_to_instr,
                                                    computation_map)    // opcode string → typed op
            id_to_instr[instr_proto.id()] = instr                      // operand_ids resolve here
        comp = builder.Build(id_to_instr[comp_proto.root_id()])        // root_id names output
        module.AddComputation(comp, is_entry = (id == entry_id))
    // 2. Attach module-level tables:
    module.set_schedule(proto.schedule())                              // HloScheduleProto
    module.set_input_output_alias(proto.input_output_alias())
    module.set_frontend_attributes(proto.frontend_attributes())
    module.set_stack_frame_index(proto.stack_frame_index())           // interned provenance
    return module

HloInstruction::CreateFromProto 是 opcode 字符串通过 StringToHloOpcode 映射回 C++ opcode 的地方，也是 union 字段选择发生的地方："dot" 读取字段 30/51，"convolution" 读取 16/15/50，"custom-call" 读取 28/43/77。operand_ids 会根据逐 computation 的 id_to_instr map 解析；该 map 在 instruction 按 id 顺序创建时构建。这就是 proto 按 id 拓扑序列化 instruction 的原因。

GOTCHA — HloModuleConfig 不在 HloModuleProto 中；它单独重建。 HloModuleProto 携带图；HloModuleConfigProto（在 xla.proto 中，与 HloModuleProtoWithConfig 成对）携带 entry-computation layout、replica/partition 数、SPMD flag、matrix_unit_operand_precision（MXU precision）、device_memory_size 以及 290 字段的 DebugOptions。CreateModuleConfigFromProto (0x1e5e0480) 会从 proto 加运行时的 DebugOptions/ExecutionOptions 构建 config。只反序列化 HloModuleProto 并默认化 config 的重新实现者，会得到一个没有 committed entry layout 且默认 precision 的模块；Phase 1 依赖的 layout-assignment 和 MXU-precision 决策来自config，而不是图 proto。

---

本页未覆盖内容

• 摄取之后运行的 HLO 优化流水线（pre-pass、sharding、layout、fusion、MSA、schedule）— 见 compile-phases.md 和 hlo-pre-passes.md。
• IR 层栈和五阶段主干总览 — 见 overview.md。
• 从 HLO 离开 的 MLIR descent（HloLegalizeToStablehloPass 和 MHLO→tpu lowering）— 见 mhlo-xtile-tpu-lowering.md。
• PjRt phase-compile C-ABI 表面（PJRT_Api extension type 9、options marshalling）— 见 ../pjrt/ext-compile-phasecompile.md。
• 此构建的精确 HloOpcodeString 助记符表。 opcode 集以字符串序列化；描述符池（正确地）不携带 enum。精确的约 200 项助记符列表必须从二进制 text/rodata 中的 HloOpcodeString jump table 提取；这里未枚举。[Confidence: the category bindings are CONFIRMED from the recovered attribute fields; the verbatim per-build mnemonic spellings are LOW.]

---

交叉引用

• overview.md — Part V 导向；IR 层栈和五个编译阶段（Phase 0 在其中命名）。
• compile-phases.md — 逐 phase 细节；Phase 1（HLO pass 流水线）会在本页生成的 HloModule 上运行。
• hlo-pre-passes.md — 首先接触摄取后模块的流水线前端 HLO pre-pass 集（custom-call expander、DynamicPadder、precision rewriter）。
• hlo-pass-registry.md — 这些 pass 派生自的 HloPassInterface 类目录。
• mhlo-xtile-tpu-lowering.md — 反向跨越：HLO 回到 StableHLO/MHLO 并下降到 tpu dialect（Phase 2a）。
• ../pjrt/ext-compile-phasecompile.md — 调用 CompilePhase0StablehloToHlo 的 PjRt phase-compile 入口。
• 二进制： extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）
• 索引条目： Part V — Compiler: Lowering & Optimization Passes / Front-end and pipeline — 返回索引