PJRT 扩展链

本页中的所有地址和偏移均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d），PJRT C-API v0.103。其他版本会有所不同。

摘要

PJRT C-API 有一个固定宽度的 vtable：140 个函数指针槽位，冻结在插件构建时所针对的版本（此处为 v0.103）。如果插件想暴露超出这个冻结表面的能力，又不能在不提升版本的情况下新增槽位，因此 API 带有第二条开放式通道：挂在 PJRT_Api.extension_start 上的 extension 节点单向链表。每个节点都以一个公共头开头：{ size_t struct_size; uint32 type; uint32 _pad; PJRT_Extension_Base* next; }。框架通过沿着 next 遍历来发现功能，直到找到 type 匹配已知 PJRT_Extension_Type id 的节点，或遇到 NULL。这与 LLVM 用于 pass-plugin registry、COM 用于 QueryInterface 的模式相同：稳定 ABI 主干加上自描述、版本容忍的旁表。

libtpu 发布 17 个扩展。其中 16 个位于 .bss，作为 pjrt::tpu_plugin::GetTpuPjrtApi @ 0xE6AA440 内零初始化的函数局部静态对象；每个对象都在 C++ __cxa_guard 保护下只填充一次，由扁平表初始化器 pjrt::Create<Name>Extension(node, next, ...) 写入 struct_size、type、next 和固定一段 fn-ptr 槽位，然后返回。第 17 个 Profiler（type 1）位于 .data，带有静态初始化重定位，并作为链的种子：RawBuffer 的 next 指向它，而它自己的 next 是 NULL。由于每个新节点的 next 指向此前构建的节点，链表按尾部优先构建；extension_start 被设为最后构建的节点，因此消费者遍历链时会以最新优先的顺序看到扩展。该顺序对消费者没有语义意义，发现过程是按 type-id 线性扫描，但它映射了 GetTpuPjrtApi 内源级声明顺序。

本页负责说明节点布局、PJRT_Extension_Type enum、完整的有序扩展清单以及遍历算法。各扩展的深入解析位于专门页面（在清单表中链接）。140 个非扩展函数指针槽位是完全独立的结构，在 API Vtable Reconstruction 中重建。

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

PJRT_Extension_Base 公共头：字段顺序、偏移、32-bit type 加 4-byte pad，以及位于 +0x10 的 next。
PJRT_Extension_Type enum：libtpu 使用哪些 id（1, 4, 6, 8, 9, 12–23），有意省略哪些 id（0, 2, 3, 5, 7, 10, 11），以及这些省略意味着什么。
17 个节点按遍历顺序的清单，包括每个节点的存储 VA、struct_size、type 和填充它的 creator。
构建顺序与遍历顺序的反转，以及为什么这无害。
消费端的功能检测循环及其前向兼容契约。


链头	`PJRT_Api.extension_start` @ `+0x08` → `0x224C3F68`（HostMemoryAllocator 节点）
节点头	`{ size_t struct_size; uint32 type; uint32 _pad; PJRT_Extension_Base* next; }` — 24 bytes，`next` 位于 `+0x10`
链长	17 个节点；终止节点是 Profiler（type 1），`next = NULL`
存在的 type id	1, 4, 6, 8, 9, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23
缺失的 type id	0, 2, 3, 5, 7, 10, 11（包括 canonical FFI、Custom_Partitioner、Stream、Triton）
Builder	`pjrt::tpu_plugin::GetTpuPjrtApi` @ `0xE6AA440`（16 个 `__cxa_guard` 块 + 第 17 个用于 `pjrt_api`）
最大节点	Megascale（type 18）@ `0x224C3D08`，248 bytes（0xF8），23 个 live + 5 个 reserved 槽位
最小节点	HostMemoryAllocator（type 23）@ `0x224C3F68`，32 bytes，1 个方法

节点布局

每个节点都是一个 PJRT_Extension_Base，后接零个或多个 8-byte 函数指针。公共头在全部 17 个扩展中不变；只有 struct_size 和 fn-ptr 尾部不同。

公共头

