Skip to content

PJRT_Api 函数指针表重建

本页所有地址均适用于 libtpu 0.0.40 (cp314, manylinux_2_31_x86_64),build-id 89edbbe81c5b328a958fe628a9f2207d。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_sizeextension_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 处编码为 qword 0x6700000000
  • 槽位到符号映射:每个已填充槽位都解析到一个 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

各区域槽位数

区域槽位数槽位范围
头部(标量)50..4
Error3(+1 个后期新增)5..7, 137
Plugin28..9
Event5(+2 个后期新增)10..14, 131..132
Client13(+约 12 个后期新增)15..27, 98, 100, 108, 115..118, 120..121, 123, 134
DeviceDescription628..33
Device6(+3 个后期新增)34..39, 126..127, 133
Memory5(+1 个后期新增)40..44, 102
Executable10(+6 个后期新增)45..54, 95..96, 99, 101, 129, 139
LoadedExecutable8(+2 个后期新增)55..62, 122, 135
Buffer19(+5 个后期新增)63..81, 97, 105, 125, 130, 136
CopyToDeviceStream582..86
TopologyDescription7(+2 个后期新增)87..93, 119, 138
Compile194
ExecuteContext1(+1 个后期新增)103..104
AsyncHostToDeviceTransferManager9(后期新增)106..107, 109..114, 124
AsyncTrackingEvent1(后期新增)128

填充路径与存储

目的

该表不是 PROGBITS section 中的静态初始化器,而是在第一次 GetPjrtApi 调用时,惰性地、仅一次地构建到零填充的 .lbss 中。理解填充路径,是返回正确表的重新实现与发生竞态或返回垃圾数据的重新实现之间的区别。

入口点

text
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;                             // 0x227BA840
c
function 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 位 = 0minor = 高 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+0x028PJRT_Error_DestroyPJRT_Error_Destroy0x0F85ECE0
6+0x030PJRT_Error_MessagePJRT_Error_Message0x0F85EDE0
7+0x038PJRT_Error_GetCodePJRT_Error_GetCode0x0F85EF40
8+0x040PJRT_Plugin_Initializetpu_plugin::PJRT_Plugin_Initialize0x0E6A9D00
9+0x048PJRT_Plugin_AttributesPJRT_Plugin_Attributes_Xla0x0F85F080
10+0x050PJRT_Event_DestroyPJRT_Event_Destroy0x0F86F920
11+0x058PJRT_Event_IsReadyPJRT_Event_IsReady0x0F86F9E0
12+0x060PJRT_Event_ErrorPJRT_Event_Error0x0F86FBA0
13+0x068PJRT_Event_AwaitPJRT_Event_Await0x0F86FA80
14+0x070PJRT_Event_OnReadyPJRT_Event_OnReady0x0F86FC60

Client(槽位 15..27)

槽位偏移字段实现符号(pjrt::…)地址
15+0x078PJRT_Client_Createtpu_plugin::PJRT_Client_Create0x0E6A8840
16+0x080PJRT_Client_DestroyPJRT_Client_Destroy0x0F85F0E0
17+0x088PJRT_Client_PlatformNamePJRT_Client_PlatformName0x0F85F4A0
18+0x090PJRT_Client_ProcessIndexPJRT_Client_ProcessIndex0x0F85F440
19+0x098PJRT_Client_PlatformVersionPJRT_Client_PlatformVersion0x0F85F500
20+0x0A0PJRT_Client_DevicesPJRT_Client_Devices0x0F85F600
21+0x0A8PJRT_Client_AddressableDevicesPJRT_Client_AddressableDevices0x0F85F660
22+0x0B0PJRT_Client_LookupDevicePJRT_Client_LookupDevice0x0F85F6C0
23+0x0B8PJRT_Client_LookupAddressableDevicePJRT_Client_LookupAddressableDevice0x0F85F880
24+0x0C0PJRT_Client_AddressableMemoriesPJRT_Client_AddressableMemories0x0F85FC60
25+0x0C8PJRT_Client_CompilePJRT_Client_Compile0x0F861820
26+0x0D0PJRT_Client_DefaultDeviceAssignmentPJRT_Client_DefaultDeviceAssignment0x0F8630C0
27+0x0D8PJRT_Client_BufferFromHostBufferPJRT_Client_BufferFromHostBuffer0x0F8644C0

DeviceDescription / Device(槽位 28..39)

