PJRT 缓冲区 ABI 与内存布局

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，745 MB，ELF x86-64，未 strip）。.text 映射在 0xe63c000；对 .text 中的函数，列出的 VA 等于文件偏移。其他 wheel 版本会不同。

摘要

本页记录 libtpu 的 PJRT 插件中的 C-ABI 缓冲区包装器：PJRT_Buffer 和 PJRT_Memory 不透明句柄、主机↔设备暂存入口点 PJRT_Client_BufferFromHostBuffer、回读路径 PJRT_Buffer_ToHostBuffer、零分配形状访问器（ElementType、Dimensions、OnDeviceSizeInBytes）、两阶段生命周期（Delete / IsDeleted / Destroy），以及内存空间表面（PJRT_Buffer_Memory、PJRT_Buffer_CopyToMemory 和 PJRT_Memory 访问器）。这些都是覆盖在 XLA C++ xla::PjRtBuffer / xla::PjRtMemorySpace 层级之上的轻量 C shim：每个 shim 都会校验调用者的 struct_size，读取其带类型参数，跳转一次虚调用进入具体的 xla::CommonPjRtBufferImpl，并把结果编组回 args 结构。熟悉上游 PJRT C API（pjrt_c_api.h）的读者会认出每个槽位；这里的价值在于精确的 libtpu 包装对象布局、vtable 偏移，以及 TPU 特定行为（HBM 不能由主机寻址、拒绝 strides layout、回读时反平铺）。

这个包装器有意作为边界，而不是缓冲区。PJRT_Buffer C 对象为 272 字节（0x110），恰好拥有一个 xla::PjRtBuffer* 加一个外部引用跟踪列表；它不携带任何设备字节。设备侧字节位于内部 C++ 缓冲区之后的 HBM 中，由 HBM 分配器和 TPU 缓冲区字节布局分配并布置，归那些页面所有，此处不重新推导。 本页负责 C-ABI 表面、主机缓冲区语义枚举和传输入口点；无类型原始字节的兄弟表面在 RawBuffer 扩展（type 8），StreamExecutor 分配器桥接在分配器集成。

本页按方面组织：先是两个 C 包装对象布局和槽位映射，然后是每个入口点（BufferFromHostBuffer、ToHostBuffer、访问器、生命周期对、内存空间表面）及其 ### 算法 和 ### 函数映射，再到 PJRT_Memory 表面、生命周期/引用计数模型和异步就绪。

对于重新实现，契约是：

两个包装对象布局：272 字节的 PJRT_Buffer（独占所有者）和 PJRT_Memory 句柄，以及每个槽位读取的字段偏移。
PJRT_Client_BufferFromHostBuffer_Args 字段布局（37 个字段，struct_size 120）和控制主机缓冲区生命周期的 PJRT_HostBufferSemantics 枚举。
每个槽位对 xla::CommonPjRtBufferImpl 的单次 vtable 跳转，以及它编组到 args 结构中的内容。
双轴生命周期（Delete 立即释放设备内存；Destroy 释放 C 包装器）以及阻塞 Delete 的外部引用 pin。


PJRT_Api 槽位范围	Buffer 槽位 63–81 + 后期槽位 97、105、125、130；`BufferFromHostBuffer` 槽位 27
`PJRT_Buffer` 包装器大小	`0x110`（272 字节）— `{ PjRtBuffer* impl@+0x00; …; ext_ref list@+0xf8/+0x100/+0x108 }`
具体内部类	`xla::CommonPjRtBufferImpl`（vtable `0x21789ec8`，vptr base `0x21789ed8`）
BufferFromHostBuffer	`pjrt::PJRT_Client_BufferFromHostBuffer` @ `0xf8644c0`，`struct_size` 当前 120 / 最少 37 个字段
ToHostBuffer	`pjrt::PJRT_Buffer_ToHostBuffer` @ `0xf86e640`（反平铺，返回 80 字节事件）
内存空间种类	`tpu_hbm`、`pinned_host`、`unpinned_host`、`device`（CPU 暂存）、跨 pod megascale
异步门控	每次传输都返回一个包装 `PjRtFuture` 的 80 字节 `PJRT_Event`

两个 C 包装对象

PJRT 的 C API 向调用者交出不透明的 PJRT_Buffer* 和 PJRT_Memory* 指针。在 libtpu 中，这些是小型堆对象，用来包装一个 C++ 内部指针；每个槽位的第一个动作都是解引用包装器以到达内部 xla::PjRtBuffer*。