struct PJRT_Extension_Base {
    /* +0x00 */ size_t                struct_size;   // per-extension total, in bytes
    /* +0x08 */ uint32_t              type;          // PJRT_Extension_Type enum value
    /* +0x0C */ uint32_t              _pad;          // padding to 8-byte align `next`
    /* +0x10 */ PJRT_Extension_Base*  next;          // chain link; NULL terminates
    /*  ...   fn-ptr slots from +0x18, one qword each ... */
};
```text

该头由每个 creator 的 store 序列直接确认。两个代表性示例：

```c
// pjrt::CreateHostMemoryAllocatorExtension @ 0xE6F5340  (type 23, 32 bytes)
function CreateHostMemoryAllocatorExtension(node, next):
    *(size_t*) (node + 0x00) = 32;                     // struct_size
    *(uint32_t*)(node + 0x08) = 23;                    // type  (32-bit store!)
    *(void**)  (node + 0x10) = next;                   // chain link
    *(void**)  (node + 0x18) = HostMemoryAllocator_Allocate;
    return node;

// pjrt::CreateShardingsExtension @ 0xF874980  (type 19, 40 bytes)
function CreateShardingsExtension(node, next):
    *(size_t*) (node + 0x00) = 40;
    *(uint32_t*)(node + 0x08) = 19;
    *(void**)  (node + 0x10) = next;
    *(void**)  (node + 0x18) = PJRT_Shardings_..._ParameterShardings;
    *(void**)  (node + 0x20) = PJRT_Shardings_..._OutputShardings;
    return node;

陷阱 — type 是位于 +0x08 的 32-bit 字段，不是 64-bit word。creator 会发出 mov dword ptr [node+8], <id>（例如 *(_DWORD*)(a1 + 8) = 23），将 +0x0C 留作 4-byte 空洞，由零初始化的 .bss/.data 后备存储保持为零。如果重新实现将 type 声明为 size_t，读取端会把 next 错放 4 bytes 并遍历到垃圾。应将 type 读为 uint32，然后跳过 4 bytes pad，再在 +0x10 读取 next。
特性 — 这些 creator 是纯表初始化器：一段扁平的 mov store 后接 ret，没有分配，也没有分支。这正是 16 个动态节点可以位于零初始化 .bss 的全部原因：没有东西需要构造，只需要填充。一次性的 __cxa_guard 是唯一同步；首次调用之后，这些节点在进程生命周期内不可变。

Fn-ptr 尾部与 Args-Struct 约定

从 +0x18 开始的槽位是方法指针。struct_size 是包含该尾部在内的完整节点大小，因此 (struct_size - 0x18) / 8 给出方法数量（例如 HostMemoryAllocator (32 - 24)/8 = 1；Shardings (40 - 24)/8 = 2；Megascale (248 - 24)/8 = 28，其中 23 个 live、5 个 reserved NULL，见 Remaining Extensions）。每个方法都接受单个 PJRT_<API>_Args*；与主 vtable 一样，每个方法体的第一个动作都是通过 pjrt::ActualStructSizeIsGreaterOrEqual("<API>_Args", min, current, args->struct_size) @ 0xF8A4EC0 执行向后兼容的 size check。包装句柄 args 将 opaque handle 放在 +0x08（没有 priv 字段）。

`PJRT_Extension_Type` Enum

type id 是发现 key。libtpu 的 id 与公开 xla/pjrt/c/pjrt_c_api.h v0.103 系列完全匹配；下表是该二进制中体现出的完整 enum，present/absent 状态来自 creator 的 type store 和 canonical header。

Type id	名称	libtpu 中的状态	记录于
0	（unused/reserved）	absent	—
1	Profiler	present（.data seed，type-1 terminator）	Profiler
2	Custom_Partitioner	absent	—（职责被吸收；见注释）
3	Stream	absent	—
4	Layouts	present	Remaining Extensions
5	FFI	absent	—（职责被吸收；见注释）
6	MemoryDescriptions	present	Remaining Extensions
7	（reserved / 某些源码树中的 Memory_Stream）	absent	—
8	RawBuffer	present	RawBuffer
9	PhaseCompile	present	PhaseCompile
10	（unused）	absent	—
11	（unused）	absent	—
12	CrossHostTransfers	present	Remaining Extensions
13	ExecutableMetadata	present（TPU factory）	Remaining Extensions
14	Callback	present（TPU slice-builder）	Remaining Extensions
15	HostAllocator	present（TPU）	Remaining Extensions
16	TopologyDescription / TpuTopology	present（TPU）	Topology Description
17	TpuExecutable	present（TPU）	Remaining Extensions
18	Megascale	present（TPU）	Remaining Extensions
19	Shardings	present	Remaining Extensions
20	AbiVersion	present	Remaining Extensions
21	Collectives	present	Remaining Extensions
22	MultiSlice	present（TPU）	Remaining Extensions
23	HostMemoryAllocator	present	Remaining Extensions

