GetPjrtApi Thunk 与 tpu_plugin 对象

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，ELF x86-64 DYN，未 stripped；反修饰后的 C++ 符号按原文引用）。其他版本会有所不同。

摘要

PJRT plugin 是一个 .so，职责只有一个：把指向 PJRT_Api 函数指针表的指针交给框架。 libtpu.so 通过一个值得拆开的双符号安排来做这件事，因为二进制采用的并不是显而易见的单函数设计。公开、导出且带版本的符号 GetPjrtApi @ 0xe6a83a0 是一个五字节 jmp thunk，本身完全不带逻辑。它 tail-call 一个内部的、未导出的引擎 pjrt::tpu_plugin::GetTpuPjrtApi @ 0xe6aa440，后者才是真正的 lazy builder：一段直线执行的十七个 Itanium-ABI __cxa_guard 保护的一次性块，负责物化 extension chain 和表，然后返回指向函数局部静态对象的指针。这个静态对象 GetTpuPjrtApi()::pjrt_api @ 0x227BA840 是 .lbss (NOBITS) 中一块 1120 字节的零填充区域；它是进程中唯一的 PJRT_Api 实例，并且不作为镜像存在于任何 PROGBITS section 中。静止状态下表本身没有任何可反汇编的内容；它的 140 个槽在首次调用时由 pjrt::CreatePjrtApi @ 0xf874160 写入。

这是把 C++ Meyers-singleton 模式应用到 C ABI：导出的 C trampoline 位于 C++ 函数局部静态对象之前，在 guard 下惰性构建，并在进程退出时泄漏。拆分有两个原因。第一，公开名称必须是规范的 lowercase-jrtGetPjrtApi（与所有其他 PJRT plugin 匹配），而内部 builder 保留描述性的 GetTpuPjrtApi 名称，并且不进入动态符号表；这样框架不会意外绑定到 builder，builder 的 TPU-specific 身份也保持私有。第二，thunk 让链接器可以把导出入口放在一个极小的桩 section 中，而 multi-kilobyte builder 可以和其余 pjrt::tpu_plugin 代码放在一起。builder 绑定到表中五个 TPU 槽的函数属于该 pjrt::tpu_plugin namespace：PJRT_Client_Create、PJRT_Plugin_Initialize、 PJRT_TopologyDescription_Create、PJRT_ExecuteContext_Create；它们是通用 XLA PJRT 表接触 TPU silicon 的唯一位置。

本页负责导出 thunk、GetTpuPjrtApi lazy-init builder（其 guard 结构、返回路径以及最终的 CreatePjrtApi 调用）、pjrt_api singleton（存储、生命周期、为什么静止时不可见），以及槽绑定到的 tpu_plugin namespace object。它不复现 140-slot field table （../pjrt/api-vtable-reconstruction.md）、16-node __cxa_guard extension 构建顺序（../pjrt/extension-chain.md），或 dlopen-time static-init 图景和 PJRT_Plugin_Initialize driver bring-up （module-init-plugin-discovery.md）。

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

Thunk 形状 — 一个导出的 GetPjrtApi@@VERS_1.0，主体为 jmp <builder>，无 prologue，无参数。
Lazy-init guard 结构 — GetTpuPjrtApi 中有 17 个 __cxa_guard 块；16 个构建 extension node，第 17 个调用 CreatePjrtApi；无论本次调用实际运行了哪些 guard，函数都返回 &pjrt_api。
Singleton — 一个 .lbss 函数局部静态对象，1120 字节，加载时为零，写入一次，随后不可变，永不释放。
tpu_plugin 对象 — CreatePjrtApi 作为参数接收的四个 TPU-specialized slot implementation 和一个 generic-XLA attributes implementation。


