PJRT 剩余扩展

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

摘要

libtpu 在 PJRT_Api.extension_start 上挂接了 17 个扩展节点。其中四个有专门页面：Profiler（type 1）、TopologyDescription/TpuTopology（type 16）、RawBuffer（type 8）和 PhaseCompile（type 9）。本页深入讨论另外十三个：每个没有自己页面的扩展的函数指针集合、结构体大小、类型 id 和行为。节点头布局、完整的 PJRT_Extension_Type 枚举、构造顺序与遍历顺序，以及消费者遍历循环由扩展链负责，这里不再重复；本页假设读者已经掌握该结构，直接说明每个节点暴露了什么。

这十三个扩展可以沿着重实现者关心的一个维度清晰划分：谁提供函数指针。其中八个是完全通用的 XLA 表面，其创建器内置 pjrt:: 命名空间包装器（Layouts、MemoryDescriptions、CrossHostTransfers、Shardings、Collectives、HostMemoryAllocator，以及 AbiVersion 的通用槽位）。三个是 TPU 专用的：每个方法都是 tpu_plugin:: 或匿名命名空间中的 TPU 实现（Megascale、TpuExecutable、MultiSlice、HostAllocator、Callback）。另外两个是混合型，创建器以参数形式接收 TPU 函数指针，并以通用方式填充其余槽位（ExecutableMetadata、AbiVersion 的两个 FromProto 工厂；PhaseCompile 是第三个混合型，但记录在自己的页面中）。每个创建器都是一个扁平表初始化器：一串 mov 存储以 ret 结束，没有分配，没有分支。因此整个节点集合都可以通过读取存储序列恢复出来，下面的大小和槽位布局正是这样获得的。

本页先按该所有权维度对十三个扩展分组，然后逐一编目：速览事实、槽位表（偏移 → 方法 → 实现符号 → 地址），以及对已反编译方法主体的 ### 行为 注释。重实现时有两个反复出现的主题：每个方法的第一个动作都是向后兼容的 ActualStructSizeIsGreaterOrEqual 检查（(min, current) 字面量对就是该方法的线兼容契约），而通用方法都遵循一次 vtable 跳转模式：在 args +0x08/+0x10 解包不透明句柄，调用 C++ 抽象基类 vtable 的一个槽位，再把结果封送回 args 结构体。

对重实现而言，契约是：

十三个节点布局：struct_size、type，以及从 +0x18 开始的有序 fn-ptr 尾部，每个槽位的实现符号和虚拟地址。
每个节点的 TPU 注入与通用划分，以及哪些槽位接收注入指针（这样重实现者就知道哪些要由插件后端提供，哪些由库提供）。
两个有意的缺席项：FFI（type 5）和 Stream（type 3）不会被公开；它们的角色由 Callback 和 TpuExecutable 取代。消费者对它们做特性检测会得到 NULL。
两个几乎相同的 host allocator（type 23 与 type 15），以及为什么匹配错误 id 是 bug。
已反编译主体的逐方法 ActualStructSizeIsGreaterOrEqual(name, min, current, args->struct_size) 契约。


本页覆盖	13 个扩展：Layouts (4)、MemoryDescriptions (6)、CrossHostTransfers (12)、ExecutableMetadata (13)、Callback (14)、HostAllocator (15)、TpuExecutable (17)、Megascale (18)、Shardings (19)、AbiVersion (20)、Collectives (21)、MultiSlice (22)、HostMemoryAllocator (23)
其他专页	Profiler (1) → ext-profiler；RawBuffer (8) → ext-rawbuffer；PhaseCompile (9) → ext-compile-phasecompile；TopologyDescription (16) → ext-topology-description
公共头部	`{ size_t struct_size; uint32 type; uint32 _pad; PJRT_Extension_Base* next; }`，fn-ptrs 从 `+0x18` 开始 — 见扩展链
构建器	`pjrt::tpu_plugin::GetTpuPjrtApi` @ `0xE6AA440`（每个节点一次 `__cxa_guard` + `Create*Extension` 调用）
逐方法兼容门控	`pjrt::ActualStructSizeIsGreaterOrEqual` @ `0xF8A4EC0`
此处最大	Megascale (18)，248 字节，23 个有效槽位 + 5 个保留 NULL 槽位
此处最小	HostMemoryAllocator (23)，32 字节，1 个方法
有意缺席	FFI (5)、Stream (3)、Custom_Partitioner (2)、Triton — 角色由 Callback (14) + TpuExecutable (17) 取代

注意 — “剩余扩展”是目录，不是算法，因此本页侧重槽位表语法（Offset | Method | Impl symbol | Addr | Confidence），而不是逐单元伪代码。两个已逐字节跟踪的方法——HostMemoryAllocator_Allocate 和 Layouts::Client_GetDefaultLayout——带有 ### Algorithm 块，作为其他每个通用方法所遵循的一次 vtable 跳转模式的规范示例。

关于两个有意缺席项

任务标题在剩余扩展中提到了 “Stream” 和 “FFI”。二者在此构建中都不存在。创建器只写入 type id 4, 6, 8, 9, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23（再加上 .data Profiler，type 1）；id 2（Custom_Partitioner）、3（Stream）、5（FFI）、7、10、11 以及任何 Triton id 都从未被存储。这一点由没有任何 Create*Extension 写入它们，以及链页面上的枚举表共同确认。这里将它们记录为带替代品的缺席项，因为进行特性检测的重实现者必须知道 libtpu 改用什么：