注释 — 这些缺失项对重新实现者具有决定意义。canonical PJRT C-API 注册了 FFI（5）、Custom_Partitioner（2）、Stream（3）以及 libtpu 未公告的 Triton id。libtpu 通过两条 TPU 专用通道交付 FFI / custom-call 扩展本会提供的内容：Callback 扩展（type 14）的 RegisterCallback 路径安装 xla::SliceBuilderCallbackState host callback，TpuExecutable 扩展（type 17）的 SetTpuCompilationEnv 路径安装序列化的 xla::CompilationEnvironmentsProto（TPU compilation-backend/XLA-flags 表面）。ExecuteContext 只通过主表槽位 103/104 暴露，而不是作为扩展暴露。要求存在 FFI 扩展的框架会在 TPU 上检测到 NULL，并且必须回退到这些通道。
陷阱 — 有两个名称几乎相同的 host-memory allocator。HostMemoryAllocator（type 23，32 bytes，一个 Allocate 方法）是链头上的通用 XLA host-staging allocator。HostAllocator（type 15，48 bytes，三个方法，均为 TPU 注入）是用于 device-host DMA staging 的 TPU pinned-host allocator。二者布局不同，服务于不同层；匹配错误的 type id 会得到错误的 allocator。

扩展清单（遍历顺序）

完整链表如下，按消费者从 extension_start 沿 next 遍历时遇到的顺序排列。遍历顺序是最新优先（构建顺序的反向）。所有存储 VA、大小和 type id 都已通过 creator 的 type/struct_size store 和二进制符号表确认。

Walk	Storage VA	Extension (type id)	Size	Methods	Creator @	深入解析
1	`0x224C3F68`	HostMemoryAllocator (23)	32	1	`CreateHostMemoryAllocatorExtension` @ `0xE6F5340`	Remaining
2	`0x224C3F20`	MultiSlice (22)	64	5	`CreateMultiSliceExtension` @ `0xE6F3C40`	Remaining
3	`0x224C3EB8`	Collectives (21)	96	9	`CreateCollectivesExtension` @ `0xE6F19A0`	Remaining
4	`0x224C3E38`	AbiVersion (20)	120	12	`CreateTpuAbiVersionExtension` @ `0xE6B7340` → `CreateAbiVersionExtension` @ `0xE6B8960`	Remaining
5	`0x224C3E08`	Shardings (19)	40	2	`CreateShardingsExtension` @ `0xF874980`	Remaining
6	`0x224C3D08`	Megascale (18)	248	23 (+5 reserved)	`CreateMegascaleExtension` @ `0xE6B97C0`	Remaining
7	`0x224C3CA8`	TpuExecutable (17)	88	7 (+1 reserved)	`CreateTpuExecutableExtension` @ `0xE6DC6E0`	Remaining
8	`0x224C3B90`	TopologyDescription / TpuTopology (16)	272	31	`CreateTpuTopologyExtension` @ `0xE6DE5E0`	Topology
9	`0x224C3B60`	Callback (14)	40	2	`CreateCallbackExtension` @ `0xE6B91E0`	Remaining
10	`0x224C3B18`	PhaseCompile (9)	64	5	`CreatePhaseCompileExtension` @ `0xE6F42A0`	PhaseCompile
11	`0x224C3AD8`	CrossHostTransfers (12)	56	4	`CreateCrossHostTransfersExtension` @ `0xF85D660`	Remaining
12	`0x224C3AA0`	HostAllocator (15)	48	3	`CreateHostAllocatorExtension` @ `0xF8A3C20`	Remaining
13	`0x224C3A70`	ExecutableMetadata (13)	40	2	`CreateExecutableMetadataExtension` @ `0xF8A3BE0`	Remaining
14	`0x224C3A40`	MemoryDescriptions (6)	40	2	`CreateMemoryDescriptionsExtension` @ `0xF874940`	Remaining
15	`0x224C39E8`	Layouts (4)	80	7	`CreateLayoutsExtension` @ `0xF8748C0`	Remaining
16	`0x224C3990`	RawBuffer (8)	80	7	`CreateRawBufferExtension` @ `0xE6F52C0`	RawBuffer
17	`0x22255B98`	Profiler (1)	40	0 (1 ptr → 8-slot vtable)	static `.data` init (no `Create*` call)	Profiler

