PJRT RawBuffer 扩展（类型 8）

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

摘要

RawBuffer 扩展（PJRT_Extension_Type id 8，struct_size 0x50 = 80 字节）是 libtpu 的无类型字节级设备内存接口。它暴露七个方法 — CreateRawAliasOfBuffer、Destroy、GetOnDeviceSizeInBytes、GetMemorySpace、CopyRawHostToDevice、CopyRawDeviceToHost、GetHostPointer — 通过 (offset, size) 元组在 HBM（或 pinned host memory）中移动并寻址原始字节，没有元素类型、没有形状、没有 tiling，也没有 de-tiling。它是有类型 PJRT_Buffer ABI 的有意设计的同级接口：后者会校验元素类型、封送维度，并在读回时 de-tile；而 RawBuffer 接口则是通向同一底层设备分配的一条扁平 void* 加长度的 DMA 通道。

这两个接口包装的是不同的 C++ 类层次。有类型 PJRT_Buffer shim 包装 xla::PjRtBuffer → xla::CommonPjRtBufferImpl；RawBuffer shim 包装 xla::PjRtRawBuffer（抽象）→ xla::CommonPjRtRawBuffer → xla::CommonPjRtRawBufferImpl（共享 copy 方法的持有者）→ xla::TpuRawBuffer（具体 TPU，vtable 0x2177cfe0）/ xla::CpuRawBuffer（具体 CPU staging，vtable 0x21789af8，顺序相同）。两个 C wrapper 对象在大小和所有权上也不同：有类型 wrapper 为 272 字节（0x110），并且独占拥有其内部 buffer；而 RawBuffer wrapper 是 16 字节（0x10），持有一个 tsl::RCReference<xla::PjRtRawBuffer> — 它是设备分配的一个共享、引用计数共同所有者。正是这一个差异，使 CreateRawAliasOfBuffer 能够在已存活的有类型 buffer 上发放零拷贝 alias。

本页涵盖该扩展 struct、其七个方法集合、16 字节 wrapper 布局、各方法 args 偏移和 struct_size 版本化、支撑这些方法的 TpuRawBuffer vtable，以及原始 host↔device DMA 语义，一直到 tpu::System::Transfer* 入口。将此扩展链接进 PJRT_Api 的 chain node 位于扩展链；跨 memory-space copy 路由和 DMA 引擎内部细节位于 DMA 与 Cross-Host Receive；最终支撑 HBM 的 StreamExecutor allocator bridge 位于 Allocator 集成。

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

扩展 struct：struct_size 0x50，type 8，七个 fn-ptr slot 位于 +0x18..+0x48，由单个扁平 table-initializer creator 填充。
16 字节 RawBuffer C wrapper（{ RCReference impl@+0x00; PJRT_Client* client@+0x08 }）及其引用计数 Destroy。
每个方法的 struct_size（min、cur）字面量、其 args 偏移，以及它执行的单次 xla::PjRtRawBuffer vtable bounce。
原始 copy 语义：一个 (offset, size) slice transfer，返回 80 字节 PJRT_Event，没有 shape/tile 转换。
两个 device-pointer/host-pointer 访问器，以及限制 host-addressability 的规则（仅 pinned_host）。


扩展类型 id	8（`PJRT_Extension_Type` RawBuffer）
扩展 struct 大小	`0x50`（80 字节）；7 个 fn-ptr slot 位于 `+0x18..+0x48`
Creator	`pjrt::CreateRawBufferExtension(PJRT_Extension_Base* next)` @ `0xe6f52c0`
`.bss` 存储	`0x224c3990`（`raw_buffer_extension`）；`next` → profiler `0x22255b98`
C wrapper 大小	`0x10`（16 字节）— `{ RCReference<PjRtRawBuffer> impl@+0x00; PJRT_Client* client@+0x08 }`
具体后端（TPU）	`xla::TpuRawBuffer`（vtable `0x2177cfe0`，vptr base `0x2177cff0`）
Copy-method holder	`xla::CommonPjRtRawBufferImpl`（vtable `+0x28`/`+0x30` 指向这里）
异步闸门	每次 raw copy 都返回一个包装 `PjRtFuture` 的 80 字节 `PJRT_Event`

扩展 Struct（类型 8，80 字节）

目的

扩展 struct 是一个扁平函数指针表，与七个追加的 raw-buffer 方法指针共享通用 PJRT_Extension_Base header（struct_size、type、_pad、next）。它驻留在 0x224c3990 的 .bss 中，并在第一次 GetTpuPjrtApi 调用时一次性初始化。

布局

