Skip to content

TpuPlatform 与 TpuNodeContext

本页所有地址均适用于 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 等于文件偏移。其他版本会不同。

摘要

TpuPlatform_*TpuNodeContext_* 是启动 TPU 进程的两个 C-ABI 生命周期名册:前者是 platform 单例表面(一个进程级 TPU platform、它的设备数量、拓扑指针、主机位置),后者是按进程的 node-context 表面(获取 TPU node 的引用、查询主机是否已关闭、将其拆除)。它们是开源 StreamExecutor TPU 后端通过 ExecutorApiFn() 函数指针表调用的 extern "C" 自由函数,用于构造并拥有运行时,有别于那些只读取已构建 tpu::TpuTopology 的拓扑访问器。两个函数簇位于 C-ABI 块中一个大致连续的区域(TpuPlatform_*0xEAB8B800xEAB8FC0TpuNodeContext_*0xEACA2600xEACA440),由二进制中对 learning/45eac/tfrc/executor/stream_executor/tpu_node_context_c_api.cc 的引用恢复而来。

最重要的结构事实,也是重新实现者默认最容易弄错的一点,是 C-ABI 的 TpuPlatform_* 函数不会触碰 tensorflow::tpu::TpuPlatform——也就是 RegisterTpuPlatform 安装到 PlatformManager 中的 0x98 字节 SE Platform 对象。它们派发到另一个更底层的对象:deepsea::executor::DeepseaPlatform,通过一个 Meyers 单例到达(GetUnderlyingDeepseaPlatform::platform,一个由 __cxa_guard 构建的函数局部 static,由 deepsea::executor::GetRegisteredDeepseaPlatform @ 0x1d4935e0 填充)。“Deepsea” 是 TPU executor 层的内部代号;TpuPlatform_* 名册是 Deepsea platform vtable 之上的薄 C wrapper,而 SE TpuPlatform C++ 类是一个独立的 StreamExecutor Platform,它自己的虚方法会通过 ExecutorApiFn() 槽位转发,而这些槽位正指向本页的函数。两个“platform”对象,一个底层 Deepsea 单例。

TpuNodeContext_* 又由另一个对象支撑——tensorflow::TPUNodeInterfaces 及其 RAII ScopedRef——它拥有按进程的驱动附着。Create 创建堆上的 ScopedRef 并运行 InitScopedRefInitialize 解析当前的 TPUNodeInterfaces*CloseTpuHost 关闭主机;Free 以两个 fatal CHECK 拆除该引用。本页负责两个名册的逐函数 impl-symbol 映射,以及 Deepsea 派发的 vtable-offset 映射;读取所得 tpu::TpuTopology 的拓扑访问器TpuTopology 与 TpuCoreLocation

对重新实现而言,契约是:

  • 两个名册——11 个 TpuPlatform_* 和 5 个 TpuNodeContext_* extern "C" 函数、它们的地址,以及每个函数支撑的 C++ 对象/方法。
  • Deepsea 单例派发——每次 TpuPlatform_* 调用都会先确保 GetUnderlyingDeepseaPlatform::platform Meyers 单例已构建,然后在其上调用固定 vtable 偏移(+16 Id、+32 VisibleDeviceCount、+48 Initialize、+72 GetExecutor、+96 ShouldRegister…、+104 GetTopologyPtr)。
  • 两种状态约定——POD 返回(Initialized1VisibleDeviceCountint)与 absl::Status-out 习惯用法(InitializeGetExecutor、所有 TpuNodeContext_* 写入输出参数 StatusRep**Unref 先前值)。
  • node-context RAII 形态——Create/Free 是成对的堆 ScopedRef 分配器/析构器,带 fatal null-CHECK;CompactionSupported 会在内部重新获取 scoped ref,而不是接收一个。