总计 = 17 个节点。Profiler（type 1）是终止节点（next = NULL）。五个 creator 以参数形式接收 TPU 函数指针的 TPU 专用节点：AbiVersion（两个 FromProto factory）、ExecutableMetadata（GetExecutableMetadata）、HostAllocator（全部三个槽位）、PhaseCompile（Get_Compiler/Destroy_Compiler），再加上完全 TPU 化的 Megascale/TpuExecutable/MultiSlice/Callback/TpuTopology，均在各自深入页面中详述；本页只关注它们在链中的位置。

链如何构建

GetTpuPjrtApi @ 0xE6AA440 是一段直线序列：16 个受 __cxa_guard 保护的 Create*Extension 调用，后接第 17 个围绕 CreatePjrtApi 的 guard。每个 Create* 调用都接收节点自身地址和此前构建节点的地址作为 next，因此链表按尾部优先生长。

// pjrt::tpu_plugin::GetTpuPjrtApi @ 0xE6AA440  (one-shot, __cxa_guard per node)
function GetTpuPjrtApi():
    // Built first; its next is the .data Profiler seed (type 1).
    once: CreateRawBufferExtension(&raw_buffer_extn,  &profiler_extension);      // -> type 8
    once: CreateLayoutsExtension(&layouts_extn,       &raw_buffer_extn);         // -> type 4
    once: CreateMemoryDescriptionsExtension(&mem_desc_extn, &layouts_extn);      // -> type 6
    once: CreateExecutableMetadataExtension(&exec_meta_extn, &mem_desc_extn,
                                            GetTpuExecutableMetadata);           // -> type 13 (TPU fn)
    once: CreateHostAllocatorExtension(&host_alloc_extn, &exec_meta_extn,
                                       TPU_HostAllocator_GetPreferredAlignment,
                                       TPU_HostAllocator_Allocate,
                                       TPU_HostAllocator_Free);                  // -> type 15 (TPU fns)
    once: CreateCrossHostTransfersExtension(&xfer_extn,  &host_alloc_extn);      // -> type 12
    once: CreatePhaseCompileExtension(&phase_extn, &xfer_extn,
                                      GetTpuPhaseCompiler, DestroyTpuPhaseCompiler); // -> type 9 (TPU fns)
    once: CreateCallbackExtension(&callback_extn,  &phase_extn);                 // -> type 14
    once: CreateTpuTopologyExtension(&topo_extn,   &callback_extn);             // -> type 16
    once: CreateTpuExecutableExtension(&tpu_exec_extn, &topo_extn);             // -> type 17
    once: CreateMegascaleExtension(&megascale_extn, &tpu_exec_extn);            // -> type 18
    once: CreateShardingsExtension(&shardings_extn, &megascale_extn);           // -> type 19
    once: CreateTpuAbiVersionExtension(&abi_extn,   &shardings_extn);           // -> type 20 (thunk)
    once: CreateCollectivesExtension(&collectives_extn, &abi_extn);            // -> type 21
    once: CreateMultiSliceExtension(&multi_slice_extn,  &collectives_extn);     // -> type 22
    // Built last -> becomes the chain head (extension_start).
    once: CreateHostMemoryAllocatorExtension(&host_mem_alloc_extn, &multi_slice_extn); // -> type 23

    once: CreatePjrtApi(&pjrt_api,
                        PJRT_Client_Create, PJRT_ExecuteContext_Create,
                        PJRT_TopologyDescription_Create, PJRT_Plugin_Initialize,
                        &host_mem_alloc_extn,            // -> PJRT_Api.extension_start
                        PJRT_Plugin_Attributes_Xla);
    return &pjrt_api;                                    // 0x227BA840
