Internal Wiki

主题

TpuTopology 与 TpuCoreLocation ​

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，ELF x86-64 DYN，未剥离；IDA 恢复的 C 名称和 demangled C++ 符号按原样引用）。.text VMA 等于文件偏移。其他版本会有所不同。

摘要 ​

TpuTopology_* 和 TpuCoreLocation_* 是 C-ABI accessor 名册，将 TPU pod 的物理 torus 几何暴露给开源 StreamExecutor 后端。它们是拓扑栈中最薄的一层：每个都是一两行的 extern "C" 自由函数，从 tpu::TpuTopology（又名 TpuTopologyExternal）或 tpu::TpuCoreLocation 对象读取单个字段 —— chip-bounds 范围、host/chip 计数、按 core type 的 device 计数、chip 坐标三元组、core index、global id。它们从二进制中的 learning/45eac/tfrc/executor/stream_executor/tpu_util_c_api.cc 引用恢复，是设备端运行时发现、AOT compiler 消费的几何信息的 C 表面。

这些 accessor 位于本 wiki 其他页面记录的两个更丰富拓扑表面之下，这种对比正是重点。PJRT TopologyDescription extension (type 16) 是现代、面向框架的拓扑对象：一个 272 字节结构体，包含 31 个函数指针（ChipBounds、CoreCountPerChip、ProcessBounds、slice configs），包装 xla::PjRtTopologyDescription 并经由其 vtable 跳转。这里的 TpuTopology_* 名册是底层的旧式 SE C-ABI：没有 args-struct size 协商，没有 PJRT_Error 返回，没有 vtable 跳转 —— 只是 return *(uint32*)(topology + offset)。PJRT extension 和 SE 名册暴露重叠的几何信息（两者都有 chip bounds 和 cores-per-chip 的概念），但它们是相互独立的 C 表面，通过不同路径到达同一组硅片事实。本页负责 SE 名册和几何字段映射；PJRT 页面负责 extension。

该名册像每个其他 SE C-ABI cluster 一样通过 ExecutorApiFn() 函数指针表到达（accessor 模式见 shim overview）。host 侧 C++ wrapper tensorflow::tpu::TpuTopologyExternal::LogicalDevicesPerHost（0x20818F80）字面上就是 ExecutorApiFn()[+632](handle, core_type) —— 它加载表、索引槽位，并调用 plugin 侧 TpuTopology_LogicalDevicesPerHost（0xEABBFC0），后者执行字段读取。重新实现者需要复现两半：索引槽位的 wrapper，以及读取偏移的 C-ABI impl。

对于重新实现，契约是：

  • 名册 —— 17 个 TpuTopology_* 和 4 个 TpuCoreLocation_* extern "C" 函数、它们的地址，以及各自支撑的 C++ accessor 或结构体字段。
  • 几何字段映射 —— 每个 accessor 读取 tpu::TpuTopology / tpu::TpuCoreLocation 的哪个字节偏移（ChipBounds_X → +88，HostCount → +108，Index → +52，等等），以及 +124/+128 处按 core type 的步进数组。
  • core-type 重映射怪异点 —— LogicalDevicesPerHost / PerChip 先通过 if t != 1: t = 2*(t==2) 折叠 TpuCoreType enum 再分发；Version 将内部 codename ordinal {0,1,2,3} → {1,2,3,4}，其他值返回 0。
  • 两种输出约定 —— 标量返回（return *(uint32*)(...)）与坐标三元组 out-param 形式（ChipCoordinates / HostCoordinates 写入三个 _DWORD* out-param）。