struct PJRT_RawBuffer_Extension {            // struct_size 0x50 (80 bytes)
    PJRT_Extension_Base base;                // +0x00 struct_size; +0x08 type=8; +0x0c _pad; +0x10 next
    /* +0x18 */ PJRT_Error* (*CreateRawAliasOfBuffer)(PJRT_RawBuffer_CreateRawAliasOfBuffer_Args*);
    /* +0x20 */ PJRT_Error* (*Destroy)               (PJRT_RawBuffer_Destroy_Args*);
    /* +0x28 */ PJRT_Error* (*GetOnDeviceSizeInBytes)(PJRT_RawBuffer_GetOnDeviceSizeInBytes_Args*);
    /* +0x30 */ PJRT_Error* (*GetMemorySpace)        (PJRT_RawBuffer_GetMemorySpace_Args*);
    /* +0x38 */ PJRT_Error* (*CopyRawHostToDevice)   (PJRT_RawBuffer_CopyRawHostToDevice_Args*);
    /* +0x40 */ PJRT_Error* (*CopyRawDeviceToHost)   (PJRT_RawBuffer_CopyRawDeviceToHost_Args*);
    /* +0x48 */ PJRT_Error* (*GetHostPointer)        (PJRT_RawBuffer_GetHostPointer_Args*);
};
```text

### Creator

`pjrt::CreateRawBufferExtension` @ `0xe6f52c0` 是一个纯 table initializer — 无分配、无分支，单个 `ret`。反编译结果是字面量式的：

```c
function CreateRawBufferExtension(slot, next):       // 0xe6f52c0
    *(u64*)(slot + 0x00) = 80                          // struct_size
    *(u32*)(slot + 0x08) = 8                           // type
    *(u64*)(slot + 0x10) = next                        // chain link (arg)
    *(u64*)(slot + 0x18) = &PJRT_RawBuffer_CreateRawAliasOfBuffer
    *(u64*)(slot + 0x20) = &PJRT_RawBuffer_Destroy
    *(u64*)(slot + 0x28) = &PJRT_RawBuffer_GetOnDeviceSizeInBytes
    *(u64*)(slot + 0x30) = &PJRT_RawBuffer_GetMemorySpace
    *(u64*)(slot + 0x38) = &PJRT_RawBuffer_CopyRawHostToDevice
    *(u64*)(slot + 0x40) = &PJRT_RawBuffer_CopyRawDeviceToHost
    *(u64*)(slot + 0x48) = &PJRT_RawBuffer_GetHostPointer
    return slot

由于该表完全静态，struct 可以位于零初始化的 .bss 中，并且只需要一次由 __cxa_guard 保护的 creator 调用。next 参数是此前构建好的 node；RawBuffer 是 GetTpuPjrtApi 内部构造的第一个 .bss node，因此其 next 被设为驻留在 .data 中的 Profiler 扩展 0x22255b98 — 也就是 chain terminator。完整 17-node 遍历以及 RawBuffer 尽管最先构建却最终位于遍历位置 16 的原因，见扩展链。

16 字节 C wrapper

