TopologyDescription 扩展（type 16）

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

摘要

PJRT 拓扑描述是在没有活跃客户端时的设备互连几何。JAX 通过 PJRT_TopologyDescription_Create("tpu_v4", {chip_bounds=[2,4,4]}) 构建它，从而可以在连接任何 TPU 之前、甚至在完全没有加速器的主机上，提前编译 XLA 程序。编译器需要环面形状（每轴芯片数、每芯片核心数、进程布局）来布置 collective 和张量分片；拓扑接口就是它在不启动运行时的情况下取得该形状的方式。这与 Client, Device & Topology 上的活跃路径形成对比：那里 PJRT_Client_TopologyDescription（槽位 100）返回已经附着到运行中 xla::TpuClient 的拓扑，即一个调用方不得销毁的借用句柄。本页负责说明独立的 AOT 拓扑：它如何从名称字符串和选项映射创建，其 C-ABI 包装器如何缓存设备描述，以及 TPU 专属扩展如何暴露环面坐标。

该接口是两层结构。第一层是主 140 槽 PJRT_Api 表中的七个通用 PJRT_TopologyDescription_* 槽位（Create、Destroy、PlatformName、PlatformVersion、GetDeviceDescriptions、Serialize、Attributes），再加三个 v0.103 后期新增项（槽位 119 的 Deserialize、槽位 138 的 Fingerprint、槽位 100 的 Client_TopologyDescription）。这些暴露 XLA 可移植抽象：平台名、不透明 protobuf 序列化、设备描述列表、属性键值列表，以及 64 位指纹；其形状与 CPU 和 GPU 插件相同。第二层是挂在扩展链上的 TPU 专属扩展（PJRT_Extension_Base.type == 16）：一个 272 字节结构体，包含 31 个函数指针，承载通用接口无法表达的环面几何（ChipBounds、CoreCountPerChip、ProcessBounds、slice 配置、ICI 路由标志）。

决定性的结构事实是堆包装器中的急切缓存。PJRT_TopologyDescription 是一个 112 字节对象，拥有一个 xla::PjRtTopologyDescription*（持有两次：+0x00 处的 owned copy 由析构函数虚析构删除，+0x08 处的 working copy 由每个读取器解引用），并在构造时预先物化三个查找 span。因此 GetDeviceDescriptions 和 Attributes 会退化为两个 mov 字段读取，不发生分配；只有 TPU 扩展上的标量/几何方法才会经由实现 vtable 分派。如果重新实现者在每次 GetDeviceDescriptions 调用时都惰性计算设备列表，就会偏离二进制的行为及其生命周期保证。

对重新实现而言，契约如下：

• 创建路径在 PJRT_TopologyDescription_Create（槽位 87，TPU 覆盖）中：校验结构体大小，解析 PJRT_NamedValue[] 选项，经由 CustomTpuTopologyNameOverride 解析拓扑名称，通过 GetTpuTopologyDescription 构建抽象 xla::PjRtTopologyDescription，并用 CreateWrapperDeviceTopology 包装。
• 包装器布局（112 字节）：impl 指针（+0x00 处 owned copy，+0x08 处 working copy）加上三个预构建 span（设备 shared-ptr、PJRT_DeviceDescription 数组、设备指针 span）以及一个缓存的 PJRT_NamedValue 属性数组，因此读取槽位从不分配。
• TPU 扩展调用约定：非标准 args 布局（包装器在 +0x08，没有 priv 字段）以及统一的 vtable-bouncer 函数体：从 wrapper+0x08 加载 wrapper->impl，调用固定 vtable 偏移，并将标量或经过边界检查的向量复制回 args 结构体。
• 生命周期拆分：Create / Deserialize 返回拥有型包装器，必须通过 Destroy（槽位 88）恰好销毁一次；Client_TopologyDescription 返回借用型包装器，其生命周期由 client 拥有。

通用接口	7 个槽位 @ PJRT_Api+0x2B8..+0x2E8（槽位 87-93）+ 后期槽位 100 / 119 / 138
Create（TPU）	pjrt::tpu_plugin::PJRT_TopologyDescription_Create @ 0xE6A9B20（槽位 87）
包装器构建	pjrt::CreateWrapperDeviceTopology(unique_ptr<>) @ 0xF870E60（函数体 @ 0xF873260）
包装器大小	112 字节（0x70），拥有 xla::PjRtTopologyDescription*（copy 在 +0x00，working ptr 在 +0x08）
TPU 扩展	type == 16，272 字节（0x110），31 个 fn-ptr @ 0x224C3B90（.bss）
扩展填充	pjrt::CreateTpuTopologyExtension @ 0xE6DE5E0（35 槽存储，无调用）
序列化	xla::PjRtTopologyDescriptionProto 的 protobuf 二进制（proto2）
C-API 版本	v0.103（{major=0, minor=103}）

---

两层接口

用途

