StreamExecutor → PJRT 适配器

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

摘要

在上游 XLA 中，从 StreamExecutor LocalClient / Platform 到现代 PjRtClient 的桥接是 xla::PjRtStreamExecutorClient（以及 TPU 专用的 TfrtTpuClient）：它是一个 PjRtClient 子类，每个设备持有一个 LocalDeviceState，每个状态拥有一个 stream_executor::StreamExecutor 和一个小型 stream_executor::Stream 池，并通过 Stream::ThenLaunch / Stream::ThenMemcpy 转发执行。到 libtpu.so 中寻找该类的重新实现者找不到它。这个二进制中没有 xla::PjRtStreamExecutorClient —— 该符号未定义（已确认：不存在已定义的 PjRtStreamExecutorClient 符号；唯一残留的 SE-PjRt 符号是 PjRtStreamExecutorDeviceDescription，属于死掉的 GPU 路径代码）。“StreamExecutor → PJRT 适配器”实际上是 xla::TpuClient，它派生自 xla::CommonPjRtClient : xla::PjRtClient，并通过 tsl::AsyncValue 事件直接位于 TFRT 原生 TPU 设备运行时 tpu::System 之上 —— 完全绕过旧的 stream_executor::tpu::TpuExecutor / TpuStream / TpuTransferManager 栈。

本页记录这个桥接层：GetTpuPjRtClient 构造路径如何把解析后的配置变成已接线的 TpuClient，TpuClient : CommonPjRtClient : PjRtClient 的类形状及其 3 参数再接 vector 的 TFRT 构造函数，设备枚举循环如何把每个 tpu::TpuCoreLocation 包装成一个 xla::TpuDevice，以及一个 SE Stream/Executor 本来会如何绑定到每个 PjRtDevice 下面 —— 只是在这里，TpuDevice 绑定的是一个共享的 tpu::System* 加一个 TpuCoreLocation 和一个节流 Semaphore，而不是每设备一个 executor。其上的 C-ABI 表面（PJRT_Client_Create、208 字节的 PJRT_Client 包装器、设备/内存访问器槽位）由客户端、设备与拓扑负责；其下的 SE 平台注册和宿主解释器由 StreamExecutor 与宿主解释器负责；每次执行的热路径由在 Stream 上异步执行负责。本页负责中间这条接缝：GetTpuPjRtClient → TpuClient → TpuDevice-包装-tpu::System。

对于重新实现，契约是：

构造链。 GetTpuPjRtClient(config) → 一次性 HostContext 构建 → TpuStatesManager / TpuSystemState 单例 → GetTpuPjRtClientInternal 设备枚举 → make_unique<TpuClient>(…) → 可选 MegaScale / Tf 装饰器 → CreateWrapperClient。
类层次和 ctor 形状。 TpuClient : CommonPjRtClient : PjRtClient；八参数 TFRT 构造函数接收设备 vector、host context、host allocator、system state 和 states manager。
设备包装规则。 每个来自 TpuTopology::logical_devices 的 tpu::TpuCoreLocation 对应一个 xla::TpuDevice，各自持有一个共享的 tpu::System*（不是私有 executor），并通过构造后的 SetClient 反向边绑定到 client，该反向边还会安装 AsyncWorkRunner。
SE→TPU 概念映射。 哪个 TPU-runtime 对象替代哪个 SE 抽象，使携带上游模型的读者能够向前翻译。


桥接入口	`xla::GetTpuPjRtClient(const PjRtTpuClientConfig&)` @ `0xF8008C0` (1226 B)
设备枚举	`xla::GetTpuPjRtClientInternal(…)` @ `0xF800DA0`
适配器客户端类	`xla::TpuClient` — ctor `0xF801980`, vtable `0x2177B598`, typeinfo `0x2177BA40`
客户端基础框架	`xla::CommonPjRtClient` — vtable `0x2178A108`, typeinfo `0x21789EB0`
抽象 PJRT 基类	`xla::PjRtClient` — vtable `0x21CA9C98`, typeinfo `0x21CA9E28`
适配器设备类	`xla::TpuDevice` — ctor `0xF7FDC40` (object `0x1E0` = 480 B), vtable `0x2177B4D0`
设备→客户端绑定	`xla::TpuDevice::SetClient(TpuClient, AsyncWorkRunner)` @ `0xF7FE520`
设备运行时后端	`tpu::System` — `Initialize` @ `0x1D0AE420`（SE `StreamExecutor` 的对应物）
C-ABI 包装	`pjrt::CreateWrapperClient(unique_ptr<PjRtClient>)` @ `0xF872060`

SE → TPU-Runtime 概念映射

目的

