PJRT_Api 函数指针表重建
本页所有地址均适用于 libtpu 0.0.40 (
cp314,manylinux_2_31_x86_64),build-id89edbbe81c5b328a958fe628a9f2207d。PJRT C-API 表面为 v0.103。其他 libtpu 版本会固定到不同的 minor,并会重新编号后期新增的槽位。
摘要
libtpu.so 是一个 PJRT 插件:它导出一个符号 GetPjrtApi,宿主框架(JAX、PyTorch/XLA)实际调用的所有内容,都会通过该符号返回的单个 PJRT_Api 结构体到达 TPU。这个结构体是一个 C ABI vtable,即由 140 个 qword 槽位组成的扁平数组:前五个标量头字段,后面跟着 135 个函数指针,每个 PJRT_* C-API 入口点对应一个。宿主读取 api->PJRT_LoadedExecutable_Execute(args),调用就落到 libtpu 的 text section 中。这里没有 C++ 继承,没有分发对象,也没有按版本区分的 trampoline table,只有这一个结构体:一次填充,进程生命周期内不可变。
本页是 PJRT 章节的骨干参考:对该表按完整顺序、逐槽位进行重建。它负责结构体头部解码(struct_size、extension_start、嵌入的 PJRT_Api_Version)、完整的 140 槽位映射(每个槽位的字段名、它指向的 libtpu 实现符号,以及该符号的虚拟地址),以及“已填充与已注入”映射(哪些五个槽位是 TPU 专用的,哪些 130 个是标准 XLA wrapper)。槽位顺序与公开的 xla/pjrt/c/pjrt_c_api.h v0.103 完全一致,包括按特性新增顺序追加到原始 v0.40 表面之后的后期新增集群。
vtable 由一个函数组装:pjrt::CreatePjrtApi @ 0xf874160。该函数只做一件事:向调用方提供的缓冲区执行 140 次 store。因此重建异常清晰:槽位索引就是初始化器中的数组下标,字段名就是被赋值的 pjrt::PJRT_* 符号。下面所有内容都锚定到这些 store。各区域行为(每个实现实际如何工作)放在底部链接的区域页面中;挂在 extension_start 上的扩展链表见 extension-chain.md。本页是它们共同依附的索引。
对重新实现而言,契约是:
- 结构体布局:1120 字节 = 140 × 8;五个标量头槽位(0..4),然后是 135 个函数指针(5..139),顺序严格如下。
- 版本固定:
PJRT_Api_Version = {struct_size=24, priv=NULL, major=0, minor=103},在偏移+0x20处编码为 qword0x6700000000。 - 槽位到符号映射:每个已填充槽位都解析到一个 text-section 符号;该构建中没有槽位保持为 null。
- 五个注入点:唯一不是编译期固定的槽位;由调用方作为参数传给
CreatePjrtApi。 - 逐槽位向后兼容保护:每个 wrapper 的第一步都是
ActualStructSizeIsGreaterOrEqual,这是让 older-minor 宿主调用这个 v0.103 插件的机制。
| 导出入口 | GetPjrtApi @ 0xE6A83A0 (thunk → GetTpuPjrtApi) |
| 构建器 | pjrt::tpu_plugin::GetTpuPjrtApi @ 0xE6AA440 |
| 初始化器 | pjrt::CreatePjrtApi @ 0xF874160(140 次 store,无逻辑) |
| 存储位置 | _ZZN4pjrt10tpu_plugin13GetTpuPjrtApiEvE8pjrt_api @ 0x227BA840 |
| 节区 | [47] .lbss(NOBITS,加载时零填充,large-model) |
| 大小 | 1120 字节 = 140 个 qword 槽位 |
| C-API 版本 | v0.103 (major=0, minor=103) |
| 已填充槽位 | 135 / 135 个函数指针;0 个 null |
| TPU 专用槽位 | 5 (slots 8, 9, 15, 87, 103) |
| 向后兼容保护 | pjrt::ActualStructSizeIsGreaterOrEqual @ 0xF8A4EC0 |
各区域槽位数
| 区域 | 槽位数 | 槽位范围 |
|---|---|---|
| 头部(标量) | 5 | 0..4 |
| Error | 3(+1 个后期新增) | 5..7, 137 |
| Plugin | 2 | 8..9 |
| Event | 5(+2 个后期新增) | 10..14, 131..132 |
| Client | 13(+约 12 个后期新增) | 15..27, 98, 100, 108, 115..118, 120..121, 123, 134 |
| DeviceDescription | 6 | 28..33 |
| Device | 6(+3 个后期新增) | 34..39, 126..127, 133 |
| Memory | 5(+1 个后期新增) | 40..44, 102 |
| Executable | 10(+6 个后期新增) | 45..54, 95..96, 99, 101, 129, 139 |
| LoadedExecutable | 8(+2 个后期新增) | 55..62, 122, 135 |
| Buffer | 19(+5 个后期新增) | 63..81, 97, 105, 125, 130, 136 |
| CopyToDeviceStream | 5 | 82..86 |
| TopologyDescription | 7(+2 个后期新增) | 87..93, 119, 138 |
| Compile | 1 | 94 |
| ExecuteContext | 1(+1 个后期新增) | 103..104 |
| AsyncHostToDeviceTransferManager | 9(后期新增) | 106..107, 109..114, 124 |
| AsyncTrackingEvent | 1(后期新增) | 128 |
填充路径与存储
目的
该表不是 PROGBITS section 中的静态初始化器,而是在第一次 GetPjrtApi 调用时,惰性地、仅一次地构建到零填充的 .lbss 中。理解填充路径,是返回正确表的重新实现与发生竞态或返回垃圾数据的重新实现之间的区别。
入口点
dlsym("GetPjrtApi")
GetPjrtApi @ 0xE6A83A0 ── thunk; tail-calls the builder
└─ pjrt::tpu_plugin::GetTpuPjrtApi @ 0xE6AA440
── 16 __cxa_guard blocks build the .bss extensions (newest→oldest)
── final __cxa_guard wraps:
└─ pjrt::CreatePjrtApi @ 0xF874160 ── 140 stores into &pjrt_api
returns &pjrt_api = 0x227BA840
```text
导出符号是一个纯 thunk。它的完整反编译体是 `return pjrt::tpu_plugin::GetTpuPjrtApi(a1);`(在 `0xE6A83A0` 确认,被 IDA 标记为 `thunk`)。`GetTpuPjrtApi` 是编排器:它为每个扩展运行一个受 Itanium-ABI `__cxa_guard` 保护的构造器,以便在 `.bss` 中构建扩展链;随后在最后一个 guard 中调用 `CreatePjrtApi`,写入每个槽位。该函数返回 `&pjrt_api`,也就是一个函数局部静态对象的地址。
### 算法
```c
function GetTpuPjrtApi(): // 0xE6AA440
// one-shot init of the extension chain (covered on extension-chain.md);
// each extension is a __cxa_guard-protected function-local static in .bss
build_extensions_if_needed(); // 16 guards, newest→oldest
if (guard_for(pjrt_api) not yet set): // _ZGV... byte == 0
CreatePjrtApi(&pjrt_api, // 0xF874160 — writes all 140 slots
Client_Create_impl, // -> slot 15 (a2)
ExecuteContext_Create_impl, // -> slot 103 (a3)
TopologyDescription_Create_impl, // -> slot 87 (a4)
Plugin_Initialize_impl, // -> slot 8 (a5)
extension_start, // -> slot 1 (a6)
Plugin_Attributes_Xla_impl); // -> slot 9 (a7)
mark_guard_set(pjrt_api);
return &pjrt_api; // 0x227BA840function CreatePjrtApi(a1, a2, a3, a4, a5, a6, a7): // 0xF874160
*a1 = 1120; // slot 0 struct_size
a1[1] = a6; // slot 1 extension_start (injected: chain head)
a1[2] = 24; // slot 2 pjrt_api_version.struct_size
a1[3] = 0; // slot 3 pjrt_api_version.priv
a1[4] = 0x6700000000LL; // slot 4 {major=0, minor=103}
a1[5] = pjrt::PJRT_Error_Destroy; // slot 5 ... 130 fixed wrappers ...
a1[8] = a5; // slot 8 Plugin_Initialize (injected)
a1[9] = a7; // slot 9 Plugin_Attributes (injected)
a1[15] = a2; // slot 15 Client_Create (injected)
a1[87] = a4; // slot 87 TopologyDescription_Create (injected)
a1[103] = a3; // slot 103 ExecuteContext_Create (injected)
// ... all other slots assigned a compile-time-fixed pjrt::PJRT_* symbol ...
a1[139] = pjrt::PJRT_Executable_ParameterMemoryKinds; // slot 139
return a1;
```text
> **注意 —** `CreatePjrtApi` 没有分支,也没有循环:它是 140 次直线 store。槽位索引就是 C 数组下标,这也是为什么下面的重建可以是精确的,而不是推断出来的。五个注入槽位(`a2`..`a7`)是该函数唯一不硬编码的值;其他所有值都是链接时烘入的 `pjrt::PJRT_*` relocation。这一点通过完整读取 `0xF874160` 处的初始化器得到确认。
### 注意事项
该结构体位于 `.lbss`(section `[47]`,NOBITS,large-model BSS)中,地址为函数局部静态对象 `_ZZN4pjrt10tpu_plugin13GetTpuPjrtApiEvE8pjrt_api @ 0x227BA840`。这对重新实现者以及任何静态检查二进制的工具有两个后果:
- **静态反汇编无法显示已填充的表。** 这 1120 字节在磁盘上为零,加载时也被零填充;只有第一次 `GetPjrtApi` 调用运行 `CreatePjrtApi` 后,函数指针才存在。下面的槽位映射是从*初始化器中的 store*重建出来的,而不是来自数据 dump:运行前没有东西可 dump。运行时确认需要在第一次调用后,在一个存活的 JAX worker 内读取 `0x227BA840` 处的 `/proc/self/mem`。
- **并发由 `__cxa_guard` 处理,而不是由读取者持有锁。** 首批调用者通过 Itanium ABI guard 串行化(一个线程运行构造器,其余线程在其 mutex 上阻塞)。第一次调用后结构体不可变,所以稳态读取者不加锁。每次调用都重建表,或省略一次性 guard 的重新实现,都会偏离这个契约。
> **注意 —** 该结构体位于具体 VA `0x227BA840` 处的 `.lbss` 中,而不是 `.bss`/`.data.rel.ro`;反汇编器中的任何 `&stru_… + _GLOBAL_OFFSET_TABLE_` 形式都是 PIC relocation artifact,不是真实地址。槽位 9 *不是* TPU-plugin override;它来自 `CreatePjrtApi` 的 `a7` 参数,即 `pjrt::PJRT_Plugin_Attributes_Xla`,是标准 XLA 实现。
---
## 结构体头部(槽位 0..4)
### 目的
前五个 qword 不是函数指针。它们是 ABI 握手:总大小(用于向前/向后兼容)、扩展链头,以及一个嵌入的 `PJRT_Api_Version` 子结构体。宿主会在触碰任何函数指针之前读取这些字段。
### 布局
| 槽位 | 偏移 | 字段 | 类型 | 值 | 来源 |
|---|---|---|---|---|---|
| 0 | `+0x000` | `struct_size` | `size_t` | `1120` | `*a1 = 1120` |
| 1 | `+0x008` | `extension_start` | `PJRT_Extension_Base*` | `0x224C3F68`(链头) | `a1[1] = a6`(注入) |
| 2 | `+0x010` | `pjrt_api_version.struct_size` | `size_t` | `24` | `a1[2] = 24` |
| 3 | `+0x018` | `pjrt_api_version.priv` | `void*` | `NULL` | `a1[3] = 0` |
| 4 | `+0x020` | `pjrt_api_version.{major,minor}` | `int,int` | `{0, 103}` | `a1[4] = 0x6700000000` |
```c
struct PJRT_Api_Version { // embedded at PJRT_Api +0x010, 24 bytes
size_t struct_size; // = 24
void* priv; // = NULL
int major_version; // = 0
int minor_version; // = 103 (0x67)
};版本 qword 按 little-endian 字节布局解码:+0x20 处的字节是 00 00 00 00 67 00 00 00,即 major = 低 32 位 = 0,minor = 高 32 位 = 0x67 = 103。struct_size = 1120 让针对不同 minor 编译的宿主精确知道它可以安全读取多少个尾部槽位。
怪癖 —
extension_start(槽位 1)指向的是最新扩展,而不是最旧扩展。该链按__cxa_guard顺序构建并以头插法链接,因此沿.next遍历会从最新走到最旧,并在.data中的 profiler extension 处终止。假设“type 0 first”或任何类型顺序的宿主都是错的;它必须一直迭代到 NULL。链本身记录在extension-chain.md。
140 槽位映射
目的
完整的有序表。槽位 0..4 是上面的标量头部。槽位 5..139 是函数指针;下面每个单元格都列出字段(按 pjrt_c_api.h v0.103)以及该槽位指向的 libtpu 符号,并给出该符号的虚拟地址。该构建中所有 135 个函数指针槽位都已填充,没有 null/未实现槽位。 对每个槽位,如果它的 store 和目标符号都已在反编译初始化器中读取并与符号表确认,则置信度为 CERTAIN;12 个抽查确认的地址已标记。
槽位顺序说明
到槽位 94 为止,顺序按区域分组(Error/Plugin/Event/Client/DeviceDescription/Device/Memory/Executable/LoadedExecutable/Buffer/CopyStream/Topology/Compile);随后是一个扁平的后期新增集群(槽位 95..139),按 minor 0.40→0.103 中的特性新增顺序追加。后期集群没有按区域重新分组:Buffer_*、Client_*、Executable_* 和 Event_* 条目按照 XLA 添加它们的顺序交错出现。“区域”列把每个后期槽位标回它的概念归属;区域深挖页面负责说明行为。
Error / Plugin / Event(槽位 5..14)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 5 | +0x028 | PJRT_Error_Destroy | PJRT_Error_Destroy | 0x0F85ECE0 |
| 6 | +0x030 | PJRT_Error_Message | PJRT_Error_Message | 0x0F85EDE0 |
| 7 | +0x038 | PJRT_Error_GetCode | PJRT_Error_GetCode | 0x0F85EF40 |
| 8 | +0x040 | PJRT_Plugin_Initialize | tpu_plugin::PJRT_Plugin_Initialize | 0x0E6A9D00 |
| 9 | +0x048 | PJRT_Plugin_Attributes | PJRT_Plugin_Attributes_Xla | 0x0F85F080 |
| 10 | +0x050 | PJRT_Event_Destroy | PJRT_Event_Destroy | 0x0F86F920 |
| 11 | +0x058 | PJRT_Event_IsReady | PJRT_Event_IsReady | 0x0F86F9E0 |
| 12 | +0x060 | PJRT_Event_Error | PJRT_Event_Error | 0x0F86FBA0 |
| 13 | +0x068 | PJRT_Event_Await | PJRT_Event_Await | 0x0F86FA80 |
| 14 | +0x070 | PJRT_Event_OnReady | PJRT_Event_OnReady | 0x0F86FC60 |
Client(槽位 15..27)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 15 | +0x078 | PJRT_Client_Create | tpu_plugin::PJRT_Client_Create | 0x0E6A8840 |
| 16 | +0x080 | PJRT_Client_Destroy | PJRT_Client_Destroy | 0x0F85F0E0 |
| 17 | +0x088 | PJRT_Client_PlatformName | PJRT_Client_PlatformName | 0x0F85F4A0 |
| 18 | +0x090 | PJRT_Client_ProcessIndex | PJRT_Client_ProcessIndex | 0x0F85F440 |
| 19 | +0x098 | PJRT_Client_PlatformVersion | PJRT_Client_PlatformVersion | 0x0F85F500 |
| 20 | +0x0A0 | PJRT_Client_Devices | PJRT_Client_Devices | 0x0F85F600 |
| 21 | +0x0A8 | PJRT_Client_AddressableDevices | PJRT_Client_AddressableDevices | 0x0F85F660 |
| 22 | +0x0B0 | PJRT_Client_LookupDevice | PJRT_Client_LookupDevice | 0x0F85F6C0 |
| 23 | +0x0B8 | PJRT_Client_LookupAddressableDevice | PJRT_Client_LookupAddressableDevice | 0x0F85F880 |
| 24 | +0x0C0 | PJRT_Client_AddressableMemories | PJRT_Client_AddressableMemories | 0x0F85FC60 |
| 25 | +0x0C8 | PJRT_Client_Compile | PJRT_Client_Compile | 0x0F861820 |
| 26 | +0x0D0 | PJRT_Client_DefaultDeviceAssignment | PJRT_Client_DefaultDeviceAssignment | 0x0F8630C0 |
| 27 | +0x0D8 | PJRT_Client_BufferFromHostBuffer | PJRT_Client_BufferFromHostBuffer | 0x0F8644C0 |
DeviceDescription / Device(槽位 28..39)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 28 | +0x0E0 | PJRT_DeviceDescription_Id | PJRT_DeviceDescription_Id | 0x0F865360 |
| 29 | +0x0E8 | PJRT_DeviceDescription_ProcessIndex | PJRT_DeviceDescription_ProcessIndex | 0x0F8653C0 |
| 30 | +0x0F0 | PJRT_DeviceDescription_Attributes | PJRT_DeviceDescription_Attributes | 0x0F865420 |
| 31 | +0x0F8 | PJRT_DeviceDescription_Kind | PJRT_DeviceDescription_Kind | 0x0F865480 |
| 32 | +0x100 | PJRT_DeviceDescription_DebugString | PJRT_DeviceDescription_DebugString | 0x0F865500 |
| 33 | +0x108 | PJRT_DeviceDescription_ToString | PJRT_DeviceDescription_ToString | 0x0F8658A0 |
| 34 | +0x110 | PJRT_Device_GetDescription | PJRT_Device_GetDescription | 0x0F8659A0 |
| 35 | +0x118 | PJRT_Device_IsAddressable | PJRT_Device_IsAddressable | 0x0F865A00 |
| 36 | +0x120 | PJRT_Device_LocalHardwareId | PJRT_Device_LocalHardwareId | 0x0F865A60 |
| 37 | +0x128 | PJRT_Device_AddressableMemories | PJRT_Device_AddressableMemories | 0x0F865AC0 |
| 38 | +0x130 | PJRT_Device_DefaultMemory | PJRT_Device_DefaultMemory | 0x0F865B20 |
| 39 | +0x138 | PJRT_Device_MemoryStats | PJRT_Device_MemoryStats | 0x0F865CE0 |
Memory(槽位 40..44)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 40 | +0x140 | PJRT_Memory_Id | PJRT_Memory_Id | 0x0F865E80 |
| 41 | +0x148 | PJRT_Memory_Kind | PJRT_Memory_Kind | 0x0F865EE0 |
| 42 | +0x150 | PJRT_Memory_DebugString | PJRT_Memory_DebugString | 0x0F865FC0 |
| 43 | +0x158 | PJRT_Memory_ToString | PJRT_Memory_ToString | 0x0F866040 |
| 44 | +0x160 | PJRT_Memory_AddressableByDevices | PJRT_Memory_AddressableByDevices | 0x0F8660C0 |
Executable(槽位 45..54)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 45 | +0x168 | PJRT_Executable_Destroy | PJRT_Executable_Destroy | 0x0F8661C0 |
| 46 | +0x170 | PJRT_Executable_Name | PJRT_Executable_Name | 0x0F866860 |
| 47 | +0x178 | PJRT_Executable_NumReplicas | PJRT_Executable_NumReplicas | 0x0F8668C0 |
| 48 | +0x180 | PJRT_Executable_NumPartitions | PJRT_Executable_NumPartitions | 0x0F866920 |
| 49 | +0x188 | PJRT_Executable_NumOutputs | PJRT_Executable_NumOutputs | 0x0F866A40 |
| 50 | +0x190 | PJRT_Executable_SizeOfGeneratedCodeInBytes | PJRT_Executable_SizeOfGeneratedCodeInBytes | 0x0F867240 |
| 51 | +0x198 | PJRT_Executable_GetCostAnalysis | PJRT_Executable_GetCostAnalysis | 0x0F867B80 |
| 52 | +0x1A0 | PJRT_Executable_OutputMemoryKinds | PJRT_Executable_OutputMemoryKinds | 0x0F869520 |
| 53 | +0x1A8 | PJRT_Executable_OptimizedProgram | PJRT_Executable_OptimizedProgram | 0x0F8672A0 |
| 54 | +0x1B0 | PJRT_Executable_Serialize | PJRT_Executable_Serialize | 0x0F86C5A0 |
LoadedExecutable(槽位 55..62)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 55 | +0x1B8 | PJRT_LoadedExecutable_Destroy | PJRT_LoadedExecutable_Destroy | 0x0F866780 |
| 56 | +0x1C0 | PJRT_LoadedExecutable_GetExecutable | PJRT_LoadedExecutable_GetExecutable | 0x0F86CFA0 |
| 57 | +0x1C8 | PJRT_LoadedExecutable_AddressableDevices | PJRT_LoadedExecutable_AddressableDevices | 0x0F866980 |
| 58 | +0x1D0 | PJRT_LoadedExecutable_Delete | PJRT_LoadedExecutable_Delete | 0x0F869A80 |
| 59 | +0x1D8 | PJRT_LoadedExecutable_IsDeleted | PJRT_LoadedExecutable_IsDeleted | 0x0F869AE0 |
| 60 | +0x1E0 | PJRT_LoadedExecutable_Execute | PJRT_LoadedExecutable_Execute | 0x0F869B40 |
| 61 | +0x1E8 | PJRT_Executable_DeserializeAndLoad | PJRT_Executable_DeserializeAndLoad | 0x0F86CC40 |
| 62 | +0x1F0 | PJRT_LoadedExecutable_Fingerprint | PJRT_LoadedExecutable_Fingerprint | 0x0F85FBE0 |
注意 — 槽位 62
PJRT_LoadedExecutable_Fingerprint是已弃用的 fingerprint 入口;v0.103 宿主改用槽位 99PJRT_Executable_Fingerprint。两者都已填充;弃用项保留给 older-minor 调用方。
Buffer(槽位 63..81)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 63 | +0x1F8 | PJRT_Buffer_Destroy | PJRT_Buffer_Destroy | 0x0F86D020 |
| 64 | +0x200 | PJRT_Buffer_ElementType | PJRT_Buffer_ElementType | 0x0F86D220 |
| 65 | +0x208 | PJRT_Buffer_Dimensions | PJRT_Buffer_Dimensions | 0x0F86D280 |
| 66 | +0x210 | PJRT_Buffer_UnpaddedDimensions | PJRT_Buffer_UnpaddedDimensions | 0x0F86D300 |
| 67 | +0x218 | PJRT_Buffer_DynamicDimensionIndices | PJRT_Buffer_DynamicDimensionIndices | 0x0F86D4C0 |
| 68 | +0x220 | PJRT_Buffer_GetMemoryLayout | PJRT_Buffer_GetMemoryLayout | 0x0F86D5E0 |
| 69 | +0x228 | PJRT_Buffer_OnDeviceSizeInBytes | PJRT_Buffer_OnDeviceSizeInBytes | 0x0F86DA80 |
| 70 | +0x230 | PJRT_Buffer_Device | PJRT_Buffer_Device | 0x0F86DB40 |
| 71 | +0x238 | PJRT_Buffer_Memory | PJRT_Buffer_Memory | 0x0F86DC60 |
| 72 | +0x240 | PJRT_Buffer_Delete | PJRT_Buffer_Delete | 0x0F86DD80 |
| 73 | +0x248 | PJRT_Buffer_IsDeleted | PJRT_Buffer_IsDeleted | 0x0F86DDE0 |
| 74 | +0x250 | PJRT_Buffer_CopyToDevice | PJRT_Buffer_CopyToDevice | 0x0F86E360 |
| 75 | +0x258 | PJRT_Buffer_ToHostBuffer | PJRT_Buffer_ToHostBuffer | 0x0F86E640 |
| 76 | +0x260 | PJRT_Buffer_IsOnCpu | PJRT_Buffer_IsOnCpu | 0x0F86ECC0 |
| 77 | +0x268 | PJRT_Buffer_ReadyEvent | PJRT_Buffer_ReadyEvent | 0x0F86ED20 |
| 78 | +0x270 | PJRT_Buffer_UnsafePointer | PJRT_Buffer_UnsafePointer | 0x0F86EE60 |
| 79 | +0x278 | PJRT_Buffer_IncreaseExternalReferenceCount | PJRT_Buffer_IncreaseExternalReferenceCount | 0x0F86EF20 |
| 80 | +0x280 | PJRT_Buffer_DecreaseExternalReferenceCount | PJRT_Buffer_DecreaseExternalReferenceCount | 0x0F86F100 |
| 81 | +0x288 | PJRT_Buffer_OpaqueDeviceMemoryDataPointer | PJRT_Buffer_OpaqueDeviceMemoryDataPointer | 0x0F86F200 |
CopyToDeviceStream(槽位 82..86)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 82 | +0x290 | PJRT_CopyToDeviceStream_Destroy | PJRT_CopyToDeviceStream_Destroy | 0x0F86F5E0 |
| 83 | +0x298 | PJRT_CopyToDeviceStream_AddChunk | PJRT_CopyToDeviceStream_AddChunk | 0x0F86F660 |
| 84 | +0x2A0 | PJRT_CopyToDeviceStream_TotalBytes | PJRT_CopyToDeviceStream_TotalBytes | 0x0F86F7E0 |
| 85 | +0x2A8 | PJRT_CopyToDeviceStream_GranuleSize | PJRT_CopyToDeviceStream_GranuleSize | 0x0F86F840 |
| 86 | +0x2B0 | PJRT_CopyToDeviceStream_CurrentBytes | PJRT_CopyToDeviceStream_CurrentBytes | 0x0F86F8A0 |
TopologyDescription / Compile(槽位 87..94)
| 槽位 | 偏移 | 字段 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|
| 87 | +0x2B8 | PJRT_TopologyDescription_Create | tpu_plugin::PJRT_TopologyDescription_Create | 0x0E6A9B20 |
| 88 | +0x2C0 | PJRT_TopologyDescription_Destroy | PJRT_TopologyDescription_Destroy | 0x0F870040 |
| 89 | +0x2C8 | PJRT_TopologyDescription_PlatformName | PJRT_TopologyDescription_PlatformName | 0x0F870200 |
| 90 | +0x2D0 | PJRT_TopologyDescription_PlatformVersion | PJRT_TopologyDescription_PlatformVersion | 0x0F870260 |
| 91 | +0x2D8 | PJRT_TopologyDescription_GetDeviceDescriptions | PJRT_TopologyDescription_GetDeviceDescriptions | 0x0F8702C0 |
| 92 | +0x2E0 | PJRT_TopologyDescription_Serialize | PJRT_TopologyDescription_Serialize | 0x0F870320 |
| 93 | +0x2E8 | PJRT_TopologyDescription_Attributes | PJRT_TopologyDescription_Attributes | 0x0F8705E0 |
| 94 | +0x2F0 | PJRT_Compile | PJRT_Compile | 0x0F870640 |
后期新增集群(槽位 95..139,新增顺序)
这些槽位是在 minor 0.40→0.103 之间追加的,顺序就是 XLA 添加它们的顺序,而不是按区域重新分组。“区域”列把每个槽位标回它的概念归属。
| 槽位 | 偏移 | 字段 | 区域 | 实现符号(pjrt::…) | 地址 |
|---|---|---|---|---|---|
| 95 | +0x2F8 | PJRT_Executable_OutputElementTypes | Executable | PJRT_Executable_OutputElementTypes | 0x0F868560 |
| 96 | +0x300 | PJRT_Executable_OutputDimensions | Executable | PJRT_Executable_OutputDimensions | 0x0F8689E0 |
| 97 | +0x308 | PJRT_Buffer_CopyToMemory | Buffer | PJRT_Buffer_CopyToMemory | 0x0F86E500 |
| 98 | +0x310 | PJRT_Client_CreateViewOfDeviceBuffer | Client | PJRT_Client_CreateViewOfDeviceBuffer | 0x0F865040 |
| 99 | +0x318 | PJRT_Executable_Fingerprint | Executable | PJRT_Executable_Fingerprint | 0x0F867AC0 |
| 100 | +0x320 | PJRT_Client_TopologyDescription | Client | PJRT_Client_TopologyDescription | 0x0F85F560 |
| 101 | +0x328 | PJRT_Executable_GetCompiledMemoryStats | Executable | PJRT_Executable_GetCompiledMemoryStats | 0x0F86CAC0 |
| 102 | +0x330 | PJRT_Memory_Kind_Id | Memory | PJRT_Memory_Kind_Id | 0x0F865F60 |
| 103 | +0x338 | PJRT_ExecuteContext_Create | ExecuteContext | tpu_plugin::PJRT_ExecuteContext_Create | 0x0E6A9A80 |
| 104 | +0x340 | PJRT_ExecuteContext_Destroy | ExecuteContext | PJRT_ExecuteContext_Destroy | 0x0F866120 |
| 105 | +0x348 | PJRT_Buffer_CopyRawToHost | Buffer | PJRT_Buffer_CopyRawToHost | 0x0F86DE40 |
| 106 | +0x350 | PJRT_AsyncHostToDeviceTransferManager_Destroy | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_Destroy | 0x0F860620 |
| 107 | +0x358 | PJRT_AsyncHostToDeviceTransferManager_TransferData | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_TransferData | 0x0F8606A0 |
| 108 | +0x360 | PJRT_Client_CreateBuffersForAsyncHostToDevice | Client | PJRT_Client_CreateBuffersForAsyncHostToDevice | 0x0F85FCC0 |
| 109 | +0x368 | PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer | 0x0F8611A0 |
| 110 | +0x370 | PJRT_AsyncHostToDeviceTransferManager_Device | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_Device | 0x0F861260 |
| 111 | +0x378 | PJRT_AsyncHostToDeviceTransferManager_BufferCount | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_BufferCount | 0x0F861380 |
| 112 | +0x380 | PJRT_AsyncHostToDeviceTransferManager_BufferSize | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_BufferSize | 0x0F8613E0 |
| 113 | +0x388 | PJRT_AsyncHostToDeviceTransferManager_SetBufferError | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_SetBufferError | 0x0F861440 |
| 114 | +0x390 | PJRT_AsyncHostToDeviceTransferManager_AddMetadata | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_AddMetadata | 0x0F861500 |
| 115 | +0x398 | PJRT_Client_DmaMap | Client | PJRT_Client_DmaMap | 0x0F860500 |
| 116 | +0x3A0 | PJRT_Client_DmaUnmap | Client | PJRT_Client_DmaUnmap | 0x0F860580 |
| 117 | +0x3A8 | PJRT_Client_CreateUninitializedBuffer | Client | PJRT_Client_CreateUninitializedBuffer | 0x0F863660 |
| 118 | +0x3B0 | PJRT_Client_UpdateGlobalProcessInfo | Client | PJRT_Client_UpdateGlobalProcessInfo | 0x0F85F940 |
| 119 | +0x3B8 | PJRT_TopologyDescription_Deserialize | Topology | PJRT_TopologyDescription_Deserialize | 0x0F870B80 |
| 120 | +0x3C0 | PJRT_Client_CreateAliasBuffer | Client | PJRT_Client_CreateAliasBuffer | 0x0F863D60 |
| 121 | +0x3C8 | PJRT_Client_FulfillAliasBuffer | Client | PJRT_Client_FulfillAliasBuffer | 0x0F8641A0 |
| 122 | +0x3D0 | PJRT_LoadedExecutable_GetDeviceAssignment | LoadedExec | PJRT_LoadedExecutable_GetDeviceAssignment | 0x0F870EA0 |
| 123 | +0x3D8 | PJRT_Client_CreateErrorBuffer | Client | PJRT_Client_CreateErrorBuffer | 0x0F8638A0 |
| 124 | +0x3E0 | PJRT_AsyncHostToDeviceTransferManager_TransferLiteral | AsyncH2D | PJRT_AsyncHostToDeviceTransferManager_TransferLiteral | 0x0F860960 |
| 125 | +0x3E8 | PJRT_Buffer_CopyRawToHostFuture | Buffer | PJRT_Buffer_CopyRawToHostFuture | 0x0F86DFE0 |
| 126 | +0x3F0 | PJRT_Device_PoisonExecution | Device | PJRT_Device_PoisonExecution | 0x0F860D00 |
| 127 | +0x3F8 | PJRT_Device_CreateAsyncTrackingEvent | Device | PJRT_Device_CreateAsyncTrackingEvent | 0x0F861080 |
| 128 | +0x400 | PJRT_AsyncTrackingEvent_Destroy | AsyncTracking | PJRT_AsyncTrackingEvent_Destroy | 0x0F861120 |
| 129 | +0x408 | PJRT_Executable_GetCompileOptions | Executable | PJRT_Executable_GetCompileOptions | 0x0F86C6E0 |
| 130 | +0x410 | PJRT_Buffer_DonateWithControlDependency | Buffer | PJRT_Buffer_DonateWithControlDependency | 0x0F86F2E0 |
| 131 | +0x418 | PJRT_Event_Create | Event | PJRT_Event_Create | 0x0F86FE00 |
| 132 | +0x420 | PJRT_Event_Set | Event | PJRT_Event_Set | 0x0F86FFA0 |
| 133 | +0x428 | PJRT_Device_GetAttributes | Device | PJRT_Device_GetAttributes | 0x0F873AC0 |
| 134 | +0x430 | PJRT_Client_Load | Client | PJRT_Client_Load | 0x0F8627E0 |
| 135 | +0x438 | PJRT_LoadedExecutable_AddressableDeviceLogicalIds | LoadedExec | PJRT_LoadedExecutable_AddressableDeviceLogicalIds | 0x0F8669E0 |
| 136 | +0x440 | PJRT_Buffer_Bitcast | Buffer | PJRT_Buffer_Bitcast | 0x0F862D00 |
| 137 | +0x448 | PJRT_Error_ForEachPayload | Error | PJRT_Error_ForEachPayload | 0x0F85EFC0 |
| 138 | +0x450 | PJRT_TopologyDescription_Fingerprint | Topology | PJRT_TopologyDescription_Fingerprint | 0x0F870520 |
| 139 | +0x458 | PJRT_Executable_ParameterMemoryKinds | Executable | PJRT_Executable_ParameterMemoryKinds | 0x0F868FC0 |
已填充与已注入映射
目的
135 个函数指针槽位每一个都已填充;该构建中没有任何一个为 null。但这些槽位分成两类来源,这个分野正是插件模式的核心:130 个槽位是链接时烘入的通用 XLA C-API wrapper,5 个则由 CreatePjrtApi 的调用方在运行时注入。
五个注入点
CreatePjrtApi 唯一不硬编码的槽位是其 a2..a7 参数。GetTpuPjrtApi 会把 TPU 专用实现传入这些精确位置;其他所有内容都是固定 relocation。
| 槽位 | 字段 | 参数 | 注入的实现 | 地址 | 为什么是 TPU 专用 |
|---|---|---|---|---|---|
| 15 | PJRT_Client_Create | a2 | tpu_plugin::PJRT_Client_Create | 0x0E6A8840 | 构造 TPU PjRtClient(设备发现、运行时初始化) |
| 103 | PJRT_ExecuteContext_Create | a3 | tpu_plugin::PJRT_ExecuteContext_Create | 0x0E6A9A80 | TPU 专用 execute context |
| 87 | PJRT_TopologyDescription_Create | a4 | tpu_plugin::PJRT_TopologyDescription_Create | 0x0E6A9B20 | TPU pod/slice 拓扑 |
| 8 | PJRT_Plugin_Initialize | a5 | tpu_plugin::PJRT_Plugin_Initialize | 0x0E6A9D00 | TPU runtime 启动 |
| 1 | extension_start | a6 | (扩展链头) | 0x224C3F68 | TPU 扩展集合 |
| 9 | PJRT_Plugin_Attributes | a7 | PJRT_Plugin_Attributes_Xla | 0x0F85F080 | 插件属性表(不在 tpu_plugin:: 中,但仍是注入的) |
四个 tpu_plugin:: 命名空间下的实现(槽位 8、15、87、103)是真正的 TPU 专用实现。槽位 9 是注入的,但解析到标准的 pjrt::PJRT_Plugin_Attributes_Xla,即调用方传入的是 XLA 实现,而不是 TPU override。槽位 1 是扩展链头,同样是注入的(见 extension-chain.md)。剩余 130 个函数指针都是来自 XLA 的 pjrt_c_api_wrapper_impl.cc 的编译期固定 pjrt::PJRT_* 符号。
怪癖 — “TPU 专用”的集合比重新实现者可能预期的要小得多。只有四个槽位承载 TPU 专用代码;另外 131 个函数指针与任何 PJRT 插件都会携带的通用 XLA wrapper 逐字节相同。TPU 特性几乎完全存在于
Client_Create背后(它构建PjRtClient,随后通用 wrapper 通过它分发)以及扩展链中,而不在逐调用 wrapper 中。
向后兼容保护
目的
针对较旧 PJRT minor 编译的宿主,会传入比这个 v0.103 插件预期更小的 _Args 结构体。插件必须接受它,并且只读取调用方实际提供的字段。其机制是每个 wrapper 顶部的逐槽位大小检查。
算法
function PJRT_Error_Destroy(args): // 0xF85ECE0
rc = pjrt::ActualStructSizeIsGreaterOrEqual( // 0xF8A4EC0
"PJRT_Error_Destroy_Args", // API name (for the log line)
23, // min accepted struct_size
24, // current (v0.103) struct_size
args->struct_size); // what the caller passed
if (rc != ok):
LogMessage(".../pjrt_c_api_wrapper_impl.cc", 545) << status; // diagnostic
if (args->struct_size < 0x18): return; // too small: bail
// ... read only fields within args->struct_size ...
```text
每个 wrapper 的第一条指令都是这个 `ActualStructSizeIsGreaterOrEqual` 调用。硬编码的 `(min, current)` 对界定了该插件接受的每个 `_Args` 版本范围;对 `PJRT_Error_Destroy_Args` 来说是 `(23, 24)`。调用方 `struct_size` 之外的字段永远不会被读取。这就是一个 v0.103 插件无需按版本设置 trampoline,也能服务一组宿主 minor 的原因:兼容逻辑由每个 args 结构体上的大小前缀承载,而不是由版本分支承载。
> **注意 —** 源路径字符串 `third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc` 被烘入 `0xF85ECE0` 处的诊断中。它确认这 130 个固定 wrapper 是 XLA 的通用 `pjrt_c_api_wrapper_impl.cc` 代码,在 PJRT 插件之间相同;这是二进制自身的 provenance 字符串,而不是外部断言。
---
## 热路径槽位
JAX/PyTorch-XLA step 最常触碰的槽位。该排名是一个*语义估计*:没有运行 call-count instrumentation,所以顺序的置信度为 LOW,尽管槽位身份为 CERTAIN。
| 排名 | 槽位 | 字段 | 每步角色 |
|---|---|---|---|
| 1 | 60 | PJRT_LoadedExecutable_Execute | 将程序启动到 TPU;主要吞吐槽位 |
| 2 | 11 | PJRT_Event_IsReady | 在紧密循环中轮询异步完成 |
| 3 | 13 | PJRT_Event_Await | 对异步结果做同步等待 |
| 4 | 14 | PJRT_Event_OnReady | 按 buffer 注册完成回调 |
| 5 | 10 | PJRT_Event_Destroy | batch await 后批量销毁 |
| 6 | 63 | PJRT_Buffer_Destroy | 按 buffer 释放;图输出清理时成批出现 |
| 7 | 27 | PJRT_Client_BufferFromHostBuffer | Host→TPU 输入上传,每个 step 的每个参数一次 |
| 8 | 75 | PJRT_Buffer_ToHostBuffer | TPU→host 下载(指标、探针) |
| 9 | 81 | PJRT_Buffer_OpaqueDeviceMemoryDataPointer | 为外部库互操作获取 zero-copy DMA 指针 |
| 10 | 79/80 | PJRT_Buffer_{In,De}creaseExternalReferenceCount | DLPack 风格的引用计数增加/减少 |
> **陷阱 —** 来源发现中的 Top-20 排名是一种排序*启发式*,不是实测数据。优化分发路径的重新实现者应当把槽位 60(Execute)和 Event 槽位(10..14)视为确定的热路径,但在没有 instrumentation 的情况下,不应信任较冷槽位的精确排名。
---
## 交叉引用
- [PJRT 概览](overview.md) — 本表在插件生命周期中的位置;上下文中的 `GetPjrtApi` → `GetTpuPjrtApi` → `CreatePjrtApi` 链
- [扩展链](extension-chain.md) — 挂在 `extension_start`(槽位 1)上的链表;17 个扩展、type-ID、逐扩展布局
- [Client 与 Device](client-and-device.md) — 槽位 15..44:`Client_Create`(关键 TPU 注入点)以及 device/memory 查询
- [Buffer 与 Memory](buffer-and-memory.md) — 槽位 63..81、97、105、125、130、136:buffer 生命周期、传输、DMA
- [Executable 执行](executable-execution.md) — 槽位 45..62、122、135,以及热路径槽位 60 `LoadedExecutable_Execute`
- [Events 与 Async](events-and-async.md) — 槽位 10..14、131..132:`PJRT_Event` 模型
- [Callbacks](callbacks.md) — `OnReady` / async-tracking-event 槽位以及 callback extension
- [Collectives Communicator](collectives-communicator.md) — 通过链到达而不是通过这些槽位到达的 collectives extension surface