拓扑可以通过两种刻意设计为不同形状的方式访问。通用层是每个 PJRT 插件都实现的可移植 XLA 抽象；只知道 pjrt_c_api.h 的框架可通过这十个槽位驱动整个拓扑，完全无需知道自己正在与 TPU 对话。TPU 扩展层是抽象 xla::PjRtTopologyDescription vtable 无法表达的几何信息的逃生口，即环面边界、每芯片核心数、ICI 路由；只有沿扩展链寻找 type == 16 的消费者才能到达。

通用槽位映射

七个核心槽位连续位于 PJRT_Api+0x2B8..+0x2E8；三个 v0.103 后期新增项散布在追加到 ABI 时分配的槽位位置。

槽位	字段	libtpu 实现	地址	最小/当前大小
87	_Create	pjrt::tpu_plugin::PJRT_TopologyDescription_Create	0xE6A9B20	0x24 / 0x38
88	_Destroy	pjrt::PJRT_TopologyDescription_Destroy	0xF870040	0x25 / 0x18
89	_PlatformName	pjrt::PJRT_TopologyDescription_PlatformName	0xF870200	0x2A / 0x28
90	_PlatformVersion	pjrt::PJRT_TopologyDescription_PlatformVersion	0xF870260	0x2D / 0x28
91	_GetDeviceDescriptions	pjrt::PJRT_TopologyDescription_GetDeviceDescriptions	0xF8702C0	0x33 / 0x28
92	_Serialize	pjrt::PJRT_TopologyDescription_Serialize	0xF870320	0x27 / 0x38
93	_Attributes	pjrt::PJRT_TopologyDescription_Attributes	0xF8705E0	0x28 / 0x28
100	Client_TopologyDescription	pjrt::PJRT_Client_TopologyDescription	0xF85F560	—
119	_Deserialize	pjrt::PJRT_TopologyDescription_Deserialize	0xF870B80	0x28 / 0x28
138	_Fingerprint	pjrt::PJRT_TopologyDescription_Fingerprint	0xF870520	0x29 / 0x20

最小大小 是传给 ActualStructSizeIsGreaterOrEqual 的 min_size 常量，即可接受的最旧历史 args 结构体；当前大小 是此构建自己的调用方传入的大小。

怪异点 — 只有槽位 87（_Create）由 TPU 覆盖，它指向 pjrt::tpu_plugin::PJRT_TopologyDescription_Create。其他六个通用槽位是 CPU 和 GPU 插件也使用的共享 pjrt::PJRT_TopologyDescription_* 包装器；TPU 专门化完全发生在这些包装器调用的 xla::PjRtTopologyDescription（一个 tpu::TpuTopology 子类）内部，以及 type-16 扩展内部。Client_TopologyDescription（槽位 100，活跃路径）记录在 Client, Device & Topology；此处列出它只是为了补全构造路径图景。

陷阱 — Destroy 槽位的 (min=0x25, current=0x18) 对具有 min > current。0x25（37）是烘入调用中的原始 min_size 字面量；当前调用方把 args 结构体截短为 24 字节。运行时测试是 actual >= min，因此来自当前调用方的 24 字节 args 结构体若按字面 >= 37 检查会失败；但反编译出的 ActualStructSizeIsGreaterOrEqual 会把两者中较小者当作向后兼容的下限。请逐槽复制精确的 (min, current) 常量，而不是假设有统一规则。

---

PJRT_TopologyDescription_Create — 独立拓扑

用途

PJRT_TopologyDescription_Create 是 AOT 入口点：它把平台名字符串加选项映射转换为一个完整构建的拓扑包装器，不需要活跃客户端，也不探测硬件。JAX 会为每个提前编译目标调用它一次。它是唯一由 TPU 覆盖的通用拓扑槽位，因为名称解析（"v4:2x2x1" → 规范配置）和 chip-bounds 查找属于 TPU 私有逻辑。

入口点

PJRT_Api slot 87 (0x2B8)                                ── pjrt::tpu_plugin::PJRT_TopologyDescription_Create (0xE6A9B20)
  ├─ pjrt::ActualStructSizeIsGreaterOrEqual              ── "PJRT_TopologyDescription_Create_Args", min=0x24, cur=0x38
  ├─ pjrt::ConvertFromPjRtNamedValueList (0xF8A43C0)     ── create_options[] → flat_hash_map<string,variant>
  ├─ pjrt::CustomTpuTopologyNameOverride (0xE6AFCE0)     ── platform_name + opts → canonical topology name (StatusOr)
  ├─ pjrt::GetTpuTopologyDescription (0xE6ADB60)         ── name + opts → xla::PjRtTopologyDescription (StatusOr)
  └─ pjrt::CreateWrapperDeviceTopology (0xF870E60)       ── wrap into 112-byte PJRT_TopologyDescription, store at args[6]

算法