struct PJRT_RawBuffer {                  // sizeof = 0x10 (16 bytes)
    /* +0x00 */ tsl::RCReference<xla::PjRtRawBuffer> impl;   // SHARED, ref-counted co-owner
    /* +0x08 */ PJRT_Client*                         client; // borrowed (not owned)
};
```text

两个字段都已按字节确认：`Destroy`（`0xe6f4e40`）读取 wrapper`+0x00` 处的内部 `RCReference`，并在 refcount 路径之后 `free` wrapper；`CreateRawAliasOfBuffer` 和 `GetMemorySpace` 读取 wrapper`+0x08` 处借用的 `client`。`impl` 是 `tsl::RCReference` — wrapper 参与对底层 `xla::PjRtRawBuffer` 的*共享*所有权，这与有类型 `PJRT_Buffer` wrapper 独占拥有其内部 `PjRtBuffer*` 形成对比。这是 raw aliasing 的结构性前提：两个 RawBuffer wrapper 可以通过共享 `RCReference` 共同拥有同一个设备分配。

> **注意 —** 不要混淆两个 buffer wrapper。有类型接口（[buffer-and-memory.md](buffer-and-memory.md)）使用 272 字节的独占所有 wrapper；本接口使用覆盖*不同* C++ base（`xla::PjRtRawBuffer`，不是 `xla::PjRtBuffer`）的 16 字节引用计数 wrapper。重新实现者如果复用同一个 wrapper 布局给两者，会错误计算 `free` 大小并错误建模所有权。

### Args 约定

每个方法的 args struct 都遵循规范的 `{ size_t struct_size; void* priv; <handle>; ... }` 形态。RawBuffer wrapper handle 位于 **args+0x10**（`priv` 占用 `+0x08`）；每个 scalar/pointer 输出写入 **args+0x18** 及其之后。每个 wrapper 的第一个动作都是调用 `pjrt::ActualStructSizeIsGreaterOrEqual(name, min_fields, cur_bytes, args->struct_size)` @ `0xf8a4ec0`；不匹配时，wrapper 会通过 `operator new(8)` 分配一个携带 size status 的 `PJRT_Error` 并返回，不触碰 buffer。

---

## 方法集合

### Slot 映射

全部七个方法都已对照反编译按字节确认：args-name 字符串、`(min, cur)` `struct_size` 字面量、args-output 偏移，以及各自 bounce 经过的单个 `xla::PjRtRawBuffer` vtable 偏移。“vtable bounce”列是内部对象 vtable 中的偏移（`TpuRawBuffer` 的 object vptr base 为 `0x2177cff0`）。

| Off | 方法 | C symbol | Addr | min/cur | vtable bounce / 后端 |
|---|---|---|---|---|---|
| `0x18` | CreateRawAliasOfBuffer | `pjrt::PJRT_RawBuffer_CreateRawAliasOfBuffer` | `0xe6f4d40` | 42 / 32 | static `xla::PjRtRawBuffer::CreateRawAliasOfBuffer` @ `0xf93f540` |
| `0x20` | Destroy | `pjrt::PJRT_RawBuffer_Destroy` | `0xe6f4e40` | 27 / 24 | `RCReference` dec-ref + vtable+0x08（deleting dtor）+ `free(0x10)` |
| `0x28` | GetOnDeviceSizeInBytes | `pjrt::PJRT_RawBuffer_GetOnDeviceSizeInBytes` | `0xe6f4f20` | 42 / 32 | vtable+0x20 `GetOnDeviceSizeInBytes()` |
| `0x30` | GetMemorySpace | `pjrt::PJRT_RawBuffer_GetMemorySpace` | `0xe6f4f80` | 34 / 32 | vtable+0x10 `memory_space()` + `PJRT_Client_FindMemoryWrapper` @ `0xf8605e0` |
| `0x38` | CopyRawHostToDevice | `pjrt::PJRT_RawBuffer_CopyRawHostToDevice` | `0xe6f5040` | 39 / 56 | vtable+0x28 `CommonPjRtRawBufferImpl::CopyRawHostToDevice` @ `0xf91c640` |
| `0x40` | CopyRawDeviceToHost | `pjrt::PJRT_RawBuffer_CopyRawDeviceToHost` | `0xe6f5180` | 39 / 56 | vtable+0x30 `CommonPjRtRawBufferImpl::CopyRawDeviceToHost` @ `0xf91c780` |
| `0x48` | GetHostPointer | `pjrt::PJRT_RawBuffer_GetHostPointer` | `0xe6f4ec0` | 34 / 32 | vtable+0x18 `TpuRawBuffer::GetHostPointer`（仅 pinned_host） |

args-name 字符串原样存在于 `.rodata`（`"PJRT_RawBuffer_Destroy_Args"`、`"PJRT_RawBuffer_CopyRawHostToDevice_Args"` 等），确认了公开 API 名称。

> **特性 —** 两个 copy 方法的 `(min, cur)` 对是 `(39, 56)` — 以字段计数来说，`min`（接受的最小字段数）比其他方法更大，因为 copy args 携带 host pointer、offset、size 和 out-event。其他五个方法接受 `min` 27–42，`cur` 为 24 或 32 字节。validator 会把调用方的 `struct_size` 与 `cur`（当前字节大小）以及字段计数 `min` 比较；只要达到 `min`，基于更旧、更小 header 构建的调用方也会被接受。

---

### CreateRawAliasOfBuffer（slot 0x18）

#### 目的

在现有*有类型* `PJRT_Buffer` 上创建一个**原始、无类型 alias**。该 alias 是一个新的 16 字节 RawBuffer wrapper，它通过 `RCReference` 共享 donor buffer 的底层设备分配。没有字节被复制；alias 是同一块 HBM 的第二个共同所有者。这就是调用方如何获得一个由有类型 `BufferFromHostBuffer` 路径创建的 buffer 的字节级 handle。

#### 算法

```c
function PJRT_RawBuffer_CreateRawAliasOfBuffer(args):    // 0xe6f4d40
    if !ActualStructSizeIsGreaterOrEqual("PJRT_RawBuffer_CreateRawAliasOfBuffer_Args", 42, 32, args.struct_size):
        return new PJRT_Error{ size_status }              // operator new(8)

    // args.buffer @ +0x10 is a typed PJRT_Buffer wrapper; **(a1+0x10) = inner xla::PjRtBuffer*
    inner_typed = *(*(args + 0x10))                       // typed wrapper -> PjRtBuffer*
    sor = xla::PjRtRawBuffer::CreateRawAliasOfBuffer(inner_typed)   // 0xf93f540 -> StatusOr<RCReference>

    if sor.is_error:                                      // discriminant low bit set
        err = new PJRT_Error(8); *err = sor.status        // and Unref the StatusRep
        return err

    // success: build the 16-byte raw wrapper
    raw_wrapper = operator new(0x10)
    raw_wrapper[+0x00] = sor.value (the RCReference, move-out)
    raw_wrapper[+0x08] = *(args[+0x10] + 0x08)            // copy client from DONOR wrapper+0x08
    args[+0x18] = raw_wrapper                             // OUT
    return NULL

