PJRT 回调与 Pre-Fatal Hook

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节）。该镜像未剥离符号；反混淆后的 C++ 符号名按原文逐字引用。PJRT C-API 版本 0.103。VA == 文件分析地址。其他版本会有所不同。

摘要

本页描述 libtpu PJRT 插件的 C-ABI 回调注册层：也就是框架（JAX、PyTorch-XLA）交给 libtpu 的函数指针接口，使正在运行的 TPU 程序或即将退出的进程能够回调到宿主代码。这里有两类互不相同的回调族，并且它们通过完全不同的路径接入插件。第一类是通过 PJRT_ExecuteOptions 携带的 send/recv 宿主回调对：PJRT_SendCallbackInfo / PJRT_RecvCallbackInfo 数组，PJRT_LoadedExecutable_Execute（slot 60，0xf869b40）会将其转换为 C++ xla::SendCallback / xla::RecvCallback 闭包，然后按 HLO channel id 作为键注册到每次启动专用的 xla::TpuHostTransferManager。第二类是 pre-fatal 错误 hook：框架通过 callback_extension（扩展类型 14）安装的回调；当 libtpu 将因不可恢复错误而 abort 之前，每个已注册闭包都会携带失败的 absl::Status 运行，这是框架刷新 trace buffer、转储遥测或给崩溃添加注解的最后机会。

参考框架是上游 XLA 的 xla/pjrt/c/pjrt_c_api.h。在公开头文件中，PJRT_ExecuteOptions 以 PJRT_SendCallbackInfo**（按 [device][channel] 索引）携带 send_callbacks / recv_callbacks，每个 PJRT_SendCallbackInfo 是 { int64_t channel_id; void* user_arg; PJRT_SendCallback send_callback; }，也就是一个普通 C trampoline。libtpu 的 pjrt::CSendCallbackToCpp(PJRT_SendCallbackInfo const&)（0xf876680）和 pjrt::CRecvCallbackToCpp(PJRT_RecvCallbackInfo const&) 会把这些 C trampoline 包装成 std::function，其错误通道是 PJRT_Error_Code + 消息字符串，而不是抛出的异常；这个边界是纯 C。pre-fatal hook 没有上游稳定 slot；它完全位于 TPU callback_extension（learning/45eac/research/pjrt/extensions/callback/callback_extension.cc）内部，这是一个 40 字节结构体，只暴露两个入口点：RegisterCallback 和 InvokeCallback。

本页负责说明 C-ABI 形状和注册管线：PJRT_SendCallbackInfo / PJRT_RecvCallbackInfo 结构体、PJRT_Chunk wire 结构体、PJRT_Callback_RegisterCallback_Args / PJRT_Callback_PrefatalArgs 结构体、跨越 C 边界时的 channel-id 键控方式，以及每个回调由哪个线程触发。下层机制，也就是按 channel 键控的 rendezvous、TpuHostTransferManager 内部的两个 flat_hash_map、miss 时 fatal、设备侧 sync flag，归 Host Callbacks 所有；本页会链接过去而不重复展开。buffer-ready PJRT_Event_OnReady 回调（slot 14）见 Events & Async。跨宿主 send/recv buffer 传输见 DMA & Cross-Host Recv。

对于重新实现，契约如下：

这两个回调族彼此独立。 Send/recv 通过 PJRT_ExecuteOptions 传递并馈入 SetUpHostCallbacksForDevice → TpuHostTransferManager；pre-fatal / slice-builder hook 通过 callback_extension 传递并馈入进程全局的 …CallbackState 注册表。两条路径互不接触。
C↔C++ shim 形状。 PJRT_SendCallbackInfo 的 C 函数指针由 CSendCallbackToCpp 包装成 std::function<Status(PjRtTransferMetadata const&, PjRtChunk, size_t, bool)>；其返回值是 PJRT_Error_Code（enum）+ 消息，libtpu 在 pjrt_c_api_wrapper_impl.cc:2190 将其重新编码为 absl::Status。跨边界没有 C++ 异常。
PJRT_Chunk wire 结构体。 32 字节：{ void* data; size_t size; void(*deleter)(void* data, void* deleter_arg); void* deleter_arg; }。pjrt::ConvertToCppChunk（0xf8a5280）会复制全部 32 字节，并用一个 C++ 闭包替换 C deleter，该闭包在 drop 时调用原始 C deleter。
pre-fatal hook ABI。 callback_extension（type 14，struct @ 0x224c3b60，ctor CreateCallbackExtension @ 0xe6b91e0）有两个 slot：RegisterCallback(PJRT_Callback_RegisterCallback_Args*) @ 0xe6b9220 和 InvokeCallback(PJRT_Callback_InvokeCallback_Args*) @ 0xe6b94c0。callback_type 判别器（1=SliceBuilder，2=Prefatal）路由到 xla::SliceBuilderCallbackState / xla::PreFatalErrorCallbackState。