function PJRT_TopologyDescription_Create(args):          // 0xE6A9B20
    // args[0]=struct_size, args[2]=platform_name.data, args[3]=platform_name.len,
    // args[4]=create_options (PJRT_NamedValue*), args[5]=num_options,
    // args[6]=PJRT_TopologyDescription** out  (offset +0x30).
    st = ActualStructSizeIsGreaterOrEqual(               // 0xE6A9B20:+0x18
            "PJRT_TopologyDescription_Create_Args", min=0x24, cur=0x38, args[0])
    if st != 1:
        return new PJRT_Error{st}                        // operator new(8)

    // (1) untyped option ingest — same parser as PJRT_Client_Create.
    opts = ConvertFromPjRtNamedValueList(args[4], args[5])   // -> flat_hash_map<string,variant<str,bool,long,vector<long>,float>>
    if args[3] < 0: BUG()                                // platform_name length sanity

    // (2) resolve the canonical topology name. Returns a Status as the
    //     first slot of the StatusOr; an error here aborts before build.
    name_or = CustomTpuTopologyNameOverride(args[2], args[3], opts)   // 0xE6AFCE0
    if name_or is error:
        return new PJRT_Error{name_or.status}            // refcount the StatusRep, then return

    // (3) build the abstract topology: looks up chip_bounds / host_bounds /
    //     slice_config for the named platform, constructs the TpuTopology subclass.
    topo_or = GetTpuTopologyDescription(name_or, opts)   // 0xE6ADB60 -> StatusOr<unique_ptr<xla::PjRtTopologyDescription>>
    if topo_or is error:
        return new PJRT_Error{topo_or.status}

    // (4) wrap for the C ABI; out-param is args[6] (+0x30). Transfers
    //     ownership of the unique_ptr into the wrapper's +0x00 slot.
    args[6] = CreateWrapperDeviceTopology(topo_or.release())   // 0xF870E60
    return NULL                                          // success

陷阱 — 如果 platform_name 为空且 create_options 非空，CustomTpuTopologyNameOverride 会返回错误字符串 "TPU PJRT_TopologyDescription_Create does not support extra create_options if no topology_name is given."（在 .rodata 中）。重新实现者必须在这里拒绝无名称但有选项的情况；选项表以解析后的名称为键，因此没有默认配置可应用这些选项。

注意 — out-param 索引不同于 Client_Create。这里拓扑写入 args[6]（args 结构体内字节偏移 +0x30），由反编译的 a1[6] = pjrt::CreateWrapperDeviceTopology(...) 确认。_Deserialize 路径（槽位 119）则把其包装器写入 args[4]（+0x20）；这些 args 布局不可互换。

注意事项

GetTpuTopologyDescription（0xE6ADB60）会把平台名字符串映射为环面配置。支持的名称会通过同一个 kPlatformConfigs 单例（见 Slice 与平台配置）往返，该单例也被 GetDefaultPlatformConfig 使用；在二进制中观察到的规范集合为 "tpu"、"tpu_v4"、"tpu_v5e"、"tpu_v6e" 以及类似的逐代 key（中等置信度：完整 key 列表需要遍历 0x224C4108 处的 kPlatformConfigs map 存储）。构造出的对象在构建后是 const：TPU 拓扑由不可变系统查询装配而成，之后永不变异，这正是缓存读取槽位可并发调用的原因。

---

包装器 — PJRT_TopologyDescription（112 字节）

用途

CreateWrapperDeviceTopology（0xF870E60）取得 unique_ptr<xla::PjRtTopologyDescription> 的所有权，并产生 C ABI 分发的 112 字节堆对象。它的决定性行为与 client 包装器完全相同，是急切缓存：构造时它遍历一次 impl 拓扑的设备列表，并构建三个 span 加一个缓存属性数组，因此通用 GetDeviceDescriptions 和 Attributes 槽位会变成字段读取。

包装器布局（PJRT_TopologyDescription，112 字节）

由析构函数（~PJRT_TopologyDescription @ 0xF8700C0）和七个通用读取器重建。偏移位于 operator new(0x70) 块内。

字段	偏移	类型	含义
impl_owned	+0x00	xla::PjRtTopologyDescription*（owned）	拥有的实现拓扑；析构函数运行其虚析构函数（vtable +0x08）。由 unique_ptr/const-ref 调用方在构造函数体返回后设置；函数体自身将其保留为 NULL
impl	+0x08	xla::PjRtTopologyDescription*（working ptr）	每个读取器和每个 type-16 扩展方法都会解引用的指针（*(wrapper+0x08)）；其 vtable 驱动每个未缓存槽位。地址与 impl_owned 相同
cached_device_shared_ptrs	+0x10	vector<shared_ptr<xla::PjRtDeviceDescription>>	begin/end/cap；拥有 impl 设备描述的 refcount
cached_device_descriptions	+0x28	vector<PJRT_DeviceDescription>	begin/end/cap；每个设备一个 0x20 字节（32 字节）包装器
cached_device_pointer_span	+0x40	vector<PJRT_DeviceDescription*>	begin/end/cap；由 GetDeviceDescriptions 原样返回的 span
attributes_data	+0x58	PJRT_NamedValue*	缓存属性数组
attributes_size	+0x60	size_t	由 Attributes 返回
attributes_capacity	+0x68	size_t	后备容量

算法