factory xla::PjRtRawBuffer::CreateRawAliasOfBuffer(PjRtBuffer*) @ 0xf93f540 会遍历全局 per-platform factory registry xla::GetFactoryFuncs()::funcs @ 0x224c70c8（guard 0x224c70d0）— 每个 entry 是一个 bool(*)(StatusOr<RCReference>&, PjRtBuffer*) — 并返回第一个接受该 buffer 平台的已注册 factory。如果没有任何 factory 接受，它会构造一个 StrCat 错误。新 alias 的 client pointer 从donor 的 wrapper +0x08 复制，因此 alias 共享 donor 的 PJRT_Client。

特性 — alias 通过 RCReference 共享 donor 的设备分配；它不是一次 copy，也不会 pin 第二块 HBM 区域。销毁 alias 会递减共享 refcount；只有最后一个共同所有者释放时，设备内存才会被释放。重新实现者必须把 alias 的 Destroy 路由到相同的引用计数 decrement，而不是对设备 buffer 做扁平 free。

函数映射

函数	Addr	角色
`pjrt::PJRT_RawBuffer_CreateRawAliasOfBuffer`	`0xe6f4d40`	C wrapper；构建 raw alias wrapper
`xla::PjRtRawBuffer::CreateRawAliasOfBuffer`	`0xf93f540`	static factory；遍历 platform registry
`xla::GetFactoryFuncs()::funcs`	`0x224c70c8`	per-platform raw-alias factory 的 registry

Destroy（slot 0x20）

目的

释放调用方对 RawBuffer 的引用。因为 wrapper 持有一个引用计数 RCReference，Destroy 不是扁平 free — 它递减共享 refcount，并且仅在计数到零时运行内部 deleting destructor。然后它无条件释放 16 字节 wrapper。

算法

function PJRT_RawBuffer_Destroy(args):                   // 0xe6f4e40
    if !ActualStructSizeIsGreaterOrEqual("PJRT_RawBuffer_Destroy_Args", 27, 24, args.struct_size):
        return new PJRT_Error{ size_status }

    wrapper = args[+0x10]                                 // a1[2]
    if wrapper != NULL:
        inner = wrapper[+0x00]                            // the RCReference's pointer
        if inner != NULL:
            // canonical TSL refcount: count at inner+0x08 (4-byte)
            if inner.refcount == 1 || atomic_dec(inner+0x08) == 0:
                inner.vtable[+0x08](inner)                // deleting dtor (~PjRtRawBuffer/D0)
        free(wrapper)                                     // 16-byte wrapper, unconditional
    return NULL
```text

fast-path 检查 `refcount == 1` 会在调用方持有唯一引用时跳过 atomic；否则用带 `lock` 前缀的 decrement 判断是否释放设备 buffer。deleting destructor 位于内部对象的 vtable`+0x08`（`xla::TpuRawBuffer::~TpuRawBuffer` [D0] @ `0xf83c5e0`）。

> **陷阱 —** 每次调用 `Destroy` 都会释放 wrapper，即使共享 refcount 还未归零（另一个 alias 仍共同拥有设备内存）也是如此。wrapper free 和 device-memory free 是解耦的：16 字节 C handle 会立即消失；HBM 会存活到最后一个 `RCReference` 释放。应把 `Destroy` 理解为“丢弃我的 handle”，而不是“释放 buffer”。

---

### GetOnDeviceSizeInBytes（slot 0x28）和 GetHostPointer（slot 0x48）

这两个是最小化的单次 bounce 访问器，结构相同：校验 `struct_size`，三重解引用 wrapper 到达内部 vtable，调用一个 virtual method，将结果写入 args`+0x18`。成功路径上它们从不分配。

```c
function PJRT_RawBuffer_GetOnDeviceSizeInBytes(args):    // 0xe6f4f20
    // min 42, cur 32
    inner = *(args[+0x10])                                // **(a1+0x10)
    args[+0x18] = inner.vtable[+0x20].GetOnDeviceSizeInBytes()   // int64