槽位偏移字段实现符号(pjrt::…)地址
28+0x0E0PJRT_DeviceDescription_IdPJRT_DeviceDescription_Id0x0F865360
29+0x0E8PJRT_DeviceDescription_ProcessIndexPJRT_DeviceDescription_ProcessIndex0x0F8653C0
30+0x0F0PJRT_DeviceDescription_AttributesPJRT_DeviceDescription_Attributes0x0F865420
31+0x0F8PJRT_DeviceDescription_KindPJRT_DeviceDescription_Kind0x0F865480
32+0x100PJRT_DeviceDescription_DebugStringPJRT_DeviceDescription_DebugString0x0F865500
33+0x108PJRT_DeviceDescription_ToStringPJRT_DeviceDescription_ToString0x0F8658A0
34+0x110PJRT_Device_GetDescriptionPJRT_Device_GetDescription0x0F8659A0
35+0x118PJRT_Device_IsAddressablePJRT_Device_IsAddressable0x0F865A00
36+0x120PJRT_Device_LocalHardwareIdPJRT_Device_LocalHardwareId0x0F865A60
37+0x128PJRT_Device_AddressableMemoriesPJRT_Device_AddressableMemories0x0F865AC0
38+0x130PJRT_Device_DefaultMemoryPJRT_Device_DefaultMemory0x0F865B20
39+0x138PJRT_Device_MemoryStatsPJRT_Device_MemoryStats0x0F865CE0

Memory(槽位 40..44)

槽位偏移字段实现符号(pjrt::…)地址
40+0x140PJRT_Memory_IdPJRT_Memory_Id0x0F865E80
41+0x148PJRT_Memory_KindPJRT_Memory_Kind0x0F865EE0
42+0x150PJRT_Memory_DebugStringPJRT_Memory_DebugString0x0F865FC0
43+0x158PJRT_Memory_ToStringPJRT_Memory_ToString0x0F866040
44+0x160PJRT_Memory_AddressableByDevicesPJRT_Memory_AddressableByDevices0x0F8660C0

Executable(槽位 45..54)

槽位偏移字段实现符号(pjrt::…)地址
45+0x168PJRT_Executable_DestroyPJRT_Executable_Destroy0x0F8661C0
46+0x170PJRT_Executable_NamePJRT_Executable_Name0x0F866860
47+0x178PJRT_Executable_NumReplicasPJRT_Executable_NumReplicas0x0F8668C0
48+0x180PJRT_Executable_NumPartitionsPJRT_Executable_NumPartitions0x0F866920
49+0x188PJRT_Executable_NumOutputsPJRT_Executable_NumOutputs0x0F866A40
50+0x190PJRT_Executable_SizeOfGeneratedCodeInBytesPJRT_Executable_SizeOfGeneratedCodeInBytes0x0F867240
51+0x198PJRT_Executable_GetCostAnalysisPJRT_Executable_GetCostAnalysis0x0F867B80
52+0x1A0PJRT_Executable_OutputMemoryKindsPJRT_Executable_OutputMemoryKinds0x0F869520
53+0x1A8PJRT_Executable_OptimizedProgramPJRT_Executable_OptimizedProgram0x0F8672A0
54+0x1B0PJRT_Executable_SerializePJRT_Executable_Serialize0x0F86C5A0

LoadedExecutable(槽位 55..62)

槽位偏移字段实现符号(pjrt::…)地址
55+0x1B8PJRT_LoadedExecutable_DestroyPJRT_LoadedExecutable_Destroy0x0F866780
56+0x1C0PJRT_LoadedExecutable_GetExecutablePJRT_LoadedExecutable_GetExecutable0x0F86CFA0
57+0x1C8PJRT_LoadedExecutable_AddressableDevicesPJRT_LoadedExecutable_AddressableDevices0x0F866980
58+0x1D0PJRT_LoadedExecutable_DeletePJRT_LoadedExecutable_Delete0x0F869A80
59+0x1D8PJRT_LoadedExecutable_IsDeletedPJRT_LoadedExecutable_IsDeleted0x0F869AE0
60+0x1E0PJRT_LoadedExecutable_ExecutePJRT_LoadedExecutable_Execute0x0F869B40
61+0x1E8PJRT_Executable_DeserializeAndLoadPJRT_Executable_DeserializeAndLoad0x0F86CC40
62+0x1F0PJRT_LoadedExecutable_FingerprintPJRT_LoadedExecutable_Fingerprint0x0F85FBE0