ExecuteOptions 上的 send/recv	`PJRT_LoadedExecutable_Execute` slot 60 @ `0xf869b40` → `SetUpHostCallbacksForDevice` → `xla::TpuHostTransferManager`
Send C→C++ shim	`pjrt::CSendCallbackToCpp(PJRT_SendCallbackInfo const&)`; inner error lambda @ `0xf876680`
Recv C→C++ shim	`pjrt::CRecvCallbackToCpp(PJRT_RecvCallbackInfo const&)`（驱动 `xla::CopyToDeviceStream`）
Chunk 结构体解码	`pjrt::ConvertToCppChunk(PJRT_Chunk const&)` @ `0xf8a5280`（32 字节结构体）
CopyToDeviceStream AddChunk	`pjrt::PJRT_CopyToDeviceStream_AddChunk` @ `0xf86f660`（slot 83）→ `xla::(anon)::TpuCopyToDeviceStream::AddChunk` @ `0xf8374e0`
Callback extension（type 14）	struct @ `0x224c3b60`，size 40；ctor `pjrt::CreateCallbackExtension` @ `0xe6b91e0`
Register 入口	`pjrt::(anon)::PJRT_Callback_RegisterCallback(PJRT_Callback_RegisterCallback_Args*)` @ `0xe6b9220`（Args min=35，current=40）
Invoke 入口	`pjrt::(anon)::PJRT_Callback_InvokeCallback(PJRT_Callback_InvokeCallback_Args*)` @ `0xe6b94c0`（Args min=33）
Pre-fatal 注册表	`xla::PreFatalErrorCallbackState` — `AddCallback` @ `0xf95dc00`，`InvokeCallbacks` @ `0xf95dc80`，ctor @ `0xf95dbe0`
Slice-builder 注册表	`xla::SliceBuilderCallbackState` — `AddCallback` @ `0xf95df80`，`InvokeCallbacks` @ `0xf95e000`
Prefatal args 结构体	`PJRT_Callback_PrefatalArgs`（min=26）：`{struct_size, PJRT_Error_Code code @+8, char* msg @+16, size_t msg_len @+24}`
源码根位置	`tpu_pjrt_client.cc`（send/recv），`extensions/callback/callback_extension.cc`（hooks）
证据等级	可重新实现级 / 已按 IDA 反编译逐字节确认（两类回调均已追踪；结构体、判别器、struct-size 包络以及 C↔C++ 错误映射均已确认）

1. 两个回调族，两条路径

PJRT 插件是一个 140-slot vtable 加上一条扩展链（见 PJRT API Vtable Reconstruction）。回调通过其中两个接口进入 libtpu，重新实现者必须将它们分开处理。

方面	Send/Recv 宿主回调	Pre-fatal / slice-builder hook
框架传入位置	`PJRT_ExecuteOptions`（`PJRT_LoadedExecutable_Execute` 的参数，slot 60）	`callback_extension`（extension type 14）→ `RegisterCallback`
C 结构体	`PJRT_SendCallbackInfo` / `PJRT_RecvCallbackInfo`	`PJRT_Callback_RegisterCallback_Args`
C→C++ shim	`CSendCallbackToCpp` / `CRecvCallbackToCpp`	内联于 `PJRT_Callback_RegisterCallback`
键控方式	HLO channel id（`int64`）	callback type 判别器（`1`/`2`）— 无键，append-only 列表
存储位置	每次启动专用 `TpuHostTransferManager` 的两个 `flat_hash_map`	进程全局 `…CallbackState` `std::vector`（按 `TpuClient`）
生命周期	单次 execute 启动	进程 / client 生命周期
触发时机	设备在执行中命中 `Send`/`Recv` op	libtpu abort 之前（通过 `InvokeCallback` 或内部错误路径触发）
触发线程	`TpuHostTransferManager` drain 线程（脱离 device-completion 线程）	检测到 fatal error 的线程
归属	Host Callbacks（机制）；本页（C ABI）	本页