function PJRT_RawBuffer_GetHostPointer(args):            // 0xe6f4ec0
    // min 34, cur 32
    inner = *(args[+0x10])
    args[+0x18] = inner.vtable[+0x18].GetHostPointer()   // void* or NULL

GetOnDeviceSizeInBytes 返回具体 HBM 字节数（xla::TpuRawBuffer::GetOnDeviceSizeInBytes @ 0xf838880，读取已解析的 TpuBufferBase+0x50）。这包括任何设备侧 tile padding — 它是物理分配大小，而不是 product(dims) * elem_size。

GetHostPointer（xla::TpuRawBuffer::GetHostPointer @ 0xf837b60）只会对 pinned_host buffer 返回 host 可解引用的 pointer；对于 tpu_hbm 则返回 NULL。实现会根据 buffer 的 memory-space kind 字符串进行 gate（对 "pinned_host" magic 做 11 字节 overlapping compare），并且在 pinned buffer 上沿着 AsyncValue indirect-node chain 返回已解析的 TpuBufferBase+0x48。

特性 — GetHostPointer 返回 NULL 不是错误 — 对任何 HBM-resident buffer 来说这是正确答案，因为 HBM 不能被 host addressable。调用方必须检查 NULL 并回退到 CopyRawDeviceToHost；把 NULL 当作 failure code 会破坏每一个 device buffer。这与有类型接口的 OpaqueDeviceMemoryDataPointer（slot 81）相呼应，后者返回一个只在所属 device/core 上有效的原始 HBM 虚拟地址 — 见 buffer-and-memory.md 上的特性。

GetMemorySpace（slot 0x30）

目的

返回 raw buffer 所在 memory space 的 C PJRT_Memory* wrapper。内部 buffer 的 memory_space() 产生一个 C++ xla::PjRtMemorySpace*；wrapper 通过 client 的 memory-wrapper cache 往返解析它，以取得 C handle。

算法

function PJRT_RawBuffer_GetMemorySpace(args):            // 0xe6f4f80
    if !ActualStructSizeIsGreaterOrEqual("PJRT_RawBuffer_GetMemorySpace_Args", 34, 32, args.struct_size):
        return new PJRT_Error{ size_status }

    inner   = *(args[+0x10])                              // **(a1+0x10)
    mem_cpp = inner.vtable[+0x10].memory_space()          // xla::PjRtMemorySpace*
    mem_c   = PJRT_Client_FindMemoryWrapper(mem_cpp, args[+0x10].client)  // 0xf8605e0; client @ wrapper+0x08
    args[+0x18] = mem_c
    if mem_c == NULL:
        return new PJRT_Error{ MakeErrorImpl<12>("Could find memory_space() for RawBuffer") }
    return NULL
```text

错误路径已按字节确认到字面量：`absl::status_internal::MakeErrorImpl<12>`（code 12 = `Internal`），message 为 `"Could find memory_space() for RawBuffer"`，来源于 `pjrt_c_api_raw_buffer_internal.cc`。finder 使用的 `client` 是 wrapper`+0x08` 处*借用的* client — 也就是 `CreateRawAliasOfBuffer` 复制进 raw alias 的同一字段。五个 memory-space 类及其 `kind` 字符串（`tpu_hbm`、`pinned_host`、`unpinned_host`、`device`、cross-pod megascale）记录在 [buffer-and-memory.md](buffer-and-memory.md#pjrt_memory-handle)；本方法只是解析 wrapper，并不分类。

---

### CopyRawHostToDevice（slot 0x38）和 CopyRawDeviceToHost（slot 0x40）

#### 目的

字节粒度的传输接口。每个方法都在 host `void*` 和 device buffer 的一个 `(offset, size)` slice 之间复制 `size` 字节，并返回一个 80 字节 `PJRT_Event`，该 event 会在 DMA 完成时触发。这里**没有元素类型、没有形状，也没有 tiling 转换** — 字节会原样移动。这是它与有类型 `ToHostBuffer`（slot 75）的定义性差异，后者在 transfer 前会通过 `ShapeUtil::DeviceShapeToHostShape` de-tile。

#### 算法

两个方向逐字节互为镜像；只有 vtable 偏移（`+0x28` vs `+0x30`）和 host-pointer 语义（source vs destination）不同。

```c
function PJRT_RawBuffer_CopyRawHostToDevice(args):       // 0xe6f5040  (mirror: 0xe6f5180, vtable+0x30)
    if !ActualStructSizeIsGreaterOrEqual("PJRT_RawBuffer_CopyRawHostToDevice_Args", 39, 56, args.struct_size):
        return new PJRT_Error{ size_status }

    inner  = *(args[+0x10])                               // **(a1+0x10)
    host   = args[+0x18]                                  // host_src (or host_dst for D2H)
    offset = args[+0x20]                                  // device byte offset
    size   = args[+0x28]                                  // byte count

    // single bounce into the shared copy-method holder:
    future = inner.vtable[+0x28].CopyRawHostToDevice(host, offset, size)   // CommonPjRtRawBufferImpl @ 0xf91c640
                                                          // returns PjRtFuture on stack (-0x58)

    // wrap the future as an 80-byte PJRT_Event:
    event = operator new(0x50)
    event[+0x00] = future.async_value (move)             // tsl::AsyncValue*
    event[+0x08..0x18] = profiling closure 0             // AnyInvocable (NULL-policy patched if empty)
    event[+0x28..0x38] = profiling closure 1             // AnyInvocable
    event[+0x48] = 0
    args[+0x30] = event                                  // OUT: PJRT_Event*

    // tear down the temp future: run any non-empty closure dtors,
    // then AsyncValue::Destroy on refcount->0
    return NULL