注意 — 槽位 62 PJRT_LoadedExecutable_Fingerprint 是已弃用的 fingerprint 入口;v0.103 宿主改用槽位 99 PJRT_Executable_Fingerprint。两者都已填充;弃用项保留给 older-minor 调用方。

Buffer(槽位 63..81)

槽位偏移字段实现符号(pjrt::…)地址
63+0x1F8PJRT_Buffer_DestroyPJRT_Buffer_Destroy0x0F86D020
64+0x200PJRT_Buffer_ElementTypePJRT_Buffer_ElementType0x0F86D220
65+0x208PJRT_Buffer_DimensionsPJRT_Buffer_Dimensions0x0F86D280
66+0x210PJRT_Buffer_UnpaddedDimensionsPJRT_Buffer_UnpaddedDimensions0x0F86D300
67+0x218PJRT_Buffer_DynamicDimensionIndicesPJRT_Buffer_DynamicDimensionIndices0x0F86D4C0
68+0x220PJRT_Buffer_GetMemoryLayoutPJRT_Buffer_GetMemoryLayout0x0F86D5E0
69+0x228PJRT_Buffer_OnDeviceSizeInBytesPJRT_Buffer_OnDeviceSizeInBytes0x0F86DA80
70+0x230PJRT_Buffer_DevicePJRT_Buffer_Device0x0F86DB40
71+0x238PJRT_Buffer_MemoryPJRT_Buffer_Memory0x0F86DC60
72+0x240PJRT_Buffer_DeletePJRT_Buffer_Delete0x0F86DD80
73+0x248PJRT_Buffer_IsDeletedPJRT_Buffer_IsDeleted0x0F86DDE0
74+0x250PJRT_Buffer_CopyToDevicePJRT_Buffer_CopyToDevice0x0F86E360
75+0x258PJRT_Buffer_ToHostBufferPJRT_Buffer_ToHostBuffer0x0F86E640
76+0x260PJRT_Buffer_IsOnCpuPJRT_Buffer_IsOnCpu0x0F86ECC0
77+0x268PJRT_Buffer_ReadyEventPJRT_Buffer_ReadyEvent0x0F86ED20
78+0x270PJRT_Buffer_UnsafePointerPJRT_Buffer_UnsafePointer0x0F86EE60
79+0x278PJRT_Buffer_IncreaseExternalReferenceCountPJRT_Buffer_IncreaseExternalReferenceCount0x0F86EF20
80+0x280PJRT_Buffer_DecreaseExternalReferenceCountPJRT_Buffer_DecreaseExternalReferenceCount0x0F86F100
81+0x288PJRT_Buffer_OpaqueDeviceMemoryDataPointerPJRT_Buffer_OpaqueDeviceMemoryDataPointer0x0F86F200

CopyToDeviceStream(槽位 82..86)

槽位偏移字段实现符号(pjrt::…)地址
82+0x290PJRT_CopyToDeviceStream_DestroyPJRT_CopyToDeviceStream_Destroy0x0F86F5E0
83+0x298PJRT_CopyToDeviceStream_AddChunkPJRT_CopyToDeviceStream_AddChunk0x0F86F660
84+0x2A0PJRT_CopyToDeviceStream_TotalBytesPJRT_CopyToDeviceStream_TotalBytes0x0F86F7E0
85+0x2A8PJRT_CopyToDeviceStream_GranuleSizePJRT_CopyToDeviceStream_GranuleSize0x0F86F840
86+0x2B0PJRT_CopyToDeviceStream_CurrentBytesPJRT_CopyToDeviceStream_CurrentBytes0x0F86F8A0

TopologyDescription / Compile(槽位 87..94)

槽位偏移字段实现符号(pjrt::…)地址
87+0x2B8PJRT_TopologyDescription_Createtpu_plugin::PJRT_TopologyDescription_Create0x0E6A9B20
88+0x2C0PJRT_TopologyDescription_DestroyPJRT_TopologyDescription_Destroy0x0F870040
89+0x2C8PJRT_TopologyDescription_PlatformNamePJRT_TopologyDescription_PlatformName0x0F870200
90+0x2D0PJRT_TopologyDescription_PlatformVersionPJRT_TopologyDescription_PlatformVersion0x0F870260
91+0x2D8PJRT_TopologyDescription_GetDeviceDescriptionsPJRT_TopologyDescription_GetDeviceDescriptions0x0F8702C0
92+0x2E0PJRT_TopologyDescription_SerializePJRT_TopologyDescription_Serialize0x0F870320
93+0x2E8PJRT_TopologyDescription_AttributesPJRT_TopologyDescription_Attributes0x0F8705E0
94+0x2F0PJRT_CompilePJRT_Compile0x0F870640