注意 — libtpu 中有两个互不相关的缓冲区包装器。此处记录的 272 字节 PJRT_Buffer 是规范的带类型/带形状表面。另一个单独的 16 字节 原始字节包装器（RCReference<PjRtRawBuffer> + 借用 client）支撑 RawBuffer 扩展，并记录在 ext-rawbuffer.md。它们包装不同的 C++ 类层级（xla::PjRtBuffer 与 xla::PjRtRawBuffer），并且有不同的所有权模型。不要混淆。

PJRT_Buffer 布局（0x110 = 272 字节）

struct PJRT_Buffer {                  // sizeof = 0x110, freed as free(wrapper, 0x110)
    /* +0x00  */ xla::PjRtBuffer* impl;              // owned EXCLUSIVELY (CommonPjRtBufferImpl)
    /* +0x08  */ PJRT_Client*     client;            // borrowed; set at construction
    /* +0x10  */ uint8_t          ...flags;          // construction-zeroed bytes
    /* ...    */ /* cached element-type / device / memory wrappers, error/status slots */
    /* +0xf8  */ ExternalReferenceHold* ext_ref_data;   // external-reference tracking list
    /* +0x100 */ size_t                 ext_ref_size;
    /* +0x108 */ size_t                 ext_ref_capacity;
};
```text

大小和字段偏移已经按字节确认：`PJRT_Buffer_Destroy`（`0xf86d020`）调用 `~PJRT_Buffer()`，然后 `free(wrapper, 0x110)`；`BufferFromHostBuffer` 构造函数执行 `operator new(0x110)`，并清零 `+0xf0..+0x110` 区域（外部引用列表）以及 `+0x10`、`+0xa8`、`+0xb0`、`+0xc8`、`+0xd0`、`+0xe8` 处的构造标志字节；`Increase/DecreaseExternalReferenceCount` 读取/跟踪 `+0xf8/+0x100/+0x108` 处的列表。包装器*独占*拥有 `impl`：C 包装器本身没有引用计数；`Destroy` 会直接运行内部析构函数。

### PJRT_Memory 句柄

`PJRT_Memory` 是 `xla::PjRtMemorySpace*`（HBM / pinned-host / CPU-device）的 C 句柄。C 包装器由 client 生成并缓存，而不是由缓冲区生成：`PJRT_Client_FindMemoryWrapper(PjRtMemorySpace*, PJRT_Client*)` @ `0xf8605e0` 会线性扫描 client 的已缓存内存包装器数组（计数在 `client+0x90`，数组基址在 `client+0x88`），寻找内部指针等于所请求 C++ 内存空间的包装器，未命中则返回 NULL。`PJRT_Buffer_Memory`（槽位 71，`0xf86dc60`）读取内部缓冲区的内存空间（内部 vtable+0x58 `memory_space()`），并*内联*运行同一个扫描；如果没有匹配的包装器，它返回 `xla::Unimplemented("PJRT_Buffer_Memory not implemented for platform '%s'")`（absl code 12），而不是成功的 NULL。五个内存空间类及其 `kind` 字符串：

| 类 | `kind` 字符串 | `kind()` 地址 | vtable |
|---|---|---|---|
| `xla::TpuHbmMemorySpace` | `tpu_hbm` | `0xf817100` | `0x2177b478` |
| `xla::PinnedHostMemorySpace` | `pinned_host` | `0xf90c8e0` | `0x21789978` |
| `xla::UnpinnedHostMemorySpace` | `unpinned_host` | `0xf90c700` | — |
| `xla::CpuDeviceMemorySpace` | `device`（CPU 暂存） | `0xf90cac0` | — |
| `xla::MegaScalePjRtMemorySpace` | cross-pod / DCN | `0xe6eb460` | — |

> **怪癖 —** `tpu_hbm` **不能由主机寻址**。HBM 中的缓冲区没有有效的主机指针；主机指针访问器（在 RawBuffer 表面上）会返回 NULL，除非缓冲区的内存空间 kind 恰好是 `"pinned_host"`。假设 `OpaqueDeviceMemoryDataPointer`（槽位 81）给出可由 CPU 解引用地址的代码，在 TPU 上是错误的：该指针是原始 HBM 虚拟地址，只在所属设备/核心上有效，只能由外部设备侧消费者（DMA 引擎、自定义内核）使用，绝不能由 `memcpy` 使用。

### 槽位映射

本页覆盖的 18 个槽位，全部位于 140 槽位的 `PJRT_Api` 中。每个包装器都会通过 `pjrt::ActualStructSizeIsGreaterOrEqual(name, min_fields, cur_bytes, caller_size)` @ `0xf8a4ec0` 校验 `struct_size`，然后读取 `args+0x10` 以到达包装器，再读取 `wrapper+0x00` 以到达内部 `PjRtBuffer*`。

| 槽位 | 偏移 | 方法 | C 符号 | 地址 | vtable 跳转 / 后端 |
|---|---|---|---|---|---|
| 27 | — | BufferFromHostBuffer | `PJRT_Client_BufferFromHostBuffer` | `0xf8644c0` | memory-space vtable+0x120（分配+暂存） |
| 63 | 0x1f8 | Destroy | `PJRT_Buffer_Destroy` | `0xf86d020` | `~PJRT_Buffer()` + `free(0x110)` |
| 64 | 0x200 | ElementType | `PJRT_Buffer_ElementType` | `0xf86d220` | inner vtable+0x10 `element_type()` + `ConvertToPjRtBufferType` |
| 65 | 0x208 | Dimensions | `PJRT_Buffer_Dimensions` | `0xf86d280` | inner vtable+0x18 `dimensions()` → `{ptr,count}` |
| 69 | — | OnDeviceSizeInBytes | `PJRT_Buffer_OnDeviceSizeInBytes` | `0xf86da80` | inner vtable+0x88 `GetOnDeviceSizeInBytes()` |
| 71 | 0x238 | Memory | `PJRT_Buffer_Memory` | `0xf86dc60` | inner vtable+0x58 `memory_space()` + 内联 client 侧包装器查找 |
| 72 | 0x240 | Delete | `PJRT_Buffer_Delete` | `0xf86dd80` | inner vtable+0xa0 `Delete()`（立即释放 HBM） |
| 73 | 0x248 | IsDeleted | `PJRT_Buffer_IsDeleted` | `0xf86dde0` | inner vtable+0xb0 `IsDeleted()` |
| 74 | 0x250 | CopyToDevice | `PJRT_Buffer_CopyToDevice` | `0xf86e360` | dst-device vtable+0x98（默认 mem）+ src vtable+0xb8 |
| 75 | 0x258 | ToHostBuffer | `PJRT_Buffer_ToHostBuffer` | `0xf86e640` | 反平铺 + inner vtable+0x78 `ToLiteral()` |
| 76 | 0x260 | IsOnCpu | `PJRT_Buffer_IsOnCpu` | `0xf86ecc0` | inner vtable+0xe8 `IsOnCpu()` |
| 77 | 0x268 | ReadyEvent | `PJRT_Buffer_ReadyEvent` | `0xf86ed20` | inner vtable+0xe0 `GetReadyFuture()` → 0x50 event |
| 79 | 0x278 | IncreaseExternalReferenceCount | `PJRT_Buffer_IncreaseExternalReferenceCount` | `0xf86ef20` | inner vtable+0x70 `AcquireExternalReference` |
| 80 | 0x280 | DecreaseExternalReferenceCount | `PJRT_Buffer_DecreaseExternalReferenceCount` | `0xf86f100` | 释放已跟踪的 hold |
| 81 | 0x288 | OpaqueDeviceMemoryDataPointer | `PJRT_Buffer_OpaqueDeviceMemoryDataPointer` | `0xf86f200` | inner vtable+0x70 + 读取 hold+0x08（原始 ptr） |
| 97 | 0x308 | CopyToMemory | `PJRT_Buffer_CopyToMemory` | `0xf86e500` | inner vtable+0xb8 `CopyToMemorySpace(PjRtMemorySpace*)` |
| 105 | 0x350 | CopyRawToHost | `PJRT_Buffer_CopyRawToHost` | `0xf86de40` | inner vtable+0x90 |
| 130 | 0x410 | DonateWithControlDependency | `PJRT_Buffer_DonateWithControlDependency` | `0xf86f2e0` | inner vtable+0xd8 |

继承的形状访问器 `element_type()`/`dimensions()` 指向抽象 `xla::PjRtBuffer` 基类（`0xe6eaac0` / `0xe6eaae0`），而不是 `CommonPjRtBufferImpl`，具体 impl 不会覆盖它们。

---

## PJRT_Client_BufferFromHostBuffer（槽位 27）

### 用途

主机→设备暂存入口点。JAX/IFRT 调用它把主机数组上传到新的设备缓冲区，并提供主机数据指针、公共元素类型、维度（以及可选的字节 strides）、目标内存空间、可选设备布局，以及一个 `PJRT_HostBufferSemantics` 值，用于决定 libtpu 是否可以别名主机缓冲区，还是必须复制它。它返回两个输出：新的 `PJRT_Buffer*`，以及 `done_with_host_buffer` `PJRT_Event`，通知调用者何时可以复用或释放主机缓冲区。

### Args 布局

包装器校验 `ActualStructSizeIsGreaterOrEqual("PJRT_Client_BufferFromHostBuffer_Args", 37, 120, caller_size)` — 37 个命名字段，当前 120 字节。直接来自反编译的字节偏移如下：

| 偏移 | 字段 | 说明 |
|---|---|---|
| `+0x00` | `struct_size` | 首先校验 |
| `+0x10` | `client` | `PJRT_Client*`；内部 XLA client 位于 `*(*client)` |
| `+0x18` | `data` | 主机源指针 |
| `+0x20` | `type` | `PJRT_Buffer_Type`；`ConvertFromPjRtBufferType(*(u32*)(a1+0x20))` |
| `+0x28` / `+0x30` | `dims` ptr / `num_dims` | 形状范围数组 |
| `+0x38` / `+0x40` | `byte_strides` ptr / `num_byte_strides` | 可选；当且仅当 `+0x38 != 0` 时存在 |
| `+0x48` | `host_buffer_semantics` | `ConvertFromPjRtHostBufferSemantics(*(u32*)(a1+0x48))` |
| `+0x50` | `memory` | 目标 `PJRT_Memory*`；分配+暂存 vtable 通过它到达 |
| `+0x58` | （内存布局相关输入） | 在 stage 调用中作为 `**(a1+0x58)` 解引用 |
| `+0x60` | `device_layout` | `PJRT_Buffer_MemoryLayout*`；其 `type` 位于该结构的 `+0x48` |
| `+0x68` | `done_with_host_buffer` | OUT：`PJRT_Event*`（`operator new(0x50)`） |
| `+0x70` | `buffer` | OUT：新的 `PJRT_Buffer*`（`operator new(0x110)`） |

`byte_strides` 存在性测试是反编译中的 `if (*(a1+0x38))` 分支：非空 strides 指针会翻转一个内部“has layout”标志（`v75`），该标志稍后会在 stage 调用前检查。

> **陷阱 —** `device_layout` 字段（`+0x60`）会在任何分配之前针对 TPU 平台校验。如果其 `type` 是 `PJRT_Buffer_MemoryLayout_Type_Strides`（枚举值 1），调用会失败并报 `"PJRT_Buffer_MemoryLayout_Type_Strides in device_layout is not supported in PJRT_Client_BufferFromHostBuffer for platform <name>"`；任何其他意外 type 值都会失败并报 `"Unexpected PJRT_Buffer_MemoryLayout_Type type: <n>"`。TPU 只接受*平铺*设备布局（type 0）。传入 strides 布局的重新实现会得到清晰错误，而不是静默错放的缓冲区。

### PJRT_HostBufferSemantics 枚举

`+0x48` 处的公共枚举值由 `pjrt::ConvertFromPjRtHostBufferSemantics` @ `0xf8a3f20` 映射到 XLA 的内部枚举（在反编译中调用了四处，每个内存空间分派分支一处）。该枚举遵循上游 `pjrt_c_api.h` 的顺序；其含义控制主机缓冲区生命周期：

| 值 | 名称 | 对暂存复制的含义 |
|---|---|---|
| 0 | `kImmutableOnlyDuringCall` | libtpu 必须在返回前完成对 `data` 的读取；调用后主机缓冲区可立即变更。强制同步复制。 |
| 1 | `kImmutableUntilTransferCompletes` | 主机缓冲区必须保持有效，直到 `done_with_host_buffer` 触发；libtpu 异步复制。 |
| 2 | `kImmutableZeroCopy` | 主机缓冲区在缓冲区整个生命周期内不可变；在布局允许时 libtpu 可为其建立别名（无复制）。 |
| 3 | `kMutableZeroCopy` | 调用者捐赠主机分配；libtpu 可接管所有权并建立别名。 |

> **注意（高置信度）—** *偏移*（`a1+0x48`）和转换器调用（`ConvertFromPjRtHostBufferSemantics`）已按字节确认；四个枚举项名称/值镜像上游 PJRT C API 枚举（`PJRT_HostBufferSemantics`），转换器必须与之匹配，ABI 才能互操作。逐值复制与别名行为是已记录的 PJRT 契约；TPU stage 调用中*作用于*每个值的确切分支未逐字节跟踪（stage 是一次 vtable+0x120 调用，将转换后的值作为参数接收）。

### 算法

```c
function PJRT_Client_BufferFromHostBuffer(args):          // 0xf8644c0
    if !ActualStructSizeIsGreaterOrEqual(..., 37, 120, args.struct_size):
        return new PJRT_Error{ size_error }                // operator new(8)

    has_strides = (args.byte_strides != NULL)              // a1+0x38; sets internal flag v75
    layout = args.device_layout                            // a1+0x60
    if layout != NULL:
        switch layout.type:                                // layout+0x48
            case 1 /*Strides*/: return Error("...Strides... not supported... for platform " + name)
            case 0 /*Tiled*/:   break
            default:            return Error("Unexpected PJRT_Buffer_MemoryLayout_Type type: " + n)

    // Build a PjRtFuture promise for done_with_host_buffer (PromiseMaker<void>::Make)
    promise = PromiseMaker<void>::Make()                   // tsl::internal, lines 210-227

    xla_type      = ConvertFromPjRtBufferType(args.type)              // a1+0x20
    semantics     = ConvertFromPjRtHostBufferSemantics(args.host_buffer_semantics)  // a1+0x48
    mem           = args.memory                            // a1+0x50, the target memory space

    // Single virtual call into the memory-space / client allocate+stage path:
    //   vtable+0x120 (offset 288) on the inner client object.
    //   This allocates the device buffer in `mem`, schedules the host->device DMA
    //   under `semantics`, and yields a StatusOr<unique_ptr<PjRtBuffer>>.
    status_or = (*inner_vtable[288])(inner_client, xla_type, dims, num_dims,
                                     byte_strides, semantics, layout_info, mem, promise)
    if !status_or.ok():
        return new PJRT_Error{ status_or.status }

    buf_wrapper = operator new(0x110)                      // LABEL_46
    buf_wrapper.impl   = status_or.value                   // +0x00
    buf_wrapper.client = args.client                       // +0x08
    zero(buf_wrapper + 0xf0 .. 0x110)                      // external-ref list + flags
    args.buffer = buf_wrapper                              // a1+0x70

    done_event = operator new(0x50)                        // 80-byte PJRT_Event
    move promise.future.async_value -> done_event+0x08
    move 2 profiling closures           -> done_event+0x18..0x40
    done_event+0x48 = 0
    args.done_with_host_buffer = done_event                // a1+0x68
    return NULL                                            // success

