PJRT 客户端、设备与拓扑

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

摘要

libtpu.so 是一个 PJRT 插件：它导出单个 C 符号 GetPjrtApi，返回一个包含 140 个槽位的函数指针 vtable（PJRT_Api 结构体，见 API 与 vtable 重建）。框架（JAX、PyTorch/XLA）通过这个扁平的 C ABI 驱动整个 TPU 运行时。本页记录客户端、设备、设备描述、内存空间和拓扑槽位背后的行为：每个槽位如何处理它包装的堆对象，以及重新实现者如何从 PJRT_Client_Create 一路重建调用图，直到得到一个完整连线的 PJRT_Client 对象，其中持有设备、内存空间和拓扑句柄。

其形态精确对应上游 pjrt_c_api_wrapper_impl.cc，二进制中仍带有该路径字符串（从 pjrt::CreateWrapperClient 的 0xf872060 引用）。C ABI 背后有三类对象：wrapper 结构体（PJRT_Client、PJRT_Device、PJRT_Memory、PJRT_DeviceDescription、PJRT_TopologyDescription），它们是用 operator new 分配的小型 POD，持有指向真实 C++ 实现的裸指针；implementation 对象（xla::PjRtClient 及其 TPU 子类 xla::TpuClient、xla::PjRtDevice、xla::PjRtMemorySpace、xla::PjRtDeviceDescription）；以及把扁平 PJRT_NamedValue[] 数组转换为带类型 PjRtTpuClientConfig 的 options 管线。140 个槽位中只有 5 个是 TPU 专用的（tpu_plugin::*）；设备/内存/拓扑访问器全都是通用的 pjrt::PJRT_* wrapper，与 CPU 和 GPU 插件共享。TPU 专用化完全发生在这些 wrapper 调用的实现对象内部。

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

PJRT_Client_Create 中的 options-kv 摄取：解析 PJRT_NamedValue[] → flat-hash-map，合并 10 项默认选项表，验证，解析成 PjRtTpuClientConfig，然后通过 GetTpuPjRtClient 以及 MegaScale / TfPjRtClient 客户端选择分支分发。
CreateWrapperClient 中的 wrapper-client 构造：一个 208 字节的 PJRT_Client，会急切物化设备向量、可寻址设备子集、内存空间向量，以及两个 PjRt*→PJRT_* 查找哈希表，因此之后每个访问器都是 O(1) 字段读取。
访问器模式：每个槽位都以 ActualStructSizeIsGreaterOrEqual(name, min, current, args->struct_size) 开头；成功后要么从 wrapper 读取字段，要么对实现对象做一次虚调用。


插件入口	`GetPjrtApi` @ `0xE6A83A0` → `pjrt::tpu_plugin::GetTpuPjrtApi` @ `0xE6AA440`
客户端创建	`pjrt::tpu_plugin::PJRT_Client_Create` @ `0xE6A8840`（槽位 15，TPU 专用）
客户端后端	`xla::GetTpuPjRtClient(const PjRtTpuClientConfig&)` @ `0xF8008C0` → `xla::TpuClient`
Wrapper 构建	`pjrt::CreateWrapperClient(unique_ptr<xla::PjRtClient>)` @ `0xF872060`
`PJRT_Client` 大小	208 字节（`0xD0`），`operator new`
C-API 版本	v0.103（`{major=0, minor=103}`）
Wrapper impl 路径	`third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc`（二进制内字符串）

PJRT_Client_Create — 选项摄取与后端选择

目的

PJRT_Client_Create 是插件中最重的单个函数：它把框架传入的无类型选项数组转换为带类型的 TPU 客户端配置，启动底层 xla::TpuClient（探测硬件、构建设备列表并连接 TPU 运行时），可选地在其上叠加 MegaScale 或 TF-PjRt 客户端，最后把结果包装给 C ABI。它是仅有的五个 TPU 覆盖槽位之一；通用的 pjrt::PJRT_Client_Create 不会安装，槽位 15 指向 tpu_plugin::PJRT_Client_Create（0xE6A8840）。

入口点

text