两个方向的 args 布局：

Offset	字段	方向
`+0x00`	`struct_size`	首先校验
`+0x10`	`buffer`	`PJRT_RawBuffer*` wrapper
`+0x18`	`host`	source（H2D）/ destination（D2H）host pointer
`+0x20`	`offset`	slice 的 device byte offset
`+0x28`	`size`	要传输的字节数
`+0x30`	`event`	OUT：`PJRT_Event*`（`operator new(0x50)`）

vtable+0x28/+0x30 slot 都指向 xla::CommonPjRtRawBufferImpl（共享 copy-method holder），而不是直接指向 TpuRawBuffer — 具体 TPU/CPU 类继承这些 copy 方法。中间层随后调用 AndReturnEvent 变体（TpuRawBuffer::CopyRawHostToDeviceAndReturnEvent @ 0xf8388c0 位于 vtable+0x40，CopyRawDeviceToHostAndReturnEvent @ 0xf838ea0 位于 vtable+0x48），它们会对 slice 做 bounds-check（ValidateSlice @ 0xf837be0）、取 sub-view（SliceBuffer @ 0xf837d80）、用 RAII tpu::WithTransferRequirements 包装 DMA，并通过 tpu::System::TransferToDevice @ 0x1d0afa20 / TransferFromDevice @ 0x1d0b0160 驱动硬件 byte-mover。完整三层 DMA pipeline 和 tpu::TpuPxcDriver byte-mover 记录在 DMA 与 Cross-Host Receive；本页停在 C-ABI wrapper 和第一次 vtable bounce。

陷阱 — (offset, size) slice 的 bounds-check 位于 AndReturnEvent 层内部（ValidateSlice），而不是 C wrapper 中。C wrapper 不会根据 buffer 的 on-device size 对 offset/size 做任何校验 — 它会把它们直接转发给内部方法。重新实现者不能假定 wrapper 会防护越界 slice；防护在下一层，失败会以返回的 event/future 内部错误呈现，而不是 wrapper 同步返回的 PJRT_Error。

函数映射

函数	Addr	角色
`pjrt::PJRT_RawBuffer_CopyRawHostToDevice`	`0xe6f5040`	C wrapper H2D；把 future 包装为 event
`pjrt::PJRT_RawBuffer_CopyRawDeviceToHost`	`0xe6f5180`	C wrapper D2H（镜像）
`xla::CommonPjRtRawBufferImpl::CopyRawHostToDevice`	`0xf91c640`	vtable+0x28；返回 PjRtFuture
`xla::CommonPjRtRawBufferImpl::CopyRawDeviceToHost`	`0xf91c780`	vtable+0x30
`xla::TpuRawBuffer::CopyRawHostToDeviceAndReturnEvent`	`0xf8388c0`	vtable+0x40；bounds-check + DMA
`xla::TpuRawBuffer::CopyRawDeviceToHostAndReturnEvent`	`0xf838ea0`	vtable+0x48
`tpu::System::TransferToDevice`	`0x1d0afa20`	硬件 DMA（host→HBM）
`tpu::System::TransferFromDevice`	`0x1d0b0160`	硬件 DMA（HBM→host）

TpuRawBuffer 后端 Vtable