阅读这一层最快的方式，是把它看作读者已经掌握的上游 PjRtStreamExecutorClient 模型的差异。每个 StreamExecutor 抽象都有一个 TPU-runtime 对应物，但形状不同：SE 是同步 stream 模型（每个 Stream 上有一个操作 FIFO，用 Event 做跨 stream 排序）；TPU runtime 是 async-value DAG 模型（通过 tpu::TpuEventIssuer 序列点发出的 tpu::TpuEvent 依赖图）。该适配器不会在运行时把一种模型翻译成另一种模型 —— 它根本不构造 SE 对象。它原生构建在异步模型之上。

映射表

对于携带上游模型的读者来说，这是最重要的一张表。每一行都是“你本来会寻找的东西” → “libtpu 改用的东西” → “底层 TPU-runtime 对象”。

上游 SE / PjRt 概念	libtpu 适配器对象	TPU-runtime 后端
`PjRtStreamExecutorClient` / `TfrtTpuClient`	`xla::TpuClient` (`: CommonPjRtClient : PjRtClient`)	`tpu::System`（单例，通过 `TpuSystemState` / `TpuStatesManager`）
`PjRtStreamExecutorDevice`	`xla::TpuDevice`	`tpu::System*` + `tpu::TpuCoreLocation` + `xla::Semaphore`
`LocalDeviceState` + `StreamExecutor`（每设备）	(无 —— 没有每设备 executor)	共享 `tpu::System` + `AsyncWorkRunner` + `tfrt::ConcurrentWorkQueue`
`stream_executor::StreamExecutor`	(无)	`tpu::System`（一个，在所有设备间共享）
`Stream::ThenLaunch`	(没有 `Stream`)	`tpu::System::Execute(AsyncValueRef<ProgramHandle>, …)` @ `0x1D0B33E0`
`Stream::ThenMemcpy`	(没有 `Stream`)	`tpu::System::TransferTo/FromDevice` @ `0x1D0AFA20` / `0x1D0B0160`
`stream_executor::Event`	`xla::TpuTrackedDeviceEvent`	`tpu::TpuEvent`（一个 `tsl::AsyncValueRef`）
`Stream` 排序 FIFO	(没有 FIFO)	`tpu::TpuEventIssuer` 序列点 + `tsl::AsyncValue` 依赖
`DeviceMemoryBase`	`tpu::TpuSharedMemoryLocation`	TPU HBM 地址
`PjRtStreamExecutorBuffer`	`xla::TpuRawBuffer` (`: CommonPjRtRawBuffer`)	位于 `TpuSharedMemoryLocation` 的 `tpu::TpuBuffer`
`TransferManager`	(内联在 `TpuRawBuffer` / `TpuClient` 中)	`tpu::System::TransferTo/FromDevice`
SE `Compiler`（通过 factory 注册）	`xla::TpuCompiler` (`: PjRtCompiler`)	jellyfish JIT (`TpuJitCompileHloWithOptions`)

怪癖 — 空单元格正是这张表的重点。若重新实现者为每个设备构建一个 LocalDeviceState 风格的对象，每个对象拥有一个 StreamExecutor 和一个 stream 池，那么他复刻的是上游 XLA，不是 libtpu。TPU 适配器在所有 TpuDevice 间共享一个 tpu::System；没有每设备 executor，也完全没有 Stream。确实存在的每设备状态很小：一个 TpuCoreLocation（这个设备映射到哪个物理 core）和一个 Semaphore（进行中 launch 的节流器）。SE 后端本来会从每 stream FIFO 得到的排序，在这里来自进程级的 TpuEventIssuer DAG。

为什么 SE 栈是缺席，而不是被适配

旧的 SE TPU 栈确实存在于二进制中 —— stream_executor::tpu::TpuExecutor、tensorflow::tpu::TpuStream、tensorflow::tpu::TpuTransferManager，以及 Tpu*_* / TpuExecutor_* / TpuStream_* C-ABI 导出 —— 但它是一条并行路径，不是 PJRT 下面的一层。它支撑 xla::LocalClient / xla::Service 和 TF-TPU op kernels，通过模块初始化时注册的 SE TpuPlatform 进入（RegisterTpuPlatform → PlatformManager，记录在 StreamExecutor 与宿主解释器）。PJRT TpuClient 适配器从不触碰它：在整个 TpuClient / TpuDevice / TpuRawBuffer / TpuLoadedExecutable 代码范围（0xF800DA0–0xF816EC0）中，没有对 ExecutorApiFn、SE_StreamExecutor 或 TpuExecutor 的引用。这两个 host 抽象共享同一个 TPU driver core，但使用不同的 host 侧调度：TFRT async-values（PJRT）对 SE Streams（旧路径）。本页只记录 PJRT 分支。