GetPjrtApi (0xE6A83A0)                                  ── 导出的 thunk，5 字节 JMP
  └─ pjrt::tpu_plugin::GetTpuPjrtApi (0xE6AA440)         ── 对 140 槽位表做一次性 __cxa_guard 初始化
       └─ slot 15 = pjrt::tpu_plugin::PJRT_Client_Create (0xE6A8840)
            ├─ pjrt::ConvertFromPjRtNamedValueList        ── PJRT_NamedValue[] → flat_hash_map<string,variant>
            ├─ pjrt::ValidateCreateOptions                ── 根据 10 项默认表检查键
            ├─ libtpu::telemetry::UptimeMetric::MaybeRefreshInstance  ── 如果存在 ml_framework_name/version
            ├─ pjrt::tpu_plugin::ParseTpuClientConfig     ── 带类型的 PjRtTpuClientConfig
            ├─ xla::GetTpuPjRtClient (0xF8008C0)          ── 构建 xla::TpuClient
            ├─ xla::MegaScalePjRtClient::CreateMegaScalePjRtClient   ── 多 slice 覆盖层（条件式）
            ├─ xla::TfPjRtClient::CreateTfPjRtClient       ── 如果 use_tf_pjrt_client == 1
            └─ pjrt::CreateWrapperClient (0xF872060)       ── 包装成 208 字节 PJRT_Client，存到 args[8]
```text

### 算法

```c
function PJRT_Client_Create(args):                       // 0xE6A8840
    // args[0]=struct_size, args[2]=PJRT_NamedValue* options,
    // args[3]=num_options, args[8]=PJRT_Client** out.
    st = ActualStructSizeIsGreaterOrEqual(               // 0xE6A8840:+0
            "PJRT_Client_Create_Args", min=23, current=88, args[0])
    if st != 1:
        return new PJRT_Error{st}                        // caller compiled too old

    // (1) untyped ingest: caller's flat array -> typed map keyed by string.
    user_opts = ConvertFromPjRtNamedValueList(args[2], args[3])  // -> flat_hash_map<string,variant>

    // (2) build the DEFAULT-OPTIONS table inline on the stack: 10 entries,
    //     each a {key, PJRT_NamedValue_Type, default} record. Keys are
    //     pjrt::tpu_configs::k* string constants. See "Default Options" below.
    default_opts = { kMaxInflightComputations:i64=1, kUseTfPjrtClient:i64=1,
                     kMlFrameworkName:str="", kMlFrameworkVersion:str="",
                     kUseGlobalTpuSystem:bool, kTpuAllowAsyncAllocations:bool,
                     kExecutableCompatibilityCheckOnDeserialization:bool,
                     kThrottleLowPriorityHostTransfers:bool,
                     kPinnedHostAllocationMode:str, kPremappedBufferSize:i64,
                     kMaximumPremappedBufferSizeForTransfersInBytes:i64,
                     kNumPremappedPartitions:i64, kSkipMegascalePjrtClient:bool }
    merged = insert_range(default_opts, user_opts)       // user values win on key collision

    // (3) validate: reject unknown keys / wrong types against default table.
    vstatus = ValidateCreateOptions(merged, default_table)   // 0xE6A8840:LABEL_15
    if vstatus != 1:
        return new PJRT_Error{vstatus}

    // (4) telemetry: only if ml_framework_name AND ml_framework_version set.
    if FLAGS_enable_runtime_uptime_telemetry and merged has both keys:
        UptimeMetric::MaybeRefreshInstance(framework_name, framework_version)

    // (5) typed parse: produce a PjRtTpuClientConfig from the merged map.
    cfg = ParseTpuClientConfig(merged)                   // LABEL_69
    if cfg is error:
        return new PJRT_Error{cfg.status}

    // (6) bring up the underlying TpuClient. This is the heavy step:
    //     hardware enumeration, runtime attach, device/core construction.
    client = GetTpuPjRtClient(cfg)                       // 0xF8008C0 -> xla::TpuClient
    if client is error:
        return new PJRT_Error{client.status}

    // (7) MegaScale overlay: skip if kSkipMegascalePjrtClient,
    //     env SKIP_MEGASCALE_PJRT_CLIENT, or FLAGS_megascale_port < 0.
    skip_ms = merged[kSkipMegascalePjrtClient] | (getenv("SKIP_MEGASCALE_PJRT_CLIENT") != NULL)
    if FLAGS_megascale_port >= 0 and not skip_ms:
        client = MegaScalePjRtClient::CreateMegaScalePjRtClient(client)

    // (8) optional TF-PjRt wrapping (legacy client). use_tf_pjrt_client default = 1.
    if merged[kUseTfPjrtClient] == bool && value == 1:
        client = TfPjRtClient::CreateTfPjRtClient(client)
    // (variant type-mismatch on any lookup -> __throw_bad_variant_access -> LABEL_125)

    // (9) wrap for the C ABI and hand back through the out-param.
    args[8] = CreateWrapperClient(client)                // 0xF872060
    return NULL                                          // success