导出入口符号	`GetPjrtApi @ 0xe6a83a0`（5 字节 `jmp` thunk，`GetPjrtApi@@VERS_1.0`）
Thunk 目标 / builder	`pjrt::tpu_plugin::GetTpuPjrtApi @ 0xe6aa440`（1336 B，未导出）
表构造器	`pjrt::CreatePjrtApi @ 0xf874160`（1872 B，仅第 17 个 guard）
Singleton 存储	`GetTpuPjrtApi()::pjrt_api @ 0x227BA840`，`.lbss` (NOBITS)，1120 B = 140 x 8
Guard 数量	17 个 `__cxa_guard` 块（16 个 extension builder + 1 个 `CreatePjrtApi`）
Guard 实现	libtpu 自己的 libc++abi `__cxa_guard_acquire/release/abort @ 0x213e9ac0 / 0x213e9be0 / 0x213e9c20`
传给 slot 1 的 chain head	`&host_memory_allocator_extension @ 0x224c3f68`
TPU 注入槽	8, 9, 15, 87, 103（Plugin_Initialize, Plugin_Attributes, Client_Create, TopologyDescription_Create, ExecuteContext_Create）
C-API 版本	v0.103 — version qword `0x6700000000` -> `{major=0, minor=0x67=103}`
置信度	CONFIRMED（以字节锚定，对照 decompile），除非某行或 callout 另有说明

导出 Thunk — `GetPjrtApi`

目的

GetPjrtApi 是每个 XLA front-end 用 dlsym 解析的单一 rendezvous 符号。它存在的目的，是给这个名称一个稳定、带版本、导出的地址，同时把所有实际逻辑保留在内部函数中。该 thunk 本身没有行为；它只做转发。

入口点

text

framework (JAX / TF / PyTorch-XLA)
  dlsym(handle, "GetPjrtApi")           ── the ONLY exported name matching /Pjrt/
    │
    └─ GetPjrtApi  0xe6a83a0  (5 bytes)  ── e9 9b 20 00 00  jmp 0xe6aa440
         └─ pjrt::tpu_plugin::GetTpuPjrtApi  0xe6aa440  (the real builder)
```text

### 算法

该 thunk 是一个 tail-call。IDA 将它渲染为单行；原始字节是一个相对 `jmp`。

```c
// GetPjrtApi @ 0xe6a83a0  (5 bytes, attribute: thunk)
//   machine code:  e9 9b 20 00 00   jmp 0xe6aa440
const PJRT_Api* GetPjrtApi(void) {
    return pjrt::tpu_plugin::GetTpuPjrtApi();   // tail call — no own frame
}

因为它是 jmp 而不是 call，GetPjrtApi 不消耗自己的 stack frame，也不留下自己的 return address：builder 的 ret 直接返回到框架。重新实现也可以直接在公开名称下发出 builder；该 thunk 是 linker/visibility 便利设施，而非语义要求。

函数图

函数	地址	大小	作用
`GetPjrtApi`	`0xe6a83a0`	5 B	导出的 `jmp` thunk -> builder
`pjrt::tpu_plugin::GetTpuPjrtApi`	`0xe6aa440`	1336 B	它转发到的 builder（内部）

GOTCHA — 拼写和大小写是 ABI 的一部分。导出符号是 GetPjrtApi（lowercase jrt）；GetTpuPjrtApi 是内部符号，且不在动态符号表中。对 GetTpuPjrtApi 执行 dlsym 的加载器，或只导出 Tpu 大小写名称的构建，都会发现失败。194 个确实导出的 Tpu*_* 符号（FUNC GLOBAL，全部 @@VERS_1.0）是 legacy StreamExecutor C-ABI，从不经由 PJRT 到达；见 tftpu-initialize-bootstrap.md。
NOTE — GetPjrtApi 是唯一匹配 /Pjrt/ 的 GLOBAL FUNC 导出，版本为 GetPjrtApi@@VERS_1.0。签名是规范 PJRT plugin 入口：const PJRT_Api* GetPjrtApi(void)，无参数，返回表指针。

Lazy Builder — `GetTpuPjrtApi`

目的

GetTpuPjrtApi 是整个 PJRT surface 的构造引擎。首次调用时，它构建 16 个驻留在 .bss 的 extension node，并写入 140-slot 表；之后每次调用都会跳过全部 17 个已满足的 guard 并返回同一个指针。它是 pjrt_api 的函数局部静态生命周期拥有者。

入口点

text