// Construction body (const* overload) @ 0xF873260.
function CreateWrapperDeviceTopology(impl):              // 0xF870E60 (thin) -> 0xF873260 (body)
    w = operator new(0x70)                                // 112-byte PJRT_TopologyDescription
    w[+0x00] = 0                                          // owned slot left NULL by the body
    w->impl = impl                                        // working pointer at +0x08

    // (cached at construction by the body @ 0xF873260)
    // (1) DEVICE DESCRIPTIONS: walk impl->DeviceDescriptions() (vtable +0x30),
    //     a vector<shared_ptr<PjRtDeviceDescription>>; for each, allocate a
    //     PJRT_DeviceDescription wrapper (0x20 bytes) and push into cached_device_descriptions.
    for sp in impl->DeviceDescriptions():                 // vtable +0x30
        w.cached_device_shared_ptrs.push_back(sp)         // keep the refcount alive
        w.cached_device_descriptions.push_back(PJRT_DeviceDescription{ sp.get() })
        w.cached_device_pointer_span.push_back(&w.cached_device_descriptions.back())

    // (2) ATTRIBUTES: snapshot impl->Attributes() (vtable +0xE8) into a PJRT_NamedValue[] once.
    w.attributes_data, w.attributes_size = PopulatePjrtAttributes(impl->Attributes())
    return w

// Thin overloads (unique_ptr @ 0xF870E60, const-ref deserialize path) then take ownership:
//   prev = w[+0x00]; w[+0x00] = impl;  if (prev) prev->vtable[+0x08](prev)
// so the OWNED copy lands at +0x00 and the destructor virtual-deletes it exactly once.

怪异点 — 设备描述 span（+0x40）和属性数组（+0x58）只构建一次，并由槽位 91 和 93 直接读取。反编译的 GetDeviceDescriptions 字面上就是 args[3] = *(wrapper+0x40); args[4] = *(wrapper+0x48)，即两个 mov，不调用 impl 拓扑，也不分配。若重新实现每次调用 GetDeviceDescriptions 都调用 impl->DeviceDescriptions()，就会重新遍历并重新分配，也会破坏生命周期契约：调用方可以在包装器存活期间持有返回的 PJRT_DeviceDescription* 数组，因此它必须是稳定存储，而不是每次调用的临时对象。

通用读取槽位

function PJRT_TopologyDescription_GetDeviceDescriptions(args):  // 0xF8702C0, slot 91
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TopologyDescription_GetDeviceDescriptions_Args", 0x33, 0x28, args[0]) != 1:
        return new PJRT_Error{...}
    w = args[2]                                           // PJRT_TopologyDescription*
    args[3] = *(w + 0x40)                                 // cached_device_pointer_span.data
    args[4] = *(w + 0x48)                                 // cached_device_pointer_span.size
    return NULL                                           // never touches w->impl

function PJRT_TopologyDescription_Attributes(args):       // 0xF8705E0, slot 93
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TopologyDescription_Attributes_Args", 0x28, 0x28, args[0]) != 1:
        return new PJRT_Error{...}
    w = args[2]
    args[3] = *(w + 0x58)                                 // attributes_data
    args[4] = *(w + 0x60)                                 // attributes_size
    return NULL

PlatformName（槽位 89）和 PlatformVersion（槽位 90）不缓存，而是跳转到 impl vtable：PlatformName 调用 vtable +0x18，PlatformVersion 调用 vtable +0x20，两者各返回一个 std::string_view，其 {data, size} 对会存入 args+0x18 / args+0x20。这些是廉价虚调用，返回 const 拓扑内部的指针，因此缓存没有收益。

注意 — AOT 拓扑属性是 TPU 几何暴露给通用消费者的位置。 即使框架不遍历扩展链，仍会在 Attributes（槽位 93）列表中以命名条目看到 slice 形状、芯片坐标、每芯片核心数，这与活跃路径上的 PJRT_DeviceDescription_Attributes 使用同一机制。重新实现 AOT 拓扑时，除了连接 type-16 扩展方法外，还必须把 coords / core_on_chip / slice_index / cores-per-chip 填充为 PJRT_NamedValue 属性条目。

---

TPU 扩展（type 16，31 个方法）

用途

抽象 xla::PjRtTopologyDescription vtable 只暴露每个后端共有的内容（平台名、设备描述、序列化、指纹）。TPU 环面几何（芯片边界、每芯片核心数、进程布局、ICI 路由策略、slice 配置）在该 vtable 中没有槽位。type-16 扩展就是承载者：它是一个包含 31 个函数指针的 272 字节结构体，链挂在 extension_start 之后；TPU 感知的消费者会通过沿链寻找 base->type == 16 找到它。

存储与填充

扩展结构体是 pjrt::tpu_plugin::GetTpuPjrtApi::tpu_topology_extension @ 0x224C3B90（.bss）处的静态对象，guard @ 0x224C3CA0，在受 __cxa_guard 保护的 GetTpuPjrtApi 初始化期间，由 pjrt::CreateTpuTopologyExtension(PJRT_Extension_Base*) @ 0xE6DE5E0 填充一次。该函数是 34 次字段存储加一个 ret：三个 header 槽位，然后 31 个函数指针（后者为 31 对 lea/store），唯一的 rsi 参数会写入 +0x10 作为链上的 next（指向 callback_extension @ 0x224C3B60）。