怪癖 — 这里有两个正在流转的事件，而且含义不同。done_with_host_buffer（+0x68）表示主机缓冲区何时可触碰（由 host_buffer_semantics 控制）。设备缓冲区自身的就绪性是另一个事件，之后通过 PJRT_Buffer_ReadyEvent（槽位 77）获取。重新实现者不能把二者混用：缓冲区可能早在设备 DMA 完成前就“done with host”，反之亦然。

函数映射

函数	地址	角色
`pjrt::PJRT_Client_BufferFromHostBuffer`	`0xf8644c0`	C 包装器，args 校验 + 编组
`pjrt::ConvertFromPjRtBufferType`	`0xf8a3e60`	`PJRT_Buffer_Type` → `xla::PrimitiveType`
`pjrt::ConvertFromPjRtHostBufferSemantics`	`0xf8a3f20`	public → XLA semantics enum
`tsl::internal::PromiseMaker<void>::Make`	（内联）	构建 `done_with_host_buffer` promise
inner client vtable+0x120	（按平台）	分配 + 调度 host→device DMA

PJRT_Buffer_ToHostBuffer（槽位 75）

用途

回读路径：把设备缓冲区内容复制到调用者提供的主机缓冲区。在 TPU 上这并不简单，因为设备上数据是平铺的（填充到 TPU 原生 tile 形状），而主机期望密集的行主序线性布局。ToHostBuffer 会在复制前执行反平铺形状转换。

