Internal Wiki

主题

PJRT 事件与异步跟踪 ​

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d）。该镜像未 strip；demangled C++ symbol name 均逐字引用。.text VMA 等于文件偏移。PJRT C-API 版本为 v0.103。其他版本会不同。

摘要 ​

PJRT_Event 是 PJRT plugin 用来报告异步操作完成的 C-ABI 句柄，例如 program launch、host-to-device upload、device-to-host copy。它是不透明指针，调用者从 PJRT_LoadedExecutable_Execute、PJRT_Buffer_ReadyEvent 和 transfer 路径取得，然后轮询（PJRT_Event_IsReady）、阻塞等待（PJRT_Event_Await），或按惯用方式附加 done-callback（PJRT_Event_OnReady）。完成后用 PJRT_Event_Destroy 释放。C-API 占据五个连续 vtable 槽位（10-14）以及两个后加项 PJRT_Event_Create/PJRT_Event_Set（槽位 131-132），后者允许调用者显式创建并 resolve 一个 event。本页负责该 C-ABI 接口及其背后的对象。

PJRT_Event 本身不是完成 primitive。每个 event 背后都有一个 xla::PjRtFuture<void>，具体是包裹 refcounted、single-assignment tsl::AsyncValue 的 tsl::internal::FutureBase<absl::Status, false>。wrapper 函数很薄：每个函数在版本检查后展开 C args 结构，在固定偏移处到达 future，然后转发到一个 FutureBase 方法。IsReady 读取 async value 的 state byte；Await 调用 FutureBase::Await；OnReady 调用 future 的 AndThen，它要么内联运行 callback（value 已可用），要么创建一个 AsyncValue waiter node 并把它串到 value 的 waiter list。这里没有 poll loop；completion 是推送式的：runtime fulfil async value，触发其 waiter list，并在执行 fulfilment 的线程上运行已注册的 done-callback。

本页是 C-ABI event wrapper 层。fulfil 这些 event 的 runtime-internal 机制，包括 linked promise pair、device-side TpuTrackedDeviceEventPromise、tpu::System::Execute 的 define-event registration，以及在 device retirement 时 resolve 它们的 TpuEventIssuer，位于下一层 Completion Loop & AsyncTrackingEvent。放置这些函数的 140 槽表在 API Vtable Reconstruction。阅读那些页面可了解什么会产生 event；阅读本页可了解调用者如何使用 event。

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

  • PJRT_Event 对象布局：0x50 字节 heap struct，持有 { AsyncValue* future_av, two profiling callbacks, PromiseBase<absl::Status> promise }，由 PromiseMaker<void>::Make 构建，并由 PJRT_Event_Destroy 拆除。
  • future/async-value 后端：tsl::internal::FutureBase<absl::Status, false> 包裹 tsl::AsyncValue，其 state byte（& 2 = concrete/available）控制 readiness，且 OnReady 必须遍历 indirect-value chain。
  • 五方法 dispatch 契约：每个入口的 args-struct min/current size、future 偏移、IsValid()/IsReady() 前置条件，以及 error-wrapping 返回约定。
  • OnReady dispatch：ready 时内联运行，否则 enqueue waiter node；node 携带 base::Context，因此 callback 在调用者的 trace context 中运行，而不是 fulfiller 的 context。
Backing futurexla::PjRtFuture<void> = tsl::internal::FutureBase<absl::Status, false> over a tsl::AsyncValue
Event object0x50 字节 heap struct { av@+0, profiling_cb@+8/+40, PromiseBase<absl::Status>@+72 }
PJRT_Event_Destroyslot 10 · pjrt::PJRT_Event_Destroy @ 0xf86f920（args min 18 / cur 24）
PJRT_Event_IsReadyslot 11 · pjrt::PJRT_Event_IsReady @ 0xf86f9e0（args min 18 / cur 25）
PJRT_Event_Errorslot 12 · pjrt::PJRT_Event_Error @ 0xf86fba0（args min 16 / cur 24）
PJRT_Event_Awaitslot 13 · pjrt::PJRT_Event_Await @ 0xf86fa80（args min 16 / cur 24）
PJRT_Event_OnReadyslot 14 · pjrt::PJRT_Event_OnReady @ 0xf86fc60（args min 18 / cur 40）
PJRT_Event_Createslot 131 · pjrt::PJRT_Event_Create @ 0xf86fe00（args min 17 / cur 24）
PJRT_Event_Setslot 132 · pjrt::PJRT_Event_Set @ 0xf86ffa0（args min 14 / cur 48）
Readiness flagav+8 处的 AsyncValue state qword，bit & 2 = concrete（value available）
Allocated flagstatus-rep / async-value low bit（& 1）或 byte[+4] & 8 = heap-owned；控制 refcount drop
Evidence gradeReimplementation-grade / 已对 IDA decompile 做 byte-confirmed