struct PJRT_Extension_Base {                              // 24-byte common header
    /* +0x00 */ size_t                struct_size;         // = 0x110
    /* +0x08 */ uint32_t              type;                // = 16
    /* +0x0c */ uint32_t              _pad0;
    /* +0x10 */ PJRT_Extension_Base*  next;                // → callback_extension @ 0x224C3B60
    /* +0x18 .. +0x108 : 31 function pointers (see method table) */
};
_Static_assert(sizeof(PJRT_TpuTopology_Extension) == 0x110, "3*8 header + 31*8 fn-ptrs");

陷阱 — TPU 扩展 args 结构体使用与主 PJRT_Api 不同的布局。每个结构体都把包装器指针放在 +0x08，紧跟在 struct_size 之后，没有 priv 字段。主 PJRT_Api args 遵循 pjrt_c_api.h 的 {size_t struct_size; void* priv; ...} 约定，payload 在 +0x10。这在反编译中表现为 mov 0x8(%rbx),%rax（extension；a1[1]）对比 mov 0x10(%rbx),%rax（main；a1[2]）。如果重新实现者把 main-API 偏移复制到扩展方法里，就会读取错误字段并解引用垃圾。

调用约定 — 统一函数体

31 个方法分为三类，但每个方法开头都相同：校验 struct_size，然后加载 wrapper = args[1]（偏移 +0x08）和 topo = wrapper->impl（*(wrapper+0x08)，即 wrapper +0x08 处的 xla::PjRtTopologyDescription* working 指针；注意内层解引用）。占多数的 vtable-bouncer 随后调用固定 vtable 偏移并复制结果。

PJRT_Error* method(PJRT_TpuTopology_<Name>_Args* args) {  // uniform template
    if (ActualStructSizeIsGreaterOrEqual(
            "PJRT_TpuTopology_<Name>_Args", <min>, <cur>, args->struct_size) != 1)
        return new PJRT_Error{returned_status};            // operator new(8)

    PJRT_TopologyDescription* wrapper = args[1];            // args+0x08 (NO priv)
    xla::PjRtTopologyDescription* topo = wrapper->impl;     // wrapper+0x08
    StatusOr<T> r = topo->vtable[<offset>]();               // fixed per-method offset

    if (r.ok()) { args->output = move(r.value()); return NULL; }
    return new PJRT_Error{r.status()};                      // refcount StatusRep, then wrap
}

两个已确认的函数体锚定了该模板：

function CoreCountPerChip(args):                          // 0xE6DF540, ext slot 9
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TpuTopology_CoreCountPerChip_Args", 0x26, 0x14, args[0]) != 1:
        return new PJRT_Error{...}
    topo = *(args[1] + 8)                                  // wrapper->impl
    r = topo->vtable[+0x80]()                              // StatusOr<int32>
    if r.ok():
        *(int32*)(args + 0x10) = r.value                   // scalar out at +0x10
        return NULL
    return new PJRT_Error{r.status}

function ChipBounds(args):                                 // 0xE6E0AE0, ext slot 25
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TpuTopology_ChipBounds_Args", 0x20, 0x28, args[0]) != 1:
        return new PJRT_Error{...}
    topo = *(args[1] + 8)
    sv = topo->vtable[+0xD0]()                             // StatusOr<vector<int32>>  (X,Y,Z bounds)
    if not sv.ok(): return new PJRT_Error{sv.status}
    n = sv.size()
    args[4] = n                                            // actual_count at +0x20
    if args[2] >= n:                                       // capacity at +0x10
        for i in 0..n: ((int32*)args[3])[i] = sv[i]        // buffer at +0x18
        return NULL
    // (caller buffer too small: vector path returns InvalidArgument "needed N, provided M")

方法表（31 个函数）

偏移位于 272 字节扩展结构体内；“vtable” 是 bouncer 调用的抽象 xla::PjRtTopologyDescription 上的偏移。类别：VB = vtable-bouncer，HLP = 下转到 tpu::TpuTopology 的 helper，CFG = kPlatformConfigs 查找。