GetTpuPjrtApi  0xe6aa440
  ├─ guard 1..16:  CreateXxxExtension(&node, &prev_node[, tpu_fns])  ── 16 .bss nodes
  │                  (RawBuffer → … → HostMemoryAllocator; build order owned by
  │                   ../pjrt/extension-chain.md)
  └─ guard 17:     CreatePjrtApi(&pjrt_api, Client_Create, ExecuteContext_Create,
                                 TopologyDescription_Create, Plugin_Initialize,
                                 &host_memory_allocator_extension, Plugin_Attributes_Xla)
                     └─ writes all 140 slots into 0x227BA840
       return &pjrt_api  =  0x227BA840   (.lbss)
```text

### 算法

函数体由十七个相同的 `__cxa_guard` gated block 组成，随后是单个 `return`。每个 block 测试 guard byte，若未设置则 acquire，运行一次性工作，然后 release。decompile 逐字确认了这一点；下面压缩为第一个 builder、结构不变量和最终 `CreatePjrtApi` block。

```c
// pjrt::tpu_plugin::GetTpuPjrtApi @ 0xe6aa440  (1336 B)
// Each of the 17 blocks is: if (!guard_byte && __cxa_guard_acquire(&guard)) { work; __cxa_guard_release(&guard); }
function GetTpuPjrtApi():

    // --- guards 1..16: build the 16 .bss extension nodes (construction order) ---
    once(raw_buffer_extension):                          // built FIRST
        CreateRawBufferExtension(&raw_buffer_extension,
                                 &profiler_extension);    // next = .data Profiler seed (type 1)
    once(layouts_extension):
        CreateLayoutsExtension(&layouts_extension, &raw_buffer_extension);
    // … 13 more builders, each next = previously-built node …
    once(host_memory_allocator_extension):               // built LAST → becomes chain head
        CreateHostMemoryAllocatorExtension(&host_memory_allocator_extension,
                                           &multi_slice_extension);

    // --- guard 17: materialize the 140-slot table ---
    once(pjrt_api):
        CreatePjrtApi(&pjrt_api,                          // 0x227BA840 (.lbss)
                      PJRT_Client_Create,                 // a2 → slot 15
                      PJRT_ExecuteContext_Create,         // a3 → slot 103
                      PJRT_TopologyDescription_Create,    // a4 → slot 87
                      PJRT_Plugin_Initialize,             // a5 → slot 8
                      &host_memory_allocator_extension,   // a6 → slot 1 (extension_start)
                      PJRT_Plugin_Attributes_Xla);        // a7 → slot 9

    return &pjrt_api;                                     // 0x227BA840 — unconditional

16 个 extension builder 及其精确构造顺序、next 链接和 creator 地址归 ../pjrt/extension-chain.md 所有；本页只复现第一个和最后一个，以固定形状以及 seed/head 关系。这里要说明的是控制结构：返回是无条件的，不依赖本次调用运行了哪些 guard。首次调用时全部 17 个 block 都会执行；之后每次调用中，guard-byte 快路径（!(_BYTE)guard 为 false）会跳过每个 block，函数直接落到 return &pjrt_api。

QUIRK — 表构造器 CreatePjrtApi 由与 16 个 extension node 同类的 guard 控制；它是第 17 个 __cxa_guard block，与 pjrt_api guard byte 共享，而不是独立机制。因此“构建 extensions”和“构建表”是一个函数里的十七个同级项，不是两个阶段。如果重新实现者急切构建表（在 dlopen 时或在 guard 外），就会丢失框架依赖的 lazy-on-first-call 和 concurrent-serialization 语义。
GOTCHA — guard 变量使用的是 libtpu 自己的 libc++abi __cxa_guard_acquire/release/abort @ 0x213e9ac0 / 0x213e9be0 / 0x213e9c20，不是 glibc 的。它们是 .bss 中 17 个不同的 guard byte（例如 raw_buffer_extension guard @ 0x224c39e0，其 node static 位于 0x224c3990）。并发 first-caller 通过 Itanium-ABI guard 语义串行化：失败者阻塞直到胜者 release，然后看到已满足 byte 并跳过。一次性过程之后，这些 byte 在进程生命周期内保持 set，reader 不再取锁。

函数图