陷阱 — merged[kUseTfPjrtClient] 和 merged[kSkipMegascalePjrtClient] 查找会在读取值之前先读取 variant 类型标签（*(byte)(slot+48)）。如果调用者用正确的键传入错误的 PJRT_NamedValue 类型，代码会落到 LABEL_125 → std::__throw_bad_variant_access，从而中止进程，而不是返回 PJRT_Error。步骤 (3) 的验证是唯一防护；重新实现必须在 ValidateCreateOptions 期间做类型检查，否则会继承这个崩溃。

默认选项表

默认表在栈上逐字段构建（没有静态数据段；它使用来自 pjrt::tpu_configs::k* 字符串常量和每项类型字节的 vmovaps 内联构造）。这些是 libtpu 接受的唯一选项键；ValidateCreateOptions 会拒绝其他所有键。

Key (`pjrt::tpu_configs::`)	类型	默认值	作用
`kMaxInflightComputations`	i64	1	每个设备的并发 in-flight 执行数
`kUseTfPjrtClient`	i64	1	包装进 `TfPjRtClient`（旧版兼容 shim）
`kMlFrameworkName`	str	""	遥测标签（例如 `"JAX"`）
`kMlFrameworkVersion`	str	""	遥测标签；与名称配对以触发 `UptimeMetric`
`kUseGlobalTpuSystem`	bool	—	附着到共享/全局 TPU 系统，而不是私有系统
`kTpuAllowAsyncAllocations`	bool	—	允许分配器中的异步设备分配
`kExecutableCompatibilityCheckOnDeserialization`	bool	—	在反序列化并加载时验证可执行文件兼容性
`kThrottleLowPriorityHostTransfers`	bool	—	对低优先级 H2D/D2H 传输限速
`kPinnedHostAllocationMode`	str	—	pinned-host 分配器策略
`kPremappedBufferSize`	i64	—	DMA 预映射暂存区域大小
`kMaximumPremappedBufferSizeForTransfersInBytes`	i64	—	预映射传输缓冲区上限
`kNumPremappedPartitions`	i64	—	预映射区域分区数
`kSkipMegascalePjrtClient`	bool	—	强制禁用 MegaScale 多 slice 覆盖层

说明 — 与每个默认键一起存储的类型字节（反编译中 1=bool? 4=i64 等）是 PJRT_NamedValue_Type 枚举。i64 类型的布尔标志（kUseTfPjrtClient 是类型 4、值 1）是有意为之：此构建把“使用 TF client”编码为整数标志，而不是 PJRT_NamedValue_kBool，因此调用者如果在这里传 bool，会触发上面的 bad_variant_access 路径。

后端选择 — 三个客户端层

GetTpuPjRtClient（0xF8008C0）是真正进入 TPU 运行时的入口；它委托给静态 GetTpuPjRtClientInternal，后者拥有一个 tfrt::HostContext 和一个 TpuSystemState，然后构造 xla::TpuClient。TpuClient 构造函数签名（0xF801980）是：

TpuClient(const PjRtTpuClientConfig& cfg,
          int process_index,
          std::string platform_version,
          std::vector<unique_ptr<TpuDevice>> devices,
          ...)
