PhaseCompile 扩展（类型 9）
本页中的所有地址均适用于
libtpu-0.0.40-cp314wheel（libtpu_lts_20260413_b_RC00，BuildID89edbbe81c5b328a958fe628a9f2207d，ELF x86-64 DYN，约 745 MB）里的libtpu.so。PJRT C-API 表面是 v0.103。其他构建会重新编号槽位并偏移地址。
摘要
PhaseCompile 扩展（PJRT_Extension_Type == 9，一个 64 字节结构体）是 libtpu 的增量式编译 ABI。一次性 PJRT_Client_Compile（槽位 25）会把整个程序和一个 CompileOptionsProto 交给客户端，并返回完成的、已加载到设备的可执行体；而 PhaseCompile 扩展允许宿主框架一次驱动一个具名阶段来运行同一个 xla::TpuCompiler，例如先运行 phase0_stablehlo_to_hlo，缓存中间 HLO blob，之后再运行 phase1_hlo_opts 以及后续阶段。阶段边界是序列化的 PjRtPartialProgramProto blob，因此部分结果可以被缓存、检查，或拆分到不同宿主上执行。两个表面都由同一个 xla::TpuCompiler 对象和同一个 CompileOptionsProto 反序列化关口支撑；它们的差异只在于中间产物是否暴露出来。
该扩展是一个扁平的函数指针表，附加在 PJRT_Api 扩展链上。五个槽位中有三个是标准通用 XLA C-ABI 驱动（pjrt::PJRT_PhaseCompile_*，逐字节等同于上游 pjrt_c_api_phase_compile_internal.cc 代码）；另两个 Get_Compiler 和 Destroy_Compiler 由 TPU 注入（tpu_plugin::*TpuPhaseCompiler），也是该表面上唯一的 TPU 专用代码。注入的 Get_Compiler 会创建一个 xla::TpuCompiler，并在其上注册恰好六个具名阶段；随后通用的 Run_Phase / Get_Phase_Names 驱动通过已注册编译器上的两个虚槽（+48 GetPhaseNames，+56 RunPhases）分发。
本页负责说明PhaseCompile 扩展结构体、五槽位阶段 API，以及编译选项 proto 摄入路径，也就是每一次 PhaseCompile 调用与整体式路径共享的部分。编译生成的最终可执行体，包括它的 C-ABI 盒子、serialize / deserialize 往返以及 Execute，见可执行体加载与执行。这些阶段实际驱动的 HLO pass 管线（phase1_hlo_opts 的主体和 lowering 阶段）见 HLO Pass 注册表；缓存部分结果和完整结果的内容寻址存储见编译缓存。
对于重新实现，契约是：
- 64 字节扩展结构体。 一个
PJRT_Extension_Base头（size=64，type=9，next），后跟五个函数指针，位于+0x18 … +0x38。创建器从参数写入两个 TPU 注入的指针，并硬编码三个通用驱动。 - 编译器句柄。
Get_Compiler通过operator new分配一个0x40字节的xla::TpuCompiler（vtableoff_2177D660），调用RegisterAllPhases，把它包装进一个0x10字节的{ base, owner }holder，然后把 holder 存入 args 输出字段。宿主持有的不透明句柄是 holder，而不是编译器；Run_Phase/Get_Phase_Names解引用holder.base，Destroy_Compiler释放holder.owner。 - 六个已注册阶段、它们的顺序，以及
phase3_linking/phase3_linking_test_only是同一个函数的两次注册，仅绑定的bool不同。 - 共享的编译选项关口。
Run_Phase（和每个编译入口一样）执行CompileOptionsProto::ParseFromString→CompileOptions::FromProto，解析失败时返回完全相同的"PJRT_Client_Compile: failed to deserialize CompileOptionsProto"。 - char-buffer 编组。 阶段名和部分程序 blob 以并行的
(char**, size_t*, count)span 跨越 C-ABI；Convert*辅助函数在它们和 C++string/PjRtPartialProgramProtovector 之间转换，C_Buffers_Destroy释放返回的数组。
| 扩展类型 | 9 (PJRT_Extension_Type::PhaseCompile) |
| 扩展结构体大小 | 0x40 (64 B); 5 个 fn-ptr 位于 +0x18..+0x38 |
| 创建器 | pjrt::CreatePhaseCompileExtension @ 0x0E6F42A0 |
| 存储位置 | .bss slot GetTpuPjrtApi::phase_compile_extension @ 0x224C3B18 |
Get_Compiler (+0x18, TPU) | pjrt::tpu_plugin::GetTpuPhaseCompiler @ 0x0E6AA320 |
Destroy_Compiler (+0x20, TPU) | pjrt::tpu_plugin::DestroyTpuPhaseCompiler @ 0x0E6AA400 |
Run_Phase (+0x28, generic) | pjrt::PJRT_PhaseCompile_Run_Phase @ 0x0E6F42E0 |
Get_Phase_Names (+0x30, generic) | pjrt::PJRT_PhaseCompile_Get_Phase_Names @ 0x0E6F4A60 |
C_Buffers_Destroy (+0x38, generic) | pjrt::PJRT_PhaseCompile_C_Buffers_Destroy @ 0x0E6F4CC0 |
| 编译器对象 | xla::TpuCompiler (0x40 B, vtable off_2177D660) |
| 阶段数量 | 6（5 个不同函数 + 1 个 test_only 别名） |
| 来源（断言） | …/pjrt/c/pjrt_c_api_phase_compile_internal.cc; …/research/pjrt/tpu_pjrt_compiler.cc |
扩展结构体
目的
PhaseCompile 扩展是 PJRT_Api 扩展链中的一个链节（见扩展链）：宿主遍历 PJRT_Extension_Base 的 next 列表，匹配 type == 9，然后把该链节转换为 PJRT_PhaseCompile_Extension 以访问五个阶段函数。该结构体在插件初始化时由 CreatePhaseCompileExtension 组装一次，放入固定的 .bss 槽位（0x224C3B18），并由 GetTpuPjrtApi 使用作为参数传入的两个 TPU 注入函数指针进行填充。
布局
struct PJRT_PhaseCompile_Extension { // struct_size 0x40 (64 B), type 9
PJRT_Extension_Base base; // +0x00 size=64; +0x08 type=9; +0x10 next
/* +0x18 */ Get_Compiler* get_compiler; // TPU-injected (arg a3)
/* +0x20 */ Destroy_Compiler* destroy_compiler; // TPU-injected (arg a4)
/* +0x28 */ Run_Phase* run_phase; // generic pjrt:: driver
/* +0x30 */ Get_Phase_Names* get_phase_names; // generic pjrt:: driver
/* +0x38 */ C_Buffers_Destroy* c_buffers_destroy; // generic pjrt:: driver
};算法 — 创建器
function CreatePhaseCompileExtension(ext, next, get_fn, destroy_fn): // 0xe6f42a0
ext[+0x00] = 64 // struct_size
ext[+0x08] = 9 // PJRT_Extension_Type::PhaseCompile
ext[+0x10] = next // chain link (arg a2)
ext[+0x18] = get_fn // arg a3 (TPU GetTpuPhaseCompiler)
ext[+0x20] = destroy_fn // arg a4 (TPU DestroyTpuPhaseCompiler)
ext[+0x28] = pjrt::PJRT_PhaseCompile_Run_Phase // generic driver
ext[+0x30] = pjrt::PJRT_PhaseCompile_Get_Phase_Names // generic driver
ext[+0x38] = pjrt::PJRT_PhaseCompile_C_Buffers_Destroy // generic driver
return ext创建器是纯粹的扁平表写入器，没有分配也没有分支；在 0x0E6F42A0 可验证为八次 store 和一次 return。注入（get/destroy）和通用（run/names/destroy-buffers）之间的拆分就是整个设计：后端只提供编译器工厂和销毁逻辑，上游 C-ABI 驱动负责针对工厂返回的任何编译器对象完成所有编组。
注意 — 两个注入指针以调用参数（
a3、a4）形式到达，而不是来自创建器直接引用的 TPU 符号。这是 libtpu 扩展链中共享的通用创建器模式：CreatePhaseCompileExtension本身是标准 XLA 代码，TPUGetTpuPjrtApi会用 TPU 函数指针调用它。将 libtpu 移植到另一个后端的重新实现者只需重写Get_Compiler/Destroy_Compiler；创建器和三个驱动保持不变。
函数映射
| 槽位 | 字段 | 符号 | 地址 | 来源 |
|---|---|---|---|---|
+0x18 | get_compiler | tpu_plugin::GetTpuPhaseCompiler | 0x0E6AA320 | TPU 注入 |
+0x20 | destroy_compiler | tpu_plugin::DestroyTpuPhaseCompiler | 0x0E6AA400 | TPU 注入 |
+0x28 | run_phase | pjrt::PJRT_PhaseCompile_Run_Phase | 0x0E6F42E0 | 通用 |
+0x30 | get_phase_names | pjrt::PJRT_PhaseCompile_Get_Phase_Names | 0x0E6F4A60 | 通用 |
+0x38 | c_buffers_destroy | pjrt::PJRT_PhaseCompile_C_Buffers_Destroy | 0x0E6F4CC0 | 通用 |
| — | 创建器 | pjrt::CreatePhaseCompileExtension | 0x0E6F42A0 | 通用 |
Get_Compiler — 创建阶段编译器
目的
Get_Compiler 是宿主首先调用的入口：它生成每个后续 Run_Phase / Get_Phase_Names 调用都会携带的不透明编译器句柄。对于 TPU 后端，这是 tpu_plugin::GetTpuPhaseCompiler（0x0E6AA320），它构建一个新的 xla::TpuCompiler，在其上注册六个阶段，并把它装箱到一个很小的所有权 holder 中。
入口点
GetTpuPhaseCompiler (0xe6aa320, ext +0x18)
── ActualStructSizeIsGreaterOrEqual("PJRT_PhaseCompile_Get_Compiler_Args", 35, 24)
── operator new(0x40) ── xla::TpuCompiler, vtable off_2177D660, body zeroed
── xla::TpuCompiler::RegisterAllPhases (0xf849ec0) ── registers the 6 phases
└─ operator new(0x10) ── holder { base, owner } → args[+0x10]算法
function GetTpuPhaseCompiler(args): // 0xe6aa320
if !ActualStructSizeIsGreaterOrEqual( // line 14
"PJRT_PhaseCompile_Get_Compiler_Args", min=35, cur=24, args->struct_size):
return new PJRT_Error{status} // guard miss
compiler = operator new(0x40) // xla::TpuCompiler
zero(compiler+0x08 .. compiler+0x3F) // vmovups ymm 32B x2
compiler->vtable = off_2177D660 // line 24
status = xla::TpuCompiler::RegisterAllPhases(compiler) // 0xf849ec0, line 26
if status == OK: // (== 1)
holder = operator new(0x10) // line 29
holder->base = compiler // +0x00 — Run/Names deref this
holder->owner = compiler // +0x08 — Destroy frees this
args[+0x10] = holder // args[2], line 32
return NULL // success
else:
err = new PJRT_Error{status} // line 38
compiler->vtable[+8](compiler) // TpuCompiler dtor — tear down
return err注意事项
0x10holder 是句柄，不是编译器。 宿主接收的是holder，即{ TpuCompiler* base; TpuCompiler* owner; }这一层间接引用。Run_Phase/Get_Phase_Names读取holder->base并通过它的 vtable 分发；Destroy_Compiler置空并释放holder->owner。创建时两个指针相同，但 holder 的存在使所有权（要删除什么）与访问权（要调用什么）可以分离，这是上游 PJRT 用于可转移但可借用句柄的模式。- 失败路径会干净销毁。 如果
RegisterAllPhases失败，半构造的编译器会在返回错误前通过vtable+8（其 dtor）销毁，不会泄漏这个0x40字节对象。错误是堆上的PJRT_Error，持有absl::Status，形状与每个 PJRT wrapper 返回的错误相同。 struct_size包络。 args 结构体 guard 是min=35, cur=24。与该表面其他位置一样（见 API Vtable 重建），这是ActualStructSizeIsGreaterOrEqual(name, expected_size_for_min_version, current_size, host_supplied_size)；结构体小于最小值的宿主会被PJRT_Error拒绝，主体不会读取任何内容。
Destroy_Compiler
function DestroyTpuPhaseCompiler(args): // 0xe6aa400
holder = args[+0x10] // args[2]
if holder:
owner = holder->owner // +0x08
holder->owner = NULL
if owner: owner->vtable[+8](owner) // TpuCompiler dtor
free(holder) // free the 0x10 boxDestroy_Compiler 没有 struct_size guard，返回 void，它是纯清理例程。它通过 Get_Compiler 在失败路径上使用的同一个 vtable+8 槽位释放编译器，然后释放 holder。在 dtor 调用前置空 owner，可以防止宿主错误地销毁两次时出现 double-free。
Run_Phase — 驱动一个具名阶段
目的
Run_Phase（pjrt::PJRT_PhaseCompile_Run_Phase @ 0x0E6F42E0）是通用驱动，用来针对编译器句柄执行一个有序的具名阶段列表。它是增量 ABI 的核心：它从 C-ABI char-buffer span 反编组输入的部分程序 blob 和阶段名列表，反序列化 CompileOptionsProto（共享关口），分发到 RunPhases（编译器 vtable+56），再把输出的部分程序 vector 重新编组回 char-buffer span。
入口点
PJRT_PhaseCompile_Run_Phase (0xe6f42e0, ext +0x28)
── ActualStructSizeIsGreaterOrEqual("PJRT_PhaseCompile_Run_Phase_Args", 32, 120)
── ConvertCharBuffersToCppStrings (phase names → vector<string>)
── ConvertCharBuffersToPjRtPartialProgramProtos (input progs → vector<PjRtPartialProgramProto>)
── CompileOptionsProto::ParseFromString → CompileOptions::FromProto [shared chokepoint]
── null-check holder (args[+0x10]) ── "…Run_Phase: phase compiler is null"
── (*holder->base->vtable[+56])(…) ── PjRtPhaseCompiler::RunPhases
└─ ConvertPjRtPartialProgramProtosToCharBuffers (output vec → char buffers, written out)算法
function PJRT_PhaseCompile_Run_Phase(args): // 0xe6f42e0
if !ActualStructSizeIsGreaterOrEqual( // line 76
"PJRT_PhaseCompile_Run_Phase_Args", min=32, cur=120, args->struct_size):
return new PJRT_Error{status}
// --- unmarshal the two parallel char-buffer spans ---
phase_names = ConvertCharBuffersToCppStrings( // line 84
ptrs=args[+0x30], count=args[+0x40],
sizes=args[+0x38]) // → vector<string>
input_programs = ConvertCharBuffersToPjRtPartialProgramProtos( // line 85
ptrs=args[+0x18], count=args[+0x28],
sizes=args[+0x20]) // → vector<PjRtPartialProgramProto>
// --- the compile-options chokepoint (see §Compile-Options Ingest) ---
proto = CompileOptionsProto{} // line 134
if !proto.ParseFromString(args[+0x48], args[+0x50]): // line 135 — ptr +0x48, len +0x50
return new PJRT_Error{MakeErrorImpl<3>(
"PJRT_Client_Compile: failed to deserialize CompileOptionsProto",
…, "…/pjrt_c_api_phase_compile_internal.cc")} // :46 (NOT wrapper_impl:1113)
opts = CompileOptions::FromProto(proto) // line 138
holder = args[+0x10] // args[2]
if !holder: // line ~230
return new PJRT_Error{MakeErrorImpl<13>(
"PJRT_PhaseCompile_Run_Phase: phase compiler is null", // :77
…, "…/pjrt_c_api_phase_compile_internal.cc")}
// --- drive the pipeline: RunPhases via compiler vtable +56 ---
out_programs = (*holder->base->vtable[+56])( // line 238
holder->base, &opts, input_programs,
topology = *(args[+0x88]) + 8, // PjRtTopologyDescription
phase_names) // → StatusOr<vector<PjRtPartialProgramProto>>
if !out_programs.ok(): return new PJRT_Error{out_programs.status}
// --- remarshal results into the out char-buffer span ---
ConvertPjRtPartialProgramProtosToCharBuffers(out_programs, args+0x60..)
return NULLargs 布局
反编译出的主体会读取 PJRT_PhaseCompile_Run_Phase_Args 内的下列偏移。两个并行 char-buffer span（程序和阶段名）各自携带一个 (ptrs, sizes, count) 三元组；这两个 span 的语义标签是从 Convert* 调用参数顺序推断出来的。
| 偏移 | 字段（恢复） | 含义 |
|---|---|---|
+0x00 | struct_size | guard 输入；min 32，cur 120 |
+0x10 | phase_compiler | 来自 Get_Compiler 的 {base,owner} holder（args[2]） |
+0x18 | input_programs | char** — 序列化的 PjRtPartialProgramProto blob |
+0x20 | input_programs_sizes | size_t* — 每个 blob 的字节长度 |
+0x28 | input_programs_count | 元素数量 |
+0x30 | phase_names | char** — 阶段名字符串 |
+0x38 | phase_names_sizes | size_t* — 每个名称的长度 |
+0x40 | phase_names_count | 元素数量 |
+0x60 | out_programs | char** out — 成功时写入 |
+0x68 | out_programs_sizes | size_t* out |
+0x70 | out_programs_count | count out |
+0x88 | topology | *(args+0x88)+8 → PjRtTopologyDescription |
+0x48 | compile_options | char* — 序列化的 CompileOptionsProto，送入 ParseFromString |
+0x50 | compile_options_size | size_t — proto 字节长度 |
注意 — 上面的偏移直接来自
Run_Phase主体0xe6f42e0（第 84-85、134-135 行）：两个 char-buffer span 是不同的(ptr, size, count)三元组，程序位于(+0x18, +0x20, +0x28)，阶段名位于(+0x30, +0x38, +0x40)；选项 proto 是(ptr=+0x48, len=+0x50)。反编译器把这些渲染为十进制（*(a1+72)/*(a1+80)= 十六进制+0x48/+0x50）；不要把十进制误读为十六进制。确切的逐字段名称取决于公开的pjrt_c_api_phase_compile.h头文件顺序，而它不在二进制中；偏移具有高置信度。
注意事项
- 分发是
vtable+56。 反编译主体调用(*(holder->base->vtable + 56))(…)，即PjRtPhaseCompiler::RunPhases @ 0x1D16AF20。这是整体式PJRT_Client_Compile驱动的同一个编译器对象，只是通过不同 vtable 槽位到达；分阶段 ABI 暴露整体式路径会丢弃的中间PjRtPartialProgramProto产物。 - 错误字符串有误导性，它即使在这里也命名为
PJRT_Client_Compile。CompileOptionsProto解析失败会返回"PJRT_Client_Compile: failed to deserialize CompileOptionsProto"，虽然这里是Run_Phase。解析逻辑是每个编译入口都会运行的同一个双调用序列，但Run_Phase携带其自己的内联副本（发射在pjrt_c_api_phase_compile_internal.cc:46，MakeErrorImpl<3>），不同于整体式入口共享的pjrt_c_api_wrapper_impl.cc:1113位置；这个字符串只是逐字复制（和DeserializeAndLoad的怪癖相同，见可执行体执行）。重新实现者不应基于该字符串进行错误处理。 - 每一步的错误都会以
PJRT_Error*暴露。 反编组失败、选项解析失败、编译器为空，或RunPhases失败，都会生成一个堆上的PJRT_Error，其中持有absl::Status；成功返回NULL。
Get_Phase_Names — 枚举已注册阶段
目的
Get_Phase_Names（pjrt::PJRT_PhaseCompile_Get_Phase_Names @ 0x0E6F4A60）返回编译器注册的有序阶段名列表，让宿主可以在驱动管线前发现它。它是 Run_Phase 的只读配套入口。
算法
function PJRT_PhaseCompile_Get_Phase_Names(args): // 0xe6f4a60
if !ActualStructSizeIsGreaterOrEqual( // line 25
"PJRT_PhaseCompile_Get_Phase_Names_Args", min=38, cur=48, args->struct_size):
return new PJRT_Error{status}
holder = args[+0x10] // args[2]
if !holder: // line 34
return new PJRT_Error{MakeErrorImpl<13>(
"PJRT_PhaseCompile_Get_Phase_Names: phase compiler is null", // :100
…, "…/pjrt_c_api_phase_compile_internal.cc")}
names = (*holder->base->vtable[+48])(holder->base) // PjRtPhaseCompiler::GetPhaseNames
// 0x1d16ae20, line 40
args[+0x18] = ConvertCppStringsToCharBuffers(names, &args[+0x20]) // ptrs / sizes
args[+0x28] = names.count // line 71
free(names) // line 60-114: release the vector
return NULL注意事项
- 分发是
vtable+48。GetPhaseNames @ 0x1D16AE20返回已注册阶段名 vector（即RegisterAllPhases追加它们的顺序）。随后驱动把它们通过ConvertCppStringsToCharBuffers转为args[+0x18]/args[+0x20]（指针数组 / sizes），并把 count 放在args[+0x28]。 - args 结构体名称有两种拼写。
struct_size断言字符串是"PJRT_PhaseCompile_Get_Phase_Names_Args"（0xe6f4a60，第 25 行），但 IDA 恢复出的函数 mangled symbol 接收PJRT_PhaseCompile_Get_PhaseNames_Args*（Phase和Names之间没有下划线）。两种拼写都出现在二进制中；线上结构体使用断言命名的带下划线形式。 - 返回的 vector 会就地释放。 转换成 char buffer 后，主体会遍历
vector<string>，释放每个长字符串分配（libc++ SSO 符号位检查(char)x < 0），再释放 backing array，也就是标准的0xAAAAAAAAAAAAAAABreciprocal-of-24 stride math。宿主只拥有转换后的 char buffer，之后通过C_Buffers_Destroy释放它们。
C_Buffers_Destroy — 释放返回的数组
目的
C_Buffers_Destroy（pjrt::PJRT_PhaseCompile_C_Buffers_Destroy @ 0x0E6F4CC0）释放 Run_Phase（out_programs）和 Get_Phase_Names（阶段名）返回给宿主的 char-buffer 数组。它是宿主针对这些分配必须调用的清理钩子。
算法
function PJRT_PhaseCompile_C_Buffers_Destroy(args): // 0xe6f4cc0
ptrs = args[+0x10] // args[2] — char** array
count = args[+0x20] // args[4] — element count
if ptrs:
for i in 0 .. count-1: free(ptrs[i]) // each buffer
free(ptrs) // the array
sizes = args[+0x18] // args[3] — size_t* array
if sizes: free(sizes)
// no status returned — pure cleanup它不返回状态；与 Destroy_Compiler 一样，它是 void 清理例程。重新实现必须为每个产生 buffer 的 Run_Phase / Get_Phase_Names 配对调用一次 C_Buffers_Destroy，否则会泄漏返回的数组。
六个已注册阶段
目的
xla::TpuCompiler::RegisterAllPhases @ 0x0F849EC0 是 TPU 后端声明其分阶段管线的位置。它调用 PjRtPhaseCompiler::RegisterPhase @ 0x1D16AB20 六次，每次传入一个名称、一个编译函数，以及共享的 CommonPhaseValidator。RegisterPhase 会把 {name → CompilationPhaseFunctions{compile_fn, validator_fn}} 插入 hash map，并把名称追加到 GetPhaseNames 返回的有序 vector 中。
阶段表
这六个注册在 RegisterAllPhases（0xf849ec0，六次 RegisterPhase 调用）中得到验证；名称字符串在 .rodata 中逐字存在。
| 顺序 | 阶段名 | 编译函数 | 地址 | 输出格式 |
|---|---|---|---|---|
| 1 | phase0_stablehlo_to_hlo | CompilePhase0StablehloToHlo | 0x0F84DE60 | unopt_hlo |
| 2 | phase1_hlo_opts | CompilePhase1HloOptimizations | 0x0F84EE00 | optimized HLO |
| 3 | phase2a_tlp_lowering | CompilePhase2aTlpLowering | 0x0F850840 | TLP-lowered |
| 4 | phase2b_deduped_lowering | CompilePhase2bDedupedLowering | 0x0F852180 | deduped-lowered |
| 5 | phase3_linking | CompilePhase3Linking(test=false) | 0x0F852F40 | device program |
| 6 | phase3_linking_test_only | CompilePhase3Linking(test=true) | 0x0F852F40 | device program |
算法 — 注册
function TpuCompiler::RegisterAllPhases(this): // 0xf849ec0
RegisterPhase(this, "phase0_stablehlo_to_hlo", // line 74-80
CompilePhase0StablehloToHlo, CommonPhaseValidator)
RegisterPhase(this, "phase1_hlo_opts", // line 101-108
CompilePhase1HloOptimizations, CommonPhaseValidator)
RegisterPhase(this, "phase2a_tlp_lowering", // line 137-143
CompilePhase2aTlpLowering, CommonPhaseValidator)
RegisterPhase(this, "phase2b_deduped_lowering", // line 175-181
CompilePhase2bDedupedLowering, CommonPhaseValidator)
RegisterPhase(this, "phase3_linking", // line 204-212
bind_front(CompilePhase3Linking, false), CommonPhaseValidator)
RegisterPhase(this, "phase3_linking_test_only", // line 237-244
bind_front(CompilePhase3Linking, true), CommonPhaseValidator)
// each returns Status; a failure aborts and propagates (Get_Compiler tears down)每个阶段函数具有如下签名
StatusOr<vector<PjRtPartialProgramProto>>
CompilePhaseN(CompileOptions,
Span<const PjRtPartialProgramProto> /* input */,
const PjRtTopologyDescription&);共享 validator 是 Status CommonPhaseValidator(Span<const PjRtPartialProgramProto>)，在每个阶段前针对输入程序运行。RegisterPhase 将两者存入一个 CompilationPhaseFunctions 记录；RunPhases 在该记录上通过 vtable+72 调用 validator，并通过 vtable+40 调用编译函数。
怪癖 —
phase3_linking和phase3_linking_test_only是同一个函数（CompilePhase3Linking @ 0xf852f40）的两次注册，仅绑定的bool test参数不同（false生产，true测试）。它们不是五阶段管线加第六阶段；它们是五阶段管线（phase0→phase1→phase2a→phase2b→phase3）以及终端 linker 的一个测试别名。通过Get_Phase_Names枚举阶段的宿主会看到六个名称；正常编译驱动五个。
驱动 — RunPhases
function PjRtPhaseCompiler::RunPhases(opts, programs, topology, phase_names): // 0x1d16af20
current = programs // vector<PjRtPartialProgramProto>
for name in phase_names: // host-supplied order
rec = phase_map.find(name)
if !rec:
return Error("No phase compiler/validator registered with phase name \"%s\"", name) // :281
rec->validator(current) // CompilationPhaseFunctions vtable +72
current = rec->compile_fn(opts, current, topology) // vtable +40 → new vector
return current // final artifacts每个阶段的输出会成为下一阶段的输入（vector 被重新赋值），因此宿主可以单独运行 ["phase0_stablehlo_to_hlo"]，持久化 unopt_hlo blob，之后再从该 blob 恢复并运行 ["phase1_hlo_opts", "phase2a_tlp_lowering", …]。RegisterPhase 会用 "A phase compiler/validator with Phase name \"%s\" already exists"（pjrt_compiler.cc:259）拒绝重复名称，并用第 250/253/256 行的 guard 拒绝空名称或空函数。
部分程序产物格式
跨越每个阶段边界的 blob 是一个 PjRtPartialProgramProto（proto3，package xla，pjrt_partial_program.proto）：
message PjRtPartialProgramProto {
bytes program = 1; // serialized stage output (e.g. unopt HLO)
string program_format = 2; // e.g. "unopt_hlo"
string producer_phase = 3; // e.g. "phase0_stablehlo_to_hlo"
repeated string consumer_phases = 4; // e.g. ["phase1_hlo_opts"]
string version = 5;
string program_name = 6;
}CompilePhase0StablehloToHlo（0xf84de60）是一个完整示例：它对输入 StableHLO module 执行 ParseMlirModuleString，使用 getDefaultChloToHighLevelMhloOptions 通过 MlirToXlaComputation 生成 HloModuleProto，把 arg/output layout-mode 和 memory-space 属性填入 HLO 的 frontend_attributes（键为 kArgLayoutModesAttr / kOutLayoutModesAttr / kArgMemorySpacesAttr / kOutMemorySpacesAttr），并发出 PjRtPartialProgramProto{ program=serialized HLO, program_format="unopt_hlo", producer_phase="phase0_stablehlo_to_hlo", consumer_phases=["phase1_hlo_opts"] }。下游阶段的 HLO-pass 内容见 HLO Pass 注册表。
注意 —
phase1_hlo_opts、phase2a_tlp_lowering、phase2b_deduped_lowering和phase3_linking的内部主体在此只按输入/输出契约层面刻画（每个都把vector<PjRtPartialProgramProto>映射到vector<PjRtPartialProgramProto>）；已确认 phase 3 会构建BufferAssignmentProto并发出可加载的设备程序。它们的内部 pass 管线很大，位于编译器页面，而不在这里。
编译选项摄入
目的
libtpu 中的每个编译入口，包括 PJRT_Client_Compile、PJRT_Compile、PJRT_Client_Load、PJRT_Executable_DeserializeAndLoad 以及 PhaseCompile Run_Phase，都会通过同一个双调用序列（CompileOptionsProto::ParseFromString → CompileOptions::FromProto）反序列化同一个 CompileOptionsProto，并发出逐字节相同的失败字符串。前四个共享一个发射位置（pjrt_c_api_wrapper_impl.cc:1113）；Run_Phase 用自己的位置（pjrt_c_api_phase_compile_internal.cc:46）内联重新实现该序列，逻辑相同、诊断相同、源码行不同。这是线上选项变成活的 C++ xla::CompileOptions 的汇聚点。TPU 编译器配置（拥有 1,121 个字段的 xla.jellyfish.TpuCompilationEnvironment）以及 xla::DebugOptions / XLA flags 全部都位于这个 proto 内部。
关口
// PJRT_Client_Compile / PJRT_Compile / PJRT_Client_Load / DeserializeAndLoad → wrapper_impl.cc:1113
// PhaseCompile Run_Phase (line 134) → its own inline copy at phase_compile_internal.cc:46
proto = CompileOptionsProto{} // arena = 0
if !proto2::MessageLite::ParseFromString(&proto, ptr, len):
return new PJRT_Error{MakeErrorImpl<3>(
"PJRT_Client_Compile: failed to deserialize CompileOptionsProto", // verbatim, all 5
…, /* wrapper_impl.cc:1113 | Run_Phase: phase_compile_internal.cc:46 */)}
opts = xla::CompileOptions::FromProto(proto) // wire → live C++ options线格式 schema
CompileOptionsProto（proto3，package xla，…/pjrt/proto/compile_options.proto）是载体。TPU 编译环境在其中嵌套三层：
CompileOptionsProto {
repeated ShapeProto argument_layouts = 1;
bool parameter_is_tupled_arguments = 2;
ExecutableBuildOptionsProto executable_build_options = 3; // ← env carrier
bool compile_portable_executable = 4;
int64 profile_version = 5;
bytes serialized_multi_slice_config = 6;
repeated EnvOptionOverridesEntry env_option_overrides = 7; // ← per-flag overrides
GpuTargetConfigProto target_config = 8;
bool allow_in_place_mlir_modification = 9;
PrecisionConfig.Precision matrix_unit_operand_precision = 10;
string compiler_variant = 11;
}
ExecutableBuildOptionsProto {
int64 device_ordinal = 1;
DebugOptions debug_options = 3; // ← XLA-standard flags
CompilationEnvironmentsProto comp_envs = 13; // ← THE env nest
int64 num_replicas = 4; num_partitions = 5;
bool use_spmd_partitioning = 6; use_auto_spmd_partitioning = 7;
bytes fdo_profile = 14; // feedback-directed-opt
...
}
CompilationEnvironmentsProto { repeated google.protobuf.Any environments = 1; }三种注入通道
重新实现者需要知道，TPU 编译器配置到达编译器有三种同心方式，从最窄到最宽：
| 通道 | 位置 | 作用域 |
|---|---|---|
| 单 flag 覆盖 | CompileOptionsProto.env_option_overrides（字段 7） | 一个 flag，本次编译；map<name, {string|bool|int|double}> |
| 完整 env，按编译 | executable_build_options.comp_envs.environments[Any → TpuCompilationEnvironment]（3→13→1） | 整个 1,121 字段表，本次编译 |
| 进程全局 | TpuExecutable ext SetTpuCompilationEnv（类型 17，槽位 +0x40，0xe6dd400） | 一个 CompilationEnvironments singleton，之后的编译会继承 |
前两个位于 CompileOptionsProto 内部，因此和其他每个编译入口一样，以完全相同的方式通过 Run_Phase 的关口。第三个是 TpuExecutable 扩展（类型 17）上的独立带外表面，不是 PhaseCompile 扩展；它在那个扩展处记录，这里只作提示，让重新实现者知道 PhaseCompile 路径并不是唯一的 env 通道。
嵌套环境的解码
FromProto 运行时，CompilationEnvironments::CreateFromProto @ 0x1E63E5A0 会遍历 comp_envs.environments 中的每个 Any，通过 FindMessageTypeByName 在生成的 descriptor pool 中解析其 type_url（例如 xla.jellyfish.TpuCompilationEnvironment），用 MessageFactory::New 创建 prototype，对 Any.value 字节执行 InternalUnpackTo，并用 AddEnvImpl 把活的 message 加入 FlatHashMap<Descriptor*, unique_ptr<Message>>。格式错误的 Any 会产生 compilation_environments.cc:139..160 处的诊断（"Invalid/Unknown/Unsupported CompilationEnvironment message type"、"Unable to unpack…"）。因此 TpuCompilationEnvironment 是在编译时通过反射解码的，即按 protobuf 字段名查询，而不是按固定偏移查询。完整字段表见编译器页面。
怪癖 — libtpu 的编译完全在进程内进行：
CompileOptionsProto在同一地址空间中被消费并通过反射解码；没有编译器子进程。xla::CommonPjRtClient::supports_two_phase_launch @ 0xe6edbc0硬返回1，但 "two-phase" 指的是编译/加载是可分离的 API 调用，而不是编译会 fork 到进程外。（这与某些其他 PJRT 后端相反，那些后端的Compile会 fork 一个外部编译器二进制，并以 argv/临时文件传递选项。libtpu 的 TPU 编译器是静态链接的。）
与整体式路径的关系
PhaseCompile 扩展和一次性 PJRT_Client_Compile（槽位 25）是同一个编译器之上的两个 ABI：
| 方面 | 整体式（PJRT_Client_Compile，槽位 25） | 分阶段（PhaseCompile ext，类型 9） |
|---|---|---|
| 调用者提供 | 整个程序 + CompileOptionsProto | 逐阶段 PjRtPartialProgramProto blob + 阶段名列表 + 选项 |
| 返回 | 完成的、已加载到设备的 PJRT_LoadedExecutable | 阶段之间的中间 PjRtPartialProgramProto 产物 |
| 底层对象 | xla::TpuCompiler（通过 client vtable） | xla::TpuCompiler（通过 ext Get_Compiler holder） |
| 编译器分发 | client Compile 虚函数 | 在已注册阶段上执行 RunPhases（vtable+56） |
| 中间产物 | 丢弃 | 暴露（可缓存、可检查、可跨宿主拆分） |
| 选项关口 | CompileOptionsProto::ParseFromString → FromProto | 相同 |
分阶段 ABI 提供了整体式 ABI 无法提供的三件事：（a）缓存部分编译的程序（在 phase0 后持久化 unopt_hlo blob），（b）跨宿主拆分阶段，以及（c）当只有后期阶段选项改变时，只重新编译下游阶段。任一路径生成的最终可执行体，包括其 serialize/deserialize 往返和 Execute，见可执行体加载与执行。
相关组件
| 名称 | 关系 |
|---|---|
xla::TpuCompiler | Get_Compiler 创建并由 RegisterAllPhases 配置的 0x40 字节编译器对象 |
xla::PjRtPhaseCompiler | 提供 RegisterPhase / RunPhases / GetPhaseNames 的基类（vtable +40/+48/+56/+72） |
CompilationPhaseFunctions | 按阶段名存储的 {compile_fn, validator_fn} 记录 |
pjrt::ActualStructSizeIsGreaterOrEqual | 每个入口的 struct_size 版本 guard，与每个 PJRT wrapper 共享 |
xla::CompilationEnvironments | 在 FromProto 时解码嵌套的 TpuCompilationEnvironment Any |
| TpuExecutable 扩展（类型 17） | 另一个 env 通道（SetTpuCompilationEnv），相对于编译是带外的 |
交叉引用
- 扩展链 — 宿主如何遍历
PJRT_Extension_Base.next并匹配type == 9以到达该结构体 - API Vtable 重建 —
PJRT_Api槽位表，以及每个入口共享的struct_size向后兼容 guard - 可执行体加载与执行 — 编译生成的最终可执行体：
PJRT_Client_Compile、serialize / deserialize 往返，以及Execute - 客户端与设备 —
PJRT_Client_Create，构建TpuClient并安装 env 解码所用 descriptor pool 的注入槽位 - PJRT 概览 — 编译表面位于插件生命周期中的位置
- HLO Pass 注册表 —
phase1/phase2/phase3主体驱动的 HLO pass 管线 - 编译缓存 — 缓存部分（
PjRtPartialProgramProto）和完整结果的内容寻址存储