规范扩展	Type id	libtpu 状态	替代提供角色者
FFI / custom-call host callbacks	5	缺席	Callback（type 14）`RegisterCallback` — 安装 `xla::SliceBuilderCallbackState` host callbacks
FFI / compilation-env injection	5	缺席	TpuExecutable（type 17）`SetTpuCompilationEnv` — 安装 `xla::CompilationEnvironmentsProto`
Stream（原始 device-stream 句柄）	3	缺席	不公开；device-async 通过主表上的 `PJRT_Event` 呈现
Custom_Partitioner	2	缺席	不公开；SPMD partitioning 通过 Shardings（type 19）查询
ExecuteContext	（主表）	不是扩展	主表槽位 103/104（`PJRT_ExecuteContext_Create/Destroy`）

易错点 — 如果框架要求 FFI 扩展（type 5）来注册 custom call，它会在 TPU 上特性检测到 NULL，并且必须回退到 Callback / TpuExecutable 通道。扩展链上的遍历循环会对缺席 id 返回 NULL；缺席是有效答案，不是错误。

第 1 组 — 通用 XLA 表面

八个节点的创建器只内置 pjrt:: 命名空间包装器，没有 TPU 注入指针。它们包装由 TPU 支撑的 C++ 类，但线表面是规范 XLA 表面。八者都遵循一次 vtable 跳转模板：解包句柄，调用一个 vtable 槽位，再封送返回。

Layouts — type 4，80 字节

规范的 PJRT_Layouts 表面：不透明 PjRtLayout 句柄、序列化、来自 client 或 topology 的默认 layout 查询，以及逐 buffer 和逐 executable 的 layout 获取。创建器 pjrt::CreateLayoutsExtension @ 0xF8748C0。


存储 VA	`0x224C39E8`
struct_size	80（`(80-24)/8 = 7` 个方法）
TPU 注入	无 — 通用，由 TPU layout-assignment 支撑

偏移	方法	实现符号	地址
`+0x18`	`MemoryLayout_Destroy`	`pjrt::PJRT_Layouts_MemoryLayout_Destroy`	`0xF871360`
`+0x20`	`MemoryLayout_Serialize`	`pjrt::PJRT_Layouts_MemoryLayout_Serialize`	`0xF871400`
`+0x28`	`Client_GetDefaultLayout`	`pjrt::PJRT_Layouts_PJRT_Client_GetDefaultLayout`	`0xF8714A0`
`+0x30`	`Buffer_MemoryLayout`	`pjrt::PJRT_Layouts_PJRT_Buffer_MemoryLayout`	`0xF871620`
`+0x38`	`Topology_GetDefaultLayout`	`pjrt::PJRT_Layouts_PJRT_Topology_GetDefaultLayout`	`0xF8716A0`
`+0x40`	`Executable_GetOutputLayouts`	`pjrt::PJRT_Layouts_PJRT_Executable_GetOutputLayouts`	`0xF871C40`
`+0x48`	`Executable_GetParameterLayouts`	`pjrt::PJRT_Layouts_PJRT_Executable_GetParameterLayouts`	`0xF871820`

算法

Client_GetDefaultLayout @ 0xF8714A0 是跳转模式的展开示例。args 布局：client 句柄在 +0x10，元素类型在 +0x18，dims 指针在 +0x20，dims 数量在 +0x28，输出句柄写入 +0x30。

function Layouts_Client_GetDefaultLayout(args):                 // 0xF8714A0
    // (min 0x2E=46, cur 0x38=56) — wire-compat literal pair
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Layouts_PJRT_Client_GetDefaultLayout_Args",
            46, 56, args->struct_size):
        return new PJRT_Error(status);                          // size too small
    client    = *(void**)(args + 0x10);                         // unwrap PJRT_Client
    elem_type = ConvertFromPjRtBufferType(*(u32*)(args + 0x18)); // 0xF8A3E60: C enum -> xla::PrimitiveType
    // client vtable +0x98 == GetDefaultLayout(elem_type, dims_ptr, dims_count)
    status = (*client->vtable[0x98])(client, elem_type,
                 *(void**)(args + 0x20), *(u64*)(args + 0x28));  // -> StatusOr<Layout>
    if status.ok():
        layout = new xla::Layout(...);                          // 0xF0-byte PjRtLayout wrapper
        *(void**)(args + 0x30) = box(layout);                   // opaque handle, 2-qword box
    return status_to_PJRT_Error(status);

注意 — client vtable 槽位 +0x98 是 TPU PjRtClient::GetDefaultLayout。整个 Layouts 扩展从不直接看到 TPU 代码；每次调用都通过实时 client/buffer/executable/topology 句柄的 vtable 跳转。返回的 PjRtLayout 是一个堆对象（new(0xF0)），在 +0x18 持有一个 xla::Layout；调用者通过 MemoryLayout_Destroy（+0x18）释放它。

MemoryDescriptions — type 6，40 字节

枚举附加到 PJRT_DeviceDescription 的 memory-space 描述：在 topology-query / pre-compile 时可用、任何实时 PJRT_Memory 对象存在之前的静态内存分类（"device" / "pinned_host" / "unpinned_host" kinds）。它不同于 Buffer and Memory 上的实时 PJRT_Memory 表面。创建器 pjrt::CreateMemoryDescriptionsExtension @ 0xF874940。


存储 VA	`0x224C3A40`
struct_size	40（2 个方法）
TPU 注入	无 — 通用

偏移	方法	实现符号	地址
`+0x18`	`DeviceDescription_MemoryDescriptions`	`pjrt::PJRT_DeviceDescription_MemoryDescriptions`	`0xF865580`
`+0x20`	`MemoryDescription_Kind`	`pjrt::PJRT_MemoryDescription_Kind`	`0xF865920`