GOTCHA — pre-fatal hook 不是上游 PJRT slot。 它不会出现在 140-slot PJRT_Api 表的任何位置。只能通过遍历扩展链到 callback_extension（type 14）并调用它的 RegisterCallback 函数指针来访问它。重新实现者如果在 vtable 中搜索 "fatal" 或 "abort" slot，将找不到任何结果。必须迭代链（extension_start → 跟随 .next，直到 type==14 或 NULL）；见 Extension Chain。置信度：CONFIRMED — CreateCallbackExtension（0xe6b91e0）写入 type=14 和两个函数指针；没有 PJRT_Api slot 引用任一函数。

2. `PJRT_ExecuteOptions` 上的 Send / Recv 回调

目的

执行 outside-compilation 或跨宿主 send/recv 的已编译 XLA 程序，会发出按 channel_id 键控的成对 Send/Recv HLO op。执行时，框架必须为每个 channel 提供一个宿主回调：Send op（设备→宿主）把数据块向上传给 SendCallback；Recv op（宿主→设备）从 RecvCallback 向下拉取数据块。这些回调不是通过 vtable slot 注册的，而是在每次 PJRT_LoadedExecutable_Execute 调用中作为 PJRT_ExecuteOptions 结构体的一部分传递。

C-ABI 结构体

公开 pjrt_c_api.h v0.103 布局如下，并由 libtpu 链接的 wrapper 符号确认（pjrt::CSendCallbackToCpp(PJRT_SendCallbackInfo const&)、pjrt::CRecvCallbackToCpp(PJRT_RecvCallbackInfo const&)）：

// PJRT_Chunk — the host-side data buffer. 32 bytes; confirmed by
// pjrt::ConvertToCppChunk @ 0xf8a5280 (copies 32 bytes, re-wraps deleter).
typedef struct {
    void*  data;                                    // +0   raw bytes
    size_t size;                                    // +8   byte count
    void  (*deleter)(void* data, void* deleter_arg);// +16  C free hook
    void*  deleter_arg;                             // +24  passed to deleter
} PJRT_Chunk;

// PJRT_SendCallback fires when the device 'Send's a chunk to the host.
typedef PJRT_Error* (*PJRT_SendCallback)(
    PJRT_Chunk* chunk, PJRT_CallbackError* callback_error,
    size_t total_size_in_bytes, bool done, void* user_arg);

typedef struct {
    int64_t            channel_id;     // the HLO Send op's channel id
    void*              user_arg;       // opaque, threaded back to the callback
    PJRT_SendCallback  send_callback;  // the C trampoline
} PJRT_SendCallbackInfo;

// PJRT_RecvCallback fires when the device 'Recv's: the host pushes a
// chunk down a PJRT_CopyToDeviceStream rather than receiving one.
typedef void (*PJRT_RecvCallback)(
    PJRT_CopyToDeviceStream* stream, void* user_arg);

typedef struct {
    int64_t            channel_id;
    void*              user_arg;
    PJRT_RecvCallback  recv_callback;
} PJRT_RecvCallbackInfo;

// In PJRT_ExecuteOptions:
//   PJRT_SendCallbackInfo** send_callbacks;  // [num_devices][num_send_ops]
//   PJRT_RecvCallbackInfo** recv_callbacks;
//   size_t num_send_ops; size_t num_recv_ops;
```text

> **QUIRK — "Send" 是设备→宿主，"Recv" 是宿主→设备。** 命名来自*设备程序*的视角，这与宿主工程师的直觉相反。`SendCallback` 会*接收* `PJRT_Chunk`；`RecvCallback` 则通过馈入 `PJRT_CopyToDeviceStream` 来*产生*数据。这个反转由 [Host Callbacks §2.1](../runtime/host-callbacks.md) 负责并详细解释；这里重复它只是因为 C 结构体名称把这种语义跨 ABI 携带了出来。置信度：CONFIRMED。

### 算法 — C→C++ send shim

`CSendCallbackToCpp` 构造一个 `std::function<absl::Status(PjRtTransferMetadata const&, PjRtChunk, size_t, bool)>`。运行时，它会把 C++ `PjRtChunk` 编组回 `PJRT_Chunk`，调用用户的 C `send_callback`，并把其 `PJRT_Error_Code` 结果转换为 `absl::Status`。错误映射 lambda 是逐字节确认的叶子函数：

```c
// pjrt::CSendCallbackToCpp(PJRT_SendCallbackInfo const&)::$_0
//   ::{lambda(PJRT_Error_Code, char const*, size_t)}::__invoke   sub_F876680
// a1 = PJRT_Error_Code (the C callback's status code)
// a2 = char* message, a3 = message length
function send_error_to_status(code, msg, msg_len):       // sub_F876680
    rep = operator new(8)                                 // the StatusOr<...> slot
    if msg_len > 0x7FFF...F6: throw length_error           // SSO bound guard
    // copy msg into a fresh std::string (SSO if <= 22 bytes, else heap):
    buf = (msg_len > 0x16) ? operator new(round_up(msg_len)) : inline_buf
    memcpy(buf, msg, msg_len); buf[msg_len] = 0
    *rep = absl::Status::MakeRep(4*code + 1, buf, msg_len, /*line*/2190,
                "third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc")
    if heap_allocated(buf): free(buf)
    return rep                                            // StatusOr<chunk-ack> holding the Status