函数	地址	大小	作用
`pjrt::tpu_plugin::GetTpuPjrtApi`	`0xe6aa440`	1336 B	17-guard lazy builder
`pjrt::CreatePjrtApi`	`0xf874160`	1872 B	第 17 个 guard 中的 slot-fill constructor
`__cxa_guard_acquire`	`0x213e9ac0`	—	libtpu 自己的一次性 acquire
`__cxa_guard_release`	`0x213e9be0`	—	libtpu 自己的一次性 release
`__cxa_guard_abort`	`0x213e9c20`	—	libtpu 自己的一次性 abort

注意事项

builder 只运行 PJRT-table 和 extension-chain 构造。它不初始化 TPU driver、不扫描 silicon，也不运行 GoogleInitializer module DAG；那些发生在之后，即框架调用 PJRT_Plugin_Initialize（slot 8）和 PJRT_Client_Create（slot 15）时。GetTpuPjrtApi 返回时，表已经填充，但硬件尚未触碰。这个分离让框架可以在承诺 driver bring-up 之前检查 struct_size、版本和 extension chain（纯 metadata）。deep init path 归 module-init-plugin-discovery.md 所有。

`pjrt_api` Singleton

目的

pjrt_api 是进程 ever produces 的唯一 PJRT_Api 实例。它是 GetTpuPjrtApi 的函数局部静态对象，这就是为什么其构造受 guard 保护、生命周期覆盖整个进程。理解它的存储对重新实现者至关重要：二进制中没有可复制的静态表镜像，因此 slot 值必须从 CreatePjrtApi 的函数体重建，而不是从 section dump 读取。

存储与布局

text

.lbss (section [47], NOBITS — no file bytes)
  0x227BA840 ┌─────────────────────────────────────────────┐
             │ GetTpuPjrtApi()::pjrt_api  (1120 B = 140×8)  │
             │   slot 0  +0x000  struct_size       = 1120   │  written by CreatePjrtApi
             │   slot 1  +0x008  extension_start    → chain  │  (guard 17), zero at load
             │   slots 2..4      pjrt_api_version   {0,103}  │
             │   slots 5..139    135 fn-ptrs                 │
             └─────────────────────────────────────────────┘