```text

因此设备列表和进程索引由运行时启动流程内部决定，而不是由 C-API 层决定。随后 C-API 层选择要堆叠多少个 wrapper 层：

| 条件 | 结果客户端 | 符号 |
|---|---|---|
| 始终 | 基础 `xla::TpuClient` | `xla::GetTpuPjRtClient` @ `0xF8008C0` |
| `megascale_port >= 0` 且未跳过 | 包装基础客户端的 `MegaScalePjRtClient` | `CreateMegaScalePjRtClient` @ `0xE6EA680` |
| `use_tf_pjrt_client == 1`（默认） | 包装步骤 (8) 产物的 `TfPjRtClient` | `TfPjRtClient::CreateTfPjRtClient` |

MegaScale 层是多 slice（跨 pod）覆盖层；它消耗基础 `unique_ptr<TpuClient>` 和一个可选 `MultiSliceConfig`。TF 层是旧版兼容客户端。因为两者都会消耗并替换 `client`，最终包装对象可以有一层、两层或三层，但 C ABI 只看到最外层的 `xla::PjRtClient*`。

---

## CreateWrapperClient — 构建 PJRT_Client

### 目的

`CreateWrapperClient`（`0xF872060`）接管 `unique_ptr<xla::PjRtClient>` 的所有权，并物化 C ABI 发放的 208 字节 `PJRT_Client` wrapper。它的定义性行为是**急切缓存**：它只遍历一次实现客户端的设备和内存列表，为每个对象分配 `PJRT_Device` / `PJRT_Memory` wrapper 对象，并构建两个 flat-hash-map（`xla::PjRtDevice* → PJRT_Device*` 和 `xla::PjRtMemorySpace* → PJRT_Memory*`）以及可寻址设备子集向量。因此之后每个设备/内存访问器都只是指针追踪，不分配内存。

### Wrapper 布局（`PJRT_Client`，208 字节）

根据 `CreateWrapperClient` 中的字段写入重建；偏移量相对于 `operator new(0xD0)` 块。反编译器把对象索引为 `_QWORD[]`，因此 qword 索引 `n` = 字节偏移 `8n`。

| 字段 | 偏移 | 类型 | 含义 |
|---|---|---|---|
| `client` | `+0x00` | `xla::PjRtClient*`（owned） | 实现对象；vtable 驱动所有行为 |
| (reserved) | `+0x08` | `_QWORD[3]`（qw 1..3） | 构造期间写入的临时/辅助字段；设备访问器不读取 |
| `devices` | `+0x20` | `vector<PJRT_Device>`（begin/end/cap 位于 qw 4..6） | 所有 wrapper 设备；每项 9 个 qword（72 B）。`PJRT_Client_Devices` 读取 begin=qw4、end=qw5 |
| `addressable_devices` | `+0x38` | `vector<PJRT_Device*>`（begin/end/cap 位于 qw 7..9） | `device->IsAddressable()` 为 true 的子集。`PJRT_Client_AddressableDevices` 读取 begin=qw7、end=qw8 |
| `device_map` | `+0x50` | `flat_hash_map<PjRtDevice*, PJRT_Device*>`（qw 10..13） | O(1) impl→wrapper；由 `GetCDevice` 在 wrapper+0x50 读取 |
| `memories` | `+0x70` | `vector<PJRT_Memory>`（qw 14..16） | 所有 wrapper 内存空间；每项 5 个 qword（40 B） |
| `addressable_memories` | `+0x88` | `vector<PJRT_Memory*>`（qw 17..19） | 每设备可寻址子集 |
| `memory_map` | `+0xA0` | `flat_hash_map<PjRtMemorySpace*, PJRT_Memory*>`（qw 20..23） | O(1) impl→wrapper；由 `GetCMemory` 在 wrapper+0xA0 读取 |

### 算法

```c
function CreateWrapperClient(client_uptr):               // 0xF872060
    w = operator new(0xD0)                                // 208-byte PJRT_Client
    PJRT_Client::PJRT_Client(w, client_uptr)              // store owned client at +0x00

    // (1) DEVICES: walk client->devices() (vtable+40), reserve, wrap each.
    devs = w.client->devices()                            // vtable +40
    w.devices.reserve(devs.size())
    for d in devs:
        attrs = d->description()->Attributes()            // PjRtDevice vtable+32 -> desc vtable+56
        pjrt_dev = PJRT_Device{ device=d, description_attrs=PopulatePjrtAttributes(attrs) }
        w.devices.push_back(pjrt_dev)
        if d->IsAddressable():                            // PjRtDevice vtable+24
            w.addressable_devices.push_back(&w.devices.back())
        w.device_map[d] = &w.devices.back()               // SwissTable insert

    // (2) INVARIANT CHECK (fatal on mismatch):
    CHECK(w.addressable_devices.size() == w.client->addressable_device_count())
        // vtable+32, file pjrt_c_api_wrapper_impl.cc:3435, msg quoted below

    // (3) MEMORIES: walk client->memory_spaces() (vtable+80), wrap each.
    mems = w.client->memory_spaces()                      // vtable +80
    w.memories.reserve(mems.size())
    for m in mems:
        w.memories.push_back(PJRT_Memory{ memory=m })
        w.memory_map[m] = &w.memories.back()              // CHECK iter != c_memory_map.end()

    // (4) cross-link addressable memories per addressable device
    //     (client->addressable_memory_spaces, vtable+144) and
    //     fix up each device's addressable-memory list via the device_map
    //     (CHECK "iter != c_device_map.end()", line 101 / 110).
    for d in w.addressable_devices:
        for m in d->memory_spaces():                      // vtable+144
            d_wrapper.addressable_memories.push_back(memory_map[m])
    return w