NOTE — 规范状态编码是 4*code + 1。 absl::Status::MakeRep 调用参数中使用 4*code + 1；这是 absl tagged-pointer 约定（低位中的 code，bit 0 中的 inline-vs-heap flag）。反向方向，即 absl::Status → PJRT_Error_Code，使用 pjrt::StatusCodeToPjrtErrorCode（0xf8a3cc0，见于 prefatal trampoline，§3）及其逆函数 pjrt::PjrtErrorCodeToStatusCode（0xf8a3ca0，见于 InvokeCallback）。两者在有效 0..16 范围内都是恒等映射，也就是说 PJRT_Error_Code enum 与 absl::StatusCode 共享相同整数值，因此数值确实一致。转换器存在只是为了守住边界：PjrtErrorCodeToStatusCode 是裸 return a1，而 StatusCodeToPjrtErrorCode 是带有守卫的 return a1，会在 INT_MIN/INT_MAX/DO_NOT_USE sentinel code 上 LOG(FATAL)（pjrt_c_api_helpers.cc:251-256）。重新实现者可以原样传递 code，但必须拒绝这些 sentinel。

算法 — recv shim 与 `CopyToDeviceStream`

RecvCallback 不接收数据块；它会拿到一个 PJRT_CopyToDeviceStream 并向其中推送数据块。CRecvCallbackToCpp 将 C recv_callback 包装成 std::function<void(PjRtTransferMetadata const&, unique_ptr<CopyToDeviceStream>)>。宿主通过 slot 83，PJRT_CopyToDeviceStream_AddChunk，馈入该 stream：

// pjrt::PJRT_CopyToDeviceStream_AddChunk   sub_F86F660   (PJRT_Api slot 83)
// a1 = PJRT_CopyToDeviceStream_AddChunk_Args*  (Args min=37, current=40)
function PJRT_CopyToDeviceStream_AddChunk(args):
    if !ActualStructSizeIsGreaterOrEqual("..._AddChunk_Args", 37, 40, args->struct_size):
        return error
    stream = *args->stream                               // xla::CopyToDeviceStream*
    cpp_chunk = ConvertToCppChunk(args->chunk)            // 32-byte struct -> C++ PjRtChunk
    av = stream->vtable[2](stream, cpp_chunk)             // TpuCopyToDeviceStream::AddChunk @ 0xf8374e0
    // returns a PJRT_Event wrapping the AddChunk completion AsyncValue