DeviceDescription_MemoryDescriptions 返回设备描述的 opaque PJRT_MemoryDescription* 数组；MemoryDescription_Kind 返回某个描述的 kind 字符串。配对方式是标准的“先列出再查询”惯用法。

CrossHostTransfers — type 12，56 字节

跨主机（DCN）buffer 传输表面：带 descriptor 握手的 receive-buffer 分配、点对点 send/receive，以及 copy-to-remote-device。用于 pipeline-parallel 和跨 pod buffer 移动；规范的 PjRtCrossHostTransfers，由 TPU 支撑。创建器 pjrt::CreateCrossHostTransfersExtension @ 0xF85D660。配套的主表 DMA 路径见 DMA and Cross-Host Recv。


存储 VA	`0x224C3AD8`
struct_size	56（4 个方法）
TPU 注入	无 — 通用

偏移	方法	实现符号	地址
`+0x18`	`Client_MakeCrossHostReceiveBuffers`	`pjrt::PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers`	`0xF85C9A0`
`+0x20`	`Buffer_CopyToRemoteDevice`	`pjrt::PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice`	`0xF85CE20`
`+0x28`	`Client_CrossHostReceiveBuffers`	`pjrt::PJRT_Transfers_PJRT_Client_CrossHostReceiveBuffers`	`0xF85BBA0`
`+0x30`	`Client_CrossHostSendBuffers`	`pjrt::PJRT_Transfers_PJRT_Client_CrossHostSendBuffers`	`0xF85C2A0`

注意 — MakeCrossHostReceiveBuffers（返回 descriptor）和 CrossHostReceiveBuffers（调用者提供 descriptor）是两种不同的接收惯用法：前者将 descriptor 带回给 sender 以便带外传递；后者消费已经交换好的 descriptor。重实现必须同时连通两者。

Shardings — type 19，40 字节

公开已编译 executable 的逐参数和逐输出 xla::OpSharding / HloSharding，以 OpSharding proto 序列化，用于 SPMD partitioning 查询。创建器 pjrt::CreateShardingsExtension @ 0xF874980。


存储 VA	`0x224C3E08`
struct_size	40（2 个方法）
TPU 注入	无 — 通用

偏移	方法	实现符号	地址
`+0x18`	`PJRT_Executable_ParameterShardings`	`pjrt::PJRT_Shardings_PJRT_Executable_ParameterShardings`	`0xF868000`
`+0x20`	`PJRT_Executable_OutputShardings`	`pjrt::PJRT_Shardings_PJRT_Executable_OutputShardings`	`0xF868A60`

这是缺席的 Custom_Partitioner（type 2）的替代品：消费者不能注册 custom partitioner，但可以读回编译器选择的 shardings。