```text

`.lbss` 是 *large* BSS，即 NOBITS，不占用文件空间，由加载器零填充。在 `dlopen` 时，整个 1120 字节 slab 为零。直到第一次 `GetPjrtApi` 调用运行 guard 17 且 `CreatePjrtApi` 写入每个槽之前，它都保持为零。header 写入已按字节对照 `CreatePjrtApi @ 0xf874160` 确认：`*a1 = 1120; a1[1] = chain_head; a1[2] = 24; a1[3] = 0; a1[4] = 0x6700000000`。

| 字段 | 值 | 来源 |
|---|---|---|
| Storage VA | `0x227BA840` | `.lbss` static，builder 中的 `mov rbx, offset` |
| Section | `.lbss` [47], NOBITS | section table |
| 大小 | 1120 B (= 140 x 8) | `*a1 = 1120` in `CreatePjrtApi` |
| 加载时的值 | 全零 | NOBITS，零填充 |
| 写入者 | `CreatePjrtApi @ 0xf874160`（guard 17） | builder call site |
| 生命周期 | 进程（退出时泄漏） | 函数局部静态对象，无 atexit dtor |

> **NOTE —** 逐槽重建，即 135 个 fn-ptr 分别映射到哪个 `pjrt::PJRT_*` wrapper，见 [../pjrt/api-vtable-reconstruction.md](../pjrt/api-vtable-reconstruction.md)。本页固定的是*表在哪里以及如何出现*：一个零 `.lbss` slab，在 guard 17 下写入一次。静态分析者如果 grep `PROGBITS` 中已填充的 `PJRT_Api`，会找不到；该表是运行时 artifact。
>
> **QUIRK —** singleton 永不释放。没有为 `pjrt_api` 或 16 个 `.bss` extension node 注册 `atexit`/`__cxa_thread_atexit`；它们是 plugin `.so` 的正常 leaked Meyers-singleton lifetime。进程的 `FINI_ARRAY` 只 teardown 一个 trivial `__do_fini` 桩和 BoringSSL per-thread RNG state，而不是 PJRT 表。Client、executable 和 buffer 通过其显式 `PJRT_*_Destroy` C-API 调用 teardown，而不是通过表的生命周期。

---

## `tpu_plugin` 对象

### 目的

`pjrt::tpu_plugin` 是拥有通用 XLA PJRT 表的 TPU specialization 的 namespace。在 140 个槽中，`CreatePjrtApi` 将 130 个硬编码为编译期固定的 `pjrt::PJRT_*` wrapper（与 generic XLA layer 共享），只从调用者接收**五个**函数指针。这五个中的四个是 `pjrt::tpu_plugin` 成员，即真正的 TPU hook；第五个是 generic XLA attributes implementation。知道哪五个槽由调用者提供，就能准确告诉重新实现者 TPU backend 在哪里挂接到通用表上。

### 注入槽

`CreatePjrtApi` 的 argument-to-slot 映射已按字节确认：在 `GetTpuPjrtApi` 中，调用传入
`(Client_Create, ExecuteContext_Create, TopologyDescription_Create, Plugin_Initialize, &host_memory_allocator_extension, Plugin_Attributes_Xla)`；
在 `CreatePjrtApi` 中它们落为 `a1[15]=a2`、`a1[103]=a3`、`a1[87]=a4`、`a1[8]=a5`、`a1[1]=a6`、`a1[9]=a7`。
`CreatePjrtApi` 的 mangled symbol 编码了这个精确参数顺序（`Client_Create_Args`、`ExecuteContext_Create_Args`、
`TopologyDescription_Create_Args`、`Plugin_Initialize_Args`、`PJRT_Extension_Base*`、`Plugin_Attributes_Args`），
从而佐证 register-to-slot trace。

| Slot | Field | Implementation | Address | Namespace |
|---|---|---|---|---|
| 8 | `PJRT_Plugin_Initialize` | `tpu_plugin::PJRT_Plugin_Initialize` | `0xe6a9d00` | `pjrt::tpu_plugin` (TPU) |
| 9 | `PJRT_Plugin_Attributes` | `pjrt::PJRT_Plugin_Attributes_Xla` | `0xf85f080` | `pjrt` (generic XLA) |
| 15 | `PJRT_Client_Create` | `tpu_plugin::PJRT_Client_Create` | `0xe6a8840` | `pjrt::tpu_plugin` (TPU) |
| 87 | `PJRT_TopologyDescription_Create` | `tpu_plugin::PJRT_TopologyDescription_Create` | `0xe6a9b20` | `pjrt::tpu_plugin` (TPU) |
| 103 | `PJRT_ExecuteContext_Create` | `tpu_plugin::PJRT_ExecuteContext_Create` | `0xe6a9a80` | `pjrt::tpu_plugin` (TPU) |

Slot 1（`extension_start`）也是 caller-supplied，即参数 `a6`、chain head `&host_memory_allocator_extension @ 0x224c3f68`，但它是数据指针，不是 `tpu_plugin` 函数。其余 130 个函数指针槽是 `CreatePjrtApi` 函数体中的 `lea`-loaded constants：编译期固定的 `pjrt::PJRT_*` wrapper（例如 `a1[5]=PJRT_Error_Destroy`、`a1[16]=PJRT_Client_Destroy`），与 generic XLA PJRT layer 共享，而非 TPU-specialized。

> **NOTE —** slot 9 是刻意的例外。虽然它是一个*注入*参数（`a7`），但 `CreatePjrtApi` 的调用者传入的值是 `pjrt::PJRT_Plugin_Attributes_Xla @ 0xf85f080`，也就是 generic XLA attributes implementation（它通告 `xla_version`、`supported_devices`、serialization metadata），**不是** `tpu_plugin` override。因此五个 caller-supplied slot 中，只有四个（8/15/87/103）是真正 TPU-specific。重新实现者必须注入 TPU 的 `Client_Create`/`Plugin_Initialize`/`TopologyDescription_Create`/`ExecuteContext_Create`，但可以复用 stock XLA attributes 函数。

### 函数图

| 函数 | 地址 | 作用 |
|---|---|---|
| `pjrt::tpu_plugin::PJRT_Client_Create` | `0xe6a8840` | Silicon scan + live client construction（slot 15） |
| `pjrt::tpu_plugin::PJRT_Plugin_Initialize` | `0xe6a9d00` | One-time TPU driver bring-up（slot 8） |
| `pjrt::tpu_plugin::PJRT_TopologyDescription_Create` | `0xe6a9b20` | AOT pod topology，无 client（slot 87） |
| `pjrt::tpu_plugin::PJRT_ExecuteContext_Create` | `0xe6a9a80` | Per-execution context（slot 103） |
| `pjrt::PJRT_Plugin_Attributes_Xla` | `0xf85f080` | Generic XLA attribute table（slot 9） |

### 注意事项

四个 TPU 槽在这里被绑定，但不在这里*运行*。`PJRT_Plugin_Initialize` 是发现之后框架的第一个 deep call：它获取跨进程 TPU lock、读取 `LIBTPU_INIT_ARGS`，并运行 GoogleInitializer module DAG，后者注册 per-`TpuVersion` HAL factory。`PJRT_Client_Create` 是实际发生 silicon scan，并在 StreamExecutor `TpuPlatform` 之上构建 live `xla::PjRtClient` 的地方。二者均归 [module-init-plugin-discovery.md](module-init-plugin-discovery.md) 所有；本页只关心 `GetTpuPjrtApi` 通过 `CreatePjrtApi` 把它们的地址接到 slot 8 和 15 这一事实。

---

## 相关组件

| 组件 | 关系 |
|---|---|
| `GetPjrtApi @ 0xe6a83a0` | 导出的 `jmp` thunk；唯一 PJRT 入口符号 |
| `pjrt::tpu_plugin::GetTpuPjrtApi @ 0xe6aa440` | 它转发到的 17-guard lazy builder |
| `pjrt::CreatePjrtApi @ 0xf874160` | 填充 140 个槽的第 17 个 guard constructor |
| `GetTpuPjrtApi()::pjrt_api @ 0x227BA840` | builder 返回的 `.lbss` singleton |
| `pjrt::tpu_plugin` object（5 个注入槽） | 接入 slot 8/9/15/87/103 的 TPU specialization |
| 16-node `__cxa_guard` extension chain | 同一个 builder 中的 guard 1..16；归 `../pjrt/extension-chain.md` 所有 |
| `Tpu*_*` C-ABI（194 个导出） | 共享该二进制但不通过 PJRT 到达的 legacy StreamExecutor surface |

## 交叉引用

- [overview.md](overview.md) — 本页所在的生命周期 section map
- [module-init-plugin-discovery.md](module-init-plugin-discovery.md) — dlopen-time static-init 图景，以及这些槽会进入的 `PJRT_Plugin_Initialize` -> `Client_Create` driver bring-up
- [tftpu-initialize-bootstrap.md](tftpu-initialize-bootstrap.md) — 共享该二进制的 legacy `Tpu*_*` StreamExecutor C-ABI bootstrap
- [../pjrt/overview.md](../pjrt/overview.md) — PJRT C-ABI map：handshake、struct shape，以及此 builder 返回的表的 region index
- [../pjrt/api-vtable-reconstruction.md](../pjrt/api-vtable-reconstruction.md) — `pjrt_api` singleton 的完整 140-slot field-by-field reconstruction
- [../pjrt/extension-chain.md](../pjrt/extension-chain.md) — 16 个 `__cxa_guard` extension builder（guard 1..16）及其组装的 newest-first chain

GetPjrtApi Thunk 与 tpu_plugin 对象 ​

摘要 ​

导出 Thunk — GetPjrtApi ​

目的 ​

入口点 ​

函数图 ​

Lazy Builder — GetTpuPjrtApi ​

目的 ​

入口点 ​

函数图 ​

注意事项 ​

pjrt_api Singleton ​

目的 ​

存储与布局 ​

GetPjrtApi Thunk 与 tpu_plugin 对象

摘要

导出 Thunk — `GetPjrtApi`

目的

入口点

函数图

Lazy Builder — `GetTpuPjrtApi`

目的

入口点

函数图

注意事项

`pjrt_api` Singleton

目的

存储与布局