```text

`xla::(anon)::TpuCopyToDeviceStream::AddChunk(PjRtChunk)`（`0xf8374e0`）把数据块交给设备侧的 `tpu::host_commands::CopyToDeviceStream::AddChunk`（`0x1d0a6320`），这是 transfer manager 内部宿主→设备路径的底部。粒度/大小计数 slot（`TotalBytes` `0xf86f7e0`、`GranuleSize` `0xf86f840`、`CurrentBytes` `0xf86f8a0`）让宿主回调按设备要求的粒度切分写入。

### 函数映射

| 函数 | 地址 | 角色 |
|---|---|---|
| `PJRT_LoadedExecutable_Execute` | `0xf869b40` | Slot 60；接收带有回调数组的 `PJRT_ExecuteOptions` |
| `xla::(anon)::SetUpHostCallbacksForDevice` | (`tpu_pjrt_client.cc`) | 将 `Span<SendCallback>`/`Span<RecvCallback>` 转换为每次启动专用的 `TpuHostTransferManager` map |
| `pjrt::CSendCallbackToCpp` | inner @ `0xf876680` | 将 C `send_callback` 包装成 C++ `std::function`；映射 `PJRT_Error_Code`→`Status` |
| `pjrt::CRecvCallbackToCpp` | — | 将 C `recv_callback` 包装成驱动 `CopyToDeviceStream` 的 `std::function` |
| `pjrt::ConvertToCppChunk` | `0xf8a5280` | 解码 32 字节的 `PJRT_Chunk`，并把 C deleter 重新包装为 C++ 闭包 |
| `pjrt::PJRT_CopyToDeviceStream_AddChunk` | `0xf86f660` | Slot 83；宿主把数据块推入 recv stream |
| `xla::(anon)::TpuCopyToDeviceStream::AddChunk` | `0xf8374e0` | stream 侧向 transfer manager 执行 AddChunk |

> **GOTCHA — 未注册的 channel 会导致 fatal crash。** 一旦 `SetUpHostCallbacksForDevice` 填充了 `TpuHostTransferManager` map，设备 `Send`/`Recv` 的 channel id 如果没有匹配回调，就会触发 `LOG(FATAL)`，而不是静默丢弃。这是在下一层 `HandleSendChunk`/`HandleRecvChunk` 中强制执行的；见 [Host Callbacks §2.3](../runtime/host-callbacks.md)。C-ABI 层的职责是确保程序需要的每个 `PJRT_SendCallbackInfo`/`PJRT_RecvCallbackInfo` 在启动前都存在于 `PJRT_ExecuteOptions` 中。置信度：CONFIRMED。

---

## 3. Pre-Fatal 错误 Hook（`callback_extension`，type 14）

### 目的

当 libtpu 检测到不可恢复条件（SDC checksum mismatch、slice-builder failure、driver-level fault）时，框架希望在进程死亡*之前*执行清理，例如刷新 profiler trace、写入 crash annotation、快照设备状态。PJRT 没有用于此目的的上游 slot，因此 TPU 插件通过私有扩展暴露它。`callback_extension`（type 14）允许框架注册 `std::function<void(absl::Status const&)>` 闭包，这些闭包会在 abort 之前携带失败状态触发。

### 入口点

```text
extension_start (PJRT_Api +0x08)  ──> walk .next to type == 14
  callback_extension @ 0x224c3b60  (size 40, built by CreateCallbackExtension @ 0xe6b91e0)
    +0x00  struct_size = 40
    +0x08  type        = 14
    +0x10  next         ──> phase_compile_extension (type 9)
    +0x18  RegisterCallback ──> pjrt::(anon)::PJRT_Callback_RegisterCallback  @ 0xe6b9220
    +0x20  InvokeCallback   ──> pjrt::(anon)::PJRT_Callback_InvokeCallback    @ 0xe6b94c0

register / invoke ABI

RegisterCallback 和 InvokeCallback 在它们 args 结构体的 +16 位置共享一个 callback_type 判别器：

`callback_type`	含义	注册表	C++ 状态对象
`1`	Slice-builder failure	追加到 `SliceBuilderCallbackState`	callback `void(SliceFailureType)`
`2`	Pre-fatal error	追加到 `PreFatalErrorCallbackState`	callback `void(absl::Status const&)`
other	被拒绝	—	`MakeErrorImpl<12>("Callback type not supported.")`

// pjrt::(anon)::PJRT_Callback_RegisterCallback   sub_E6B9220
// a1 = PJRT_Callback_RegisterCallback_Args*  (min=35, current=40)
//   +0x00 struct_size   +0x08 PJRT_Client**   +0x10 callback_type (int)
//   +0x18 callback fn ptr   +0x20 user_arg
function PJRT_Callback_RegisterCallback(args):                 // sub_E6B9220
    if !ActualStructSizeIsGreaterOrEqual("..._RegisterCallback_Args", 35, 40, args->struct_size):
        return error
    client = **args->client                                    // unwrap PJRT_Client -> TpuClient
    if client->vtable.tpu_id() != 0x83D71ADBA77968AA:          // xla::TpuId() guard
        return null                                            // wrong backend -> no-op
    if !args->callback_fn: return null
    state_base = client[81]                                    // the client's callback-state block
    switch args->callback_type:
      case 2:  // Prefatal
        st = state_base + 280                                  // PreFatalErrorCallbackState mutex
        closure = { thunk = RegisterPrefatalCallback::$_0,     // C->C++ status trampoline
                    storage = {args->callback_fn, args->user_arg} }
        xla::PreFatalErrorCallbackState::AddCallback(st, closure)   // sub_F95DC00
      case 1:  // SliceBuilder
        st = state_base + 264
        xla::SliceBuilderCallbackState::AddCallback(st, closure)    // sub_F95DF80
      default:
        return MakeErrorImpl<12>("Callback type not supported.")    // callback_extension.cc:96
    return null   // success