GetTpuPjRtClient — 构造桥接

目的

GetTpuPjRtClient 是 C-ABI 选项层交给运行时的位置。调用它时，PJRT_Client_Create 已经验证了 option map，并将其解析成类型化的 PjRtTpuClientConfig（该路径由客户端、设备与拓扑负责）。GetTpuPjRtClient 只启动一次进程级基础设施 —— TFRT HostContext 和 tpu::System 单例 —— 然后枚举设备并构造 TpuClient。它返回一个 StatusOr<unique_ptr<TpuClient>>，C-ABI 层会可选地装饰它，然后再包装它。

入口点

text

pjrt::tpu_plugin::PJRT_Client_Create (0xE6A8840)        ── slot 15; owns option parse (see client-and-device.md)
  └─ xla::GetTpuPjRtClient (0xF8008C0)                   ── THIS PAGE: brings up HostContext + tpu::System, builds TpuClient
       ├─ xla::CreateDefaultHostContext (lambda, __cxa_guard one-shot)
       │    ├─ tfrt::CreateMallocAllocator
       │    └─ tfrt::CreateMultiThreadedWorkQueue(DefaultThreadPoolSize(), …)
       ├─ xla::GetSingletonTpuStatesManager(bool)        (0xF958360)  ── mutex+guard singleton
       │    └─ TpuStatesManager::GetOrCreateTpuSystemState (0xF956E40)
       │         └─ xla::CreateTpuSystemState (0xF95A4A0)  ── lazily brings up tpu::System
       └─ xla::GetTpuPjRtClientInternal(config, host_ctx, system_state, states_mgr) (0xF800DA0)
            ├─ tpu::System::topology() / hal_location()
            ├─ for each TpuCoreLocation in TpuTopology::logical_devices:
            │    ├─ xla::TpuDeviceDescription(core_loc, topology, "TPU", "TpuDevice")
            │    └─ operator new(0x1E0) + xla::TpuDevice::TpuDevice(desc, name, ordinal, tpu::System*)
            ├─ make_unique<xla::TpuClient>(config, process_index, platform_version, devices,
            │                              host_ctx, host_allocator, system_state, states_mgr)
            └─ for each device: TpuDevice::SetClient(TpuClient*, AsyncWorkRunner*)  (0xF7FE520)
```text

### 算法

```c
function GetTpuPjRtClient(config):                       // 0xF8008C0
    // (1) One-shot process-wide TFRT HostContext. Guarded by the
    //     CreateDefaultHostContext::kInstance __cxa_guard so every client
    //     in the process shares one HostContext / thread pool.
    if !kInstance:                                        // guard var, lines 70-83
        if cxa_guard_acquire(&kInstance_guard):
            alloc = tfrt::CreateMallocAllocator()         // line 85
            queue = tfrt::CreateMultiThreadedWorkQueue(   // line 92
                        DefaultThreadPoolSize(config), …) //   thread count from config
            kInstance = make_shared<tfrt::HostContext>(alloc, queue)   // line 94
            cxa_guard_release(&kInstance_guard)
    host_ctx = kInstance                                  // shared_ptr; refcount bumped (line 116)

    // (2) The tpu::System singleton, reached through the states manager.
    //     GetOrCreateTpuSystemState lazily runs CreateTpuSystemState the
    //     first time, which initialises tpu::System (Initialize @ 0x1D0AE420).
    states_mgr  = GetSingletonTpuStatesManager(use_global_tpu_system)  // 0xF958360, line 212
    system_state = states_mgr.GetOrCreateTpuSystemState(opts)          // 0xF956E40, line 214

    // (3) Tail into device enumeration + client construction.
    //     The 'a6' boolean selects an own-vs-borrow path for system_state.
    return GetTpuPjRtClientInternal(config, host_ctx,    // 0xF800DA0, lines 179/236
                                    system_state, states_mgr)

注意 — HostContext 是一个进程级单例，不是每客户端一个。CreateDefaultHostContext::kInstance 上的 __cxa_guard 意味着同一进程中的第二次 PJRT_Client_Create 会复用同一个 malloc allocator、同一个多线程 work queue，因而也复用同一个 AsyncWorkRunner 线程池。重新实现者不得为每个 client 分配新的线程池，否则执行路径中的节流/调度假设会被破坏。

设备枚举 — GetTpuPjRtClientInternal

GetTpuPjRtClientInternal（0xF800DA0）是把物理 TPU core 转换成 PjRt devices 的循环。它最接近上游“为每个设备构建一个 LocalDeviceState”的对应物，只是不会创建 executor 或 stream 池 —— 只创建一个 description 和一个 TpuDevice，每个都指向同一个共享的 tpu::System。