名册前缀TpuTopology_*（17）、TpuCoreLocation_*（4）
C-ABI 块TpuTopology_* @ 0xEABBFC0–0xEABC2A0；TpuCoreLocation_* @ 0xEABC2C0–0xEABC340
可用性附加项TpuTopology_AvailableCoreCount/CoresPerChip/MaybeAvailableSparseCoresPerLogicalDevice @ 0xF6A1CA0–0xF6A1EA0
源文件（已恢复）…/stream_executor/tpu_util_c_api.cc（来自 .rodata 日志字符串）
到达方式ExecutorApiFn() 槽（例如 LogicalDevicesPerHost = ExecutorApiFn()[+632]）
后备 C++ 对象tpu::TpuTopology（== tensorflow::tpu::TpuTopologyExternal）/ tpu::TpuCoreLocation
输出形式标量 uint32 返回；坐标三元组 _DWORD* out-param
证据等级重新实现级 / 已按 IDA 反编译逐字节确认

范围 — *ApiFn() accessor 模式、不透明句柄约定和名册映射见 shim overview（链接，不在此重新推导）。构造拓扑的 TpuPlatform_* / TpuNodeContext_* 生命周期见 TpuPlatform & TpuNodeContext。现代 PJRT 拓扑对象见 TopologyDescription Extension。设备端 torus 形状发现见 ICI Topology Discovery。完整 tpu::TpuTopology 结构布局（这些 accessor 索引的 Target+0x3b8 按 codename 的几何）由硅片页面负责 —— SparseCore Architecture —— 本页只摘要这些 accessor 触及的偏移。

1. 专用化的 Accessor 模式 ​

目的 ​

本页每个函数都是一次 SE shim 调用的 plugin 侧一半。host 侧一半是 TpuTopologyExternal::<Method> C++ wrapper，它加载 ExecutorApiFn() 表并调用固定槽位。两半在 init 时连接（bootstrap 用匹配的 TpuTopology_* impl 地址填充槽位），只有 TfTpu_ExecutorApiFn* 指针跨越 ABI 边界。拓扑对象自身从不作为 C++ 类型跨越边界 —— host 将它作为不透明 void* 持有，并在每次调用时传回。

入口点 ​

LogicalDevicesPerHost 的工作路径，已在反编译中确认：