七个 C wrapper bounce 进入具体 xla::TpuRawBuffer vtable @ 0x2177cfe0（object vptr base 0x2177cff0），该 vtable 从 R_X86_64_RELATIVE relocation 重建而来。xla::CpuRawBuffer（vtable 0x21789af8）具有相同顺序 — 抽象 xla::PjRtRawBuffer base 固定了 slot 布局。这里只复现 C wrapper 实际使用的 slot；完整 vtable 还会继续包含跨 memory-space 路径使用的 copy/slice/async helper。

Vtbl off	Target	方法	使用者
`+0x00`	`0xf83c5a0`	`~TpuRawBuffer()` [complete, D2]	—
`+0x08`	`0xf83c5e0`	`~TpuRawBuffer()` [deleting, D0]	`Destroy`
`+0x10`	`0xf83c640`	`memory_space()` → `this+0x10`	`GetMemorySpace`
`+0x18`	`0xf837b60`	`GetHostPointer()` → 仅 pinned_host，否则 0	`GetHostPointer`
`+0x20`	`0xf838880`	`GetOnDeviceSizeInBytes()` → `TpuBufferBase+0x50`	`GetOnDeviceSizeInBytes`
`+0x28`	`0xf91c640`	`CommonPjRtRawBufferImpl::CopyRawHostToDevice(const void*, l, l)`	`CopyRawHostToDevice`
`+0x30`	`0xf91c780`	`CommonPjRtRawBufferImpl::CopyRawDeviceToHost(void*, l, l)`	`CopyRawDeviceToHost`
`+0x38`	`0xf83c660`	`OpaqueDeviceMemoryDataPointer()` → `TpuBufferBase+0x48`	（typed slot 81）
`+0x40`	`0xf8388c0`	`CopyRawHostToDeviceAndReturnEvent(const void*, l, l)`	copy 中间层
`+0x48`	`0xf838ea0`	`CopyRawDeviceToHostAndReturnEvent(void*, l, l)`	copy 中间层

相关 TpuRawBuffer / TpuBufferBase 字段偏移（从 accessor body 解码）：

text

TpuRawBuffer+0x10 = PjRtMemorySpace*                       (returned by memory_space())
TpuRawBuffer+0x18 = tsl::AsyncValueRef<tpu::TpuBufferBase>  (the device handle; accessors
                    BlockUntilReady then follow the indirect-node chain
                    while (state & 3) node = node->[+0x10])