1. 后端对象：PjRtFuture 与 PJRT_Event 结构 ​

目的 ​

每个 C-ABI event 方法都是围绕单个 C++ 对象的一页 wrapper。先理解该对象，五个方法就很简单。PJRT_Event*（args->event 中的值）指向一个由 PJRT_Event_Create 或 runtime execute 路径构建的 heap struct；该 struct 嵌入用户可见的 future（xla::PjRtFuture<void>）以及 runtime fulfil 的 promise。future 是 tsl::internal::FutureBase<absl::Status, false>：absl::Status payload 携带错误信息，false 模板参数表示不可复制，底层是 tsl::AsyncValue（completion primitive）。

Event 对象布局 ​

PJRT_Event_Create（0xf86fe00）是该结构最清楚的字节级证据。它 operator new 0x50 bytes 并填充，PJRT_Event_Destroy 按同一布局逐字段拆除。交叉引用二者：

FieldOffsetTypeMeaning
future_av+0x00tsl::AsyncValue*future 观察的 async value（readiness / error 位于此处）
profiling_cb_a+0x08..+0x20std::function-style policy pairProfilingKeys() open callback（__policy_func + policy ptr）
profiling_cb_b+0x28..+0x40std::function-style policy pairvoid(ProfilingKeys) close callback
promise+0x48tsl::internal::PromiseBase<absl::Status>promise half；qword 术语中的 obj[9]

两个 profiling callback 是由 tsl::FutureHelpers::ProfilingKeys 标识的 TraceMe span open/close 对，即 runtime 在 TpuClient::CreateProfiledFuture 中应用的同一个 profiling wrapper（completion loop §5）。在 C-API mint 路径（PJRT_Event_Create）上，二者都默认为空策略（__create_empty），因此显式创建的 event 不携带 profiling span。

对象如何创建 ​