算法

function PJRT_Buffer_ToHostBuffer(args):                  // 0xf86e640
    wrapper = args[+0x10]; inner = wrapper.impl
    shape = build xla::Shape from inner.on_device_shape()
    if !shape.is_static():
        ... (dynamic-shape path)
    host_shape = ShapeUtil::DeviceShapeToHostShape(shape)  // 0x20cec000 — de-tile
    if args.host_layout != NULL:
        layout = ConvertToLayout(args.host_layout)         // 0xf8a5640, tiled-layout struct
    // layout-aware copy into the caller's host literal:
    future = inner.vtable[+0x78].ToLiteral(MutableLiteralBase{host_ptr, host_shape, layout})
    args.event = wrap_as_PJRT_Event(future)                // operator new(0x50)
    return NULL
```text

原始、无形状的回读变体是 `CopyRawToHost`（槽位 105，内部 vtable+0x90 `CopyRawToHost(void*, off, size)`）和 `CopyRawToHostFuture`（槽位 125，`0xf86dfe0`，内部 vtable+0x98 `CopyRawToHostFuture(Future<void*>, off, size)`）；它们移动字节时不反平铺，是 [RawBuffer 扩展](ext-rawbuffer.md)的 device→host copy 的带类型缓冲区镜像。当主机数组必须匹配逻辑（线性）形状时使用 `ToHostBuffer`；需要逐字节精确设备 dump 时使用 raw 变体。

> **陷阱 —** 跳过 `DeviceShapeToHostShape` 步骤并把设备上的字节直接 `memcpy` 到主机，对于任何范围尚未 tile 对齐的张量都会得到垃圾：设备字节包含主机布局不期望的 tile padding。反平铺是必需的，不是优化。

### 函数映射

| 函数 | 地址 | 角色 |
|---|---|---|
| `pjrt::PJRT_Buffer_ToHostBuffer` | `0xf86e640` | C 包装器；反平铺 + ToLiteral |
| `xla::ShapeUtil::DeviceShapeToHostShape` | `0x20cec000` | tiled → linear shape |
| `pjrt::ConvertToLayout` | `0xf8a5640` | `PJRT_Buffer_MemoryLayout_Tiled` → `xla::Layout` |
| `CommonPjRtBufferImpl::ToLiteral` | inner vtable+0x78（`0xf9295a0`） | layout 感知的 device→host copy |

### 注意事项

完整的 `PJRT_Buffer_ToHostBuffer_Args` 字段偏移（主机目标指针、目标大小、可选 `host_layout`）除了形状/layout 序言之外**未按字节跟踪**。反平铺和 `ToLiteral` 跳转已确认；精确参数偏移标记为 LOW，重新实现者应与上游 `PJRT_Buffer_ToHostBuffer_Args` 交叉核对。

---

## 形状访问器（槽位 64、65、69）

### 用途

三个零分配访问器，读取缓冲区不可变的已缓存形状，并把答案写回 args 结构。它们不加锁，成功时不分配任何东西，并且可安全地与其他只读操作并发调用。

### 算法

```c
function PJRT_Buffer_ElementType(args):                   // 0xf86d220
    inner = args[+0x10].impl
    prim  = inner.vtable[+0x10].element_type()             // xla::PrimitiveType
    args[+0x18] = ConvertToPjRtBufferType(prim)            // 0xf8a3d80 -> PJRT_Buffer_Type