function GetTpuPjRtClientInternal(config, host_ctx, system_state, states_mgr):  // 0xF800DA0
    system   = system_state->system()                    // tpu::System* (line 136: j+64)
    topology = system->topology()                         // tpu::TpuTopology* (line 138)
    hal_loc  = system->hal_location()                     // line 132

    devices = vector<unique_ptr<TpuDevice>>{}
    cores   = topology->logical_devices(TpuCoreType)      // line 140; stride 56 bytes (line 148/323)
    for core_loc in cores:                                // each TpuCoreLocation, 56 B apart
        // (a) static identity: id/process/kind/attributes/memory
        desc = TpuDeviceDescription(core_loc, topology,   // line 155
                                    type_prefix="TPU", name="TpuDevice")
        // (b) the PjRt device object — 0x1E0 (480) bytes, holds a SHARED system ptr
        dev_mem = operator new(0x1E0)                      // line 156
        device  = TpuDevice::TpuDevice(desc, name, ordinal, system)   // line 206, 0xF7FDC40
        devices.push_back(unique_ptr<TpuDevice>{device})   // grows the vector (line 260 = length-error guard)
        // record local-device-id -> TpuSharedMemoryLocation (LocalSharedMemory, line 374)

    // (c) build the client; ctor takes ownership of the device vector
    client = make_unique<TpuClient>(                       // line 545
                 config, process_index, platform_version,
                 devices, host_ctx, host_allocator,
                 system_state, states_mgr)

    // (d) close the back-edge: each device learns its client + work runner
    for device in client->devices():
        device->SetClient(client.get(), client->async_work_runner())   // 0xF7FE520
    return client
```text

> **陷阱 —** 设备→client 链接是一个*构造后*反向边（`SetClient`），不是构造函数参数。`TpuClient` ctor 接收设备 vector，但每个 `TpuDevice` 直到之后运行 `SetClient` 时才知道自己的 owner client（或它的 `AsyncWorkRunner`）。重新实现者如果试图在 `TpuClient` 构造函数内部使用 `device->client()`，会读到未初始化指针。两阶段接线存在的原因是 `AsyncWorkRunner` 是完整构造后的 client 的成员；在 client 对象存在之前无法交给设备。
>
> **怪癖 —** `TpuCoreLocation` 数组以硬编码的 56 字节步长遍历（lines 148, 323）。`TpuTopology::logical_devices` 返回一个连续块，循环执行的是 `v18 += 56`，而不是调用 iterator。这个 56 字节的 `TpuCoreLocation` 大小是一个 ABI 事实，拓扑侧的重新实现者必须精确匹配，否则设备枚举循环会读到垃圾 core。

### Internal 返回后 — 装饰器和 C 包装器

回到 `PJRT_Client_Create` 中，裸的 `unique_ptr<TpuClient>` 会被可选包装：

- 多 slice 激活时（`FLAGS_megascale_port >= 0` 且未跳过）使用 `MegaScalePjRtClient::CreateMegaScalePjRtClient`（`0xE6EA680`）。
- `use_tf_pjrt_client == 1`（默认）时使用 `TfPjRtClient::CreateTfPjRtClient`（`0x108524E0`）。

然后 `CreateWrapperClient`（`0xF872060`）把最外层的 `xla::PjRtClient*` 装箱进 208 字节的 C-ABI `PJRT_Client`。这三个步骤和 wrapper 布局由 [客户端、设备与拓扑](client-and-device.md) 负责；本页止于 `GetTpuPjRtClient` 产出的 `unique_ptr<TpuClient>`。

---

## TpuClient — 适配器客户端类

### 目的

`xla::TpuClient` 是 TPU 的具体 `PjRtClient`。它是每个 C-ABI client 槽位最终 virtual-call 进入的对象（在解包任何 MegaScale/Tf 装饰器和 `PJRT_Client` POD 之后）。它重写了约 75 个 `CommonPjRtClient` 框架方法 —— buffer 创建、executable 加载、device assignment、transfer 管线 —— 并拥有一个 StreamExecutor client 本会分散到各个 `LocalDeviceState` 中的那些进程级 runtime handle。

### 类层次

```text
xla::PjRtClient                 (abstract base; vtable 0x21CA9C98, typeinfo 0x21CA9E28)
   ▲
xla::CommonPjRtClient           (async-value framework; vtable 0x2178A108, typeinfo 0x21789EB0)
   │   BufferFromHostBuffer, DefineBuffer, PrepareArguments, CreateOutputs,
   │   TrackFuture, ShouldRetryOnOom, GetTransposePlan, … (~79 methods)
   ▲