陷阱 — CreateWrapperClient 会在三个不变量上中止（通过 pjrt_c_api_wrapper_impl.cc 中的 LogMessageFatal），而不是返回错误：addressable_devices.size() == client->addressable_device_count()（第 3435 行）、iter != c_memory_map.end()（第 110 行）和 iter != c_device_map.end()（第 101 行）。如果重新实现的 xla::PjRtClient 的 devices()、addressable_device_count() 和 memory_spaces() 彼此不一致，会在客户端创建时、任何执行之前使宿主进程崩溃。
特性 — 设备属性会在包装时快照（PopulatePjrtAttributes 把描述的属性映射复制到 PJRT_Device）。C-ABI 的 PJRT_DeviceDescription_Attributes 槽位读取这个缓存副本，而不是实时实现映射；Client_Create 之后的属性变更不会跨边界可见。

设备枚举与查找

目的

一旦 wrapper 持有急切缓存，客户端级设备槽位就很简单：PJRT_Client_Devices 和 PJRT_Client_AddressableDevices 返回缓存向量中的 span；PJRT_Client_LookupDevice / LookupAddressableDevice 按全局设备 id 索引。它们都不分配，也不接触实现对象。

算法

function PJRT_Client_Devices(args):                      // 0xF85F600, slot 20
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_Client_Devices_Args", 24, 40, args[0]) != 1:
        return new PJRT_Error{...}
    client = args[2]                                      // PJRT_Client*
    args[3] = client->devices.begin                       // +0x20 within wrapper = qw 4
    args[4] = client->devices.end                         // +0x28 = qw 5; count = (end-begin)/sizeof(PJRT_Device)
    return NULL