function PJRT_Buffer_Dimensions(args):                     // 0xf86d280
    inner = args[+0x10].impl
    span  = inner.vtable[+0x18].dimensions()               // {const int64* data, size_t count}
    args[+0x18] = span.data                                // zero-copy into cached shape
    args[+0x20] = span.count

function PJRT_Buffer_OnDeviceSizeInBytes(args):            // 0xf86da80
    if !ActualStructSizeIsGreaterOrEqual("PJRT_Buffer_OnDeviceSizeInBytes_Args", 36, 32, args.struct_size):
        return new PJRT_Error{ size_error }
    sz = inner.vtable[+0x88].GetOnDeviceSizeInBytes()       // StatusOr<int64>
    if sz.ok(): args[+0x18] = sz.value; return NULL
    else:       return new PJRT_Error{ sz.status }

OnDeviceSizeInBytes 是三者中唯一可能失败的（它返回 StatusOr，因此有 operator new(8) 错误路径）；ElementType 和 Dimensions 是对已缓存形状的不可失败读取。维度 span 是对缓冲区内部形状的借用，其生命周期绑定到缓冲区，因此调用者必须在缓冲区销毁前复制它。

注意 — 槽位 69 返回的设备上大小是包含 tile padding 的 HBM 字节数，从内部缓冲区读取（设备层的 TpuBufferBase+0x50，见 TPU 缓冲区布局）。它通常大于 product(dims) * sizeof(element_type)；不要用它来确定 ToHostBuffer 的主机缓冲区大小，后者期望的是主机（反平铺后）大小。