xla::TpuClient                  (the adapter; ctor 0xF801980, vtable 0x2177B598, typeinfo 0x2177BA40)
       overrides ~75 methods; holds devices, HostContext, HostMemoryAllocator,
       TpuSystemState, TpuStatesManager, Semaphore

两个可选装饰器位于这个层次外部，作为包装器而不是子类：

装饰器	作用	启用条件	符号
`xla::MegaScalePjRtClient`	覆盖 `unique_ptr<TpuClient>` + `MultiSliceConfig` 的多 slice 层；跨 slice 转发/增强 compile+execute	`FLAGS_megascale_port >= 0`，且未跳过	`CreateMegaScalePjRtClient` @ `0xE6EA680`; typeinfo `0x215FD068`
`xla::TfPjRtClient`	TF 线程安全装饰器（`DestroyWrappedBuffersAndClient`）；拥有被包装的 buffers	`use_tf_pjrt_client == 1`（默认）	`CreateTfPjRtClient` @ `0x108524E0`; typeinfo `0x217EFED0`

xla::AbstractTpuClient（typeinfo 0x215FD080）是一个很薄的基类，用于共享拓扑 helpers —— 例如 GetSubsliceTopologyForCompilation（0xE6EE120）—— 在 TpuClient 和 MegaScale 路径之间共享。

TFRT 构造函数

构造函数签名由 0xF801980 处的 demangled symbol 按字节确认。这是 TFRT 形状 —— 原生 async-value，接收预构建基础设施，而不是构造 executors：

TpuClient::TpuClient(                                    // 0xF801980
    const PjRtTpuClientConfig&        config,
    int                               process_index,
    std::string                       platform_version,
    std::vector<unique_ptr<TpuDevice>> devices,          // takes ownership of the enumerated devices
    std::shared_ptr<tfrt::HostContext> host_context,     // the process-wide singleton
    std::unique_ptr<HostMemoryAllocator> host_allocator, // host-staging allocator
    std::unique_ptr<TpuSystemState, MaybeOwningDeleter> system_state,  // may own or borrow
    std::shared_ptr<TpuStatesManager> states_manager)