```text

> **QUIRK — 在错误 backend 上注册会静默成为 no-op。** 存储任何内容前，`RegisterCallback` 会检查 `client->tpu_id() == 0x83D71ADBA77968AA`（魔数 `xla::TpuId()` 常量，惰性 `__cxa_guard` 初始化）。如果 `PJRT_Client` 不是 TPU client，该函数返回 `null`（成功），但*不会*注册回调。跨 backend 共享此扩展的重新实现者必须复现 magic-id gate，否则非 TPU client 会看到“注册成功”但回调永远不会触发。置信度：CONFIRMED — 同一个 `0x83D71ADBA77968AA` guard 同时出现在 `RegisterCallback` 和 `InvokeCallback` 中。

### prefatal trampoline — C → C++ status

为 type-2 回调存储的闭包是一个 `RegisterPrefatalCallback::$_0` trampoline，它将 C++ `absl::Status` 转换成用户 C 回调所期望的 C `(PJRT_Error_Code, message)` 对：

```c
// __call_func<pjrt::(anon)::RegisterPrefatalCallback(...)::$_0>   sub_E6B9700
// a1 = __policy_storage* (holds the user C fn ptr + user_arg)
// a2 = &absl::Status (the failing status)
function prefatal_trampoline(storage, status):                 // sub_E6B9700
    code   = absl::status_internal::MapToLocalCode(status)
    c_code = pjrt::StatusCodeToPjrtErrorCode(code)             // absl code -> PJRT_Error_Code
    msg, msg_len = status.message()                            // inline-or-heap string
    user_fn = storage[0]; user_arg = storage[1]
    user_fn(/*PrefatalArgs{c_code, msg, msg_len}*/ ..., user_arg)
    status.Unref()                                            // drop our ref if heap-backed

`InvokeCallback` — 框架驱动的触发路径

该扩展还允许调用方通过 C ABI 触发已注册回调（由 SDK / telemetry 侧用于触发 pre-fatal flush）。InvokeCallback 执行反向转换，也就是接收 C PrefatalArgs，重建 absl::Status，然后运行每个已注册闭包：

// pjrt::(anon)::PJRT_Callback_InvokeCallback   sub_E6B94C0
// a1 = PJRT_Callback_InvokeCallback_Args*  (min=33)
//   +0x10 callback_type   +0x18 &type-specific args (here PJRT_Callback_PrefatalArgs)
function PJRT_Callback_InvokeCallback(args):                   // sub_E6B94C0
    if !ActualStructSizeIsGreaterOrEqual("..._InvokeCallback_Args", 33, 32, args->struct_size):
        return error
    if args->callback_type != 2:
        return MakeErrorImpl<12>("Callback type can not be invoked.")   // callback_extension.cc:130
    pf = args->prefatal_args                                   // PJRT_Callback_PrefatalArgs*
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Callback_PrefatalArgs", 26, 32, pf->struct_size):
        return error
    code = pjrt::PjrtErrorCodeToStatusCode(pf->code)           // PJRT_Error_Code -> absl code
    status = absl::Status::MakeRep(4*code + 1, pf->msg, pf->msg_len,
                /*line*/111, "...callback_extension.cc")
    if client->tpu_id() == 0x83D71ADBA77968AA:
        xla::PreFatalErrorCallbackState::InvokeCallbacks(client[81]+280, status)  // sub_F95DC80
    status.Unref()
    return null
```text

```c
// PJRT_Callback_PrefatalArgs   (Args min=26, current=32)
typedef struct {
    size_t          struct_size;   // +0
    void*           priv;          // (reserved)        -- folded into +0x08 slot in this build
    PJRT_Error_Code code;          // +8   (uint32)     failing status code
    char*           message;       // +16                error message
    size_t          message_size;  // +24
} PJRT_Callback_PrefatalArgs;

`PreFatalErrorCallbackState` — 注册表

该注册表是一个由 mutex 保护的 std::vector<std::function<void(absl::Status const&)>>。它是 32 字节对象：+0 处一个 mutex，以及 vector 的 {begin, size, capacity} 三元组。每个 std::function slot 为 32 字节（因此有 32 * count 算术）。

// xla::PreFatalErrorCallbackState::AddCallback   sub_F95DC00
function AddCallback(this, fn):                                // sub_F95DC00
    this.mutex.lock()
    if this.size >= this.capacity:
        __emplace_back_slow_path(&this.vec, fn)                // grow + move
    else:
        memcpy(this.vec[this.size], fn, 32)                    // place the std::function inline
        this.size += 1
    this.mutex.unlock()