后期新增集群(槽位 95..139,新增顺序)

这些槽位是在 minor 0.40→0.103 之间追加的,顺序就是 XLA 添加它们的顺序,而不是按区域重新分组。“区域”列把每个槽位标回它的概念归属。

槽位偏移字段区域实现符号(pjrt::…)地址
95+0x2F8PJRT_Executable_OutputElementTypesExecutablePJRT_Executable_OutputElementTypes0x0F868560
96+0x300PJRT_Executable_OutputDimensionsExecutablePJRT_Executable_OutputDimensions0x0F8689E0
97+0x308PJRT_Buffer_CopyToMemoryBufferPJRT_Buffer_CopyToMemory0x0F86E500
98+0x310PJRT_Client_CreateViewOfDeviceBufferClientPJRT_Client_CreateViewOfDeviceBuffer0x0F865040
99+0x318PJRT_Executable_FingerprintExecutablePJRT_Executable_Fingerprint0x0F867AC0
100+0x320PJRT_Client_TopologyDescriptionClientPJRT_Client_TopologyDescription0x0F85F560
101+0x328PJRT_Executable_GetCompiledMemoryStatsExecutablePJRT_Executable_GetCompiledMemoryStats0x0F86CAC0
102+0x330PJRT_Memory_Kind_IdMemoryPJRT_Memory_Kind_Id0x0F865F60
103+0x338PJRT_ExecuteContext_CreateExecuteContexttpu_plugin::PJRT_ExecuteContext_Create0x0E6A9A80
104+0x340PJRT_ExecuteContext_DestroyExecuteContextPJRT_ExecuteContext_Destroy0x0F866120
105+0x348PJRT_Buffer_CopyRawToHostBufferPJRT_Buffer_CopyRawToHost0x0F86DE40
106+0x350PJRT_AsyncHostToDeviceTransferManager_DestroyAsyncH2DPJRT_AsyncHostToDeviceTransferManager_Destroy0x0F860620
107+0x358PJRT_AsyncHostToDeviceTransferManager_TransferDataAsyncH2DPJRT_AsyncHostToDeviceTransferManager_TransferData0x0F8606A0
108+0x360PJRT_Client_CreateBuffersForAsyncHostToDeviceClientPJRT_Client_CreateBuffersForAsyncHostToDevice0x0F85FCC0
109+0x368PJRT_AsyncHostToDeviceTransferManager_RetrieveBufferAsyncH2DPJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer0x0F8611A0
110+0x370PJRT_AsyncHostToDeviceTransferManager_DeviceAsyncH2DPJRT_AsyncHostToDeviceTransferManager_Device0x0F861260
111+0x378PJRT_AsyncHostToDeviceTransferManager_BufferCountAsyncH2DPJRT_AsyncHostToDeviceTransferManager_BufferCount0x0F861380
112+0x380PJRT_AsyncHostToDeviceTransferManager_BufferSizeAsyncH2DPJRT_AsyncHostToDeviceTransferManager_BufferSize0x0F8613E0
113+0x388PJRT_AsyncHostToDeviceTransferManager_SetBufferErrorAsyncH2DPJRT_AsyncHostToDeviceTransferManager_SetBufferError0x0F861440
114+0x390PJRT_AsyncHostToDeviceTransferManager_AddMetadataAsyncH2DPJRT_AsyncHostToDeviceTransferManager_AddMetadata0x0F861500
115+0x398PJRT_Client_DmaMapClientPJRT_Client_DmaMap0x0F860500
116+0x3A0PJRT_Client_DmaUnmapClientPJRT_Client_DmaUnmap0x0F860580
117+0x3A8PJRT_Client_CreateUninitializedBufferClientPJRT_Client_CreateUninitializedBuffer0x0F863660
118+0x3B0PJRT_Client_UpdateGlobalProcessInfoClientPJRT_Client_UpdateGlobalProcessInfo0x0F85F940
119+0x3B8PJRT_TopologyDescription_DeserializeTopologyPJRT_TopologyDescription_Deserialize0x0F870B80
120+0x3C0PJRT_Client_CreateAliasBufferClientPJRT_Client_CreateAliasBuffer0x0F863D60
121+0x3C8PJRT_Client_FulfillAliasBufferClientPJRT_Client_FulfillAliasBuffer0x0F8641A0
122+0x3D0PJRT_LoadedExecutable_GetDeviceAssignmentLoadedExecPJRT_LoadedExecutable_GetDeviceAssignment0x0F870EA0
123+0x3D8PJRT_Client_CreateErrorBufferClientPJRT_Client_CreateErrorBuffer0x0F8638A0
124+0x3E0PJRT_AsyncHostToDeviceTransferManager_TransferLiteralAsyncH2DPJRT_AsyncHostToDeviceTransferManager_TransferLiteral0x0F860960
125+0x3E8PJRT_Buffer_CopyRawToHostFutureBufferPJRT_Buffer_CopyRawToHostFuture0x0F86DFE0
126+0x3F0PJRT_Device_PoisonExecutionDevicePJRT_Device_PoisonExecution0x0F860D00
127+0x3F8PJRT_Device_CreateAsyncTrackingEventDevicePJRT_Device_CreateAsyncTrackingEvent0x0F861080
128+0x400PJRT_AsyncTrackingEvent_DestroyAsyncTrackingPJRT_AsyncTrackingEvent_Destroy0x0F861120
129+0x408PJRT_Executable_GetCompileOptionsExecutablePJRT_Executable_GetCompileOptions0x0F86C6E0
130+0x410PJRT_Buffer_DonateWithControlDependencyBufferPJRT_Buffer_DonateWithControlDependency0x0F86F2E0
131+0x418PJRT_Event_CreateEventPJRT_Event_Create0x0F86FE00
132+0x420PJRT_Event_SetEventPJRT_Event_Set0x0F86FFA0
133+0x428PJRT_Device_GetAttributesDevicePJRT_Device_GetAttributes0x0F873AC0
134+0x430PJRT_Client_LoadClientPJRT_Client_Load0x0F8627E0
135+0x438PJRT_LoadedExecutable_AddressableDeviceLogicalIdsLoadedExecPJRT_LoadedExecutable_AddressableDeviceLogicalIds0x0F8669E0
136+0x440PJRT_Buffer_BitcastBufferPJRT_Buffer_Bitcast0x0F862D00
137+0x448PJRT_Error_ForEachPayloadErrorPJRT_Error_ForEachPayload0x0F85EFC0
138+0x450PJRT_TopologyDescription_FingerprintTopologyPJRT_TopologyDescription_Fingerprint0x0F870520
139+0x458PJRT_Executable_ParameterMemoryKindsExecutablePJRT_Executable_ParameterMemoryKinds0x0F868FC0