```text

> **注意 —** `TpuSystemState` 上的 `MaybeOwningDeleter` 是单进程多 client 行为的铰链。当设置 `use_global_tpu_system` 时，system state 从全局单例*借用*（deleter 是 no-op）；否则由 client 拥有。这就是为什么 `a6` boolean 会从 `GetTpuPjRtClient` 流到 `GetTpuPjRtClientInternal` —— 它选择 `MaybeOwningDeleter` 的 own-vs-borrow 构造。

### 对象字段布局 — 未完全解码

构造函数体的精确字段写入没有从二进制中解引用。根据 ctor *签名*和调用点，该对象持有：设备 vector、`HostContext` shared_ptr、`HostMemoryAllocator`、`TpuSystemState` 指针（带其 maybe-owning deleter）、`TpuStatesManager` shared_ptr，以及一个每 client 的 `Semaphore`（进行中 computation 的节流器）。这些成员的具体字节偏移是**低置信度** —— 重新实现者需要完整反编译 `0xF801980` 才能放置它们。可以*确定*的是成员集合，以及 `async_work_runner()` 和 `devices()` 在构造后可立即访问（它们会在 `GetTpuPjRtClientInternal` 的 `SetClient` 循环中被读取）。

### 函数图

| 函数 | 地址 | 作用 |
|---|---|---|
| `TpuClient::TpuClient` (ctor) | `0xF801980` | 从枚举出的设备 + 基础设施构建适配器 |
| `TpuClient::LookupDevice` | `0xF8033A0` | global device id → `PjRtDevice*` |
| `TpuClient::platform_name` | `0xF816D20` | `"tpu"` |
| `TpuClient::platform_id` | `0xF816D00` | TPU platform id 常量 |
| `TpuClient::process_index` | `0xF816C40` | host process index |
| `TpuClient::memory_spaces` | `0xF816CE0` | HBM + pinned/unpinned host spaces |
| `TpuClient::AllocateBuffer` | `0xF7FC5A0` | device HBM alloc → `tpu::AllocateBuffer` |
| `TpuClient::AllocateRawBuffer` | `0xF7FB1E0` | host/CPU-space alloc dispatch |
| `TpuClient::CompileAndLoad` (XlaComputation) | `0xF804F20` | compile entry → jellyfish JIT |
| `TpuClient::CompileAndLoad` (MLIR) | `0xF8068E0` | MLIR compile entry |
| `TpuClient::TrackFuture` | `0xF7FAD60` | `tsl::AsyncValue` → `PJRT_Event` future |
| `TpuClient::CreateProfiledFuture` | `0xF7FAE80` | profiled completion future |
| `TpuClient::CreateErrorEvent` | `0xF808420` | `tpu::TpuEvent` set to error status |
| `TpuClient::CreateDeviceEventSet` | `0xF813D40` | mint wait/define event sets |

> **注意 —** 上面的 compile、execute、buffer 和 event 表面在各自负责的页面中有深入记录（[Executable 与执行](executable-execution.md)、[Buffer 与内存](buffer-and-memory.md)、[事件与异步](events-and-async.md)）。它们在这里只是为了展示*适配器客户端*暴露的内容；本页负责它们如何接入 SE→PJRT 桥接，而不是它们的槽位级细节。

---

## TpuDevice — 在每个 PjRtDevice 下包装 tpu::System

### 目的

`xla::TpuDevice` 是 C-ABI device 槽位包装的 `PjRtDevice`。在上游 SE 模型中，这个对象*本会*持有一个 `StreamExecutor` 和一个 stream 池。这里它不持有 executor：它指向*共享的* `tpu::System`，命名自己映射到的一个物理 core，并携带一个 launch 节流 semaphore。这是“SE → PJRT 适配器”这一说法的关键 —— 设备级绑定面向异步 runtime，而不是面向基于 stream 的 executor。

### TpuDevice 持有什么

`TpuDevice` 对象大小为 `0x1E0`（480）字节（枚举循环中的 `operator new(0x1E0)`）。从 ctor 签名（`0xF7FDC40`）和 `SetClient`（`0xF7FE520`）恢复出的成员如下：

| 成员 | 来源 | 作用 |
|---|---|---|
| `tpu::System*` | ctor arg 4 | **共享的**设备 runtime —— 每个 `TpuDevice` 中都是同一个指针 |
| `tpu::TpuCoreLocation` | 来自枚举循环的 `core_loc` | 此 PjRt device 映射到的物理 core |
| `xla::TpuDeviceDescription` | ctor arg 1 | 静态身份：id/process/kind/attributes/memory |
| `xla::TpuClient*` | `SetClient` arg 1 | 指向 owning client 的反向指针 |
| `xla::AsyncWorkRunner*` | `SetClient` arg 2 | 此设备 continuations 的 host scheduler |
| `xla::Semaphore` | 构造；由 `ExecutablesStart`/`Complete` 围住 | 进行中 computation 的节流器（受 `max_inflight_computations` 限制） |

所有设备的 `tpu::System*` 都*相同*，这是与 SE 的结构性差异。在上游 XLA 中，`device[i]->local_device_state()->executor()` 返回每个设备不同的 `StreamExecutor`。这里，`device[i]` 和 `device[j]` 共享一个 `tpu::System`；区分它们的是 `TpuCoreLocation`，`tpu::System::Execute` 和 `LoadProgram` 会使用它来定位正确的 core。

### SetClient 反向边

```c
function TpuDevice::SetClient(client, work_runner):      // 0xF7FE520
    this->client       = client                          // back-pointer used by C-ABI device→client lookups
    this->work_runner  = work_runner                     // the client's AsyncWorkRunner (host scheduler)

这会在 TpuClient 构造后立即对每个设备运行一次。它是 TpuDevice 得知自己 scheduler 的唯一方式。C-ABI PJRT_Device 槽位用来把 implementation device 翻译回其 wrapper 的 client() 访问器（通过 client 的 device_map）会读取这个字段 —— 因此在 SetClient 之前它无效，这就是 GetTpuPjRtClientInternal 中构造顺序（先构建 client，然后循环 SetClient）对加载至关重要的原因。

Launch 节流器 — 最接近 Stream 的东西

TpuDevice 携带一个每设备的 xla::Semaphore，由 max_inflight_computations 限制（默认 1，来自 config）。ExecutablesStart（0xF800300）在 launch 前获取它；ExecutablesComplete（0xF800740）在该 launch 的 completion event 触发时释放它。这是 SE stream 深度限制的功能性替代物：它阻止 host 在 device 前面无界排入过多 executions。它不是 FIFO —— launch 之间的排序由 TpuEventIssuer 序列点强制，而不是由 semaphore 强制。semaphore 只限制并发。

怪癖 — 当 max_inflight_computations == 1（默认）时，单个 TpuDevice 的行为很像单个 SE stream：最多一个 execution 进行中，下一个会被阻塞直到前一个完成。但这种相似很浅 —— 提高这个旋钮会让设备接受多个并发 launch，且它们之间没有 FIFO 排序，这是 SE stream 永远不会做的。重新实现者必须把节流器视为并发上限，而不是序列化队列。

设备内存空间

每个 TpuDevice 附加三个 memory space（每设备的 AttachMemorySpace 调用）：