偏移	方法	实现地址	vtable / 经由	类别
+0x18	Subslice	0xE6DE7A0	input-array, computed	VB
+0x20	IsSubsliceTopology	0xE6DEC20	+0x38 (bool)	VB
+0x28	SubsliceDeviceIdFromFullDeviceId	0xE6DEC80	computed cast	VB
+0x30	ReplaceHostBounds	0xE6DEF40	helper	HLP
+0x38	IsEnhancedBarrierEnabled	0xE6DF200	helper (TpuTopology+0x188/+0x190)	HLP
+0x40	HasLimitedIciConnectivity	0xE6DF2A0	helper	HLP
+0x48	IsReachableOverLimitedIci	0xE6DF340	helper	HLP
+0x50	ProcessCount	0xE6DF400	+0x40 (int)	VB
+0x58	ChipsPerProcess	0xE6DF4A0	+0x48 (int)	VB
+0x60	CoreCountPerChip	0xE6DF540	+0x80 (int)	VB
+0x68	ChipCount	0xE6DF5E0	+0x50 (int)	VB
+0x70	CoreCount	0xE6DF680	+0x58 (int)	VB
+0x78	LogiDeviceCountPerProcess	0xE6DF720	+0x60 (int)	VB
+0x80	LogiDeviceCount	0xE6DF7C0	+0x68 (int)	VB
+0x88	LogiDeviceCountPerChip	0xE6DF860	+0x70 (int)	VB
+0x90	CoreCountPerProcess	0xE6DF900	+0x78 (int)	VB
+0x98	ProcessIds	0xE6DF9A0	+0x88 (vector)	VB
+0xA0	LogiDeviceIdsOnProcess	0xE6DFB80	computed (vector)	VB
+0xA8	ProcIdAndIdxOnProcForChip	0xE6DFD60	+0x98	VB
+0xB0	ProcIdAndIdxOnProcForLogiDevice	0xE6DFE20	+0xA0	VB
+0xB8	ProcessCoordFromId	0xE6DFEE0	+0xA8	VB
+0xC0	ChipIdFromCoord	0xE6E00A0	computed	VB
+0xC8	LogiDeviceIdFromChipCoordAndIdx	0xE6E03C0	computed	VB
+0xD0	ChipCoordAndIdxForLogiDevice	0xE6E06E0	+0xC0	VB
+0xD8	ChipsPerProcessBounds	0xE6E0920	+0xC8 (vector)	VB
+0xE0	ChipBounds	0xE6E0AE0	+0xD0 (vector)	VB
+0xE8	ProcessBounds	0xE6E0CA0	computed (vector)	VB
+0xF0	GetRoutingStrategy	0xE6E0E60	helper	HLP
+0xF8	GetSliceConfig	0xE6E1080	kPlatformConfigs	CFG
+0x100	GetSliceConfigs	0xE6E13A0	kPlatformConfigs	CFG
+0x108	GetDefaultPlatformConfig	0xE6E16A0	kPlatformConfigs	CFG

注意 — 不存在 TpuGetMeshShape / TpuGetCoreIds 方法。Mesh shape 通过组合 ChipBounds + ProcessBounds + ChipsPerProcessBounds 恢复；core ID 通过 LogiDeviceIdFromChipCoordAndIdx ↔ ChipCoordAndIdxForLogiDevice 以及 ProcIdAndIdxOnProcForLogiDevice 恢复。一个已废弃的 args 字符串 PJRT_TpuTopology_ProcessIdAndIndexOnProcessForLogiDeviceOfDefaultType_Args（文件偏移 0x8551469）是 ProcIdAndIdxOnProcForLogiDevice 的改名前身，任何槽位都未引用它。

Args 布局

// Scalar-output (counts, IsSubsliceTopology, ProcessCoordFromId, ChipIdFromCoord, ...)
struct PJRT_TpuTopology_<Scalar>_Args {
    size_t                     struct_size;   // +0x00
    PJRT_TopologyDescription*  topology;      // +0x08 (NO priv)
    int32_t                    input_or_id;   // +0x10 (input methods only)
    int32_t /* or i64 / bool */ output;       // +0x10 or +0x14
};

// Vector-output (ChipBounds, ProcessBounds, ProcessIds, LogiDeviceIdsOnProcess,
//                ChipsPerProcessBounds, ChipCoordAndIdxForLogiDevice, Subslice, ReplaceHostBounds)
struct PJRT_TpuTopology_<Vector>_Args {
    size_t                     struct_size;       // +0x00
    PJRT_TopologyDescription*  topology;          // +0x08
    int64_t                    buffer_capacity;   // +0x10  caller-supplied cap
    int32_t*                   buffer;            // +0x18  caller-allocated
    size_t                     actual_count;      // +0x20  written by callee
};

向量函数体总是先写入 actual_count（+0x20），再按 buffer_capacity（+0x10）做边界检查：如果 actual > capacity，它会返回 xla::InvalidArgument("<class>: needed %zu, provided %d", ...)，且不写入缓冲区。因此调用方会通过 capacity == 0 先调用一次以探测所需大小，读取 actual_count，然后重新分配并再次调用；这就是经典的两遍 C-ABI 向量惯用法。

Helper 方法 — 下转到 TpuTopology

五个方法（ReplaceHostBounds、IsEnhancedBarrierEnabled、HasLimitedIciConnectivity、IsReachableOverLimitedIci、GetRoutingStrategy）读取不在抽象 vtable 中的 TPU 私有字段。它们把 vtable 调用替换为 pjrt::tpu_topology_extension::* 中的自由函数，后者通过 xla::GetTpuTopologyFromDescription(topo)（0xF84AC20）下转为 tpu::TpuTopology*，并直接读取私有偏移：

absl::StatusOr<bool> tpu_topology_extension::IsEnhancedBarrierEnabled(
        const xla::PjRtTopologyDescription& topo) {       // helper @ 0xE6E4C40
    auto t = xla::GetTpuTopologyFromDescription(&topo);    // downcast (StatusOr)
    if (!t.ok()) return t.status();
    return (*t)->_has_eb_field   /* TpuTopology+0x190 */
         ? (*t)->enhanced_barrier_enabled /* +0x188 */ : false;
}