// xla::PreFatalErrorCallbackState::InvokeCallbacks   sub_F95DC80
function InvokeCallbacks(this, status):                        // sub_F95DC80
    this.mutex.lock()
    for slot in this.vec[0 .. this.size]:                      // 32-byte stride
        slot.invoke(&status)                                   // calls the prefatal_trampoline
    this.mutex.unlock()
    if heap_backed(status): status.Unref()
```text

> **NOTE — 注册顺序就是触发顺序，并且没有 de-registration。** `AddCallback` 只会追加；该扩展没有暴露 remove/unregister 入口。`InvokeCallbacks` 在锁下从前到后遍历 vector。重新实现者会得到有序、append-only、锁串行化的触发语义，并且必须接受已注册闭包会存活整个 client 生命周期。相同形状也支撑 `SliceBuilderCallbackState`（`AddCallback` `0xf95df80`，`InvokeCallbacks` `0xf95e000`），区别只在闭包签名（`void(SliceFailureType)`）。置信度：CONFIRMED。
>
> **GOTCHA — 回调在失败线程上触发，并持有注册表锁。** `InvokeCallbacks` 在仍持有 `this.mutex` 的情况下，在命中 fatal 条件的任意线程上内联运行每个闭包。阻塞、重入 PJRT 或尝试注册另一个回调的 pre-fatal 回调会死锁或拖延 abort。让闭包保持简短：flush、annotate、return。置信度：CONFIRMED — `InvokeCallbacks`（`0xf95dc80`）在 `lock()`/`unlock()` 对内部调用每个 slot。

### 函数映射

| 函数 | 地址 | 角色 |
|---|---|---|
| `pjrt::CreateCallbackExtension` | `0xe6b91e0` | 构建 type-14 扩展结构体（size 40，两个 fn ptr） |
| `pjrt::(anon)::PJRT_Callback_RegisterCallback` | `0xe6b9220` | 注册 slice-builder（`1`）或 pre-fatal（`2`）回调 |
| `pjrt::(anon)::PJRT_Callback_InvokeCallback` | `0xe6b94c0` | 从 C ABI 触发 pre-fatal 回调 |
| `RegisterPrefatalCallback::$_0` trampoline | `0xe6b9700` | C++ `Status` → C `(code, msg)`；调用用户 C fn |
| `xla::PreFatalErrorCallbackState::AddCallback` | `0xf95dc00` | Append-only、mutex 保护的注册 |
| `xla::PreFatalErrorCallbackState::InvokeCallbacks` | `0xf95dc80` | 在锁下按顺序触发全部回调 |
| `xla::PreFatalErrorCallbackState` ctor | `0xf95dbe0` | 将 32 字节注册表置零初始化（vxorps/vmovups） |
| `xla::SliceBuilderCallbackState::AddCallback` | `0xf95df80` | Slice-builder 对应函数（`type==1`） |
| `xla::SliceBuilderCallbackState::InvokeCallbacks` | `0xf95e000` | Slice-builder 触发 |
| `pjrt::StatusCodeToPjrtErrorCode` | `0xf8a3cc0` | absl code → `PJRT_Error_Code`（对 0–16 恒等；在 `INT_MIN/INT_MAX/DO_NOT_USE` sentinel 上 `LOG(FATAL)`，`pjrt_c_api_helpers.cc:251-256`） |
| `pjrt::PjrtErrorCodeToStatusCode` | `0xf8a3ca0` | `PJRT_Error_Code` → absl code（纯恒等：`return a1`） |

---

## 4. 重新实现说明

- **两个接口面，互不重叠。** Send/recv 回调是 `Execute` 的参数（按启动）；pre-fatal hook 是扩展方法（按 client）。重新实现必须同时暴露两者：`PJRT_ExecuteOptions.send_callbacks`/`recv_callbacks` 数组，以及带有 `RegisterCallback`/`InvokeCallback` 的 type-14 扩展。
- **chunk 结构体就是 wire 契约。** `PJRT_Chunk` 是 `{data, size, deleter, deleter_arg}` = 32 字节。deleter 是 C 函数指针；libtpu 会在 C++ 闭包中包装它（`ConvertToCppChunk` @ `0xf8a5280`），并在 drop 时调用它。字段顺序和 deleter 签名 `void(void* data, void* arg)` 必须完全正确，否则 buffer 会泄漏或 double-free。
- **错误跨边界以 code 传递，而不是以异常传递。** `SendCallback` 返回 `PJRT_Error*`；libtpu 通过 `4*code + 1`（`pjrt_c_api_wrapper_impl.cc:2190`）把它的 `PJRT_Error_Code` + 消息转换为 `absl::Status`。Pre-fatal 回调通过 `StatusCodeToPjrtErrorCode`（`0xf8a3cc0`）走相反方向。两个转换器（`0xf8a3cc0` / `0xf8a3ca0`）都是恒等映射，因为 enum 值一致，但 `StatusCodeToPjrtErrorCode` 会在 `INT_MIN/INT_MAX/DO_NOT_USE` sentinel 上 `LOG(FATAL)`，因此重新实现仍必须拒绝这些值。
- **Channel id 为 send/recv 路径提供键；缺失 channel 是 fatal。** 启动前注册已编译程序使用的每个 channel。键控和 miss 时 `LOG(FATAL)` 位于 [Host Callbacks](../runtime/host-callbacks.md)；C-ABI 层的职责是在 `ExecuteOptions` 中交付完整的 `PJRT_SendCallbackInfo`/`PJRT_RecvCallbackInfo` 集合。
- **按 backend id gate 注册。** `RegisterCallback`/`InvokeCallback` 只有在 `client->tpu_id() == 0x83D71ADBA77968AA` 时才不是 no-op。复现这个 gate，否则跨 backend 注册会静默无效。
- **pre-fatal hook 是 append-only，并在 dying thread 上持锁触发。** 没有 de-registration；闭包按注册顺序触发，期间持有 registry mutex，运行在线程命中错误的线程上。让它们保持简短且不可重入。
- **`callback_type` 同时是 register 和 invoke 的判别器。** `1`=slice-builder，`2`=pre-fatal。`InvokeCallback` 会以 `"Callback type can not be invoked."` 拒绝除 `2` 外的任何值；`RegisterCallback` 会以 `"Callback type not supported."` 拒绝未知类型。为兼容旧调用方，应遵守 `ActualStructSizeIsGreaterOrEqual` 包络（`RegisterCallback_Args` min 35，`InvokeCallback_Args` min 33，`PrefatalArgs` min 26）。

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `xla::TpuHostTransferManager` | send/recv 回调注册进入的每次启动专用 manager（机制见 [Host Callbacks](../runtime/host-callbacks.md)） |
| `xla::PreFatalErrorCallbackState` | 进程全局 pre-fatal 回调注册表（本页主题） |
| `xla::SliceBuilderCallbackState` | 通过同一扩展注册的 slice-failure 对应物 |
| `callback_extension` (type 14) | 暴露 `RegisterCallback`/`InvokeCallback` 的扩展结构体 |
| `PJRT_LoadedExecutable_Execute` (slot 60) | 携带含有回调数组的 `PJRT_ExecuteOptions` 的 execute 入口 |
| `xla::CopyToDeviceStream` | `RecvCallback` 通过 `PJRT_CopyToDeviceStream_AddChunk` 向其中推送数据块的 stream |

## 交叉引用

- [Host Callbacks](../runtime/host-callbacks.md) — 本页下层按 channel 键控的 rendezvous *机制*：`TpuHostTransferManager` 中的两个 `flat_hash_map`、miss 时 fatal、设备侧 sync flag。本页是其上方的 C-ABI 注册层。
- [PJRT API Vtable Reconstruction](api-vtable-reconstruction.md) — 140-slot 表和扩展链；`callback_extension`（type 14）和 `Execute`（slot 60）如何位于接口面中
- [Extension Chain](extension-chain.md) — 如何遍历 `extension_start` 到达 type-14 `callback_extension`
- [Executable Execution](executable-execution.md) — `PJRT_LoadedExecutable_Execute`（slot 60）以及携带 send/recv 回调数组的 `PJRT_ExecuteOptions` 结构体
- [Events & Async](events-and-async.md) — `PJRT_Event_OnReady`（slot 14），buffer-completion 回调（与本页回调接口面不同）
- [DMA & Cross-Host Recv](dma-and-cross-host-recv.md) — 跨宿主 send/recv buffer 传输，即 channel send/recv 路径的多宿主对应物
- [PJRT Overview](overview.md) — 回调接口面在 libtpu PJRT 插件栈中的位置

PJRT 回调与 Pre-Fatal Hook ​

摘要 ​

1. 两个回调族，两条路径 ​

2. PJRT_ExecuteOptions 上的 Send / Recv 回调 ​

目的 ​

C-ABI 结构体 ​

算法 — recv shim 与 CopyToDeviceStream ​

register / invoke ABI ​

InvokeCallback — 框架驱动的触发路径 ​

PreFatalErrorCallbackState — 注册表 ​