已填充与已注入映射

目的

135 个函数指针槽位每一个都已填充;该构建中没有任何一个为 null。但这些槽位分成两类来源,这个分野正是插件模式的核心:130 个槽位是链接时烘入的通用 XLA C-API wrapper,5 个则由 CreatePjrtApi 的调用方在运行时注入。

五个注入点

CreatePjrtApi 唯一不硬编码的槽位是其 a2..a7 参数。GetTpuPjrtApi 会把 TPU 专用实现传入这些精确位置;其他所有内容都是固定 relocation。

槽位字段参数注入的实现地址为什么是 TPU 专用
15PJRT_Client_Createa2tpu_plugin::PJRT_Client_Create0x0E6A8840构造 TPU PjRtClient(设备发现、运行时初始化)
103PJRT_ExecuteContext_Createa3tpu_plugin::PJRT_ExecuteContext_Create0x0E6A9A80TPU 专用 execute context
87PJRT_TopologyDescription_Createa4tpu_plugin::PJRT_TopologyDescription_Create0x0E6A9B20TPU pod/slice 拓扑
8PJRT_Plugin_Initializea5tpu_plugin::PJRT_Plugin_Initialize0x0E6A9D00TPU runtime 启动
1extension_starta6(扩展链头)0x224C3F68TPU 扩展集合
9PJRT_Plugin_Attributesa7PJRT_Plugin_Attributes_Xla0x0F85F080插件属性表(不在 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 顶部的逐槽位大小检查。

算法

c
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..8197105125130136:buffer 生命周期、传输、DMA
- [Executable 执行](executable-execution.md) — 槽位 45..62122135,以及热路径槽位 60 `LoadedExecutable_Execute`
- [Events 与 Async](events-and-async.md) — 槽位 10..14131..132:`PJRT_Event` 模型
- [Callbacks](callbacks.md) — `OnReady` / async-tracking-event 槽位以及 callback extension
- [Collectives Communicator](collectives-communicator.md) — 通过链到达而不是通过这些槽位到达的 collectives extension surface