```text

`PJRT_Client_AddressableDevices`（`0xF85F660`，槽位 21）除了返回 `addressable_devices` span（qw 7..8）之外逐字节相同。调用者通过指针相减计算数量；`Devices` 的元素步长是 `sizeof(PJRT_Device)`（72 字节），`AddressableDevices` 的元素步长是 8 字节（一个指针）。这是重新实现者必须保留的细微差异，因为 PJRT 客户端会用不同元素类型迭代这两个数组。

### 函数映射

| 槽位 | 函数 | 地址 | 行为 |
|---|---|---|---|
| 20 | `PJRT_Client_Devices` | `0xF85F600` | `PJRT_Device` 的 span（begin/end） |
| 21 | `PJRT_Client_AddressableDevices` | `0xF85F660` | `PJRT_Device*` 的 span |
| 22 | `PJRT_Client_LookupDevice` | `0xF85F6C0` | id → `PJRT_Device*`，通过缓存查找 |
| 23 | `PJRT_Client_LookupAddressableDevice` | `0xF85F880` | id → 可寻址 `PJRT_Device*` |
| 18 | `PJRT_Client_ProcessIndex` | `0xF85F440` | 来自 `client->process_index()` 的 int |
| 17 | `PJRT_Client_PlatformName` | `0xF85F4A0` | string view（`"tpu"`） |
| 19 | `PJRT_Client_PlatformVersion` | `0xF85F500` | string view（TPU 代际/拓扑） |

---

## 设备与 DeviceDescription 访问器

### 目的

`PJRT_Device` 在 `PJRT_Device+0x00` 持有实现 `xla::PjRtDevice*`，并在 `PJRT_Device+0x08` 内联一个 `PJRT_DeviceDescription` 子对象；该描述又持有它分发到的 `xla::PjRtDeviceDescription*`。访问器划分对应上游 PJRT：`PJRT_Device_*` 槽位处理运行时状态（可寻址性、内存、统计信息）；`PJRT_DeviceDescription_*` 槽位处理静态身份（id、进程索引、kind、坐标）。`PJRT_Device_GetDescription`（`0xF8659A0`，槽位 34）是桥梁，它只是返回 `PJRT_Device+0x08`，也就是指向内联描述子对象的指针（`a1[3] = a1[2] + 8`）。

### 算法 — 访问器惯用形态

每个 `DeviceDescription` 访问器都遵循同一种形态：验证 args 结构体大小，然后通过实现对象的 vtable 做一次虚调用，并把结果存到输出字段。

```c
function PJRT_DeviceDescription_Id(args):                // 0xF865360, slot 28
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_DeviceDescription_Id_Args", 30, 28, args[0]) != 1:
        return new PJRT_Error{...}
    // args[2] = PJRT_DeviceDescription*; **(args+16) = xla::PjRtDeviceDescription*
    desc = *(args[2])                                     // load impl ptr
    args[3] = desc->vtable[+16]( desc )                   // virtual: id() -> int32
    return NULL

ProcessIndex（槽位 29）调用 vtable +...，Kind（槽位 31）返回 string view（设备 kind 字符串，例如 "TPU v4"），DebugString/ToString 返回格式化字符串。Coordinates 不是专用的 v0.103 槽位；torus (x,y,z) 坐标通过 PJRT_DeviceDescription_Attributes（槽位 30）作为命名属性暴露，这些属性来自描述的属性映射，并在包装时缓存（见上面的 PopulatePjrtAttributes）。重新实现 TPU 设备时，必须把 coords / core_on_chip / slice_index 暴露为属性项，而不是一等访问器槽位。

函数映射

槽位	函数	地址	返回
28	`PJRT_DeviceDescription_Id`	`0xF865360`	`int32` device id（vtable +16）
29	`PJRT_DeviceDescription_ProcessIndex`	`0xF8653C0`	`int32` process index
30	`PJRT_DeviceDescription_Attributes`	`0xF865420`	named-value list（携带 coords）
31	`PJRT_DeviceDescription_Kind`	`0xF865480`	string view（设备 kind）
32	`PJRT_DeviceDescription_DebugString`	`0xF865500`	string view
33	`PJRT_DeviceDescription_ToString`	`0xF8658A0`	string view
34	`PJRT_Device_GetDescription`	`0xF8659A0`	`PJRT_DeviceDescription*`
35	`PJRT_Device_IsAddressable`	`0xF865A00`	bool（vtable +24）
36	`PJRT_Device_LocalHardwareId`	`0xF865A60`	`int32` local hw id
37	`PJRT_Device_AddressableMemories`	`0xF865AC0`	`PJRT_Memory*` 的 span
38	`PJRT_Device_DefaultMemory`	`0xF865B20`	`PJRT_Memory*`（vtable +152 → `GetCMemory`）
39	`PJRT_Device_MemoryStats`	`0xF865CE0`	HBM stats struct