生命周期：Delete、IsDeleted、Destroy（槽位 72、73、63）

用途

libtpu 将设备内存回收与 C 包装器回收分开。Delete（槽位 72）立即释放 HBM 分配，同时让包装器保持有效；Destroy（槽位 63）释放 C 包装器本身。它们是对两个不同资源的独立操作。

算法

function PJRT_Buffer_Delete(args):                         // 0xf86dd80
    inner = args[+0x10].impl
    inner.vtable[+0xa0].Delete()        // eagerly free DEVICE memory; wrapper stays valid

function PJRT_Buffer_IsDeleted(args):                      // 0xf86dde0
    inner = args[+0x10].impl
    args[+0x18] = inner.vtable[+0xb0].IsDeleted()          // true after Delete()

function PJRT_Buffer_Destroy(args):                        // 0xf86d020
    wrapper = args[+0x10]
    ~PJRT_Buffer(wrapper)               // 0xf86d0a0: drain ext-ref list, dtor inner buffer
    free(wrapper, 0x110)                // exclusive ownership — no refcount
```text

> **怪癖 —** `Delete` 和 `Destroy` *不是*同一个调用，并且顺序很重要。预期顺序是：当设备上不再需要结果时，先 `Delete` 以立即回收 HBM（降低峰值内存），之后在丢弃主机侧句柄时再 `Destroy`。不调用 `Delete` 而直接 `Destroy` 也可以，内部析构函数同样会释放设备内存；但在 `Destroy` 之后再调用包装器槽位就是 use-after-free。`IsDeleted` 在 `Delete` 后返回 true；它*不会*告诉你 `Destroy` 是否已经运行（不存在“is destroyed”谓词，这是调用者的责任）。

### 外部引用 pin

`IncreaseExternalReferenceCount`（槽位 79）跳转到内部 vtable+0x70 `AcquireExternalReference`，并把返回的 hold 存入包装器位于 `+0xf8/+0x100/+0x108` 的列表；`DecreaseExternalReferenceCount`（槽位 80）释放它。只要持有任何外部引用，`Delete` **就不能**释放设备分配，这就是外部消费者（dlpack、NumPy 零拷贝、读取 `OpaqueDeviceMemoryDataPointer` 的自定义内核）保持 HBM 存活的方式。`~PJRT_Buffer()` 会在 `Destroy` 期间清空该列表。

> **陷阱 —** 不平衡的 `IncreaseExternalReferenceCount` 会造成设备内存泄漏，`Delete` 无法回收：HBM 会一直被 pin，直到 `Destroy` 清空列表。每个 Increase 都需要在 `Delete` 真正释放之前有一个匹配的 Decrease。

---

## 内存空间表面（槽位 71、97、74）

### 用途

缓冲区恰好位于一个内存空间中；这些槽位查询它并把缓冲区复制到另一个空间。`Memory`（槽位 71）返回缓冲区当前的 `PJRT_Memory*`。`CopyToMemory`（槽位 97）复制到调用者指定的内存空间。`CopyToDevice`（槽位 74）复制到某个设备的*默认*内存空间。

### 算法

```c
function PJRT_Buffer_Memory(args):                         // 0xf86dc60, slot 71
    inner   = args[+0x10].impl
    mem_cpp = inner.vtable[+0x58].memory_space()           // xla::PjRtMemorySpace*
    client  = args[+0x10].client                           // inner XLA client at *(client+0x08)
    // inlined equivalent of PJRT_Client_FindMemoryWrapper(mem_cpp, client):
    mem_c   = scan client cache (count @+0x90, array @+0x88) for wrapper whose *wrapper == mem_cpp
    if mem_c == NULL:
        args[+0x18] = NULL
        return Unimplemented("PJRT_Buffer_Memory not implemented for platform '%s'", platform_name)
    args[+0x18] = mem_c