ReplaceHostBounds 更重一些：它通过 tpu::TpuTopologySerdes::Distill（0x20805BC0）在受限 host bounding box 下重新渲染拓扑。这五个方法正是扩展存在的原因：它们暴露了在可移植 XLA 抽象中没有位置的 ICI/barrier 事实。

---

序列化与反序列化

Serialize（槽位 92）

PJRT_TopologyDescription_Serialize 生成一个 XLA 可移植 protobuf blob；线上的字节可由任何 PJRT 插件解析，而不仅限于 libtpu。

function PJRT_TopologyDescription_Serialize(args):       // 0xF870320, slot 92
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TopologyDescription_Serialize_Args", 0x27, 0x38, args[0]) != 1:
        return new PJRT_Error{...}
    topo = (args[2])->impl
    proto_or = topo->vtable[+0xF8]()                      // StatusOr<xla::PjRtTopologyDescriptionProto>  (topo = *(args[2]+8))
    if not proto_or.ok(): return new PJRT_Error{proto_or.status}
    h = operator new(0x18)                                // 24-byte std::libcpp::string object (the handle)
    MessageLite::SerializeToString(&proto_or.value, h)    // writes the proto bytes into *h
    args[5] = h                                           // serialized_handle_out  (+0x28)
    args[3], args[4] = h->data(), h->size()               // serialized_{data,size} out  (+0x18 / +0x20, SSO-aware)
    args[6] = &PJRT_TopologyDescription_Serialize::$_0::__invoke  // serialized_deleter out  (+0x30) @ 0xF876F20
    return NULL                                           // (on serialize failure: PJRT_Error "Failed to serialize PjRtTopologyDescriptionProto.")

handle（args[5]，+0x28）是一个普通的 24 字节堆 std::__u::string；deleter 单独在 args[6]（+0x30）返回，而不是嵌入 handle 中：

// args struct (out fields):
//   +0x18 serialized_data    : const char*   (string contents, interior pointer)
//   +0x20 serialized_size    : size_t
//   +0x28 serialized_handle  : std::string*  (24-byte heap object, opaque to caller)
//   +0x30 serialized_deleter : void (*)(std::string* handle)   -> $_0::__invoke @ 0xF876F20

陷阱 — 调用方必须调用 serialized_deleter(serialized_handle) 来回收堆上的 std::string。这些字节并不由包装器拥有；没有其他东西会释放它们。deleter lambda（...Serialize::$_0::__invoke @ 0xF876F20，仅有符号）调用 delete std::string*（低置信度：函数体由符号识别，未反编译）。

Deserialize（槽位 119）

PJRT_TopologyDescription_Deserialize 未由 TPU 覆盖；它是通用的 pjrt::PJRT_TopologyDescription_Deserialize，会按平台名经由编译器注册表路由。这就是一台主机上序列化的拓扑在另一台主机上重建的方式。

function PJRT_TopologyDescription_Deserialize(args):     // 0xF870B80, slot 119
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TopologyDescription_Deserialize_Args", 0x28, 0x28, args[0]) != 1:
        return new PJRT_Error{...}
    proto = xla::PjRtTopologyDescriptionProto{}           // stack-construct
    if not MessageLite::ParseFromString(&proto, {args->bytes, args->size}):  // 0x21057460
        return new PJRT_Error{InvalidArgument("Failed to parse PjRtTopologyDescriptionProto at the C API level, from binary string of size: %d")}
    name = proto.platform_name()                          // SSO string read directly from the proto
    comp_or = xla::GetDefaultPjRtCompiler(name)           // 0x1D169DA0 -> StatusOr<PjRtCompiler*>
    if not comp_or.ok():
        return new PJRT_Error{comp_or.status}             // "no compiler registered for <name>" status raised inside GetDefaultPjRtCompiler
    topo_or = comp_or.value->vtable[+0x20]( serialized_bytes )  // DeserializeTopology
    if not topo_or.ok(): return new PJRT_Error{topo_or.status}
    args[4] = CreateWrapperDeviceTopology(topo_or.value)  // 0xF873260 (const* overload), out at +0x20
    return NULL

注意 — 线上的字节是 XLA 可移植的，但解析后的拓扑依赖于有一个为 platform_name 注册的编译器。libtpu 会在 "tpu" 和逐代 key 下注册其编译器；若在没有 TPU compiler library 的 CPU-only 主机上反序列化 TPU 拓扑，会在 GetDefaultPjRtCompiler 内部失败，并返回 “no compiler registered” 状态（精确的 “not supported by any compiler / missing a compiler library” 文案不是此二进制中的字面量；如果出现该消息，也源自 compiler-registry library，而不是这里）。PjRtTopologyDescriptionProto schema 只恢复了一部分；platform_name 已确认（在 deserialize 路径中直接读取），完整 proto field-tag map 未追踪。

---