Collectives — type 21，96 字节 {#collectives--type-21-96-bytes}

规范的进程内 XLA collectives 表面：communicator 生命周期以及七个 collective 原语。由 CPU executor 支撑且为通用实现；megascale（跨 pod）collectives 位于 Megascale 扩展中，而不是这里。创建器 pjrt::CreateCollectivesExtension @ 0xE6F19A0。主表 communicator 表面见 Collectives Communicator。


存储 VA	`0x224C3EB8`
struct_size	96（9 个方法）
TPU 注入	无 — 通用，由 CPU executor 支撑

偏移	方法	实现符号	地址	min/cur
`+0x18`	`Collectives_Destroy`	`(anon)::CollectivesDestroy`	`0xE6F1A20`	0x1D/0x10
`+0x20`	`CreateCommunicators`	`(anon)::CollectivesCreateCommunicators`	`0xE6F1AA0`	0x29/0x68
`+0x28`	`Communicator_Destroy`	`(anon)::CommunicatorDestroy`	`0xE6F21A0`	—
`+0x30`	`Communicator_AllReduce`	`(anon)::CommunicatorAllReduce`	`0xE6F2220`	0x2C/0x58
`+0x38`	`Communicator_ReduceScatter`	`(anon)::CommunicatorReduceScatter`	`0xE6F2440`	—
`+0x40`	`Communicator_AllGather`	`(anon)::CommunicatorAllGather`	`0xE6F2660`	—
`+0x48`	`Communicator_CollectivePermute`	`(anon)::CommunicatorCollectivePermute`	`0xE6F2880`	—
`+0x50`	`Communicator_AllToAll`	`(anon)::CommunicatorAllToAll`	`0xE6F2CC0`	—
`+0x58`	`Communicator_ToString`	`(anon)::CommunicatorToString`	`0xE6F3280`	—

行为

CreateCommunicators（min 0x29，cur 0x68）把调用者的 int32 数组（数据在 args +0x10，数量在 +0x18）复制到一个新 new 出来的 buffer 中，然后构建 N 个 communicators。Communicator_AllReduce（min 0x2C，cur 0x58）读取 args +0x10/+0x20 处的输入 buffer descriptor（device-ptr + count，被打包成两个用 xmm 加载的 16 字节字段）、+0x48 处的 executor、+0x30 处的元素类型（通过 ConvertFromPjRtBufferType @ 0xF8A3E60）以及 +0x40 处的 reduction op；它分配一个 PJRT_Collectives_CpuExecutor（CreateCpuExecutor @ 0xE6F3740），并通过 communicator vtable 槽位 +0x30 跳转，返回一个 0x50 字节的 PJRT_Event。

注意 — 只有 AllReduce vtable 索引（+0x30）经过逐字节确认；ReduceScatter/AllGather/CollectivePermute/AllToAll 的 vtable 偏移是根据声明顺序推断的（MEDIUM confidence）。CpuExecutor 支撑使它成为通用实现而非 TPU 专用实现。

HostMemoryAllocator — type 23，32 字节 {#hostmemoryallocator--type-23-32-bytes}

在链头（遍历位置 1）公开的通用 XLA host-staging buffer allocator。一个方法。创建器 pjrt::CreateHostMemoryAllocatorExtension @ 0xE6F5340。


存储 VA	`0x224C3F68`（链头）
struct_size	32（1 个方法）
TPU 注入	无 — 通用

偏移	方法	实现符号	地址	min/cur
`+0x18`	`Allocate`	`(anon)::HostMemoryAllocator_Allocate`	`0xE6F5380`	0x26/0x40

算法

Allocate @ 0xE6F5380 是带错误路径的跳转惯用法的规范展开示例。args 布局：memory-space/client 句柄在 +0x10，size 在 +0x18，alignment（int）在 +0x20，输出 {data, deleter} 对从 +0x28 写入。

function HostMemoryAllocator_Allocate(args):                    // 0xE6F5380
    // (min 0x26=38, cur 0x40=64)
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_HostMemoryAllocator_Allocate_Args",
            38, 64, args->struct_size):
        return new PJRT_Error(status);
    client = *(void**)(args + 0x10);
    if client == NULL:
        return new PJRT_Error(InvalidArgument(                  // MakeErrorImpl<3>
            "Received null client in HostMemoryAllocator_Allocate"));
    // client vtable +0x108 -> the underlying host allocator (may be absent)
    allocator = (*client->vtable[0x108])(client);
    if allocator == NULL:
        return new PJRT_Error(Unimplemented(                    // MakeErrorImpl<12>
            "HostMemoryAllocator not implemented for client"));
    // allocator vtable +0x10 -> Allocate(size, &alignment)
    owned = (*allocator->vtable[0x10])(allocator,
                *(u64*)(args + 0x18), &(int){ *(u32*)(args+0x20) });
    if owned.data != NULL:
        *(pair*)(args + 0x28) = owned;                          // {void* data; deleter} 16-byte store
        *(void**)(args + 0x38) = owned.deleter;
    else:
        *(pair*)(args + 0x28) = {0, 0};                         // empty: zeroed store
    return OK;                                                  // returns 0 (no error)

易错点 — 这是通用 host allocator（type 23）。不要把它和 HostAllocator（type 15，三个 TPU 注入方法）混淆。家族名称相同，节点不同，布局不同，层次不同：type 23 是通过实时 client vtable 实现的 XLA host staging；type 15 是带插件供应函数指针的 TPU pinned-host DMA staging。

AbiVersion（通用部分）— type 20，120 字节 {#abiversion-generic-portion--type-20-120-bytes}

运行时 ABI 与序列化 executable 之间的跨版本兼容性检查：保存的 executable 能否针对当前 runtime 加载？十二个槽位中有十个是通用 pjrt:: 命名空间包装器；两个 FromProto 工厂是 TPU 注入的（见第 3 组）。创建器 pjrt::CreateAbiVersionExtension(node, runtime_fn, executable_fn, next) @ 0xE6B8960，通过 TPU thunk CreateTpuAbiVersionExtension @ 0xE6B7340 调用。


存储 VA	`0x224C3E38`
struct_size	120（12 个方法）
TPU 注入	2 个槽位（`+0x68`、`+0x70`）— 见第 3 组

偏移	方法	实现符号	地址
`+0x18`	`Client_RuntimeAbiVersion`	`(anon)::ClientRuntimeAbiVersion`	`0xE6B8A00`
`+0x20`	`Executable_GetAbiVersion`	`(anon)::ExecutableGetAbiVersion`	`0xE6B8AE0`
`+0x28`	`RuntimeAbiVersion_Destroy`	`(anon)::RuntimeAbiVersionDestroy`	`0xE6B8BC0`
`+0x30`	`RuntimeAbiVersion_IsCompatibleWithRuntime`	`(anon)::RuntimeAbiVersionIsCompatibleWithRuntime`	`0xE6B8C40`
`+0x38`	`RuntimeAbiVersion_IsCompatibleWithExecutable`	`(anon)::RuntimeAbiVersionIsCompatibleWithExecutable`	`0xE6B8CA0`
`+0x40`	`RuntimeAbiVersion_ToProto`	`(anon)::RuntimeAbiVersionToProto`	`0xE6B8D00`
`+0x48`	`RuntimeAbiVersion_PlatformId`	`(anon)::RuntimeAbiVersionPlatformId`	`0xE6B8EA0`
`+0x50`	`ExecutableAbiVersion_Destroy`	`(anon)::ExecutableAbiVersionDestroy`	`0xE6B8F00`
`+0x58`	`ExecutableAbiVersion_ToProto`	`(anon)::ExecutableAbiVersionToProto`	`0xE6B8F80`
`+0x60`	`ExecutableAbiVersion_PlatformId`	`(anon)::ExecutableAbiVersionPlatformId`	`0xE6B9120`
`+0x68`	`RuntimeAbiVersion_FromProto` (TPU)	`(anon)::TpuRuntimeAbiVersionFromProto`	`0xE6B7360`
`+0x70`	`ExecutableAbiVersion_FromProto` (TPU)	`(anon)::TpuExecutableAbiVersionFromProto`	`0xE6B7380`

怪癖 — 创建器从其 a4 参数写入 next，不是 a2：CreateAbiVersionExtension(node, runtime_fn, executable_fn, next) 将 runtime_fn 存到 +0x68，将 executable_fn 存到 +0x70，将 next 存到 +0x10。thunk CreateTpuAbiVersionExtension(node, next) 只是用两个 TPU FromProto 函数接好后尾调用它。如果重实现假设 next 总是第二个创建器参数，就会把这个节点链接错。

第 2 组 — 完全 TPU 专用表面

五个节点（Megascale、TpuExecutable、MultiSlice、Callback，加上 HostAllocator）的每个槽位都是 TPU 实现：要么是创建器中内置的匿名命名空间 TPU 函数，要么（对 HostAllocator 而言）是作为创建器参数注入的 tpu_plugin:: 指针。重实现必须为所有这些提供后端。

Megascale — type 18，248 字节 {#megascale--type-18-248-bytes}

这里最大的扩展，也是多主机 JAX 训练使用的 multi-pod 数据中心网络（DCN）运行时控制表面：client-context 生命周期、提前（AoT）配置、multi-slice 配置（MultiSlice 的配套项）、megascale collectives 工厂、device↔megascale id 转换，以及用于容错训练的完整异步错误聚合子系统。创建器 pjrt::CreateMegascaleExtension @ 0xE6B97C0。完全 TPU/megascale 专用。


存储 VA	`0x224C3D08`
struct_size	248（`(248-24)/8 = 28` 个槽位：23 个有效 + 5 个保留 NULL）
保留槽位	`+0x40`..`+0x60` — 构造时清零（`vmovups ymm0` 覆盖 `+0x40..+0x5F`，然后在 `+0x60` 执行 `movq $0`）
TPU 注入	所有 23 个有效槽位都是匿名命名空间 TPU 实现

text

+0x18  CreateClientContextFromPjRtClient   0xE6B9920    +0x88  DeviceId_To_MegascaleId            0xE6BA8E0
+0x20  CreateDefaultClientContext          0xE6B9A20    +0x90  MegascaleId_To_DeviceId            0xE6BA980
+0x28  DeleteClientContext                 0xE6B9B20    +0x98  RegisterMegascaleErrorHandler      0xE6BAA60
+0x30  CreateAoTConfig                     0xE6B9BC0    +0xA0  UnregisterMegascaleErrorHandler    0xE6BAB20
+0x38  CreateMultiSliceConfig              0xE6B9CA0    +0xA8  ErrorAggregator_Create             0xE6BAB80
+0x40  (reserved, NULL)                    —            +0xB0  ErrorAggregator_Delete             0xE6BAC00
+0x48  (reserved, NULL)                    —            +0xB8  ErrorDigest_Delete                 0xE6BACE0
+0x50  (reserved, NULL)                    —            +0xC0  ErrorAggregator_AddError           0xE6BAD60
+0x58  (reserved, NULL)                    —            +0xC8  ErrorAggregator_ProcessAndShutdown 0xE6BAE40
+0x60  (reserved, NULL)                    —            +0xD0  ErrorAggregator_LogErrorDigest     0xE6BAEC0
+0x68  ClientContext_Initialize            0xE6B9EC0    +0xD8  ErrorAggregator_Size               0xE6BAF20
+0x70  ClientContext_UnblockPendingWork    0xE6B9F20    +0xE0  ErrorAggregator_Active             0xE6BAF80
+0x78  ClientContext_MegascalePort         0xE6B9FE0    +0xE8  GetInterfaceAddressesHelper        0xE6BAFE0
+0x80  CreateMegascaleCollectives          0xE6BA080    +0xF0  GetOrCreateRuntimeError            0xE6BB620

所有符号都是 pjrt::(anonymous namespace)::<Name>。功能分组：

Client-context 生命周期（+0x18–+0x28、+0x68–+0x78）：从 PjRtClient 创建 megascale client context，或默认构造一个；删除它；初始化它；解除待处理工作的阻塞；查询 megascale 端口。
配置工厂（+0x30–+0x38）：AoT 配置和 multi-slice 配置；后者会产生 MultiSlice 扩展查询的 PJRT_MultiSlice_Config 句柄。
Collectives + id 映射（+0x80–+0x90）：megascale（跨 pod）collectives 工厂，以及 device-id↔megascale-id 双向转换。跨 pod collectives 位于这里，不在 Collectives 扩展中。
错误聚合子系统（+0x98–+0xF0）：注册/取消注册 error handler，然后是完整异步 aggregator：创建/删除 aggregator，添加逐主机错误，process-and-shutdown，记录 error digest，查询 size 和 active state，以及 interface-address discovery 和 runtime-error fetch。这是多主机训练的容错机制：收集逐主机错误，决定是否关闭 slice，发出 digest。

注意 — +0x40..+0x60 的五个保留 NULL 槽位是已声明但未实现的占位符（确认构造时清零；其预期的未来 client-context 方法未知 — LOW）。消费者不得通过它们调用；struct_size（248）覆盖这些槽位，但指针为 NULL。罕用 error-aggregator 方法的完整 args-struct 布局未逐字节跟踪（上方槽位表来自创建器存储序列；method-arg 偏移为 LOW）。

TpuExecutable — type 17，88 字节

最 TPU 专用的已编译 executable 表面：target-argument 和 HLO-module 自省、compiled-memory 与 cost analysis，以及对重实现者最重要的compilation-env 注入路径（SetTpuCompilationEnv）和 megascale 容错使用的predetermined-error 分类器（IsTpuPredeterminedError）。创建器 pjrt::CreateTpuExecutableExtension @ 0xE6DC6E0。


存储 VA	`0x224C3CA8`
struct_size	88（7 个有效槽位 + 1 个保留 NULL，位于 `+0x28`）
TPU 注入	所有 7 个有效槽位都是匿名命名空间 TPU 实现

偏移	方法	实现符号	地址	min/cur
`+0x18`	`GetTargetArguments`	`(anon)::GetTargetArguments`	`0xE6DC760`	—
`+0x20`	`GetHloModuleWithConfig`	`(anon)::GetHloModuleWithConfig`	`0xE6DC8A0`	—
`+0x28`	（保留，NULL）	—	—	—
`+0x30`	`GetCompiledMemoryStats`	`(anon)::GetCompiledMemoryStats`	`0xE6DCA40`	—
`+0x38`	`RunHloCostAnalysis`	`(anon)::RunHloCostAnalysis`	`0xE6DCC40`	—
`+0x40`	`SetTpuCompilationEnv`	`(anon)::SetTpuCompilationEnv`	`0xE6DD400`	0x2C/0x38
`+0x48`	`GetTpuCompilationEnvFieldAsString`	`(anon)::GetTpuCompilationEnvFieldAsString`	`0xE6DD620`	—
`+0x50`	`IsTpuPredeterminedError`	`(anon)::IsTpuPredeterminedError`	`0xE6DD880`	0x2F/0x19

行为

SetTpuCompilationEnv（min 0x2C，cur 0x38）通过 proto2::MessageLite::ParseFromString @ 0x21057460 解析序列化的 xla::CompilationEnvironmentsProto（字节位于 args +0x08/+0x10），通过 CreateFromProto @ 0x1E63E5A0 构建 xla::CompilationEnvironments，然后通过 tpu_executable_extension::SetTpuCompilationEnv(CompilationEnvironments*) @ 0xE6DE1C0 将其安装为进程全局配置。这是 TPU compilation-backend config / XLA-flags 注入表面，也是 libtpu 对 custom-call/FFI env 的替代。

IsTpuPredeterminedError（min 0x2F，cur 0x19）解析序列化的 tensorflow::StatusProto（args +0x08/+0x10），通过 tsl::StatusFromProto @ 0xF8BB9E0 转换为 absl::Status，调用 tpu_executable_extension::IsTpuPredeterminedError(absl::Status) @ 0xE6DE4A0，并把一个 bool 写入 args +0x18。它会为 megascale 训练中的容错重启逻辑分类某个 runtime error 是否为 “predetermined”（可在执行前检测）。

注意 — SetTpuCompilationEnv 和 IsTpuPredeterminedError 共享同一个 .rodata source-location 字符串，VA 为 0x877CD0F；它是作为 absl::SourceLocation 参数传给状态构造器的 tpu_executable_extension.cc 源文件路径，不是错误消息。两者都是扩展链上命名的 FFI 替代通道：找不到 FFI 扩展的框架会使用这些通道推送编译器配置并读取故障分类。

MultiSlice — type 22，64 字节 {#multislice--type-22-64-bytes}

查询 PJRT_MultiSlice_Config 不透明包装器：一个 8 字节堆对象，在 +0x00 持有 unique_ptr<MultiSliceConfig>。用于 multi-slice（跨 pod）训练的 slice-topology 自省；config 本身由 Megascale 的 CreateMultiSliceConfig 产生。创建器 pjrt::CreateMultiSliceExtension @ 0xE6F3C40。


存储 VA	`0x224C3F20`
struct_size	64（5 个方法）
TPU 注入	全部为匿名命名空间 TPU 实现

偏移	方法	实现符号	地址	min/cur
`+0x18`	`Config_Destroy`	`(anon)::ConfigDestroy`	`0xE6F3CA0`	0x23/0x10
`+0x20`	`Config_NumSlices`	`(anon)::ConfigNumSlices`	`0xE6F3D20`	0x25/0x14
`+0x28`	`Config_SliceId`	`(anon)::ConfigSliceId`	`0xE6F3D80`	—
`+0x30`	`Config_NumDevicesPerSlice`	`(anon)::ConfigNumDevicesPerSlice`	`0xE6F3DE0`	—
`+0x38`	`Config_Serialize`	`(anon)::ConfigSerialize`	`0xE6F3FE0`	—

行为

Config_Destroy 读取 args +0x08 处的 wrapper；如果非空，它解引用 wrapper->impl vtable +0x08（析构函数），将指针置空，并 free(wrapper, 8)。Config_NumSlices 跳转到 wrapper->impl vtable +0x10 并将 int32 写入 args +0x10。其他读取器（SliceId、NumDevicesPerSlice、Serialize）在不同 vtable 偏移处遵循相同的一次 vtable 跳转模板。

注意 — Config_Serialize @ 0xE6F3FE0 使用的 MultiSliceConfig 序列化 proto schema 未恢复（LOW）。handle-at-+0x08 约定（没有 priv）与 TpuTopology 家族一致；由 mov 0x8(%rbx),... 读取确认。

Callback — type 14，40 字节

在 TPU slice-builder 故障事件上触发的 host-side callbacks；这是 libtpu 的 host-callback / fault-handler 表面，替代通用 FFI host-callback 扩展。创建器 pjrt::CreateCallbackExtension @ 0xE6B91E0。


存储 VA	`0x224C3B60`
struct_size	40（2 个方法）
TPU 注入	全部为匿名命名空间 TPU 实现

偏移	方法	实现符号	地址	min/cur
`+0x18`	`RegisterCallback`	`(anon)::PJRT_Callback_RegisterCallback`	`0xE6B9220`	0x23/0x28
`+0x20`	`InvokeCallback`	`(anon)::PJRT_Callback_InvokeCallback`	`0xE6B94C0`	—

行为

RegisterCallback（min 0x23，cur 0x28）读取 args +0x10 处的 callback-type enum。对 type==1，它验证目标设备是 TPU：检查 device-kind id 是否等于 xla::TpuId()::kTpuId @ 0x224C3FC8（guard 0x224C3FD0）；然后把 std::function<void(accel_ssw::deepsea::slice_builder::SliceFailureType)> 注册到 xla::SliceBuilderCallbackState::AddCallback @ 0xF95DF80（通过 RegisterSliceBuilderCallback::$_0 policy thunk @ 0xE6B96C0，policy @ 0x215FB718）。type==2 是第二类 callback。

注意 — 只有 type==1（slice-builder）路径被完整跟踪；type==2 分发目标没有跟踪（LOW）。更广泛的 host-callback 表面见 Callbacks。

HostAllocator — type 15，48 字节 {#hostallocator--type-15-48-bytes}

用于 device-host DMA staging 的 TPU pinned-host memory allocator；不同于链头的通用 HostMemoryAllocator（type 23）。全部三个槽位都是作为创建器参数注入的 tpu_plugin:: TPU 实现。创建器 pjrt::CreateHostAllocatorExtension(node, next, prefalign_fn, alloc_fn, free_fn) @ 0xF8A3C20；node 是返回槽位，next（a2）进入 +0x10，创建器直接把来自 a3/a4/a5 参数的三个 TPU 指针写入 +0x18/+0x20/+0x28。


存储 VA	`0x224C3AA0`
struct_size	48（3 个方法）
TPU 注入	全部 3 个槽位（`a3`/`a4`/`a5`）

偏移	方法	实现符号	地址
`+0x18`	`GetPreferredAlignment`	`tpu_plugin::TPU_PJRT_HostAllocator_GetPreferredAlignment`	`0xE6AA060`
`+0x20`	`Allocate`	`tpu_plugin::TPU_PJRT_HostAllocator_Allocate`	`0xE6AA140`
`+0x28`	`Free`	`tpu_plugin::TPU_PJRT_HostAllocator_Free`	`0xE6AA240`

易错点 — 这两个 host allocator 是此链中的典型混淆点。Type 23（HostMemoryAllocator，32 字节，1 个方法，通用，无 priv，payload 指针在 +0x08）是通过实时 client vtable 实现的 XLA host staging。Type 15（HostAllocator，48 字节，3 个方法，全部 TPU 注入）是带插件供应指针的 TPU pinned-host DMA staging。结构体布局不同，层次不同；匹配错误 type id 会返回错误 allocator。

第 3 组 — 混合型（TPU 注入工厂） {#group-3--hybrid-tpu-injected-factory}

这些节点的创建器以参数形式接收一个或两个 TPU 函数指针，并用通用包装器填充其余槽位。重实现者只提供注入工厂；库提供周边方法。

ExecutableMetadata — type 13，40 字节

返回序列化 executable metadata（描述已编译 executable 的 TPU 专用 blob），以及返回 buffer 的 deleter。Get 路径由 TPU 提供；Destroy 是通用的。创建器 pjrt::CreateExecutableMetadataExtension(node, next, get_metadata_fn) @ 0xF8A3BE0。


存储 VA	`0x224C3A70`
struct_size	40（2 个方法）
TPU 注入	1 个槽位（`+0x18`，即 `Get` 路径）

偏移	方法	实现符号	地址
`+0x18`	`GetExecutableMetadata` (TPU)	`tpu_plugin::GetTpuExecutableMetadata`	`0xE6A9E40`
`+0x20`	`DestroySerializedMetadata`	`pjrt::DestroySerializedMetadata`	`0xF8A3BA0`

怪癖 — 注入的 TPU 函数落在第一个槽位（+0x18），通用 deleter 位于 +0x20。创建器把它的 a3（TPU get fn）存到 +0x18，并把内置的 DestroySerializedMetadata 存到 +0x20。这与 PhaseCompile（TPU 工厂在前）的排列相反，但概念上相符：生产者是 TPU，析构器是通用的。

AbiVersion FromProto 工厂 — type 20

AbiVersion 的两个 TPU 注入槽位（主体节点已在第 1 组中记录）。thunk CreateTpuAbiVersionExtension @ 0xE6B7340 提供它们：

偏移	方法	实现符号	地址	包装
`+0x68`	`RuntimeAbiVersion_FromProto`	`(anon)::TpuRuntimeAbiVersionFromProto`	`0xE6B7360`	`xla::PjRtRuntimeAbiVersionFromProto` @ `0xE6B73E0`
`+0x70`	`ExecutableAbiVersion_FromProto`	`(anon)::TpuExecutableAbiVersionFromProto`	`0xE6B7380`	`xla::PjRtExecutableAbiVersionFromProto` @ `0xE6B7660`

每个函数消费其 proto（xla::PjRtRuntimeAbiVersionProto / xla::PjRtExecutableAbiVersionProto），并通过 pjrt::Common{Runtime,Executable}AbiVersionFromProto @ 0xE6B86A0 / 0xE6B8800 产生一个 StatusOr<unique_ptr<...AbiVersion>>。它们由 TPU 提供，是因为 proto 携带 TPU platform id；AbiVersion 表面的其余部分与平台无关。

PhaseCompile（Get/Destroy compiler）— type 9

PhaseCompile 是第三个混合型：CreatePhaseCompileExtension(node, next, get_compiler_fn, destroy_compiler_fn) @ 0xE6F42A0 注入 tpu_plugin::GetTpuPhaseCompiler @ 0xE6AA320（+0x18）和 tpu_plugin::DestroyTpuPhaseCompiler @ 0xE6AA400（+0x20）；三个驱动方法（Run_Phase、Get_Phase_Names、C_Buffers_Destroy，+0x28/+0x30/+0x38）是通用的。完整的 compiler-object 内部细节在其专页中，见 PhaseCompile。

TPU 注入 vs 通用 — 重实现者速查表

对重实现者最有用的单一分类：哪些节点需要 TPU 后端，哪些由库提供。（Profiler、RawBuffer、PhaseCompile、TpuTopology 各有专页，这里只是为了划分完整性而列出。）

类别	扩展	重实现者提供什么
完全 TPU	Megascale (18)、TpuExecutable (17)、HostAllocator (15)、MultiSlice (22)、Callback (14)、TpuTopology (16)	每个方法主体
TPU 注入工厂	AbiVersion (20) [2 个 `FromProto`]、ExecutableMetadata (13) [`Get`]、PhaseCompile (9) [`Get`/`Destroy` compiler]	仅注入的 factory fn(s)；库提供其余部分
通用 XLA（TPU 支撑）	Layouts (4)、MemoryDescriptions (6)、RawBuffer (8)、CrossHostTransfers (12)、Shardings (19)、Collectives (21)、HostMemoryAllocator (23)、Profiler (1)	extension 侧无需提供；通过实时句柄 vtable 跳转

组件	关系
扩展链	负责节点头部、type-id 枚举、清单和遍历；本页是其逐节点深入说明
Buffer and Memory	MemoryDescriptions（静态）之前对应的实时 `PJRT_Memory` 表面
Collectives Communicator	与 Collectives (21) 扩展并行的主表 communicator 表面
DMA and Cross-Host Recv	与 CrossHostTransfers (12) 并行的主表 DMA 路径
Callbacks	Callback (14) 扩展馈入的更广泛 host-callback 表面

交叉引用

扩展链 — 节点布局、PJRT_Extension_Type 枚举、17 节点清单、构造顺序与遍历顺序、消费者遍历循环
PJRT Plugin Overview — dlsym("GetPjrtApi") 如何到达 GetTpuPjrtApi 以及一次性初始化
API Vtable Reconstruction — extension_start 挂接其上的 140 个主表槽位（独立结构；ExecuteContext 位于这里，而不是作为扩展）
Profiler Extension — type 1，.data 种子 / 链终止符
Topology Description Extension — type 16，最大的有效扩展（31 个方法）
RawBuffer Extension — type 8，无类型 device-memory 表面
PhaseCompile Extension — type 9，命名阶段部分编译和注入的 TPU compiler 工厂

PJRT 剩余扩展 ​

摘要 ​

关于两个有意缺席项 ​

第 1 组 — 通用 XLA 表面 ​

Layouts — type 4，80 字节 ​

算法 ​

MemoryDescriptions — type 6，40 字节 ​

CrossHostTransfers — type 12，56 字节 ​

Shardings — type 19，40 字节 ​

Collectives — type 21，96 字节 {#collectives--type-21-96-bytes} ​

行为 ​

HostMemoryAllocator — type 23，32 字节 {#hostmemoryallocator--type-23-32-bytes} ​

算法 ​

AbiVersion（通用部分）— type 20，120 字节 {#abiversion-generic-portion--type-20-120-bytes} ​

第 2 组 — 完全 TPU 专用表面 ​

Megascale — type 18，248 字节 {#megascale--type-18-248-bytes} ​

TpuExecutable — type 17，88 字节 ​

行为 ​

MultiSlice — type 22，64 字节 {#multislice--type-22-64-bytes} ​

行为 ​

Callback — type 14，40 字节 ​

行为 ​

HostAllocator — type 15，48 字节 {#hostallocator--type-15-48-bytes} ​

第 3 组 — 混合型（TPU 注入工厂） {#group-3--hybrid-tpu-injected-factory} ​

ExecutableMetadata — type 13，40 字节 ​

AbiVersion FromProto 工厂 — type 20 ​

PhaseCompile（Get/Destroy compiler）— type 9 ​

TPU 注入 vs 通用 — 重实现者速查表 ​

相关组件 ​

交叉引用 ​

PJRT 剩余扩展

摘要

关于两个有意缺席项

第 1 组 — 通用 XLA 表面

Layouts — type 4，80 字节

算法

MemoryDescriptions — type 6，40 字节

CrossHostTransfers — type 12，56 字节

Shardings — type 19，40 字节

Collectives — type 21，96 字节 {#collectives--type-21-96-bytes}

行为

HostMemoryAllocator — type 23，32 字节 {#hostmemoryallocator--type-23-32-bytes}

算法

AbiVersion（通用部分）— type 20，120 字节 {#abiversion-generic-portion--type-20-120-bytes}

第 2 组 — 完全 TPU 专用表面

Megascale — type 18，248 字节 {#megascale--type-18-248-bytes}

TpuExecutable — type 17，88 字节

行为

MultiSlice — type 22，64 字节 {#multislice--type-22-64-bytes}

行为

Callback — type 14，40 字节

行为

HostAllocator — type 15，48 字节 {#hostallocator--type-15-48-bytes}

第 3 组 — 混合型（TPU 注入工厂） {#group-3--hybrid-tpu-injected-factory}

ExecutableMetadata — type 13，40 字节

AbiVersion FromProto 工厂 — type 20

PhaseCompile（Get/Destroy compiler）— type 9

TPU 注入 vs 通用 — 重实现者速查表

相关组件

交叉引用