function PJRT_Buffer_CopyToMemory(args):                   // 0xf86e500, slot 97
    inner    = args[+0x10].impl
    dst_mem  = args.dst_memory                             // PjRtMemorySpace* directly
    future   = inner.vtable[+0xb8].CopyToMemorySpace(dst_mem)   // 0xf926c80
    args.dst_buffer = wrap(future.value); args.event = wrap_as_PJRT_Event(future)

function PJRT_Buffer_CopyToDevice(args):                   // 0xf86e360, slot 74
    dst_dev  = args.dst_device
    dst_mem  = dst_dev.vtable[+0x98].default_memory_space()     // TpuDevice::default_memory_space 0xf7feda0
    future   = src_inner.vtable[+0xb8].CopyToMemorySpace(dst_mem)
    ...

两个复制槽位都会汇入同一个 CommonPjRtBufferImpl::CopyToMemorySpace(PjRtMemorySpace*) @ 0xf926c80（内部 vtable+0xb8）；区别只在于目标内存空间如何获得（直接命名 vs. 设备默认）。该函数按源/目标形状和类路由：同设备、ICI 跨设备、跨 pod megascale 都经过它；CPU↔设备路径走专用分支。其下方的跨内存空间路由和 DMA 引擎记录在 DMA 与跨主机接收；设备侧分配器记录在 HBM 分配器。本页停留在 C-ABI 跳转层。