Slice 与平台配置 {#slice--platform-configs}

GetSliceConfig、GetSliceConfigs 和 GetDefaultPlatformConfig（扩展槽位 28-30）不调用拓扑 vtable，而是查询静态单例 kPlatformConfigs @ 0x224C4108（.bss，guard 0x224C4128）：这是一个由 GetAllPlatformConfigs()::$_0::operator() @ 0xE6E1DA0 惰性构建的 flat_hash_map<string, TpuPlatformConfig>，并通过 ConvertSliceConfig @ 0xE6E1BC0 转换为 wire-format C 结构体。

struct PJRT_TpuTopology_SliceConfig {                     // sizeof = 0x1C
    int64_t  n_dims;          // +0x00  mirrors chip_bounds rank (0..3 dims)
    int32_t  chip_bounds[4];  // +0x08  per-axis chip extent
    bool     ici_polarity[3]; // +0x18  per-axis ICI wrap polarity
};

SliceConfig proto（proto2，字段为 repeated int32 chip_bounds、repeated bool ici_polarity，rank 匹配）和更丰富的 TpuPlatformConfig proto（嵌套的 SliceConfig，加上 core_layout、sparse_cores_per_chip、host_per_chip、embedding_cores_per_chip 整数 enum）是 AOT 编译开关：它们在没有活跃设备时告知编译器逐代硅片几何。规范 key 集合（"tpu_v4"、"tpu_v5e"、"tpu_v6e"，等等）映射到 SparseCore 硅片页面中编目的逐代 core/chip 计数。

---

生命周期语义

路径	槽位	返回	销毁责任
_Create	87	args +0x30 处的拥有型包装器（112 B，new）	调用方，通过 _Destroy 恰好一次
_Deserialize	119	args +0x20 处的拥有型包装器（112 B，new）	调用方，通过 _Destroy 恰好一次
Client_TopologyDescription	100	借用句柄	活跃 PJRT_Client（不要调用 _Destroy）

function PJRT_TopologyDescription_Destroy(args):         // 0xF870040, slot 88
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_TopologyDescription_Destroy_Args", 0x25, 0x18, args[0]) != 1:
        return new PJRT_Error{...}
    if args->topology != NULL:
        ~PJRT_TopologyDescription(args->topology)         // 0xF8700C0
        free(args->topology, /*size=*/112)
    return NULL                                           // success

析构函数（0xF8700C0）按顺序释放：缓存属性 span（+0x58）、设备指针 span（+0x40）、PJRT_DeviceDescription vector（+0x28，从后向前遍历并释放每个元素拥有的子分配）、设备 shared-ptr vector（+0x10，运行每个元素的虚析构函数），然后运行拥有型拓扑的虚析构函数（*(wrapper+0x00)，vtable +0x08）。C ABI 之间没有引用计数，也没有 shared-pointer 语义；Create/Deserialize 包装器必须恰好销毁一次。

陷阱 — 对 PJRT_Client_TopologyDescription（槽位 100）返回的句柄调用 PJRT_TopologyDescription_Destroy 是 use-after-free：该拓扑由活跃 client 拥有，并由 PJRT_Client_Destroy 释放。C ABI 没有携带区分 owned 与 borrowed 句柄的标志；区别只存在于哪个槽位产生了该指针。借用路径见 Client, Device & Topology。

线程安全：底层拓扑在构造后为 const，且包装器缓存的 span 从不变异，因此并发读取调用（Attributes、Fingerprint、GetDeviceDescriptions、扩展 getter）是安全的。Destroy 不能安全地与同一包装器上的任何其他槽位并发调用。

---

相关组件

组件	关系
PJRT_Api 140 槽表	持有 10 个通用拓扑槽位（87-93、100、119、138）；slot→address 映射见 vtable 页面
xla::PjRtTopologyDescription	包装器拥有的抽象 impl 对象；在 TPU 上是 tpu::TpuTopology 子类
扩展链	type-16 扩展是挂在 extension_start 后的 17 个扩展之一
kPlatformConfigs 单例	slice-config 方法查询的逐代 TpuPlatformConfig 表
xla::PjRtCompiler 注册表	支撑 _Deserialize：GetDefaultPjRtCompiler(platform_name) → DeserializeTopology

交叉引用

• Client、Device 与 Topology — live-client 设备/拓扑访问器和 PJRT_Client_TopologyDescription（槽位 100）；对比：该页负责借用拓扑，本页负责独立 AOT 拓扑
• API 与 vtable 重建 — 完整 140 槽表、@0x227BA840 .lbss 存储，以及每个槽位开头都会使用的 ActualStructSizeIsGreaterOrEqual 向后兼容机制
• 扩展链 — extension_start 如何链接 17 个扩展；type-16 条目在链中的位置，以及消费者如何遍历到它
• 概览 — C-API 版本、扩展链概念，以及运行 CreateTpuTopologyExtension 的 GetPjrtApi 填充路径
• 可执行文件与执行 — PJRT_Compile 消费一个拓扑以进行提前编译；此处构建的 AOT 拓扑会供给该路径
• ICI 拓扑发现 — 同一环面几何的设备端运行时发现；本页为 AOT 编译静态暴露该几何
• SparseCore 架构 — Part-IV 硅片几何：CoreCountPerChip 和 TpuPlatformConfig 表报告的逐代 core/chip 计数