text
xla / SE TPU backend
  └─ tensorflow::tpu::TpuTopologyExternal::LogicalDevicesPerHost(core_type)   0x20818F80  (host-side C++ wrapper)
       └─ tbl = stream_executor::tpu::ExecutorApiFn()                          0x20819360  — singleton fn-ptr struct
       └─ tbl[+632]( *(void**)this, core_type )                                — call through the slot
              │  (slot populated at init to point at:)
              ▼
       TpuTopology_LogicalDevicesPerHost                                       0xEABBFC0   — plugin-side C-ABI impl
              └─ tpu::TpuTopology::LogicalDevicesPerHost(topo, core_type')      0x20AD3920  — after enum remap
```text

> **说明 —** wrapper 在传参前先解引用句柄一次（`*(void**)this`）：`TpuTopologyExternal` host 对象在其 `+0` 处持有真正的 `tpu::TpuTopology*`，这个内部指针就是每个 `TpuTopology_*` impl 接收的 `a1`。accessor 的偏移相对于*那个*内部对象，而不是 host wrapper。

### 两种输出约定

名册可以清晰分为两种调用形状，重新实现者必须逐一精确匹配：

```c
// (A) scalar return — the majority. Reads one uint32 field and returns it.
//     e.g. TpuTopology_HostCount @ 0xEABC000
uint32 TpuTopology_HostCount(void* topo):
    return *(uint32*)(topo + 108)

// (B) coordinate-triple out-param — TpuCoreLocation_ChipCoordinates / _HostCoordinates.
//     Three int32* out-params plus a redundant scalar return: _HostCoordinates returns
//     the x read (loc+8), _ChipCoordinates returns the last-written triple element (z).
//     e.g. TpuCoreLocation_HostCoordinates @ 0xEABC300
int32 TpuCoreLocation_HostCoordinates(void* loc, int32* x, int32* y, int32* z):
    *x = *(uint32*)(loc + 8); *y = *(uint32*)(loc + 12); *z = *(uint32*)(loc + 16)
    return *(uint32*)(loc + 8)

这个表面没有 PJRT_Error/absl::Status 管线，也没有 args-struct size 检查 —— 那些属于上层 PJRT 层。SE C-ABI 信任调用方并无条件读取字段；唯一的有效性门控位于可用性和版本 accessor 内部的 BUG()/return 0 guard（§4）。

2. TpuTopology_ 名册 ​

函数映射 ​

14 个规范 accessor 形成连续块 0xEABBFC0–0xEABC2A0；三个 "available" accessor 位于 0xF6A1CA0–0xF6A1EA0 的单独块中，因为它们先从 mesh/ops-util singleton 获取拓扑，而不是将其作为 a1 接收。"读取"列给出 tpu::TpuTopology（内部对象）上的字段偏移，或其 thunk 到的 C++ member。

函数地址读取 / 支撑输出
TpuTopology_LogicalDevicesPerHost0xEABBFC0enum 重映射后的 tpu::TpuTopology::LogicalDevicesPerHost（0x20AD3920）标量
TpuTopology_LogicalDevicesPerChip0xEABBFE0enum 重映射后的 tpu::TpuTopology::LogicalDevicesPerChip标量
TpuTopology_HostCount0xEABC000*(uint32*)(topo + 108)标量
TpuTopology_ChipsPerHost0xEABC020*(uint32*)(topo + 116)标量
TpuTopology_ChipBounds_X0xEABC040*(uint32*)(topo + 88)标量
TpuTopology_ChipBounds_Y0xEABC060*(uint32*)(topo + 92)标量
TpuTopology_ChipBounds_Z0xEABC080*(uint32*)(topo + 96)标量
TpuTopology_HasChip0xEABC0A0tpu::TpuTopology::HasChip(topo, x, y, z)bool
TpuTopology_CoreForId0xEABC0E0thunk → tpu::TpuTopology::LogicalDeviceForIdTpuCoreLocation
TpuTopology_Core0xEABC100tpu::TpuTopology::Core(topo, type, x, y, z, idx)TpuCoreLocation
TpuTopology_NumCores0xEABC140tpu::TpuTopology::logical_devices()（计数）标量
TpuTopology_Cores0xEABC160tpu::TpuTopology::logical_devices() → 填充 TpuCoreLocation*[]数组填充
TpuTopology_IdForHost0xEABC260tpu::TpuTopology::IdForHost(topo, x, y, z)标量
TpuTopology_Version0xEABC2A0**(uint32**)(topo + 8) → codename ordinal，经重映射标量
TpuTopology_AvailableCoreCount0xF6A1CA0*(uint32*)(GetTpuTopology() + 12*type + 128)标量
TpuTopology_AvailableCoresPerChip0xF6A1DE0*(uint32*)(GetTpuTopology() + 12*type + 124)标量
TpuTopology_MaybeAvailableSparseCoresPerLogicalDevice0xF6A1EA0xla::jellyfish::NumEmbeddingDevices(...)（StatusOr<int>）StatusOr

怪异点 — 名册混合了纯字段读取和 member thunk，重新实现时区分它们至关重要。ChipBounds_X/Y/Z、HostCount、ChipsPerHost 是扁平的 *(uint32*)(topo+off) 读取 —— bounds 和 counts 已在 topology 对象上预先物化为标量。但 HasChip、Core、IdForHost、CoreForId 会调用 tpu::TpuTopology member functions，遍历 chip/core 布局。只复制字段偏移能复现廉价 accessor，却不能复现查找；只复制 thunk 又会漏掉通用几何是无计算的扁平读取这一点。

几何字段映射 ​

这些 accessor 触及的偏移是 tpu::TpuTopology 布局的一个窗口（完整结构见 silicon pages）。C-ABI 函数体直接确认的内容如下：

偏移字段读取者说明
+8Target*（内嵌）Version（双重解引用 **）Target 的第一个 uint32 是 codename ordinal
+52core-on-chip indexTpuCoreLocation_Index（在 core-location 上，见 §3）—
+88chip_bounds.xChipBounds_Xtorus 范围，X 轴
+92chip_bounds.yChipBounds_Ytorus 范围，Y 轴
+96chip_bounds.zChipBounds_Ztorus 范围，Z 轴
+108host_countHostCountpod 中 host 总数
+116chips_per_hostChipsPerHost连接到一个 host 的 chip 数
+124 + 12·typeavailable cores-per-chip[type]AvailableCoresPerChip按 core type 步进的条目
+128 + 12·typeavailable core-count[type]AvailableCoreCount按 core type 步进的条目

说明 — 带 12·type 步长的 +124/+128 读取，正是硅片页面描述的同一个按 core type 的几何数组，在更大结构内位于 Target+0x3b8 偏移：每个 TpuCoreType（TensorCore / SparseCore / embedding）都有一个 12 字节记录，保存 cores-per-chip 和 core-count。C-ABI 暴露该记录三个字段中的两个。完整记录和 codename→geometry 表由 SparseCore Architecture 负责；本页只确认 AvailableCoresPerChip 和 AvailableCoreCount 按 type 用 12 字节步长索引它，并带有 type < 3 边界（越界 BUG()）。

算法 — 两个非平凡函数体 ​

c
function TpuTopology_Version(topo):                      // 0xEABC2A0
    ordinal = **(uint32**)(topo + 8)                      // first u32 of the embedded Target
    if ordinal < 4:
        return ordinal + 1                                // internal {0,1,2,3} -> public {1,2,3,4}
    return 0                                              // unknown codename -> 0

function TpuTopology_AvailableCoreCount(mesh_state, core_type):   // 0xF6A1CA0
    if mesh_state != NULL:
        topo = tensorflow::TpuMeshCommonState::tpu_topology(*mesh_state)   // VLOG "from mesh_state"
    else:
        topo = tensorflow::tpu_ops_util::GetTpuTopology()                 // VLOG "from tpu_ops_util"
    if core_type >= 3: BUG()                              // only 3 core types
    return *(uint32*)(topo + 12*core_type + 128)          // strided per-type read
```text

> **陷阱 —** `AvailableCoreCount` 接收的是 **mesh state**，而不是拓扑句柄，并根据 `mesh_state` 是否为空用两种不同方式解析拓扑（`else` 分支回退到进程全局 `tpu_ops_util::GetTpuTopology()`）。同级的 `AvailableCoresPerChip` / `MaybeAvailableSparseCoresPerLogicalDevice` *只*走 ops-util 路径，并在没有注册拓扑时返回默认值（`4`）或 `absl::Status` 错误（`"TPU system is not available"`）。若重新实现者把这三个 "available" accessor 都当作从传入 topology 做纯字段读取，就会错误处理这些函数专门为了支持的无设备情形。

### core-type 重映射

`LogicalDevicesPerHost` 和 `LogicalDevicesPerChip` 会在分发前折叠传入的 `TpuCoreType` enum：

```c
function TpuTopology_LogicalDevicesPerHost(topo, core_type):   // 0xEABBFC0
    if core_type != 1:
        core_type = 2 * (core_type == 2)                  // {0->0, 1->1, 2->2, other->0}
    return tpu::TpuTopology::LogicalDevicesPerHost(topo, core_type)   // 0x20AD3920

该重映射会将除 1 和 2 之外的每个 enum 值折叠到 0，把越界或 "default" core type 规范化为第一个槽位。这是 C-ABI 边界处的防御性 clamp，因为 host 可能传入来自更新/更旧 enum 的 TpuCoreType，plugin 无法识别；内部 member function 随后索引固定三项表，避免越界读取。

注意事项 ​

TpuTopology_Cores（0xEABC160）是唯一具有非平凡函数体的 accessor：它获取 logical_devices() span，并为每一项发出一个 TpuCoreLocation*，把 result + 56·i 指针（TpuCoreLocation 步长为 56 字节）物化到调用方缓冲区。反编译中用 AVX（vpmuludq/vpaddq）向量化了指针步长写入，但语义只是简单的 for i in 0..n: out[i] = base + 56*i。56 字节步长即已确认的 sizeof(tpu::TpuCoreLocation)，也与 §3 读取的字段偏移一致（+52 是最后一个 4 字节字段）。TpuTopology_NumCores 返回同一 span 的元素数，也是调用方传给 _Cores 的大小 —— 典型的 count-then-fill C-ABI 对。

3. TpuCoreLocation_ 名册 ​

目的 ​

tpu::TpuCoreLocation 是由 TpuTopology_Core / _CoreForId / _Cores 生成的逐 core 坐标记录：一个 logical device 位于 torus 中何处。四个 TpuCoreLocation_* accessor 将其字段暴露给 host —— 两个坐标三元组（chip 和 host）、core-on-chip index，以及扁平化的 global id。对象大小为 56 字节（由 _Cores 步长确认）；这些 accessor 读取其前部和尾部的标量字段。

函数映射 ​

函数地址读取 / 支撑输出
TpuCoreLocation_ChipCoordinates0xEABC2C0tpu::TpuCoreLocation::chip_coordinates() → (x,y,z)3× int32* out + 标量
TpuCoreLocation_HostCoordinates0xEABC300*(uint32*)(loc + 8/12/16) → (x,y,z)3× int32* out + 标量
TpuCoreLocation_Index0xEABC320*(uint32*)(loc + 52)标量
TpuCoreLocation_Id0xEABC340thunk → tpu::TpuCoreLocation::LogicalDeviceId标量

字段映射 ​

偏移字段读取者
+0chip coords / id base（通过 member fn）ChipCoordinates, Id
+8host_coord.xHostCoordinates（loc[2]）
+12host_coord.yHostCoordinates（loc[3]）
+16host_coord.zHostCoordinates（loc[4]）
+52core-on-chip indexIndex

算法 ​

c
function TpuCoreLocation_ChipCoordinates(loc, x_out, y_out, z_out):   // 0xEABC2C0
    int32 tmp[3]
    tpu::TpuCoreLocation::chip_coordinates(&tmp)        // member computes the chip-coord triple
    *x_out = tmp[0]; *y_out = tmp[1]; *z_out = tmp[2]
    return tmp[2]                                        // (redundant) scalar return = z

function TpuCoreLocation_HostCoordinates(loc, x_out, y_out, z_out):   // 0xEABC300
    *x_out = *(uint32*)(loc + 8)                         // direct field reads — no member call
    *y_out = *(uint32*)(loc + 12)
    *z_out = *(uint32*)(loc + 16)
    return *(uint32*)(loc + 8)

function TpuCoreLocation_Index(loc):                     // 0xEABC320
    return *(uint32*)(loc + 52)
```text

> **怪异点 —** `ChipCoordinates` 经由 `chip_coordinates()` member（一次计算，而不是已存储三元组），而 `HostCoordinates` 是 `+8/+12/+16` 处的三个扁平字段读取。这种不对称告诉重新实现者其存储决策：host 坐标原样存储在 core-location 记录中，但 chip 坐标是派生的（很可能来自扁平 id 加 chip bounds），必须每次调用重新计算。两者都暴露相同的 `(int32*, int32*, int32*)` out-param 形状和冗余标量返回，因此 host 签名统一，即使实现不同。
>
> **说明 —** `TpuCoreLocation_Id` thunk 到 `LogicalDeviceId`，也就是 *logical*-device id，而不是原始 physical core id —— 这与 `TpuTopology_CoreForId` thunk 到 `LogicalDeviceForId` 一致。这个 C 表面上的 "id" 是 XLA device-assignment 使用的扁平 logical-device index，与 PJRT [device descriptions](../pjrt/ext-topology-description.md) 和 `LogiDeviceIdFromChipCoordAndIdx` extension method 操作的是同一 id 空间。重新实现者必须在整个名册中保持 "id" 表示 logical-device id。

---

## 4. 有效性门控

不同于 PJRT extension（会协商 args-struct 大小并返回 `PJRT_Error`），SE C-ABI accessor 信任输入，只在少数硬性不变量上做门控：

| Accessor | Guard | 失败行为 |
|---|---|---|
| `LogicalDevicesPerHost` / `PerChip` | enum 重映射 `if t != 1: t = 2·(t==2)` | 将未知 core type clamp 到 `0` |
| `Version` | `ordinal < 4` | 对未识别 codename 返回 `0` |
| `AvailableCoreCount` | `core_type < 3` | 越界时 `BUG()`（中止） |
| `AvailableCoresPerChip` | topology 非 null；`type < 3` | 无 topology 时返回 `4`；越界时 `BUG()` |
| `MaybeAvailableSparseCoresPerLogicalDevice` | topology 非 null；`type == 2` | `absl::Status` 错误：`"TPU system is not available"` / `"Invalid core type queried"` |
| 纯字段读取器（`HostCount`, `ChipBounds_*`, …） | 无 | 无条件字段读取 |

> **陷阱 —** 纯字段读取器对 topology 指针**没有 null check**。它们假设 host 传入了一个存活且完整构造的 `tpu::TpuTopology`（由 [TpuPlatform/TpuNodeContext](tpu-platform-and-topology.md) bring-up 构建）。只有三个 "available" 函数会防御缺失拓扑，因为它们从 mesh/ops-util singleton 懒解析拓扑，所以必须应对 "no TPU attached"。重新实现者必须在任何 pure-reader 调用前构造 topology；`topo == NULL` 时 `ChipBounds_X` 等没有优雅路径 —— 它会解引用 `NULL + 88`。

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `tensorflow::tpu::TpuTopologyExternal::*` | host 侧 C++ wrapper，索引 `ExecutorApiFn()` 槽进入这些 C-ABI impl |
| `tpu::TpuTopology` / `tpu::TpuCoreLocation` | 名册读取字段的 C++ 对象 |
| `tensorflow::TpuMeshCommonState::tpu_topology` / `tpu_ops_util::GetTpuTopology` | "available" accessor 从中解析拓扑的 singleton |
| `xla::jellyfish::NumEmbeddingDevices` | 支撑 `MaybeAvailableSparseCoresPerLogicalDevice` |
| PJRT type-16 topology extension | 暴露重叠几何信息的现代并列表面，使用不同 C ABI |

## 交叉引用

- [TfTpu C-API Shim](overview.md) —— `*ApiFn()` accessor 模式、不透明句柄约定，以及本页所属的名册映射
- [TpuPlatform & TpuNodeContext](tpu-platform-and-topology.md) —— 构造这些 accessor 读取的 `tpu::TpuTopology` 的 `TpuPlatform_*` / `TpuNodeContext_*` 生命周期
- [TpuExecutor 名册](tpu-executor-roster.md) —— 通过同一个 `ExecutorApiFn()` 表到达的同级 `TpuExecutor_*` 单设备运行时 C-ABI cluster
- [TopologyDescription Extension (type 16)](../pjrt/ext-topology-description.md) —— *现代 PJRT* 拓扑对象；对比：本页是其下方的 C-accessor 层，没有 args-size 协商，没有 `PJRT_Error`
- [ICI Topology Discovery](../ici/topology-discovery.md) —— 设备端运行时发现这些 accessor 静态暴露的 torus 几何
- [SparseCore Architecture](../sparsecore/architecture.md) —— Part-IV 硅片几何：完整 `tpu::TpuTopology` 结构，以及 `AvailableCoreCount` / `AvailableCoresPerChip` 索引的 `+124/+128`（`Target+0x3b8`）处按 core type 记录