名册前缀TpuPlatform_* (11), TpuNodeContext_* (5)
C-ABI 块TpuPlatform_* @ 0xEAB8B800xEAB8FC0; TpuNodeContext_* @ 0xEACA2600xEACA440
Platform 后端对象通过 GetUnderlyingDeepseaPlatform::platform(Meyers 单例)到达的 deepsea::executor::DeepseaPlatform
单例构建器deepsea::executor::GetRegisteredDeepseaPlatform @ 0x1d4935e0
Node 后端对象tensorflow::TPUNodeInterfaces + tensorflow::TPUNodeInterfaces::ScopedRef
源文件(已恢复)…/stream_executor/tpu_node_context_c_api.cc(来自 TpuNodeContext_Free 中的 CHECK 字符串)
到达方式ExecutorApiFn() 槽位(访问器模式见 shim 概览
状态约定对易失败调用使用 absl::Status-out(StatusRep** 输出参数 + Unref);查询使用 POD 返回
证据等级可重新实现级 / 经 IDA 反编译逐字节确认

范围—— *ApiFn() 访问器模式、不透明句柄约定以及完整名册映射见 shim 概览(链接,不重新推导)。读取本页 GetTopologyPtr 返回的 tpu::TpuTopology拓扑访问器名册(TpuTopology_* / TpuCoreLocation_*)见 TpuTopology 与 TpuCoreLocation——对比:那一页读取几何信息,本页拥有 platform/node-context 生命周期。SE tensorflow::tpu::TpuPlatform 类注册(RegisterTpuPlatform @ 0xe99a3a0、0x98 字节对象、PlatformManager)见 StreamExecutor 平台与执行器模型。一次性填充指向这些函数的 ExecutorApiFn() 槽位见 TfTpu_Initialize Bootstrap

Note—— 三个看起来像应属于此名册的名称,在此构建中不是 C-ABI TpuPlatform_*/TpuNodeContext_* 函数。没有 TpuPlatform_GetTpuHostLocation,没有 TpuPlatform_TpuMemoryLimit,也完全没有 StopChipHeartbeats 符号。GetTpuHostLocationTpuMemoryLimit 只作为 SE tensorflow::tpu::TpuPlatform 类上的 C++ 方法存在(…TpuPlatform::GetTpuHostLocation() const @ 0xe999f60…TpuPlatform::TpuMemoryLimit(long*) @ 0xe99a2c0);内存限制在 C ABI 上只通过另一个名册暴露,即 TpuConfigurationApi_TpuMemoryLimit @ 0xe8cdc40(见 TpuConfigurationApi)。TpuPlatform_* 名册中与主机位置最接近的是 GetHostLocation(§2);没有 TpuPlatform_* 内存限制条目——那是 configuration-API 的职责。已验证名册就是下面记录的 11 + 5 个函数。


1. Deepsea 单例派发

目的

每个 TpuPlatform_* 函数都是无状态 C 入口,作用于一个进程全局对象:已注册的 Deepsea platform。这里没有传入 platform 句柄——不接受 platform 参数的函数(IdVisibleDeviceCountGetTopologyPtr 等)会直接到达单例。这是它与拓扑名册的决定性区别,后者的访问器都把 topo 句柄作为 a1 接收。platform 是单例;topology 是你持有的对象

入口点

整个名册的派发形态相同。以 VisibleDeviceCount 为例的路径:

text
xla / SE TPU backend
  └─ tensorflow::tpu::TpuPlatform::VisibleDeviceCount()        (host-side SE Platform method)
       └─ tbl = stream_executor::tpu::ExecutorApiFn()           0x20819360  — singleton fn-ptr struct
       └─ tbl[slot]( )                                          — call through the slot
              │  (slot populated at init to point at:)

       TpuPlatform_VisibleDeviceCount                           0xEAB8E40   — plugin-side C-ABI impl
              └─ EnsureDeepseaSingleton()                        — __cxa_guard + GetRegisteredDeepseaPlatform
              └─ (*platform->vtable[+32])(platform)              0x1d4935e0 built it; +32 = device count
```text

### 算法——单例 guard

每个 `TpuPlatform_*` 函数体都以相同的 Meyers 单例 guard 开头。IDA 呈现出两种文本上不同但语义相同的形式(编译器重排了 early-return);二者都可归约为“构建一次,然后调用 vtable”:

```c
function EnsureDeepseaSingleton():                       // inlined prologue of every TpuPlatform_*
    if !guard_byte(GetUnderlyingDeepseaPlatform::platform):
        if __cxa_guard_acquire(&guard):
            platform = deepsea::executor::GetRegisteredDeepseaPlatform()   // 0x1d4935e0
            __cxa_guard_release(&guard)
    return GetUnderlyingDeepseaPlatform::platform        // may be NULL if no TPU registered

function TpuPlatform_VisibleDeviceCount():               // 0xEAB8E40
    p = EnsureDeepseaSingleton()
    return (*(int(**)(void*))(*(void**)p + 32))(p)        // vtable +32

function TpuPlatform_New():                              // 0xEAB8B80
    p = EnsureDeepseaSingleton()
    if !p: return NULL                                    // only New / GetExecutor null-check p
    h = operator new(8); *h = p; return h                 // 8-byte handle box wrapping the singleton ptr

QUIRK—— TpuPlatform_New 并不构造任何东西。它返回一个 8 字节堆 box,其中保存共享 Deepsea 单例指针的一份拷贝——每次 New 都交回一个指向同一底层 platform 的新 box。相应地,TpuPlatform_Free0xEAB8C00)只是 if (h) free(h):它释放 8 字节 box,从不释放单例。若重新实现者让 New 分配真实 platform、让 Free 销毁它,那么第二个 New/Free 对就会 double-free 进程全局对象。这个 box 是句柄,不是 owner。