resolved TpuBufferBase+0x48 = raw HBM device pointer
resolved TpuBufferBase+0x50 = on-device size in bytes
```text

> **注意 —** `memory_space()`（vtable`+0x10`）会直接返回 `this+0x10` — memory space pointer 是 `TpuRawBuffer` 上的缓存字段，而不是计算值。相比之下，device pointer 和 size 需要解析 `+0x18` 处的 `AsyncValueRef<TpuBufferBase>`（阻塞直到 ready 并遍历 indirect chain），这就是为什么 `GetHostPointer`/`OpaqueDeviceMemoryDataPointer` 可能阻塞，而 `GetMemorySpace` 不会。

---

## 所有权和异步模型

### 共享引用计数所有权

RawBuffer wrapper 持有一个 `tsl::RCReference<xla::PjRtRawBuffer>` — wrapper 只是可能有多个共同所有者之一。`Destroy` 递减共享 refcount，并且仅在最后一次释放时删除设备 buffer；`CreateRawAliasOfBuffer` 添加一个共享同一分配的共同所有者。这与有类型 `PJRT_Buffer` wrapper 相反，后者独占拥有其内部 `PjRtBuffer*`，并在 `Destroy` 时直接销毁它。

实际后果是：raw alias 与其 donor *共享* HBM。当 raw alias 仍在读取时，通过有类型接口的 `Delete` 释放 donor 的设备内存，会在设备层面造成 use-after-free，C-ABI 层不会阻止这一点 — 所有权协调是调用方的责任，只由 `RCReference` 上的共享 refcount 调停（它保护的是 C++ 对象生命周期，而不是 eager device-memory `Delete`）。

### PJRT_Event 门控的就绪状态

两个 raw copy 都返回一个**80 字节 `PJRT_Event`**（`operator new(0x50)`），包装一个 `PjRtFuture`：一个 `AsyncValue*`、两个 profiling `AnyInvocable` closure，以及一个清零的尾部。精确字段划分与 [events-and-async.md](events-and-async.md) 上的规范 event 布局一致。因此 readiness 通过轮询（`PJRT_Event_IsReady`，slot 11）、等待（`PJRT_Event_Await`，slot 13）或注册回调（`PJRT_Event_OnReady`，slot 14）处理 — 与有类型 buffer 的 `ReadyEvent`（slot 77）使用相同机制。底层 device-side completion token 是一个 `tpu::TpuEvent`（`tpu::ReadyTpuEvent` @ `0x1d0b62e0`）；中间层通过 `CommonPjRtClient::MakeTrackedReadyFuture` @ `0xf91c2e0` 将其转换为 client-tracked `PjRtFuture`。

> **陷阱 —** 返回的 event 表示*传输完成*，不是 buffer ready。`CopyRawHostToDevice` 返回 NULL（成功）后，字节**尚未**进入 HBM — 调用方必须等待 event，然后才能假定 device buffer 反映了 host data，也才能复用/释放 host source。把同步 NULL 返回当作“transfer done”的重新实现会与 DMA 发生竞态。

---

## 注意事项

- **没有 type/shape safety。** raw 接口信任调用方的 `(offset, size)`。它面向已经知道字节布局的外部消费者 — dlpack import/export、NumPy zero-copy、自定义 device kernel — 而不是通用有类型 I/O。当元素类型和 de-tiling 重要时，请使用有类型 [`PJRT_Buffer`](buffer-and-memory.md) 接口。
- **`CpuRawBuffer` 对等性。** CPU staging 后端类共享 vtable 顺序，因此同样七个方法可用于 CPU-staging raw buffer；只有具体 DMA 路径不同（`CpuRawBuffer::CopyRaw*AndReturnEvent` vs TPU 的 `tpu::System::Transfer*`）。
- **Aliasing vs donation。** `CreateRawAliasOfBuffer` 是一个*零拷贝共同所有权* alias（共享 `RCReference`），不同于有类型接口的 `DonateWithControlDependency`（slot 130），后者会*使 donor 失效*并产生一个 HBM alias donor 的新 buffer。raw alias 会让 donor 完全保持可用。
- **未逐字节追踪（LOW）。** 80 字节 event 内部精确的 `AsyncValue`/closure 字段划分取自同级 event 页面，而不是在这里重新推导；C wrapper 的 `operator new(0x50)` 和 move-in 序列（async value + 两个 `AnyInvocable` closure + 清零的 `+0x48`）已按字节确认，但每个 closure offset 的精确语义角色由 [events-and-async.md](events-and-async.md) 负责。`AndReturnEvent`/`tpu::System::Transfer*` chain 地址已由 symbol 确认，但本页没有重新逐指令追踪 DMA body（HIGH，推迟到 [dma-and-cross-host-recv.md](dma-and-cross-host-recv.md)）。

---

## 交叉引用

- [PJRT Buffer ABI 与内存布局](buffer-and-memory.md) — 有类型/有形状的 `PJRT_Buffer` 同级接口：272 字节独占 wrapper、de-tiling readback，以及定义本页所依赖的对比
- [扩展链](extension-chain.md) — 17-node `PJRT_Api` 扩展链；RawBuffer node 如何链接到 Profiler terminator
- [PJRT API 概览](overview.md) — 140-slot `PJRT_Api` 以及扩展如何挂在 `extension_start` 上
- [API Vtable 重建](api-vtable-reconstruction.md) — 如何从 binary 恢复 `struct_size` 版本化和 vtable 偏移
- [DMA 与 Cross-Host Receive](dma-and-cross-host-recv.md) — `CopyRaw*` 之下的三层 raw DMA pipeline 和 `tpu::TpuPxcDriver` byte-mover
- [Events 与 Async](events-and-async.md) — 每个 raw copy 返回的 80 字节 `PJRT_Event` / `PjRtFuture` 布局
- [剩余扩展](ext-remaining.md) — 其他 12 个 chain extension 以及构造顺序依据
- [Allocator 集成](../runtime/allocator-integration.md) — 支撑 raw buffer 所寻址 HBM 的 StreamExecutor allocator bridge

PJRT RawBuffer 扩展（类型 8） ​

摘要 ​

扩展 Struct（类型 8，80 字节） ​

目的 ​

布局 ​

16 字节 C wrapper ​

函数映射 ​

Destroy（slot 0x20） ​

目的 ​

算法 ​

GetMemorySpace（slot 0x30） ​

目的 ​

算法 ​

函数映射 ​

TpuRawBuffer 后端 Vtable ​

PJRT RawBuffer 扩展（类型 8）

摘要

扩展 Struct（类型 8，80 字节）

目的

布局

16 字节 C wrapper

函数映射

Destroy（slot 0x20）

目的

算法

GetMemorySpace（slot 0x30）

目的

算法

函数映射

TpuRawBuffer 后端 Vtable