PJRT C-ABI 概览
本页中的所有地址均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so（buildlibtpu_lts_20260413_b_RC00，build-id md589edbbe81c5b328a958fe628a9f2207d，781,691,048 bytes，ELF x86-64 DYN，not stripped；demangled C++ symbols 按原文引用）。其他版本会有所不同。
摘要
PJRT 是稳定的 C 插件 ABI，所有 XLA 前端（JAX、TensorFlow、PyTorch-XLA）都通过它驱动硬件后端，而不需要链接后端的 C++，也不需要与后端共享 C++ ABI。该契约刻意保持最小：框架 dlopen libtpu.so，dlsym 一个导出的入口符号，无参数调用它，然后接收一个扁平 C struct（PJRT_Api）的指针，其字段都是函数指针。之后的每一次交互（创建 client、编译 HLO module、上传 buffer、启动 executable、等待 event）都是通过该 struct 某个槽位进行的间接调用。由于该 struct 是普通 C，且每次调用都做 size check，基于旧 header 构建的框架可以驱动较新的插件，反之亦然。这与 StreamExecutor 的 C-shim（TfTpu_*ApiFn）在 legacy stack 中扮演的角色相同，但 PJRT 是公开、版本化的表面，也是现代 PJRT client 接触的唯一表面。
在此构建中，入口符号是 GetPjrtApi @ 0xe6a83a0，它是一个 5-byte jmp thunk，跳入真正的引擎 pjrt::tpu_plugin::GetTpuPjrtApi @ 0xe6aa440（该 thunk 本身在 ../lifecycle/get-pjrt-api-thunk.md 中剖析）。这里没有导出的 GetTpuPjrtApi；canonical casing 的 GetPjrtApi 是唯一匹配 /Pjrt/ 的 GLOBAL FUNC 导出，版本为 GetPjrtApi@@VERS_1.0。它返回的 PJRT_Api 表不是任何 PROGBITS section 中的静态镜像，而是 .lbss 中位于 0x227BA840 的函数局部 Meyers singleton（1120 bytes = 140 × 8），加载时填零，并在首次调用时由 pjrt::CreatePjrtApi @ 0xf874160 在 __cxa_guard 保护下惰性填充。该表是 PJRT C-API v0.103：140 个 qword 槽位，包括 5 个 header scalar（struct_size、extension_start 和一个 3 字段 PJRT_Api_Version sub-struct），后接 135 个函数指针。挂在 extension_start 上的是 17 节点单向扩展链，承载不适合放入通用 vtable 的 TPU 专用表面（Megascale、MultiSlice、collectives、raw buffers、phase-compile、profiler 等）。
本页是 PJRT 章节的地图。它确定 ABI 的形状、dlsym 握手和惰性构建路径（按符号说明，深层 module-init 生命周期由其他页面负责），并按区域给出 140 槽表和 17 扩展链的一览索引，链接到负责各区域逐字段细节的同级页面。它不复述完整槽位表（api-vtable-reconstruction.md）或扩展节点布局（extension-chain.md）；它告诉你哪个同级页面负责什么。
作为定位，契约是：
- 握手：
dlsym("GetPjrtApi")→ 无参数调用 → 一个const PJRT_Api*；调用者可以假设什么，以及必须发现什么。 - 结构形状：1120-byte 扁平 C 表：5 个 header scalar + 135 个 fn-ptr，按 native order 精确匹配公开
xla/pjrt/c/pjrt_c_api.hv0.103 schema，并带有每次调用的struct_size版本化。 - 5 个 TPU 注入点：槽位 8/9/15/87/103 是
CreatePjrtApi唯一从调用者取得的槽位；其他 130 个是编译期固定的通用 XLA wrapper。 - 扩展链：一个 NULL-terminated、最新优先的 17 个 typed extension 链表，只能通过
extension_start到达，绝不能通过固定偏移到达。
| 导出入口符号 | GetPjrtApi @ 0xe6a83a0（5-byte jmp thunk，@@VERS_1.0） |
| 真实引擎 | pjrt::tpu_plugin::GetTpuPjrtApi @ 0xe6aa440（1336 B，17 个 __cxa_guard 块） |
| 表构造器 | pjrt::CreatePjrtApi @ 0xf874160（1872 B，扁平 header-write + slot-fill，无循环） |
| 表存储 | GetTpuPjrtApi()::pjrt_api @ 0x227BA840，.lbss（NOBITS），1120 B = 140 × 8 |
| C-API 版本 | v0.103 — version qword 0x6700000000 → {major=0, minor=0x67=103} |
| 槽位数 | 140（5 个 header scalar + 135 个 fn-ptr） |
| 扩展链头 | extension_start（slot 1）→ host_memory_allocator_extension @ 0x224C3F68 |
| 扩展数量 | 17（16 个驻留于 .bss 并惰性构建；1 个 .data 静态节点：profiler） |
| TPU 注入槽位 | 8, 9, 15, 87, 103（Plugin_Initialize、Plugin_Attributes、Client_Create、TopologyDescription_Create、ExecuteContext_Create） |
| 置信度 | CONFIRMED（byte-anchored vs decompile），除非某行或 callout 另有说明 |
1. 插件发现握手
目的
PJRT 存在的核心理由是提供一个单一、ABI-stable 的会合点。框架不了解 libtpu 的内部；它只知道入口符号名称，以及前几个 PJRT_Api 字段的布局。其他所有内容都在运行时通过 struct_size 字段和扩展链发现。本节固定这个会合点，使重新实现者可以产出一个 stock JAX/PyTorch-XLA 构建能够加载的 .so。
契约
// The one symbol the framework dlsym's. No arguments, returns the table.
const PJRT_Api* GetPjrtApi(void); // exported, GetPjrtApi@@VERS_1.0
```text
```text
framework (JAX / TF / PyTorch-XLA)
dlopen("libtpu.so")
dlsym(handle, "GetPjrtApi") ── the ONLY name that resolves; "GetTpuPjrtApi" is internal
│
└─ GetPjrtApi 0xe6a83a0 ── 5-byte: jmp 0xe6aa440
└─ pjrt::tpu_plugin::GetTpuPjrtApi 0xe6aa440
├─ build 16 .bss extensions (each one-shot __cxa_guard)
├─ (17th guard) pjrt::CreatePjrtApi(&pjrt_api, …) 0xf874160 ── writes all 140 slots
└─ return &pjrt_api = 0x227BA840 (.lbss)随后调用者读取 api->struct_size 来获知该插件实际提供多少槽位，读取 api->pjrt_api_version 来获知 minor version，并遍历 api->extension_start 来发现可选能力。接着它调用 api->PJRT_Plugin_Initialize（slot 8），即一次性的 TPU driver bring-up，并调用 api->PJRT_Client_Create（slot 15）来创建 live client。这两个槽位都会深入 module-init 和 silicon detection，它们属于 lifecycle 章节的范围（PJRT_Plugin_Initialize → TryAcquireTpuLock → InitializeDriver → GoogleInitializer module DAG；silicon scan 延迟到 PJRT_Client_Create）；本页只按符号引用它们，见 ../lifecycle/module-init-plugin-discovery.md 和 ../lifecycle/do-init-do-fini.md。
注释 — 该表在首次调用时构建，而不是在
dlopen时构建。GetTpuPjrtApi的pjrt_api是.lbss（NOBITS）中的函数局部静态对象，在第一次GetPjrtApi前为零。17 个受__cxa_guard保护的块（16 个 extension builder + 最后的CreatePjrtApi）恰好运行一次；并发的首次调用者通过 Itanium-ABI__cxa_guard语义串行化，一次性初始化后该 struct 在进程生命周期内不可变，读取者不加锁。因此静态反汇编无法显示已填充的槽位值；slot→impl 映射是从CreatePjrtApi的函数体重建的，而不是从填零的.lbss镜像读取的。陷阱 — 拼写和大小写是 ABI 的一部分。导出符号是
GetPjrtApi（小写jrt），匹配公开 PJRT plugin 约定；GetTpuPjrtApi是内部 helper，且不导出。若 loaderdlsymGetTpuPjrtApi，或某个构建只导出带Tpu前缀的名称，发现会失败。该二进制中共同存在的Tpu*_*导出（194 个符号，全部为FUNC GLOBAL @@VERS_1.0）是 legacy StreamExecutor C-ABI，由 TF-TPU 直接链接，永远不会通过 PJRT 到达；见 stream-executor-pjrt-adapter.md。
2. PJRT_Api 结构一览
目的
PJRT_Api 是一个包含 140 个 qword 槽位的扁平 C struct。重新实现者在陷入 135 行字段表前，需要先了解其形状：header 与函数指针区域、版本化约定，以及哪些槽位是插件专用的。本节给出该形状；逐字段重建见 api-vtable-reconstruction.md。
Header 布局
前五个槽位不是函数指针。它们是使 ABI 前向/后向兼容的自描述 header。已逐字节对照 CreatePjrtApi @ 0xf874160 确认（*a1 = 1120; a1[2] = 24; a1[4] = 0x6700000000）。
| Slot | Offset | Field | Value |
|---|---|---|---|
| 0 | +0x00 | struct_size | 1120 (= 140 × 8) |
| 1 | +0x08 | extension_start | → host_memory_allocator_extension @ 0x224C3F68 |
| 2 | +0x10 | pjrt_api_version.struct_size | 24（sub-struct 自身大小） |
| 3 | +0x18 | pjrt_api_version.priv | NULL（reserved） |
| 4 | +0x20 | pjrt_api_version.{major,minor} | {0, 103}（qword 0x6700000000） |
pjrt_api_version 是嵌入的 PJRT_Api_Version { size_t struct_size; void* priv; int major; int minor }（24 bytes），因此槽位 2..4 是一个逻辑字段。major/minor 以 little-endian 打包进 slot-4 qword：低 32 bits = major（0），高 32 = minor（0x67 = 103）。
函数指针区域
槽位 5..139 是 135 个函数指针，顺序与公开 xla/pjrt/c/pjrt_c_api.h v0.103 header 声明完全一致，包括在原始 v0.40 表面之后追加的晚期槽位（95..139），按功能落地顺序排列。这里不复述全部 135 个槽位，而是将它们分组到各同级页面负责的区域。pre-95 块是稳定的 v0.40 core；95..139 是 v0.103 late additions（Output{ElementTypes,Dimensions}、CopyToMemory、CreateViewOfDeviceBuffer、Executable_Fingerprint、AsyncHostToDeviceTransferManager_* family、DmaMap/Unmap、CreateAliasBuffer、DonateWithControlDependency、Event_Create/Set、Client_Load、Bitcast、Error_ForEachPayload、Topology_Fingerprint、ParameterMemoryKinds）。
| Slot range | 区域 | 负责页面 |
|---|---|---|
| 5–7, 137 | Error: Destroy / Message / GetCode / ForEachPayload | events-and-async.md |
| 8–9 | Plugin: Initialize, Attributes（lifecycle entry） | （lifecycle — 按符号引用） |
| 10–14, 131–132 | Event: Destroy / IsReady / Error / Await / OnReady / Create / Set | events-and-async.md |
| 15–27, 98, 100, 108, 115–123, 134 | Client: create/lookup/compile/buffer-from-host/alias/dma/load | client-and-device.md |
| 28–39, 126–127, 133 | DeviceDescription + Device | client-and-device.md |
| 40–44, 102 | Memory | buffer-and-memory.md |
| 45–54, 95–96, 99, 101, 129, 139 | Executable（program metadata, serialize, cost, fingerprint） | executable-execution.md |
| 55–62, 122, 135 | LoadedExecutable（包括 slot 60 Execute，hot path） | executable-execution.md |
| 63–81, 97, 105, 125, 130, 136 | Buffer（lifecycle, transfer, refcount, donate, bitcast） | buffer-and-memory.md |
| 82–86 | CopyToDeviceStream | dma-and-cross-host-recv.md |
| 87–93, 100, 119, 138 | TopologyDescription | ext-topology-description.md |
| 94 | PJRT_Compile（AOT, no client） | executable-execution.md |
| 103–104 | ExecuteContext | executable-execution.md |
| 106–114, 124 | AsyncHostToDeviceTransferManager | buffer-and-memory.md |
| 127–128 | AsyncTrackingEvent | events-and-async.md |
特性 — 槽位顺序不是像上表那样按对象分组。header 的区域分组只是阅读辅助；实际
pjrt_c_api.h顺序会交错各 family，因为槽位按上游功能落地顺序追加。PJRT_Compile是 slot 94，但逻辑上属于 Executable 块的Executable_OutputElementTypes是 slot 95，即在它之后，因为该功能后加入。重新实现者必须逐字复现 header schema 的 wire order；family grouping 是给人看的，slot index 才是 ABI。完整有序列表见 api-vtable-reconstruction.md。
版本化：旧 client 如何与新插件通信
header 通过两个协同机制实现前向/后向兼容。首先，每个可达的 PJRT_Api 字段本身都通过其 args struct 上的 struct_size 做版本化：几乎每个槽位的第一条指令都会调用 pjrt::ActualStructSizeIsGreaterOrEqual("<API>_Args", min, current, args->struct_size)，它接受从文档化 min 到插件 current 的任何调用者 args struct，并且不读取超出调用者 struct_size 的字段。其次，表级 struct_size（slot 0）和 pjrt_api_version（slots 2..4）让调用者在索引槽位前知道有多少槽位存在。基于 v0.95 编译的 client 看到 struct_size = 1120 和 minor 103，不会索引超过自己已知的槽位，并传入更小的 args struct，而 per-slot guard 会接受它们。反向情形，即 v0.103 client 对接旧插件，则受旧插件的 struct_size 限制。
陷阱 — 绝不能通过假设固定表大小来读取
extension_start或任何 late slot。硬编码 140 个槽位的重新实现会错误解析旧插件或新插件。应先读取struct_size；将 offset>= struct_size的任何槽位视为 absent。同一规则也通过 argsstruct_sizeguard 应用于每次调用。
3. 扩展链一览
目的
通用 PJRT_Api vtable 无法在不污染公开 schema 的情况下承载后端专用表面（TPU topology details、Megascale/MultiSlice multi-pod features、collectives、raw-buffer DMA、profiler）。PJRT 通过扩展链解决这个问题：一个 NULL-terminated 单向链表，由 typed PJRT_Extension_Base 节点组成，只能通过 extension_start（slot 1）到达。每个节点以 { size_t struct_size; PJRT_Extension_Type type; PJRT_Extension_Base* next; } 开头（.next 位于 offset +0x10），后接该扩展自己的函数指针。本节按 type 索引该链；逐节点字段图见 extension-chain.md 和各同级扩展页面。
链
extension_start = 0x224C3F68。该链以最新优先顺序遍历（构建顺序的反向），并在 profiler → NULL 处终止。总计 17 个节点：16 个驻留于 .bss，在 GetTpuPjrtApi 内通过 __cxa_guard 惰性构建；profiler 是唯一的 .data static-init 节点，并作为链尾种子。
| # | Type ID | Extension | Size (B) | 负责页面 |
|---|---|---|---|---|
| 1 | 23 | HostMemoryAllocator（chain head） | 32 | ext-remaining.md |
| 2 | 22 | MultiSlice | 64 | ext-remaining.md |
| 3 | 21 | Collectives | 96 | collectives-communicator.md |
| 4 | 20 | AbiVersion | 120 | ext-remaining.md |
| 5 | 19 | Shardings | 40 | ext-remaining.md |
| 6 | 18 | Megascale | 248 | collectives-communicator.md |
| 7 | 17 | TpuExecutable | 88 | executable-execution.md |
| 8 | 16 | TpuTopology | 272 | ext-topology-description.md |
| 9 | 14 | Callback | 40 | callbacks.md |
| 10 | 9 | PhaseCompile | 64 | ext-compile-phasecompile.md |
| 11 | 12 | CrossHostTransfers | 56 | dma-and-cross-host-recv.md |
| 12 | 15 | HostAllocator | 48 | ext-remaining.md |
| 13 | 13 | ExecutableMetadata | 40 | executable-execution.md |
| 14 | 6 | MemoryDescriptions | 40 | buffer-and-memory.md |
| 15 | 4 | Layouts | 80 | ext-remaining.md |
| 16 | 8 | RawBuffer | 80 | ext-rawbuffer.md |
| 17 | 1 | Profiler（.data, static） | 40 | ext-profiler.md |
Type ID 0, 2, 3, 5, 7, 10, 11 在此构建中未使用。公开 XLA 还注册了这里不存在的 FFI 和 Memory_Stream 扩展类型：libtpu 唯一的 custom-op 表面通过 TpuExecutable 扩展的 SetTpuCompilationEnv 汇入，而不是通过公开 FFI 扩展。
特性 — 链顺序与构建顺序相反，因此迭代会先得到最新扩展，最后得到 profiler。重新实现者必须通过遍历
.next直到 NULL 并匹配type来发现能力，绝不能按位置或按链中的固定偏移发现。构建顺序（RawBuffer → … → HostMemoryAllocator，profiler 作为 RawBuffer 的.next种入）是GetTpuPjrtApi的实现细节；只有 type tag 具有契约意义。
4. 五个 TPU 注入点
目的
CreatePjrtApi 只从其调用者接收六个函数指针性质的参数，并硬编码其余内容。其中五个成为 TPU 专用槽位；第六个是扩展链头。了解哪五个槽位由插件提供（相对于通用 pjrt::PJRT_* wrapper）能告诉重新实现者 TPU 后端究竟在哪些位置接入一个原本通用的 XLA 表。
注入槽位
已对照 CreatePjrtApi @ 0xf874160 确认：函数体是扁平 header-write 加 lea/mov slot-fill，没有循环，并且恰好五个槽位从 incoming register args 写入（a1[8]=a5、a1[9]=a7、a1[15]=a2、a1[87]=a4、a1[103]=a3）；slot 1（a1[1]=a6）是链头。
| Slot | Field | libtpu impl | Addr | 作用 |
|---|---|---|---|---|
| 8 | PJRT_Plugin_Initialize | tpu_plugin::PJRT_Plugin_Initialize | 0xE6A9D00 | 一次性 TPU driver bring-up（lifecycle） |
| 9 | PJRT_Plugin_Attributes | pjrt::PJRT_Plugin_Attributes_Xla | 0xF85F080 | Plugin attribute table：通用 XLA 实现，不是 TPU override |
| 15 | PJRT_Client_Create | tpu_plugin::PJRT_Client_Create | 0xE6A8840 | Silicon scan + live client construction |
| 87 | PJRT_TopologyDescription_Create | tpu_plugin::PJRT_TopologyDescription_Create | 0xE6A9B20 | TPU pod topology（AOT, no client） |
| 103 | PJRT_ExecuteContext_Create | tpu_plugin::PJRT_ExecuteContext_Create | 0xE6A9A80 | Per-execution context |
其余 130 个函数指针槽位是编译期固定的 pjrt::PJRT_* wrapper（在 CreatePjrtApi 中以 lea 加载的常量），与通用 XLA PJRT 层共享。调用最频繁的单个槽位是 slot 60 PJRT_LoadedExecutable_Execute @ 0xF869B40，即每步程序启动；它是通用 wrapper，最终落到 runtime 的 CommonPjRtLoadedExecutable::Execute，见 executable-execution.md 和 ../runtime/overview.md。
注释 — slot 9（
PJRT_Plugin_Attributes）虽然是注入参数，但并非 TPU 专用：CreatePjrtApi的调用者传入通用pjrt::PJRT_Plugin_Attributes_Xla。它公告标准 XLA attribute set（version metadata、supported devices、serialization info）。真正 TPU 专用的注入点只有 slots 8/15/87/103。
相关组件
| Component | Relationship |
|---|---|
GetPjrtApi / GetTpuPjrtApi | 导出入口符号及其惰性构建引擎 |
pjrt::CreatePjrtApi @ 0xf874160 | 将 140 槽表物化到 .lbss 中的构造器 |
pjrt::ActualStructSizeIsGreaterOrEqual @ 0xf8a4ec0 | 每个槽位开头使用的 per-call backward-compat size gate |
17-extension chain（0x224C3F68 head） | 通过 extension_start 到达的 typed、最新优先能力列表 |
xla::TpuClient / tpu::System | 通用槽位最终落入的 runtime（modern PJRT stack） |
Tpu*_* C-ABI（194 个导出） | 与该二进制共享但不通过 PJRT 到达的 legacy StreamExecutor 表面 |
交叉引用
- api-vtable-reconstruction.md — 完整 140 槽逐字段表（每个槽位 → impl symbol + address，hot-path ranking）
- extension-chain.md — 17 节点链表、逐节点布局和
PJRT_Extension_Base机制 - client-and-device.md —
PJRT_Client_*/PJRT_Device*槽位：create、lookup、device assignment、memory stats - buffer-and-memory.md —
PJRT_Buffer_*/PJRT_Memory_*/ AsyncHostToDevice transfer 槽位和 MemoryDescriptions 扩展 - executable-execution.md —
PJRT_Executable_*/PJRT_LoadedExecutable_*/PJRT_Compile/ ExecuteContext 以及 TpuExecutable + ExecutableMetadata 扩展 - events-and-async.md —
PJRT_Event_*、PJRT_Error_*和 AsyncTrackingEvent completion plumbing - callbacks.md — Callback 扩展（type 14）
- collectives-communicator.md — Collectives（type 21）和 Megascale（type 18）扩展
- dma-and-cross-host-recv.md — CopyToDeviceStream 槽位和 CrossHostTransfers 扩展（type 12）
- ext-profiler.md — Profiler 扩展（type 1），唯一的
.datastatic-init 链节点 - ext-topology-description.md —
PJRT_TopologyDescription_*槽位和 TpuTopology 扩展（type 16） - ext-rawbuffer.md — RawBuffer 扩展（type 8）
- ext-compile-phasecompile.md — PhaseCompile 扩展（type 9）
- ext-remaining.md — HostMemoryAllocator/HostAllocator、MultiSlice、AbiVersion、Shardings、Layouts 扩展
- stream-executor-host-interpreter.md — 通用槽位下方的 StreamExecutor host/interpreter shim
- stream-executor-pjrt-adapter.md — legacy StreamExecutor
Tpu*_*C-ABI 与 PJRT 表面的关系 - ../runtime/overview.md — hot-path 槽位（尤其 slot 60
Execute）dispatch 进入的 runtime/execution 层 - ../lifecycle/get-pjrt-api-thunk.md —
GetPjrtApithunk 和惰性表构建入口 - ../lifecycle/module-init-plugin-discovery.md —
PJRT_Plugin_Initializedriver bring-up 和 slots 8/15 背后的 module-init DAG