说明 — 每个 ActualStructSizeIsGreaterOrEqual 调用中的 (min, current) 对是向后兼容包络：PJRT_DeviceDescription_Id_Args 接受最小 30 字节，而当前是 28。注意这里 min > current，这是 v0.103 args 结构体相对记录的最小值收缩了一个保留字段所留下的产物；检查是 actual >= min。重新实现者应逐槽复制精确常量，而不是假定统一规则。

设备 ↔ 内存空间连线

目的

PJRT 内存空间（PJRT_Memory，包装 xla::PjRtMemorySpace*）建模 TPU 设备可以寻址的不同内存种类（设备 HBM、pinned host 等）。连线是双向的，并且完全由 CreateWrapperClient 预先缓存：客户端持有完整 memories 向量和一个 memory_map；每个可寻址设备持有自己的 addressable_memories 子集。

算法 — DefaultMemory

PJRT_Device_DefaultMemory 展示了通过客户端缓存的 memory_map 做 impl→wrapper 转换：

function PJRT_Device_DefaultMemory(args):                // 0xF865B20, slot 38
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_Device_DefaultMemory_Args", 30, 32, args[0]) != 1:
        return new PJRT_Error{...}
    device = *(args[2])                                   // xla::PjRtDevice*
    statusor = device->vtable[+152]()                     // default_memory_space() -> StatusOr<PjRtMemorySpace*>
    if statusor.ok():
        // translate impl memory-space ptr to its wrapper via the CLIENT's map.
        args[3] = GetCMemory( device->client (at +0x40), statusor.value )
        return NULL
    return new PJRT_Error{ statusor.status }              // refcount the StatusRep
```text

`GetCMemory` 在所属客户端的 `memory_map` 中查找 `xla::PjRtMemorySpace*`（可从设备在实现偏移 `+0x40` 处的客户端反向指针到达），并返回缓存的 `PJRT_Memory*`。这就是为什么 wrapper 必须急切构建 `memory_map`：`default_memory_space()` 的结果是一个*实现*指针，需要解析回它的 C-ABI wrapper；如果惰性执行该查找，要么每次调用都分配，要么需要 C ABI 并不携带的反向引用。

### 内存访问器（槽位 40–44）

连续的 `PJRT_Memory_*` 槽位 40–44（`Id`、`Kind`、`DebugString`、`ToString`、`AddressableByDevices`）是通用的 vtable 跳转，进入 `xla::PjRtMemorySpace`，形态与设备描述访问器相同。第六个内存访问器 `PJRT_Memory_Kind_Id`（`0xF865F60`）是*后期添加*，位于槽位 **102**，不在 40–44 这一组中。`PJRT_Memory_AddressableByDevices`（`0xF8660C0`，槽位 44）返回可寻址该内存空间的 `PJRT_Device*` span，同样通过客户端的 `device_map` 解析。这些在槽位层面记录于 [Buffer 与 Memory](buffer-and-memory.md)；本页只负责*连线*，即这些映射的存在，以及每个 memory↔device 交叉引用都是缓存查找，而不是实时扫描。

---

## 拓扑描述

### 目的

拓扑描述是在没有活动客户端时的设备 fabric 几何信息。JAX 会将其用于 ahead-of-time 编译，因为它必须在任何硬件附着之前知道 TPU torus 形状。libtpu 暴露两层接口：任意 PJRT 消费者都理解的七个通用 `PJRT_TopologyDescription_*` 槽位，以及一个携带 torus 几何的 TPU 专用扩展（type id 16）。本页覆盖客户端→拓扑链接（槽位 100）；通用槽位和扩展已在 [Topology Description Extension](ext-topology-description.md) 完整重建。

### PJRT_Client_TopologyDescription

`PJRT_Client_TopologyDescription`（`0xF85F560`，槽位 100）返回已经附着到活动客户端的拓扑，即一个*借用*句柄，调用者**不得**销毁它。

```c
function PJRT_Client_TopologyDescription(args):          // 0xF85F560, slot 100
    if ActualStructSizeIsGreaterOrEqual(
           "PJRT_Client_TopologyDescription_Args", 36, 32, args[0]) != 1:
        return new PJRT_Error{...}
    client = args[2]                                      // PJRT_Client*
    // the topology lives as a StatusOr<topology*> at client+192/+200:
    statusrep = *(client + 192)                           // StatusRep* (or the ok-sentinel)
    if (statusrep & 1) == 0:                              // heap StatusRep -> bump refcount
        atomic_inc(statusrep)
    if statusrep != ok_sentinel:                          // error case
        return new PJRT_Error{ statusrep }
    if *(client + 192) != ok_sentinel:                    // belt-and-braces ok check
        StatusOr::Helper::Crash(...)
    args[3] = *(client + 200)                             // borrowed PJRT_TopologyDescription*
    return NULL