xla::TpuHbmMemorySpace（kind "tpu_hbm"）—— device HBM。
xla::PinnedHostMemorySpace / xla::UnpinnedHostMemorySpace —— host staging。

规范的 kind 常量是 xla::kBuiltinTpuMemorySpaces / xla::kTpuHbmMemorySpaceKind。槽位级 memory accessors 和位于这些空间中的 buffer objects 由 Buffer 与内存负责。

函数图

函数	地址	作用
`TpuDevice::TpuDevice` (ctor)	`0xF7FDC40`	从 desc + name + ordinal + `tpu::System*` 构建设备
`TpuDevice::SetClient`	`0xF7FE520`	安装 client 反向指针 + `AsyncWorkRunner`
`TpuDevice::ExecutablesStart`	`0xF800300`	launch 前获取 in-flight semaphore
`TpuDevice::ExecutablesComplete`	`0xF800740`	completion 时释放 in-flight semaphore

注意 — TpuDeviceDescription（vtable 0x21787AC0，typeinfo 0x21787C38）是携带 id / process_index / device_kind / Attributes / memory 的 PjRt device description。带 SE 风格的 PjRtStreamExecutorDeviceDescription（vtable 0x2177D950，typeinfo 0x2177D9B0）在这个二进制中是仅 GPU 的死代码 —— 它的存在是唯一残留的 SE-PjRt 符号，而且在 TPU 路径上永远不会到达。阅读二进制时不要混淆两者。

Runtime 后端 — 作为 StreamExecutor 对应物的 tpu::System

目的

tpu::System（Initialize @ 0x1D0AE420）是整个适配器所在的对象 —— TFRT 原生 TPU 设备运行时。它就是本页标题中所谓 “StreamExecutor” 向前映射后的东西：提供 launch、transfer、allocate 和 event primitives 的表面，但形状是 async-value，而不是 stream。一个 tpu::System 由进程中的每个 TpuClient 和每个 TpuDevice 共享（通过 GetTpuPjRtClient 中构建的 TpuStatesManager / TpuSystemState 单例到达它）。

API 表面（等价于 SE-Stream 的操作）

操作	VA	SE 等价物
`Initialize(shared_ptr<TpuTopology>, InitOptions, ConcurrentWorkQueue*)`	`0x1D0AE420`	`StreamExecutor::Init`
`Execute(AsyncValueRef<ProgramHandle>, ExecuteOptions, Span inputs, outputs, wait, define)`	`0x1D0B33E0`	`Stream::ThenLaunch`
`LoadProgram(TpuCoreLocation, shared_ptr<TpuCoreProgram>)`	`0x1D0B2240`	`StreamExecutor::GetKernel`（加载一次）
`TransferToDevice(AsyncValueRef<TpuBuffer>, Span<uint8>, opt<TpuSyncFlagOnChip>)`	`0x1D0AFA20`	`Stream::ThenMemcpy` (H2D)
`TransferFromDevice(Span<uint8>, AsyncValueRef<TpuBuffer>, opt<…>)`	`0x1D0B0160`	`Stream::ThenMemcpy` (D2H)
`AllocateHostBuffer` / `MakeTpuHostBuffer`	`0x1D0AF180` / `0x1D0AF660`	host-pinned alloc
`EnqueueInfeed` / `DequeueOutfeed`	`0x1D0B5D00` / `0x1D0B5F00`	infeed/outfeed

Execute 接收 TpuCoreLocation（间接地，通过为该 core 加载的 ProgramHandle）—— 这就是一个共享的 tpu::System 如何在没有每设备 executor 的情况下服务多个 TpuDevice。Execute 背后的内部硬件 command-stream / ring 布局属于 tpu::System::Impl 和 TPU driver core 的私有内容；未跟踪它（它是与这个桥接层不同的关注点）。

注意 — 设备侧排序 primitive 是 tpu::TpuEventIssuer（AddDepsNoReserve、AggregateDeps、RunWhenDepsReady、NextSequencePoint、FulfillArgs、ChainScope）。它是对 Execute / Transfer / Allocate 进行相互排序的序列点引擎 —— 以 tpu::TpuEvent（每个都是一个 tsl::AsyncValueRef）DAG 的形式，功能上替代 SE 每 stream FIFO 加跨 stream Event。其精确排序算法没有按字节追踪。驱动依赖 continuations 的 host 侧 scheduler 是 xla::ThreadPoolAsyncWorkRunner（AsyncWorkRunner::ExecuteWhenReady），由 HostContext 的 tfrt::ConcurrentWorkQueue 供给。两者都通过 TpuClient::async_work_runner() 到达。

重新实现者注意事项