GOTCHA—— 单例可能为 NULL。未注册 TPU platform 时(没有设备,或注册被禁用),GetRegisteredDeepseaPlatform 会返回 null。只有 TpuPlatform_New(返回 NULL)和 TpuPlatform_GetExecutor(返回失败的 StatusOr)对此防御。查询函数(IdVisibleDeviceCountGetTopologyPtrGetHostLocationShouldRegisterTpuDeviceToDeviceCopy)会无条件解引用单例的 vtable——NULL->vtable[+N] 会崩溃。契约是:主机必须先观察到成功的 New/Initialize,再调用任何查询。

vtable-offset 映射

名册派发通过的 Deepsea platform vtable 偏移,由每个函数体中的调用表达式确认:

偏移Deepsea 方法(推断)C-ABI 调用者
+16platform idTpuPlatform_Id (0xEAB8DE0)
+32visible device countTpuPlatform_VisibleDeviceCount (0xEAB8E40)
+48initialize(返回 absl::StatusTpuPlatform_Initialize (0xEAB8C20)
+72get executor for ordinal(返回 StatusOr<exec*>TpuPlatform_GetExecutor (0xEAB8CC0)
+96should register D2D copyTpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy (0xEAB8EA0)
+104get topology pointerTpuPlatform_GetTopologyPtr (0xEAB8F00)

GetHostLocation 通过 vtable——它在单例上直接调用非虚的 deepsea::executor::DeepseaPlatform::GetHostLocation @ 0x1d0e79a0InitializedGetRuntimeVersion 完全不触碰单例 vtable(见 §2)。


2. TpuPlatform_ 名册

函数映射

11 个 extern "C" 函数,其中三个是生命周期函数,八个是查询/访问器。“支撑”列给出其派发到的 Deepsea vtable 偏移或 helper。

函数地址支撑输出 / 约定
TpuPlatform_New0xEAB8B80EnsureDeepseaSingleton;box ptrvoid* 句柄(8 字节 box)或 NULL
TpuPlatform_Free0xEAB8C00free(handle)——只释放 boxvoid
TpuPlatform_Initialize0xEAB8C20singleton vtable +48absl::Status-out (StatusRep**)
TpuPlatform_Initialized0xEAB8CA0无——常量char → 始终为 1
TpuPlatform_GetExecutor0xEAB8CC0singleton vtable +72StatusOr<executor*>(boxed)
TpuPlatform_Id0xEAB8DE0singleton vtable +16标量 id
TpuPlatform_VisibleDeviceCount0xEAB8E40singleton vtable +32int
TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy0xEAB8EA0singleton vtable +96bool
TpuPlatform_GetTopologyPtr0xEAB8F00singleton vtable +104tpu::TpuTopology*(不透明)
TpuPlatform_GetHostLocation0xEAB8F60DeepseaPlatform::GetHostLocation @ 0x1d0e79a0(非虚)host-location 句柄
TpuPlatform_GetRuntimeVersion0xEAB8FC0BuildData::Timestamp + Changelist;static kMetadata填充调用者结构体(version + string)

NOTE—— GetTopologyPtr 是本页与 TpuTopology 与 TpuCoreLocation 之间的连接点。它返回不透明 tpu::TpuTopology*(vtable +104);随后主机会把该指针作为 a1 传给每个 TpuTopology_* 访问器。platform 产生拓扑指针;topology 名册消费它。双方都不会按值传递 C++ 对象——跨越边界的只有 void*

算法——两个非派发异常项

两个函数打破单例派发模式,重新实现者必须精确复现:

c
function TpuPlatform_Initialized():                      // 0xEAB8CA0
    return 1                                              // hardcoded TRUE — no state inspected

function TpuPlatform_GetRuntimeVersion(out):             // 0xEAB8FC0
    // Meyers-singleton over a static `kMetadata` (NOT the platform singleton):
    if !guard(TpuPlatform_GetRuntimeVersion::kMetadata):
        if __cxa_guard_acquire(&guard):
            kMetadata = StrCat(BuildData::Timestamp(), " cl/", BuildData::Changelist())
            __cxa_guard_release(&guard)
    out[0]  = 0                                           // version major/minor packed
    out[8]  = 1
    out[16] = ptr-to(kMetadata string)                   // SSO-aware: inline vs heap by sign of len byte
    out[24] = kMetadata length
    return out
```text

> **QUIRK——** `TpuPlatform_Initialized` 始终返回 `1`。它不检查任何状态——不检查单例,也不检查标志。“platform 是否已初始化”这个问题在 C-ABI 边界被静态回答为 true;真正可能失败的初始化位于 `TpuPlatform_Initialize`(它可能从 vtable `+48` 返回失败的 `absl::Status`)。重新实现者不能把 `Initialized` 接到真实 readiness flag 上:SE 后端把它当作始终为真的 capability probe,实际门控是 `Initialize` 的状态。

### absl::Status-out 习惯用法

`Initialize` 和 `GetExecutor` 使用整个 shim 共享的 C-ABI 状态约定:调用者传入一个 `StatusRep**` 输出参数,其中持有*先前*的状态对象;被调用者计算一个新的 `StatusRep*`,如果不同,就将其存入输出参数并 `Unref` 旧对象(`absl::Status` 是带标记指针——`& 1` 用于区分 inline OK 与需要引用计数的堆 `StatusRep*`)。

```c
function TpuPlatform_Initialize(self, status_out):       // 0xEAB8C20
    p   = EnsureDeepseaSingleton()
    new = (*(p->vtable[+48]))(p)                          // Deepsea initialize → StatusRep* (tagged)
    old = *status_out
    if new == old:
        if (new & 1) == 0: StatusRep::Unref(new)          // both heap, same ptr → drop one ref
    else:
        *status_out = new
        if (old & 1) == 0: StatusRep::Unref(old)          // replace, release the previous status

GetExecutor0xEAB8CC0)是同一习惯用法的 StatusOr<T> 形式:Deepsea 调用(vtable +72)把一个 2 槽 {status, value} 写入栈上临时对象;如果 status 槽是 OK 哨兵(&dword_0 + 1),它就把 value box 到一个新的 operator new(8) 句柄中并返回;否则它 Unref 并返回 NULL,如果主机误用结果,则走 ThrowBadStatusOrAccess 路径。


3. TpuNodeContext_ 名册

目的

tensorflow::TPUNodeInterfaces 是按进程附着到 TPU node 的对象——它拥有当前驱动会话、node 的拓扑以及其 core location。C-ABI TpuNodeContext_* 名册管理到它的一个 ScopedRef:这是一个 RAII 引用,持有期间会保持 node 存活。主机调用 Create 一次来取得该引用,调用 Initialize 解析底层 TPUNodeInterfaces*,调用 CloseTpuHost 关闭主机,并调用 Free 释放引用。四个可失败调用都使用与 TpuPlatform_Initialize 相同的 absl::Status-out 习惯用法。从 CHECK 字符串恢复出的源文件:…/stream_executor/tpu_node_context_c_api.cc

函数映射

函数地址支撑(C++ 方法)输出 / 约定
TpuNodeContext_Create0xEACA260TPUNodeInterfaces::InitScopedRefScopedRef** 句柄 + absl::Status-out
TpuNodeContext_Free0xEACA2E0~ScopedRef + free;两个 fatal CHECKvoid
TpuNodeContext_CloseTpuHost0xEACA3C0TPUNodeInterfaces::CloseTPUHostabsl::Status-out
TpuNodeContext_Initialize0xEACA400TPUNodeInterfaces::Getabsl::Status-out(+ 解析 TPUNodeInterfaces**
TpuNodeContext_CompactionSupported0xEACA440InitScopedRefTpuChipConfig::Megacorebool

算法——RAII 对

CreateFree 是堆 ScopedRef 分配器/析构器。其不对称性——Create 分配一个保存 ScopedRef* 的 8 字节 box,Free 校验并拆除它——就是重新实现契约:

c
function TpuNodeContext_Create(node_index, status_out):  // 0xEACA260
    ref = operator new(8); *ref = NULL                    // 8-byte box for the ScopedRef
    st  = TPUNodeInterfaces::InitScopedRef(node_index, ref)   // fills *ref, returns StatusRep*
    merge_status(status_out, st)                          // same Status-out idiom as §2
    handle = operator new(8); *handle = ref               // box the box
    return handle

function TpuNodeContext_Free(handle):                    // 0xEACA2E0
    CHECK(handle != NULL,  "node_context != nullptr")            // tpu_node_context_c_api.cc:33 (FATAL)
    inner = *handle
    CHECK(inner != NULL,   "node_context->node_ref != nullptr")  // :34 (FATAL)
    *handle = NULL
    ~ScopedRef(inner); free(inner)                        // destroy + free the ScopedRef
    // (second NULL-guarded teardown of *handle handles a re-entrant edge)
    free(handle)
```text

> **GOTCHA——** `TpuNodeContext_Free` 对 null 句柄或 null 内部 ref 是 **fatal** 的——它不返回错误,而是 `LogMessageFatal` 并中止(位于 `tpu_node_context_c_api.cc:33` 和 `:34` 的两个 CHECK)。这与 `TpuPlatform_Free` 正相反,后者是静默的 `if (h) free(h)`。重新实现者必须保持 node-context teardown 对误用 fatal:SE 后端依赖该 abort 来暴露 node ref 的 double-free / use-after-free,而不是破坏驱动会话。不要“贴心地”让它容忍 null。
>
> **QUIRK——** `TpuNodeContext_CompactionSupported` **接收 node-context 句柄。它通过 `InitScopedRef` 重新获取一个一次性的 `ScopedRef`,然后——仅当该 ref 的状态为 OK 时——遍历 `TPUNodeInterfaces::tpu_topology()` → `tensor_core_location()` → `tpu::TpuChipConfig::GetUserStack(...)`,并且如果 `+40` 处的 user-stack 字节为 `1`,就返回 `tpu::TpuChipConfig::Megacore(chip_config)`。在任何失败或非 megacore 路径上,它默认返回 `1`(supported),并且总是在返回前销毁 scoped ref。因此“compaction supported”实际是“这是否为 megacore chip config”,每次调用都针对一个临时 node 引用重新计算。

### Initialize / CloseTpuHost 状态形式

`TpuNodeContext_Initialize`(`0xEACA400`)调用 `TPUNodeInterfaces::Get(node_index, &out, node_out)`——它把当前 `TPUNodeInterfaces*` 解析到调用者输出参数中,*并且*通过标准输出参数习惯用法传回 `absl::Status`。`TpuNodeContext_CloseTpuHost`(`0xEACA3C0`)调用 `TPUNodeInterfaces::CloseTPUHost(node)`,并把返回状态合并到 node 的首个 `StatusRep*` 槽中。二者都是单个 `TPUNodeInterfaces` 方法之上的纯状态管线 wrapper;唯一逻辑是对先前状态执行带标记指针的 `Unref`。

---

## 4. 有效性门控

两个名册的门控方式非常不同,这种对比正是重新实现风险:

| 函数 | Guard | 失败行为 |
|---|---|---|
| `TpuPlatform_New` | singleton non-null after build | 返回 `NULL` |
| `TpuPlatform_GetExecutor` | `StatusOr` OK sentinel | 返回 `NULL`;误用时 `ThrowBadStatusOrAccess` |
| `TpuPlatform_Initialize` / `_GetExecutor` | `absl::Status`-out | 传播失败状态,`Unref` 先前状态 |
| `TpuPlatform_Id` / `_VisibleDeviceCount` / `_GetTopologyPtr` / `_GetHostLocation` / `_ShouldRegister…` || 无条件解引用 singleton vtable——若 singleton 为 `NULL` 则崩溃 |
| `TpuPlatform_Initialized` || 始终为 `1` |
| `TpuNodeContext_Create` / `_Initialize` / `_CloseTpuHost` | `absl::Status`-out | 传播失败状态 |
| `TpuNodeContext_Free` | 两个 CHECK(`:33`, `:34`) | null 句柄/ref 上 **FATAL abort** |
| `TpuNodeContext_CompactionSupported` | scoped-ref status OK | 任何失败时默认 `1`(supported) |

> **NOTE——** 这种拆分对应拓扑名册的拆分(纯读取器 vs. 有防御的访问器)。这里,*platform 查询*函数假定单例已成功构建,正如 [拓扑纯读取器](tpu-topology.md#4-validity-gating) 假定有一个 live 的 `tpu::TpuTopology`;而*node-context* 函数处理的是可能被合法拆除的按进程资源,因此对可失败操作携带 status 输出参数,对析构器使用 fatal CHECK(这里 null 始终是调用者 bug)。把一种纪律统一套用到所有函数的重新实现者,要么会过度防御廉价 platform 查询,要么会不足防御 node-context teardown。

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `deepsea::executor::DeepseaPlatform` | 每个 `TpuPlatform_*` 派发进入的底层 TPU platform 对象(vtable `+16+104`) |
| `deepsea::executor::GetRegisteredDeepseaPlatform @ 0x1d4935e0` | 构建 `GetUnderlyingDeepseaPlatform::platform` Meyers 单例 |
| `tensorflow::tpu::TpuPlatform`(`0x98` 字节) | 由 `RegisterTpuPlatform` 注册的*独立* SE `Platform` 类;其方法通过指向此名册的 `ExecutorApiFn()` 槽位转发 |
| `tensorflow::TPUNodeInterfaces` / `…::ScopedRef` | `TpuNodeContext_*` 名册管理的按进程 node attachment + RAII ref |
| `tpu::TpuChipConfig::Megacore` | 支撑 `TpuNodeContext_CompactionSupported` 的 megacore 探针 |
| `TpuConfigurationApi_TpuMemoryLimit @ 0xe8cdc40` | C ABI 实际暴露内存限制的位置(不在 `TpuPlatform_*` 上) |

## 交叉引用

- [TfTpu C-API Shim](overview.md)——`*ApiFn()` 访问器模式、不透明句柄约定,以及本页作为其中一项的名册映射
- [TpuTopology 与 TpuCoreLocation](tpu-topology.md)——读取 `TpuPlatform_GetTopologyPtr` 返回拓扑的 `TpuTopology_*` / `TpuCoreLocation_*` 几何访问器;对比:那页读取几何信息,本页拥有 platform/node 生命周期
- [TpuExecutor 名册](tpu-executor-roster.md)——由 platform 创建并通过同一 `ExecutorApiFn()` 表到达的 `TpuExecutor_*` 按设备运行时函数簇
- [TpuCompiler 名册](tpu-compiler-roster.md)——针对来自此 platform 的执行器运行的 `TpuCompiler_*` / `TpuCompile_*` 编译 C 表面
- [StreamExecutor 平台与执行器模型](../pjrt/stream-executor-host-interpreter.md)——SE `tensorflow::tpu::TpuPlatform` 类、`RegisterTpuPlatform @ 0xe99a3a0` 和 `PlatformManager`;此 C-ABI 名册的主机侧消费者
- [TfTpu_Initialize Bootstrap](../lifecycle/tftpu-initialize-bootstrap.md)——一次性填充指向这些函数的 `ExecutorApiFn()` 槽位