```text

> **特性 —** Profiler（type 1）的构建方式不同于其他 16 个节点。它**不是**在 `GetTpuPjrtApi` 内通过 `Create*Extension` 调用创建的；它是 `.data` 中的驻留节点 `pjrt::tpu_plugin::profiler_extension` @ `0x22255B98`，在加载时由静态初始化器 / `R_X86_64_RELATIVE` 重定位填充。`GetTpuPjrtApi` 只在传给 `CreateRawBufferExtension`（第一个构建的节点）的 `next` 参数中引用它。因此 Profiler 是整个 `.bss` 链追加到其前方的种子，而链在它处*终止*，因为它自己的 `next` 是 `NULL`。

### 构建顺序与遍历顺序

```text
construction (build) order            walk order from extension_start
(GetTpuPjrtApi top -> bottom)         (consumer follows .next)
  RawBuffer(8)   --- built 1st          HostMemoryAllocator(23) <- head (built last)
  Layouts(4)                            MultiSlice(22)
  MemoryDescriptions(6)                 Collectives(21)
  ExecutableMetadata(13)                AbiVersion(20)
  HostAllocator(15)                     Shardings(19)
  CrossHostTransfers(12)                Megascale(18)
  PhaseCompile(9)                       TpuExecutable(17)
  Callback(14)                          TpuTopology(16)
  TpuTopology(16)                       Callback(14)
  TpuExecutable(17)                     PhaseCompile(9)
  Megascale(18)                         CrossHostTransfers(12)
  Shardings(19)                         HostAllocator(15)
  AbiVersion(20)                        ExecutableMetadata(13)
  Collectives(21)                       MemoryDescriptions(6)
  MultiSlice(22)                        Layouts(4)
  HostMemoryAllocator(23) - built last  RawBuffer(8)
        (profiler is the pre-existing seed)  Profiler(1) <- terminator (.next = NULL)

遍历顺序正好是构建顺序的反向，因为每个节点都会把自身 prepend 到链表中：node.next = previously_built，且 extension_start = last_built。该顺序纯粹是构建序列的产物，消费者不得依赖它。

框架如何遍历链

发现过程是 type-id 线性扫描，与 canonical xla/pjrt/c pjrt_c_api_helpers 的 FindExtension 相同。想要 Profiler 扩展的框架会这样做：

// Consumer-side feature detection. `wanted` is a PJRT_Extension_Type id.
function FindExtension(const PJRT_Api* api, PJRT_Extension_Type wanted):
    PJRT_Extension_Base* ext = api->extension_start;     // +0x08
    while (ext != NULL):
        if (ext->type == wanted):                        // +0x08, uint32 compare
            return ext;                                  // caller casts to the typed struct
        ext = ext->next;                                 // +0x10
    return NULL;                                          // feature not advertised
```text

该循环建立的契约：

- **缺失是有效答案。** 返回 `NULL` 表示插件没有公告该能力。框架必须优雅降级，例如在 TPU 上缺失 FFI 扩展（type 5）是预期情况，框架会回退到上文描述的 Callback/TpuExecutable 通道。
- **`struct_size` 约束每个方法的前向兼容性，而不是链成员资格。** 找到节点足以说明该*功能*存在；某个*具体方法*是否存在由节点的 `struct_size`（尾部是否覆盖该槽位）决定，并且在调用时由该方法自己的 `ActualStructSizeIsGreaterOrEqual` 对 `args->struct_size` 检查决定。基于较新 header 构建的消费者在解引用所需槽位前，必须验证 `ext->struct_size` 覆盖该槽位。
- **顺序无关。** 遍历不得假设任何节点位于特定位置。libtpu 的最新优先顺序是偶然的；不同插件（或未来 libtpu）可能以不同顺序排列，或增删节点。
- **终止。** 循环在 `NULL` 处终止。libtpu 的终止点是 Profiler 的 `next`；忘记将最后节点 NULL-terminate 的重新实现会遍历到相邻 `.bss`/`.data` 并读到垃圾 `type` 值。

> **注释 —** `PJRT_Api` 结构本身（`extension_start` 悬挂其上的 140 槽 vtable）是 `.lbss` 中位于 `0x227BA840` 的独立对象，由 `pjrt::CreatePjrtApi` @ `0xF874160` 填充。扩展链和 vtable 在同一个一次性 `GetTpuPjrtApi` init 中填充，但它们是不同结构；vtable 的逐槽重建见 [API Vtable Reconstruction](api-vtable-reconstruction.md)。

---

PJRT 扩展链 ​

摘要 ​

节点布局 ​

公共头 ​

Fn-ptr 尾部与 Args-Struct 约定 ​

PJRT_Extension_Type Enum ​

扩展清单（遍历顺序） ​

链如何构建 ​

框架如何遍历链 ​