注意 — CopyToDevice 访问的是设备 vtable（xla::TpuDevice，vtable 0x2177b4d0），而不是缓冲区 vtable，以找到默认内存空间（+0x98 → 0xf7feda0）。这个 +0x98 是 PjRtDevice 槽位；不要在缓冲区上找它。

异步就绪

每个传输槽位（BufferFromHostBuffer、ToHostBuffer、CopyToMemory、CopyToDevice、CopyRawToHost）都返回一个80 字节 PJRT_Event，包装一个 PjRtFuture。布局固定：

struct PJRT_Event {                  // operator new(0x50)
    /* +0x00 */ void*  vtable_or_pad;
    /* +0x08 */ tsl::AsyncValueRef<absl::Status> async_value;   // the future's status
    /* +0x18 */ AnyInvocable profiling_closure_0;               // __policy func + state
    /* +0x28 */ ...
    /* +0x38 */ AnyInvocable profiling_closure_1;
    /* +0x48 */ uint64_t zero;
};
```text

缓冲区自身的就绪性使用同一套机制：`PJRT_Buffer_ReadyEvent`（槽位 77）跳转到内部 vtable+0xe0 `GetReadyFuture()`（`0xf92aa60`），并把 future 包装成 80 字节事件。消费者会轮询 `PJRT_Event_IsReady`（槽位 11）、注册 `PJRT_Event_OnReady`（槽位 14），或通过 `PJRT_Event_Await`（槽位 13）阻塞，见[事件与异步](events-and-async.md)。其下方的设备侧完成 token 是 `tpu::TpuEvent`；中间层通过 `CommonPjRtClient::MakeTrackedReadyFuture`（`0xf91c2e0`）把它转换为 client 跟踪的 `PjRtFuture`。

---

## 捐赠（槽位 130）

`PJRT_Buffer_DonateWithControlDependency`（槽位 130，`0xf86f2e0`）跳转到内部 vtable+0xd8 `DonateWithControlDependency(tsl::Future<void>)`（`0xf92a740`）。它构建一个 `tsl::Promise`（`PromiseMaker<void>::Make`），传入控制依赖 future，并生成一个**新的**缓冲区包装器，其设备内存与捐赠者别名，捐赠者会失效。这就是 JAX 的 `donate_argnums` 为输出复用输入缓冲区 HBM 的方式，每步消除一次分配和一次复制。缓冲区捐赠/别名数据模型归[缓冲区捐赠与别名](../memory/buffer-donation-aliasing.md)所有；本页只记录 C-ABI 槽位。

---

## 交叉引用

- [PJRT API 概览](overview.md) — 140 槽位的 `PJRT_Api`，以及缓冲区槽位如何融入插件生命周期
- [API Vtable 重建](api-vtable-reconstruction.md) — 槽位表和 `struct_size` 版本化如何重建
- [Client 与 Device](client-and-device.md) — `PJRT_Client`、`PJRT_Device`，以及支撑 `PJRT_Buffer_Memory` 的内存包装器缓存
- [可执行文件执行](executable-execution.md) — 输出缓冲区产生的位置，以及输入缓冲区在执行时捐赠的位置
- [事件与异步](events-and-async.md) — 每次传输返回的 80 字节 `PJRT_Event` / `PjRtFuture` 模型
- [RawBuffer 扩展（type 8）](ext-rawbuffer.md) — 无类型 16 字节原始字节兄弟表面和主机↔设备 DMA 管线
- [DMA 与跨主机接收](dma-and-cross-host-recv.md) — `CopyToMemory` 下方的跨内存空间复制路由
- [TPU 缓冲区布局](../memory/tpu-buffer-layout.md) — 设备侧字节布局、平铺和 `TpuBufferBase` 字段（`+0x48` ptr / `+0x50` size）
- [HBM 分配器](../memory/hbm-allocator.md) — `BufferFromHostBuffer` 调用的设备侧 HBM 分配
- [缓冲区捐赠与别名](../memory/buffer-donation-aliasing.md) — 槽位 130 背后的捐赠/别名数据模型
- [分配器集成](../runtime/allocator-integration.md) — StreamExecutor 分配器桥接

PJRT 缓冲区 ABI 与内存布局 ​

摘要 ​

两个 C 包装对象 ​

PJRT_Buffer 布局（0x110 = 272 字节） ​

函数映射 ​

PJRT_Buffer_ToHostBuffer（槽位 75） ​

用途 ​

算法 ​

生命周期：Delete、IsDeleted、Destroy（槽位 72、73、63） ​