不要构建 LocalDeviceState。 最可能犯的错误是复刻上游 XLA 的每设备 StreamExecutor + stream 池。TPU 适配器没有这些。构建一个共享的 tpu::System，把它的指针交给每个 TpuDevice，并且只用 TpuCoreLocation 区分设备。
单例是进程级的，不是每客户端的。 HostContext（通过 CreateDefaultHostContext::kInstance）和 tpu::System（通过 TpuStatesManager）都是 __cxa_guard / mutex 一次性对象。进程中的第二个 client 必须复用它们，而不是重建它们。TpuSystemState 上的 MaybeOwningDeleter 编码了借用行为。
设备↔client 接线是两阶段的。 用设备 vector 构造 client，然后循环 SetClient 安装反向指针和 AsyncWorkRunner。在该循环运行之前，设备还不能完全使用。
并发来自 semaphore，排序来自 event DAG。 每设备 Semaphore 限制进行中的 launch；TpuEventIssuer 序列点对它们排序。两者正交 —— 把节流器和 FIFO 混为一谈会在 max_inflight_computations > 1 时产生错误排序。
旧 SE 栈是误导。 TpuExecutor / TpuStream / TpuTransferManager 存在于同一个二进制中，但服务 LocalClient / TF kernels，从不服务 PJRT。PJRT 路径的重新实现者应忽略每个 ExecutorApiFn / SE_* 符号。

组件	关系
`PJRT_Client_Create` (slot 15)	调用 `GetTpuPjRtClient`；负责这个桥接上方的 option parse 和 C-ABI wrapper
`xla::CommonPjRtClient`	`TpuClient` 派生自的 async-value 框架；提供 buffer/executable/event 管线
`MegaScalePjRtClient` / `TfPjRtClient`	由 `PJRT_Client_Create` 堆叠在裸 `TpuClient` 上的可选装饰器
`tpu::System`	每个 `TpuDevice` 指向的共享设备 runtime；SE-`StreamExecutor` 对应物
`tpu::TpuEventIssuer`	设备侧序列点引擎；替代 SE 每 stream FIFO + events
旧 SE `TpuPlatform`	并行的非 PJRT 设备层；从不被 `TpuClient` 引用

交叉引用

客户端、设备与拓扑 — 负责 PJRT_Client_Create、option-kv 摄入，以及装箱本页构建的 TpuClient 的 208 字节 PJRT_Client wrapper
StreamExecutor 与宿主解释器 — 旧 SE TpuPlatform 注册和 host-interpreter 路径，它与这个适配器并行运行（从不位于其下）
Executable 与执行 — PJRT_LoadedExecutable_Execute 以及适配器 client 将 programs 加载进入的 TpuExecutable / TpuLoadedExecutable / TpuRawLoadedExecutable 框架
在 Stream 上异步执行 — 下行到 tpu::System::Execute 的每执行路径，即本页从 Stream::ThenLaunch 映射而来的 launch primitive 的 runtime 侧
Buffer 与内存 — 覆盖 tpu::TpuBuffer 的 TpuRawBuffer、每个 TpuDevice 附加的 memory spaces，以及 allocator glue
事件与异步 — PJRT_Event ↔ tsl::Future<void>，以及本页为排序引用的 TpuTrackedDeviceEvent / TpuEventIssuer 模型
概览 — PJRT plugin 入口、extension chain 和 GetPjrtApi 填充

StreamExecutor → PJRT 适配器 ​

摘要 ​

SE → TPU-Runtime 概念映射 ​

目的 ​

映射表 ​

为什么 SE 栈是缺席，而不是被适配 ​

GetTpuPjRtClient — 构造桥接 ​

目的 ​

入口点 ​

设备枚举 — GetTpuPjRtClientInternal ​

TFRT 构造函数 ​

Launch 节流器 — 最接近 Stream 的东西 ​

设备内存空间 ​

函数图 ​

Runtime 后端 — 作为 StreamExecutor 对应物的 tpu::System ​

目的 ​

API 表面（等价于 SE-Stream 的操作） ​

重新实现者注意事项 ​

相关组件 ​

交叉引用 ​

StreamExecutor → PJRT 适配器

摘要

SE → TPU-Runtime 概念映射

目的

映射表

为什么 SE 栈是缺席，而不是被适配

GetTpuPjRtClient — 构造桥接

目的

入口点

设备枚举 — GetTpuPjRtClientInternal

TFRT 构造函数

Launch 节流器 — 最接近 Stream 的东西

设备内存空间

函数图

Runtime 后端 — 作为 StreamExecutor 对应物的 tpu::System

目的

API 表面（等价于 SE-Stream 的操作）

重新实现者注意事项

相关组件

交叉引用