特性 — 拓扑句柄作为 StatusOr 存在 PJRT_Client wrapper 内部（位于 +192/+200，在设备/内存缓存之后），并在创建客户端时填充。如果底层 TPU 客户端无法生成拓扑，这个槽位会返回已存储的错误，因此拓扑构建失败只会在调用 PJRT_Client_TopologyDescription 时浮现，而不是在 Client_Create 时浮现。返回的指针归客户端所有；对它调用 PJRT_TopologyDescription_Destroy 是 use-after-free。

拓扑槽位映射（链接，不重复列出）

关注点	槽位	所属页面
通用 create / destroy / serialize / 7 槽位接口	87–93	ext-topology-description
拓扑 fingerprint、deserialize、后期添加项	100, 119, 138	ext-topology-description
TPU torus 扩展（type 16，31 个方法，272 B）	extension chain	ext-topology-description
客户端→拓扑借用（本页）	100	此处

组件	关系
`PJRT_Api` 140 槽位表	这些访问器填充的 vtable；slot→address 映射位于 vtable 页面
`xla::TpuClient`	实现客户端；由 `GetTpuPjRtClient` 构建，由 wrapper 拥有
`MegaScalePjRtClient` / `TfPjRtClient`	`PJRT_Client_Create` 堆叠的可选客户端层
`PjRtTpuClientConfig`	从合并后的 options map 由 `ParseTpuClientConfig` 生成的带类型配置
Extension chain	挂在 `extension_start` 后的 17 个扩展；topology（type 16）是其中之一

交叉引用

API 与 vtable 重建 — 完整 140 槽位表：slot → libtpu impl VA、@0x227BA840 .lbss 存储、ActualStructSizeIsGreaterOrEqual 兼容机制
概览 — C-API 版本、extension-chain 思路、GetPjrtApi 填充路径
Topology Description Extension — 通用拓扑槽位 + 本页从槽位 100 链接到的 TPU torus 扩展（type 16）
Buffer 与 Memory — PJRT_Buffer_* 和 PJRT_Memory_* 槽位级细节；本页连线到设备的内存空间对象
Executable 与 Execution — PJRT_LoadedExecutable_Execute 及相关项，即消耗这里枚举设备的热路径
Runtime Overview — xla::TpuClient 运行时侧：设备构造、GetTpuPjRtClient 背后的 TPU 系统状态

PJRT 客户端、设备与拓扑 ​

摘要 ​

PJRT_Client_Create — 选项摄取与后端选择 ​

目的 ​

入口点 ​

默认选项表 ​

后端选择 — 三个客户端层 ​

设备枚举与查找 ​

目的 ​

算法 ​

函数映射 ​

设备 ↔ 内存空间连线 ​

目的 ​

算法 — DefaultMemory ​

拓扑槽位映射（链接，不重复列出） ​

相关组件 ​

交叉引用 ​

PJRT 客户端、设备与拓扑

摘要

PJRT_Client_Create — 选项摄取与后端选择

目的

入口点

默认选项表

后端选择 — 三个客户端层

设备枚举与查找

目的

算法

函数映射

设备 ↔ 内存空间连线

目的

算法 — DefaultMemory

拓扑槽位映射（链接，不重复列出）

相关组件

交叉引用