c
// pjrt::PJRT_Event_Create(PJRT_Event_Create_Args*)                    0xf86fe00
function PJRT_Event_Create(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_Create", 17, 24, args->struct_size):
        return wrap_error(...)                       // older/newer header mismatch
    // PromiseMaker<void>::Make -> { promise, async_value }   (line 35)
    PromiseMaker<void>::Make(&promise, empty_profiling_open, empty_profiling_close)
    obj = operator new(0x50)
    obj[0]  = async_value                            // +0x00 : the AsyncValue, refcount transferred in
    obj[1..8] = profiling callback pair (empty)      // +0x08..+0x40
    obj[9]  = promise                                // +0x48 : PromiseBase<absl::Status>
    args->event = obj                                // hand the PJRT_Event back
    drop_ref(async_value)                            // §discipline: & 8 allocated bit guards Destroy
    return ok
```text

`PromiseMaker<void>::Make` 是 runtime 的 [`CreateLinkedUserPromise`](../runtime/completion-loop.md#2-creating-the-per-execution-event--the-linked-promise-pair) 使用的同一个 primitive（那里第 41 行）。区别在于*谁持有 promise*：在 `PJRT_Event_Create` 中调用者持有它（装箱在同一个 event 对象中，之后通过 `PJRT_Event_Set` resolve）；在 execute 路径中 runtime 持有 device-side half 并把它链接到 user value。

> **NOTE —** `PromiseMaker<void>` 产生 void-payload promise/future，但 C-API 接口（`Await`、`Error`、`Set`）都把 status channel 类型化为 `absl::Status`。future 是 `FutureBase<absl::Status, false>`：`void` 是*成功* payload（无返回内容），`absl::Status` 是*错误* channel。无错误的 ready event resolve 为 OK status；errored event resolve 为非 OK `StatusRep`。

### Readiness / Allocated 位纪律

两个 flag bit 在所有五个方法中反复出现，重新实现若误处理任何一个都会破坏模型：

```c
// readiness test, observed identically in IsReady / Error / OnReady
is_ready = (async_value->state_byte[+8] & 2) != 0     // bit 1 of the qword at av+8 = "concrete"

// refcount drop, observed in Destroy / Create / Await / OnReady
function drop_ref(av):
    if av == nullptr: return
    if (av->byte[+4] & 8) == 0: return                // not heap-allocated -> never destroy (singletons)
    if av->refcount[+0] == 1 || AtomicDecrement(&av->refcount) == 0:
        AsyncValue::Destroy(av)

av+8 处 qword 的 & 2 是 "value is concrete / available" 标志。av+4 处 byte 的 & 8 是 "heap-allocated" 标志；静态拥有的 ready/error singleton 清除此位，永不释放。status-rep 路径使用并行的 StatusRep* 低位 tag & 1，用以区分 inline/OK status 与需要 Unref 的 heap StatusRep。二者在 completion loop page §1 中完整说明；这里重复是因为 event wrapper 会直接操作它们。

2. PJRT_Event_IsReady：非阻塞状态读取 ​

目的 ​

最便宜的查询：后端 async value 是否可用？它在 args->is_ready 中返回 bool，不阻塞。JAX/PyTorch-XLA 在防御性检查和紧密 readiness loop 中轮询此项（在 140 个槽位中排名第 2 的 hot-path slot）。

算法 ​

c
// pjrt::PJRT_Event_IsReady(PJRT_Event_IsReady_Args*)                  0xf86f9e0
function PJRT_Event_IsReady(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_IsReady", 18, 25, args->struct_size):
        return wrap_error(...)
    event = args->event                              // args+16
    future_av = *event                               // event[0] = AsyncValue*
    if future_av == nullptr:
        FATAL("IsValid()", future.h:281)             // event must hold a valid future
    args->is_ready = (future_av->state[+8] & 2) != 0 // args+24 : non-blocking concrete check
    return ok
```text

唯一逻辑是 `& 2` state read。`IsValid()` CHECK（`future.h:281`）捕获 future 已被 move-out 或从未设置的 event 误用，这是 fatal，不是 status。重新实现必须把它保持为硬不变量：后端 future 为空的 `PJRT_Event` 是编程错误，不是运行时条件。

> **GOTCHA —** `IsReady` 返回 true 表示 *async value 可用*，即 launch retired 或 transfer 的 define event resolved。它**不**表示 device→host 输出数据已经落入 host memory；这需要单独的 copy event。见 [completion loop §5 GOTCHA](../runtime/completion-loop.md#buffer-release)。把 execute event 的 "ready" 当成 "outputs are host-readable" 会把 device handle 当作 host pointer 读取。

---

## 3. `PJRT_Event_Await` 与 `PJRT_Event_Error`：阻塞与状态读取

### 目的

`Await` 是该接口中*唯一*阻塞路径：它暂停调用线程直到 future resolve，然后返回其错误 status（成功时为 null `PJRT_Error*`）。`Error` 在*不*阻塞的情况下读取已 resolved 的错误，并断言 event 已 ready。二者共享同一个 `FutureBase::Await` 后端和同一个 error-wrapping 返回约定；区别只在于是否等待。

### `PJRT_Event_Await`

```c
// pjrt::PJRT_Event_Await(PJRT_Event_Await_Args*)                      0xf86fa80
function PJRT_Event_Await(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_Await", 16, 24, args->struct_size):
        return wrap_error(...)
    status = FutureBase<absl::Status,false>::Await(args->event)   // a1[2] ; BLOCKS until available
    if (status & 1) == 0:                            // heap StatusRep, not inline-OK
        AtomicIncrement(status)                      // take a ref before returning it
    if status == OK:                                 // low-tag == &dword_0+1 sentinel
        return nullptr                               // success: no PJRT_Error
    return wrap_error(status)                        // box the non-OK status as a PJRT_Error*

FutureBase<absl::Status, false>::Await 是真正的阻塞 primitive：它注册内部 waiter 并暂停线程直到 async value 触发，然后返回 resolved absl::Status。& 1 低位 tag 区分 inline OK sentinel（&dword_0 + 1，标准 "OK" 表示）与逃逸到 C-ABI 前必须增加引用计数的 heap StatusRep*。

PJRT_Event_Error ​

c
// pjrt::PJRT_Event_Error(PJRT_Event_Error_Args*)                      0xf86fba0
function PJRT_Event_Error(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_Error", 16, 24, args->struct_size):
        return wrap_error(...)
    event = args->event                              // a1[2]
    if *event == nullptr:
        FATAL("IsValid()", future.h:281)
    if (event->future_av->state[+8] & 2) == 0:       // <-- PRECONDITION: must be ready
        FATAL("event->future.IsReady()", pjrt_c_api_wrapper_impl.cc:3032)
    status = FutureBase<absl::Status,false>::Await(event)  // already ready -> returns immediately
    ... same OK/heap-StatusRep return convention as Await ...
```text

> **GOTCHA —** `PJRT_Event_Error` **不是**阻塞调用，不能当作阻塞调用使用。如果 event 尚未 ready，它会 CHECK-fail（fatal，`pjrt_c_api_wrapper_impl.cc:3032`）。正确顺序是 `IsReady()`（或 `OnReady`/`Await`）**然后** `Error()`。虽然它内部调用相同的 `FutureBase::Await`，但 readiness 前置条件意味着该调用会立即返回而不会 park。允许调用者在 event resolve 前读取 error 的重新实现会使进程崩溃，这精确匹配 upstream PJRT 语义。

共享返回约定是整个 plugin 的 C-ABI error 习惯：方法返回 `PJRT_Error*`，成功为 `nullptr`，否则为装箱的 `absl::Status`。装箱是裸 `operator new(8)`，其中持有 `StatusRep*`；调用者用 `PJRT_Error_Destroy`（槽位 5）释放它。status 自身的 refcount 由 `& 1`/`Unref` 纪律管理，因此 rep 的生命周期长于产生它的 future。

---

## 4. `PJRT_Event_OnReady`：推送式 Done-Callback

### 目的

惯用 completion 路径，也是本页核心。调用者传入 C 函数指针（`callback`）和不透明 `user_arg`；libtpu 安排 `callback(error, user_arg)` 在 future resolve 时恰好运行一次。如果 future *已经*可用，callback 在调用线程上内联运行；否则 libtpu 创建 waiter node，将其串到 async value 的 waiter list，callback 在*fulfil* value 的线程上运行（device-completion thread、transfer thread 或调用 `PJRT_Event_Set` 的线程）。

### Args 布局

`PJRT_Event_OnReady_Args`（min 18 / current 40 bytes）：

| Field | Offset | Type | Meaning |
|---|---:|---|---|
| `struct_size` | `+0` | `size_t` | version gate |
| `event` | `+16` | `PJRT_Event*` | 要观察 completion 的 event |
| `callback` | `+24` | `void(*)(PJRT_Error*, void*)` | done-callback |
| `user_arg` | `+32` | `void*` | 不透明值，作为 `callback` 第二个实参传回 |

### 算法

```c
// pjrt::PJRT_Event_OnReady(PJRT_Event_OnReady_Args*)                  0xf86fc60
function PJRT_Event_OnReady(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_OnReady", 18, 40, args->struct_size):
        return wrap_error(...)
    av       = *(args->event)                        // event[0]
    if av == nullptr: FATAL("IsValid()", future.h:401)
    callback = args->callback                        // args+24
    user_arg = args->user_arg                        // args+32
    state    = av->qword[+8]

    if (state & 2) != 0:                             // ALREADY AVAILABLE -> run inline
        while (av->byte[+4] & 3) != 0:               // walk the indirect-value chain
            av = av->qword[+16]                       //   IndirectAsyncValue -> concrete target
        status = av->qword[+64]                       // the resolved StatusRep* (or OK sentinel)
        if (status & 1) != 0:                         // low-tag set: inline/OK-style status (no refcount)
            if status == OK:                          //   &dword_0 + 1 -> success
                callback(nullptr, user_arg)
            else:
                callback(box_error(status), user_arg) // <-- inline, on THIS thread
        else:                                         // low-tag clear: heap StatusRep -> manage refcount
            AtomicIncrement(status)                   // bump for the boxed PJRT_Error*
            err = box_error(status)
            AtomicIncrement(status)                   // bump for the callback's owned ref
            callback(err, user_arg)                   // <-- inline, on THIS thread
            StatusRep::Unref(status)                  // drop OnReady's own ref
        return ok

    else:                                            // NOT YET -> enqueue a waiter node
        node = operator new(0x80)                     // 128-byte AsyncValue waiter Node
        node.vtable = &TraceContext_node_vtable        // off_2177E068
        base::Context::Context(node + 16)              // capture caller's trace context
        node[13] = av                                  // the value to read when fired
        node[14] = callback
        node[15] = user_arg
        tsl::AsyncValue::EnqueueWaiterListNode(av, node, state)   // CAS onto waiter list
        return ok

这一函数中有三个对重新实现至关重要的机制。

(1) indirect-value 遍历。 value 可用时，av 可能是被 ForwardTo 到 concrete value 的 IndirectAsyncValue placeholder（completion loop 的 SetReady 正是把 concrete TpuEvent 拼接进这种 indirect）。while (av->byte[+4] & 3) != 0: av = av->qword[+16] 循环追踪 indirection chain 到 concrete value 后，才读取 +64 处的 status。重新实现若从 indirect placeholder 读取 status，会读到垃圾。

(2) inline vs. enqueued dispatch。 如果 value 已 concrete，callback 会同步在调用线程上运行；OnReady 不会转交给 thread pool。只有 not-yet 情况会分配 waiter node。调用者不能假定 OnReady 返回后 callback 才运行。

(3) trace context。 0x80 字节 node 嵌入一个在注册时捕获的 base::Context（TraceContext）。当 node 在 fulfiller 线程上触发时，它恢复调用者 context，使 profiling/tracing 将 callback 归因于发起请求，而不是 device-completion thread。这与函数表中按 closure type 实例化的 EnqueueWaiter<...>::Node / RunWaiterAndDeleteWaiterNode 机制相同（例如 collectives 中 tsl::AsyncValue::EnqueueWaiter 下的 FutureBase::AndThen 实例化）。

与 FutureBase::AndThen 的关系 ​

C-ABI OnReady 是手工内联的 xla::PjRtFuture<void>::OnReady，upstream 中即 FutureBase<absl::Status>::AndThen(callback)。symbol table 确认了该链接：tsl::AsyncValue::EnqueueWaiter<...FutureBase<absl::Status,false>::AndThen<...PJRT_Event_OnReady::$_0>...>::Node 类型，即 enqueued 路径中创建的 waiter node，在 0xf87a580（~Node，deleting destructor）和 0xf87a5c0（RunWaiterAndDeleteWaiterNode，value 的 waiter list 调用的 fire-and-free 入口）发出两个 out-of-line virtual method，其 vtable 位于 0x2177e058（存储的 vptr 是 off_2177E068，即 offset-to-top / typeinfo header 后的 +0x10）。反编译内联了 available-case fast path，而不是总是调用 AndThen；enqueued case 分配的正是这个 node 类型。

QUIRK — callback 接收 PJRT_Error*，不是 bool。成功时为 nullptr；错误时是 callback 拥有并必须用 PJRT_Error_Destroy 销毁的新装箱 StatusRep。inline 路径围绕调用将 status refcount 增加两次并 Unref 一次，恰好留下一个引用给 callback 释放。忘记装箱（或不增加 refcount 就把同一个 StatusRep* 交给多个 callback）的重新实现会 double-free status。

5. PJRT_Event_Set 与 PJRT_Event_Destroy：调用者侧 Resolve 与 Teardown ​

目的 ​

PJRT_Event_Create/PJRT_Event_Set（槽位 131/132）是后加的一对接口，允许调用者创建 event 并显式 resolve 它，是 runtime-driven TpuTrackedDeviceEventPromise::SetReady 的 framework-driven 对应物。调用者创建 event（§1），把它交给某个异步 producer，之后用 Set 携带 status fulfil 它。Destroy 释放任何 event，不论其创建方式。

PJRT_Event_Set ​

c
// pjrt::PJRT_Event_Set(PJRT_Event_Set_Args*)                          0xf86ffa0
function PJRT_Event_Set(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_Set", 14, 48, args->struct_size):
        return wrap_error(...)
    code   = PjrtErrorCodeToStatusCode(args->error_code)   // args+24 : PJRT_Error_Code -> absl code
    msg    = args->error_message                            // args+32 (ptr)
    msglen = args->error_message_size                       // args+40 (>= 0, else BUG())
    rep    = absl::Status::MakeRep(4*code+1, msg, msglen, 3067, "...pjrt_c_api_wrapper_impl.cc")
    promise = event + 72                                    // obj[9] = PromiseBase<absl::Status>
    PromiseBase<absl::Status>::emplace<absl::Status>(promise, &rep)   // <-- FULFILS the future
    if (rep & 1) == 0: StatusRep::Unref(rep)
    return ok
```text

`emplace` 是 resolution edge：它把 status 赋给 promise 的 async value，使其转换为 concrete（设置 `& 2`）并触发 waiter list，从而运行每个 `OnReady` callback 并解除每个 `Await` 的阻塞。`PjrtErrorCodeToStatusCode` 将 C-API error enum 映射到 `absl::StatusCode`；`4*code+1` 编码是 absl 的 tagged-status 表示。传入 OK code 会产生成功 resolution。这与 runtime 在 device 路径上内部执行的 *exact same fulfilment* 相同，区别只在触发来源（C-API 调用 vs. device retirement）。

> **NOTE —** `PJRT_Event_Set` 的 args min size 是 **14**，是所有 event method 中最小的，因为最初的 event-resolution args 早于 message-size 字段。`error_message`/`error_message_size` 字段（`+32`/`+40`）把*当前* size 推到 48；使用旧 header 编译的调用者如果只传 `{struct_size, event, error_code}`（约 24 字节），仍会 resolve event，只是没有 message。`args->error_message_size < 0` guard（`BUG()`）用于防御 caller struct 大于其实际填充内容时未初始化的尾部。

### `PJRT_Event_Destroy`

```c
// pjrt::PJRT_Event_Destroy(PJRT_Event_Destroy_Args*)                  0xf86f920
function PJRT_Event_Destroy(args):
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Event_Destroy", 18, 24, args->struct_size):
        return wrap_error(...)
    obj = args->event                                // a1[2]
    if obj == nullptr: return ok
    PromiseBase<absl::Status>::~PromiseBase(obj + 9)  // +0x48 : destroy the promise half
    destroy_policy(obj[8])(obj[5])                    // run profiling_cb_b destructor if present
    destroy_policy(obj[4])(obj[1])                    // run profiling_cb_a destructor if present
    drop_ref(obj[0])                                  // §1 discipline on the AsyncValue (& 8 gate)
    free(obj)                                         // the 0x50-byte struct itself
    return ok

Destroy 是 Create 的镜像：析构 promise（obj+9 = +0x48），运行两个 profiling-callback policy destructor（位于 obj[1]/obj[4] 与 obj[5]/obj[8] 的 std::function-style {ptr, vtable} 对），在 & 8 allocated-bit guard 下减少 async value refcount，然后释放 heap struct。销毁 event 不会取消其 operation，只会释放 handle；runtime 持有的任何 pending OnReady waiter 仍会触发（它携带自己的引用）。在 waiter 已排队时销毁 async value 的重新实现会破坏 runtime；& 8 + refcount 纪律正是为了防止这一点。

6. Event 从何而来：Producer ​

目的 ​

调用者通常不会用 PJRT_Event_Create 构造 runtime 的 event；这是显式 resolution 路径。training/inference loop 中重要的 event 由其他 C-API 调用返回，每个调用都交回一个 PJRT_Event*（或把 future 装箱为一个 event）。本节映射 producer，使重新实现者知道哪些槽位创建 event，以及每个 event 的 "ready" 含义。

Producer 映射 ​

Producer slotFunctionEvent semantics
PJRT_LoadedExecutable_Execute (60)pjrt::PJRT_LoadedExecutable_Execute @ 0xf869b40per-launch completion；ready = program retired on device
PJRT_Buffer_ReadyEvent (77)pjrt::PJRT_Buffer_ReadyEvent @ 0xf86ed20buffer definition event；ready = backing HBM is valid
PJRT_Buffer_ToHostBuffer (75)pjrt::PJRT_Buffer_ToHostBuffer @ 0xf86e640D2H copy completion；ready = host bytes landed
PJRT_Buffer_CopyRawToHostFuture (125)pjrt::PJRT_Buffer_CopyRawToHostFuture @ 0xf86dfe0raw D2H copy；ready = host bytes landed
PJRT_Client_BufferFromHostBuffer (27)pjrt::PJRT_Client_BufferFromHostBuffer @ 0xf8644c0H2D upload done-event
transfer-manager slots (106-114, 124)PJRT_AsyncHostToDeviceTransferManager_*per-chunk / per-buffer transfer events

execute event 由 CreateLinkedUserPromise 创建，并由 tpu::System::Execute 的 completion lambda 通过 TpuEventIssuer fulfil。buffer/transfer event 使用同一个 tsl::AsyncValue primitive；其 fulfilment 是 transfer-completion lambda，而不是 device-retirement lambda。每种情况下，调用者收到的 C-API event 都是本页记录的同一个 0x50 字节 wrapper；producer 不同，接口相同。

QUIRK — execute event 与 buffer 的 ReadyEvent 是不同 event，背后是不同（但 linked）的 async value。execute event resolve 为 "ready" 会使 output buffer 的 definition event 可用，后者最终翻转 PJRT_Buffer_ReadyEvent 的 future。把 "execute done" 与 "every output buffer ready" 混为一谈的重新实现会漏掉 buffer 被 donated 或 aliased 且其 readiness 由不同 define event 控制的情况。关于 buffer-side lifecycle，见 Buffer and Memory。

7. 重新实现注意事项 ​

  • Event 是 wrapper，不是 primitive。 先实现 PjRtFuture<void> / AsyncValue（completion loop）；五个 C-ABI 方法随后各自只有 10-20 行。不要把 completion 逻辑放进 C wrapper。
  • OnReady 是主路径；Await 是 fallback。 注册一个 OnReady waiter，让 runtime 推送 completion。将 Await 保留给少数同步调用者；它是唯一会 park 线程的方法。
  • 遵守 args version gate。 每个方法第一步都是 ActualStructSizeIsGreaterOrEqual("<name>", min, current, args->struct_size)。上方的 (min, current) 对已经 byte-confirmed；重新实现必须接受任何 ≥ min 的 caller struct，并且只读取 caller 声明大小以内的内容。绝不触碰 struct_size 之后的字段。
  • readiness CHECK 是 fatal，且是有意为之。 IsValid()（空后端 future）和 IsReady()（resolution 前调用 Error）是硬 LogMessageFatal，不是 status return。请复制它们；它们能早期捕获 caller bug，upstream PJRT 语义也依赖该 abort。
  • 在 OnReady 中遍历 indirect chain。 可用 value 可能是转发到 concrete value 的 IndirectAsyncValue。读取 av+64 处 status 前，需在 (av+4 & 3) != 0 时追踪 av+16，否则会读取 placeholder。
  • error return 由调用者拥有。 非 OK 方法返回装箱的 PJRT_Error*，调用者必须 PJRT_Event_Destroy/PJRT_Error_Destroy。用 & 1 low-tag + Unref 纪律管理底层 StatusRep refcount，使其生命周期长于 future。
  • Destroy 释放 handle，不释放 operation。 销毁 event 不会取消 launch 或 transfer；pending runtime waiter 持有自己的引用并仍会触发。用 & 8 allocated bit 控制 async-value free，避免释放有 live waiter 或 process-shared singleton 的 value。
  • Ready ≠ host-readable。 execute 或 buffer-definition event resolve 只表示 device-side validity。输出数据的 host readability 需要单独的 D2H copy event。不要把二者合并。

相关组件 ​

NameRelationship
xla::PjRtFuture<void> / tsl::internal::FutureBase<absl::Status,false>每个 PJRT_Event 包装的 C++ future；提供 Await/AndThen/IsReady
tsl::AsyncValuefuture 下方的 refcounted single-assignment cell；其 state byte 控制 readiness
tsl::internal::PromiseMaker<void> / PromiseBase<absl::Status>创建并 fulfil 嵌入 event 对象的 promise half
pjrt::ActualStructSizeIsGreaterOrEqual每个 event entry 首先调用的 per-method args-version gate
pjrt::PjrtErrorCodeToStatusCode / absl::Status::MakeRepSet 使用的 error-code mapping 与 status-rep construction
xla::CommonPjRtClient::CreateLinkedUserPromise产生 execute event 的 runtime mint（下一层）

交叉引用 ​