DMA 映射与跨主机接收

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

摘要

本页记录两个相关的 PJRT 数据平面接口，它们在普通 host↔device 暂存路径之外移动字节：DMA-map API 和 cross-host receive。

DMA-map API 是一对 main-table 槽位：PJRT_Client_DmaMap 与 PJRT_Client_DmaUnmap。它们固定一段主机虚拟内存，使 TPU DMA 引擎无需中间拷贝即可读写该区域。它们是 libtpu 中对应 cudaHostRegister/cudaHostUnregister 的机制：调用者传入 (void* data, size_t size)，libtpu 将该区域映射到客户端上的每个 TPU host-shared-memory 位置，之后由 DmaUnmap 拆除。不同于 buffer-and-memory.md 上的 buffer 槽位，这些槽位不携带 buffer 对象；它们直接作用于 client 和原始 host 指针。

cross-host receive 接口是 PJRT_CrossHostTransfers 扩展（type 12，struct_size 56，四个方法，由 pjrt::CreateCrossHostTransfersExtension @ 0xf85d660 创建）。它允许一个 host 上的 buffer 由另一个 host 经数据中心网络（DCN）发起的 DMA 填充。模型是描述符握手：接收 host 调用 MakeCrossHostReceiveBuffers 分配空的目标 buffer 并取得不透明的 recv descriptor；它通过带外通道把这些描述符发送给发送 host；发送方调用 Buffer_CopyToRemoteDevice，传入序列化描述符 future，将自己的 buffer 推入接收方的分配中。两个更底层的点对点变体（CrossHostReceiveBuffers / CrossHostSendBuffers）补全该扩展。这是 pipeline-parallel 和 cross-pod buffer 移动之下的 C-ABI 管道；真正承载字节的 device-side ICI/megascale 传输归 ../ici/overview.md 和 ../megascale/overview.md 所有。

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

两个 DmaMap/DmaUnmap args 结构（struct_size 字面量 + {data,size} 偏移）以及每个结构发起的一次 client-vtable 跳转（+0x160 / +0x168）。
DmaMap 会把 host 区域扇出到客户端上所有 TPU host-shared-memory 位置，而不是单个映射。
PJRT_CrossHostTransfers 扩展结构（type 12，size 56）及其四个 fn-ptr 槽位，以及每个槽位的 C wrapper。
接收侧描述符握手：MakeCrossHostReceiveBuffers 如何从 C 构建 xla::Shape，调用 client，将 C notifier callback 适配到 C++，并把返回的 buffer 包装为 typed 接口所用的同一个 272 字节 PJRT_Buffer 对象。
发送侧路径：CopyToRemoteDevice 构建序列化描述符 PjRtFuture<std::string> promise 和 on_done(absl::Status, bool) callback，然后跳转到 inner buffer 的 CopyToRemoteDevice。


DmaMap 槽位	`pjrt::PJRT_Client_DmaMap` @ `0xf860500` — args min 23 / cur 40；inner client vtable `+0x160`
DmaUnmap 槽位	`pjrt::PJRT_Client_DmaUnmap` @ `0xf860580` — args min 25 / cur 32；inner client vtable `+0x168`
DmaMap 后端	`xla::TpuClient::DmaMap(void*, size_t)` @ `0xf80ba80` → per-location `tpu::System::DmaMap` @ `0x1d0b6260`
CrossHostTransfers 扩展	type 12，size 56，4 个方法；存储 `.bss` `0x224c3ad8`；creator `0xf85d660`
Recv 分配	`pjrt::PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers` @ `0xf85c9a0` — args min 44 / cur 104
Remote send	`pjrt::PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice` @ `0xf85ce20` → inner buffer vtable `+0xc8`
Recv-buffer wrapper	标准 272 字节（`0x110`）`PJRT_Buffer` — 与 buffer-and-memory.md 相同的对象

DMA Map（main-table 槽位）

目的

PJRT_Client_DmaMap 注册一段 host 内存区域以供 device 直接 DMA。调用成功后，TPU DMA 引擎可以在 [data, data+size) 中流式读写字节，无需经由 staging buffer 反弹；该区域被固定（其物理页被锁定并赋予 device 可见的 IO 地址）。PJRT_Client_DmaUnmap 执行反向操作。这一对接口用于支持零拷贝 ingest/egress 路径（框架把大型 host 数组常驻内存并反复 DMA），也是 buffer-and-memory.md 中记录的 pinned_host memory-space buffer 可作为 DMA endpoint 使用的前置条件。

QUIRK — 这些是 client 槽位，不是 buffer 或 extension 槽位。它们不接收 PJRT_Buffer。args+0x10 处的句柄是 PJRT_Client，区域是调用者拥有的裸 (void* data, size_t size) host 指针。没有表示该映射的 PJRT 对象；唯一句柄就是原始 host 指针，调用者必须原样传回给 DmaUnmap。丢失该指针就会泄漏 pin。

入口点

text

pjrt::PJRT_Client_DmaMap (0xf860500)            ── slot; validate + bounce
  └─ inner_client.vtable[+0x160] (=xla::TpuClient::DmaMap @ 0xf80ba80)
       └─ for each TpuSharedMemoryLocation on the client's host:
            tpu::System::DmaMap(loc, data, size)  (0x1d0b6260)
              └─ TpuSharedMemory::DmaMap (0x1d4bdac0)
                   └─ {Pxc,Vxc,Jxc}DriverImpl::DmaMapImpl   ── kernel DMA-mapper ioctl
```text

### 算法

```c
function PJRT_Client_DmaMap(args):                       // 0xf860500
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Client_DmaMap_Args", 23, 40, args.struct_size):   // a1+0x00
        return new PJRT_Error{ size_error }              // operator new(8)
    client = *(args + 0x10)                              // PJRT_Client wrapper
    inner  = *client                                     // xla::PjRtClient* (TpuClient)
    data   = *(args + 0x18)                              // host pointer
    size   = *(args + 0x20)                              // byte count
    status = inner.vtable[+0x160](inner, data, size)     // DmaMap(void*, size_t)
    if status == ok: return NULL
    return new PJRT_Error{ status }

function PJRT_Client_DmaUnmap(args):                     // 0xf860580
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Client_DmaUnmap_Args", 25, 32, args.struct_size):
        return new PJRT_Error{ size_error }
    inner  = *(*(args + 0x10))                           // TpuClient
    data   = *(args + 0x18)                              // the SAME host pointer
    status = inner.vtable[+0x168](inner, data)           // DmaUnmap(void*)
    if status == ok: return NULL
    return new PJRT_Error{ status }

这两个 wrapper 是几乎相同的薄 shim：验证 struct_size，两次解引用 client wrapper 到达具体的 xla::TpuClient，读取 args，做一次 virtual call，并把任何非 OK 的 absl::Status 装箱到 operator new(8) 的 PJRT_Error。vtable 偏移（0x160 = 352，0x168 = 360）直接来自每个反编译中的 call qword ptr [rax+160h] / [rax+168h]。

TpuClient::DmaMap 扇出

xla::TpuClient::DmaMap @ 0xf80ba80 是实际工作发生的位置，并且它做了一件受 CUDA 训练的读者可能不会预期的事：它把该区域映射进每个 host-shared-memory 位置，而不是一个。

function TpuClient::DmaMap(this, data, size):            // 0xf80ba80
    // Resolve the tpu::System for this client's host. The System handle is
    // an AsyncValueRef stored at this+81*8 (+256); walk the indirect chain
    // until the value is concrete:
    node = *(this[+0x288] + 0x100)
    while (node.state & 3) != 0:                          // not-yet-resolved bits
        node = node[+0x10]
    sys = node + 0x40                                     // tpu::System*
    host_loc = tpu::System::host_location(sys)            // TpuHostLocation
    shmems   = TpuHostLocation::SharedMemories(host_loc)  // span of TpuSharedMemoryLocation
    result = OkStatus()
    for each loc (56 bytes) in shmems:                    // stride 56
        if loc.is_valid:                                  // (loc+0x30 dword == 0)
            s = tpu::System::DmaMap(sys, loc, data, size) // 0x1d0b6260
            if result is still Ok: result = s             // keep the FIRST status
            else if s is not Ok: StatusRep::Unref(s)      // drop later non-Ok dups
    return result
```text

`TpuHostLocation::SharedMemories` 返回 `TpuSharedMemoryLocation` 记录的 span（每条 56 字节，即 megascale recv 路径使用的同一结构宽度），`DmaMap` 迭代整个 span，把 `[data,size)` 映射到每个有效位置。对每个位置，它调用 `tpu::System::DmaMap(TpuSharedMemoryLocation, void*, size_t)` @ `0x1d0b6260`，后者经 `tpu::TpuSharedMemory::DmaMap` @ `0x1d4bdac0` 路由到具体 driver 的 `DmaMapImpl`（`Pxc` physical / `Vxc` virtual / `Jxc` — `0xe804460` / `0x1d12d860` / `0xe73ad60`），最终到达 `asic_sw::driver::KernelDmaMapper::MapMemory` (`0xe896dc0`) 背后的 kernel DMA-mapper ioctl。

> **GOTCHA —** status 聚合是**第一个胜出**，不是全有或全无。循环保留*第一个*位置的 status，并 `Unref` 每个后续 status，包括后续*失败*。如果 location 0 映射成功但 location 3 失败，`DmaMap` 返回 OK，部分映射静默存在。需要 all-locations-or-nothing 语义的重新实现者必须自行增加 rollback；该二进制不会在后续失败时回滚早先映射，也完全不会暴露后续失败。

### 函数映射

| Function | Addr | Role |
|---|---|---|
| `pjrt::PJRT_Client_DmaMap` | `0xf860500` | slot wrapper；验证 + 跳转 vtable `+0x160` |
| `pjrt::PJRT_Client_DmaUnmap` | `0xf860580` | slot wrapper；验证 + 跳转 vtable `+0x168` |
| `xla::TpuClient::DmaMap(void*, size_t)` | `0xf80ba80` | 扇出到所有 host shared-memory 位置 |
| `xla::TpuClient::DmaUnmap(void*)` | `0xf80bb80` | unmap 扇出（镜像） |
| `xla::PjRtClient::DmaMap` (base default) | `0xf8fff60` | abstract-base 默认实现（除非 override，否则返回 Unimplemented） |
| `xla::MegaScalePjRtClient::DmaMap` | `0xe6eda80` | cross-pod client override |
| `tpu::System::DmaMap(TpuSharedMemoryLocation, void*, size_t)` | `0x1d0b6260` | per-location map |
| `tpu::System::DmaUnmap(TpuSharedMemoryLocation, void*)` | `0x1d0b62a0` | per-location unmap |
| `tpu::TpuSharedMemory::DmaMap(Span<const uint8>)` | `0x1d4bdac0` | shared-memory map 入口 |
| `tpu::TpuSharedMemoryPxcDriverImpl::DmaMapImpl` | `0xe804460` | physical-driver kernel map |
| `asic_sw::driver::KernelDmaMapper::MapMemory` | `0xe896dc0` | kernel DMA-mapper ioctl |

> **NOTE —** 两个槽位的精确 `args` 字段数与字节大小不同（`DmaMap` min 23 fields / 40 bytes；`DmaUnmap` min 25 fields / 32 bytes）。传给 `ActualStructSizeIsGreaterOrEqual` 的 "min" 值是 wrapper 编译时对应的具名字段源码行数，"cur" 值是当前 `struct_size` 的字节数；二者*不是*同一量纲，不应被调和。二者都逐字来自 validation call 中的 `mov esi/edx` 立即数。关于这种双数字 versioning 如何贯穿整个 API，见 [api-vtable-reconstruction.md](api-vtable-reconstruction.md)。

---

## CrossHostTransfers 扩展（type 12）

### 目的

cross-host receive 接口作为 chain extension 提供，而不是 main-table 槽位，因为它是可选的：单 host plugin 不会公布它（consumer 通过遍历 type id 12 的 [extension chain](extension-chain.md) 来发现它）。存在时，它提供四个 C 函数：分配 receive buffer（`MakeCrossHostReceiveBuffers`）、把本地 buffer 推送到远端分配（`Buffer_CopyToRemoteDevice`），以及两个更底层的点对点 primitive（`CrossHostReceiveBuffers`、`CrossHostSendBuffers`）。四个函数都是 generic-XLA wrapper（标准 `PjRtClient`/`PjRtBuffer` cross-host API），背后由 TPU 实现支撑；TPU/megascale 专门化位于 C-ABI 之下的 `xla::TpuClient` 和 `xla::MegaScalePjRtBuffer`。

### 扩展结构

creator 是一个扁平表初始化器：无分支、无分配、以 `ret` 结束，正是每个 libtpu extension creator 遵循的模式（见 [ext-remaining.md](ext-remaining.md)）。

```c
struct PJRT_CrossHostTransfers_Extension {       // struct_size 56 (0x38); .bss @ 0x224c3ad8
    PJRT_Extension_Base base;                    // +0x00 size=56, +0x08 type=12, +0x10 next
    /* +0x18 */ PJRT_Error* (*Client_MakeCrossHostReceiveBuffers)(PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers_Args*);
    /* +0x20 */ PJRT_Error* (*Buffer_CopyToRemoteDevice)        (PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice_Args*);
    /* +0x28 */ PJRT_Error* (*Client_CrossHostReceiveBuffers)   (PJRT_Transfers_PJRT_Client_CrossHostReceiveBuffers_Args*);
    /* +0x30 */ PJRT_Error* (*Client_CrossHostSendBuffers)      (PJRT_Transfers_PJRT_Client_CrossHostSendBuffers_Args*);
};

四个偏移、大小以及 type=12 / struct_size=56 都直接来自 CreateCrossHostTransfersExtension @ 0xf85d660 中的 mov store：

*(a1+0x00) = 56;  *(a1+0x08) = 12;  *(a1+0x10) = next;          // header
*(a1+0x18) = PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers;
*(a1+0x20) = PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice;
*(a1+0x28) = PJRT_Transfers_PJRT_Client_CrossHostReceiveBuffers;
*(a1+0x30) = PJRT_Transfers_PJRT_Client_CrossHostSendBuffers;
```text

### 方法映射

| Off | Method | C symbol | Addr | Backing (inner client/buffer vtable) |
|---|---|---|---|---|
| `+0x18` | Client_MakeCrossHostReceiveBuffers | `PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers` | `0xf85c9a0` | client vtable `+0x150` (`MakeCrossHostReceiveBuffers`) |
| `+0x20` | Buffer_CopyToRemoteDevice | `PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice` | `0xf85ce20` | buffer vtable `+0xc8` (`CopyToRemoteDevice`) |
| `+0x28` | Client_CrossHostReceiveBuffers | `PJRT_Transfers_PJRT_Client_CrossHostReceiveBuffers` | `0xf85bba0` | 点对点 recv |
| `+0x30` | Client_CrossHostSendBuffers | `PJRT_Transfers_PJRT_Client_CrossHostSendBuffers` | `0xf85c2a0` | 点对点 send |

> **NOTE —** 标准 upstream PJRT C API 将该扩展命名为 `PJRT_CrossHostTransfers`；在 libtpu 中，实现符号带有 `pjrt::PJRT_Transfers_*` 前缀（namespace short-name），args 结构为 `PJRT_Transfers_*_Args`。它们是同一个接口；符号前缀是 libtpu 的命名产物，这一点由 `MakeCrossHostReceiveBuffers` validator string `"PJRT_Client_MakeCrossHostReceiveBuffers_Args"` 确认。

---

## MakeCrossHostReceiveBuffers（扩展 +0x18）

### 目的

在**接收** host 上调用。给定一组 shapes 和一个目标 device，它分配同等数量的空 device buffer，并通过 notifier callback 交回每个 buffer 的不透明 *recv descriptor*。这些 descriptor 是 rendezvous token：接收方通过任意带外通道把它们发送给发送 host，发送方在之后的 `CopyToRemoteDevice` 中用它们寻址接收方的分配。该函数同步返回已分配的 `PJRT_Buffer*` 数组；descriptor 通过 C notifier 异步到达。

### 算法

```c
function PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers(args):  // 0xf85c9a0
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Client_MakeCrossHostReceiveBuffers_Args", 44, 104, args.struct_size):
        return new PJRT_Error{ size_error }

    // 1. Rebuild an xla::Shape[] (320 bytes each) from the C shape arrays.
    n = args[+0x18]                                    // num_shapes (a1[3])
    shapes = vector<xla::Shape>(); shapes.reserve(n)
    for i in 0..n:
        // element_type @ a1[6]+4*i, dims @ a1[5]+8*i, num_dims @ a1[4]+8*i,
        // layout/extra @ a1[7]+8*i
        s = pjrt::BuildXlaShapeFromC(type[i], dims[i], num_dims[i], layout[i])
        if s is error: cleanup; return new PJRT_Error{ s }
        shapes.push_back(s)

    // 2. Adapt the C notifier (CrossHostRecvNotifierInfo @ a1+0x48) to a
    //    std::function<void(StatusOr<PjRtCrossHostRecvState>)>:
    notifier = CCrossHostRecvNotifierToCpp(args.notifier_info)   // 0xf85d6e0 thunk

    // 3. Call the inner client. Client = a1[2] (args+0x10, double-deref);
    //    device = a1[8] (args+0x40, single-deref). Bounce +0x150:
    inner  = *(*(args + 0x10))                          // TpuClient (client wrapper @ a1[2])
    device = *(args + 0x40)                             // PjRtDevice (a1[8])
    status_or_recv = inner.vtable[+0x150](inner, shapes.data, shapes.size,
                                          device, notifier)      // StatusOr<vector<PjRtBuffer*>>
    if status_or_recv is error: cleanup; return new PJRT_Error{ status }

    // 4. Wrap each returned PjRtBuffer* in a fresh 272-byte PJRT_Buffer.
    buffers = status_or_recv.value                     // vector<unique_ptr<PjRtBuffer>>
    args[+0x60] = buffers.size                         // num_buffers OUT (a1[12])
    for i in 0..buffers.size:
        w = operator new(0x110)                        // PJRT_Buffer wrapper
        w[+0x00] = buffers[i]                          // inner buffer (moved out)
        w[+0x08] = args.client                         // borrowed client (a1[2])
        zero w[+0x10], +0xa8, +0xb0, +0xc8, +0xd0, +0xe8   // ctor flags
        zero w[+0xf0 .. +0x110]                        // external-ref list
        args.buffers[i] = w                            // a1[11] OUT array
    return NULL

两个事实固定了这一点。第一，每个 buffer wrapper 以 operator new(0x110) 分配，并且在与 buffer-and-memory.md 上 BufferFromHostBuffer 构造器完全相同的字节偏移处清零（+0x10、+0xa8、+0xb0、+0xc8、+0xd0、+0xe8，以及 +0xf0..+0x110 external-ref 区域），因此 cross-host receive buffer 一经分配就与普通 device buffer 无法区分；C consumer 看到相同的 PJRT_Buffer 对象并使用其上的相同槽位。第二，每个 xla::Shape 以 320 字节 stride 构造（320 * v7 allocation，Shape::~Shape 以 -320 遍历），这是该 build 中标准的 xla::Shape 宽度。

QUIRK — 有两个无关的 "make receive" 入口点，不能混淆。MakeCrossHostReceiveBuffers（扩展 +0x18）是描述符握手形式：它分配 buffer 并通过 notifier 产生 recv descriptor。CrossHostReceiveBuffers（扩展 +0x28）是更底层的点对点形式，接收到已知 peer。描述符形式是 pipeline-parallel JAX 使用的接口；点对点形式是其下方的构建块。

C 到 C++ notifier 适配器

recv-descriptor 交付是一座 callback 桥。C 调用者在 args+0x48 提供 PJRT_Transfers_CrossHostRecvNotifierInfo（函数指针 + user data）；适配器 pjrt::(anon)::CCrossHostRecvNotifierToCpp 体现为 0xf85d6e0 处的 std::function __call_func invoker（__policy_func<void(StatusOr<PjRtCrossHostRecvState>)>），其函数体把 C++ PjRtCrossHostRecvState（descriptor set）封送回 C 的 PJRT_Transfers_CrossHostRecvNotifierInfo callback。0xf85dbc0（C-callback void(const char*, size_t, PJRT_Error_Code, …) 形态）和 0xf85de00（void(absl::Status) 形态）的嵌套 lambda __call_func invoker 执行内部封送步骤；对应的 std::function clone/destroy 机制位于 __large_clone @ 0xf85de80 和 __large_destroy @ 0xf85dee0。当 inner client 收集到 recv descriptor（底层 ICI/megascale endpoint 建立之后）时，它调用 C++ function，后者用接收方必须发送给发送方的 descriptor 调用原始 C callback。

函数映射

Function	Addr	Role
`pjrt::PJRT_Transfers_PJRT_Client_MakeCrossHostReceiveBuffers`	`0xf85c9a0`	C wrapper；构建 shape、调用、包装 buffer
`pjrt::BuildXlaShapeFromC`	`0xf8a59e0`	C shape tuple → `xla::Shape`（320 字节）
`pjrt::(anon)::CCrossHostRecvNotifierToCpp`	`0xf85d6e0`	C notifier → `std::function` adapter
`xla::TpuClient::MakeCrossHostReceiveBuffers`	`0xf808e60`	TPU inner impl（client vtable `+0x150`）
`xla::MegaScalePjRtClient::MakeCrossHostReceiveBuffers`	`0xe6ed9e0`	cross-pod inner impl

Buffer_CopyToRemoteDevice（扩展 +0x20）

目的

在发送 host 上调用。它把本地 buffer 内容推送到接收方先前由 MakeCrossHostReceiveBuffers 创建的远端分配。发送由序列化描述符（接收方传来的字节）索引，并以 PjRtFuture<std::string> 交付，因此发送可在描述符到达前发起，并在描述符到达时触发。调用者还提供 on_done(absl::Status, bool) callback 表示完成（bool 是 sent/aborted 标志）。

算法

function PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice(args):   // 0xf85ce20
    // 1. Build a PjRtFuture<std::string> promise for the serialized descriptor.
    promise = PromiseMaker<std::string>::Make()                 // tsl::internal

    // 2. If the caller passed a ready serialized descriptor (args+0x18 != NULL),
    //    fulfil the promise immediately by copying the descriptor string
    //    (data @ *(args+0x20), len @ *(args+0x28)) into it via small/large
    //    string emplace; otherwise the promise stays pending until AndThen fires.
    serialized = args[+0x18]
    if serialized != NULL:
        str = copy_string(*(args+0x20) /*data*/, *(args+0x28) /*len*/)
        promise.emplace(StatusOr<string>{ str })
        free the caller's descriptor buffers

    // 3. Resolve the inner buffer and the on_done callback (args+0x30 region),
    //    then bounce CopyToRemoteDevice:
    inner = *(*(args + 0x10))                                   // CommonPjRtBufferImpl
    on_done = wrap args.on_done as std::function<void(absl::Status,bool)>
    inner.vtable[+0xc8].CopyToRemoteDevice(promise.future, on_done)   // 0xf91c8c0
    return NULL   (errors delivered through on_done, not the return)
```text

inner `xla::CommonPjRtBufferImpl::CopyToRemoteDevice` @ `0xf91c8c0` 完成真正工作：

```c
function CommonPjRtBufferImpl::CopyToRemoteDevice(this, desc_future, on_done):  // 0xf91c8c0
    dev_event   = this.vtable[+0x68](this)                           // +104, device-event obj
    raw_hold    = CommonPjRtBuffer::AcquireScopedRawBuffer(this)      // pin the source bytes
    if raw_hold is error:
        on_done(raw_hold.status, false); return                      // on_done @ a3+0x10
    transfer_mgr = this.vtable[+0x58](this)                          // +88, the transfer manager
    // bounce the DEVICE-EVENT object's +0x260 slot to issue the actual DCN send;
    // transfer_mgr is passed as the first argument, not the receiver:
    dev_event.vtable[+0x260](dev_event, transfer_mgr, raw_buffer, desc_future, on_done, ...)

AcquireScopedRawBuffer 在发送期间固定 source buffer 的 device 字节（与 RawBuffer extension 使用的 external-reference 机制相同），因此 source 不能在飞行中被 Delete。最终跳转位于 device-event 对象的 vtable 槽位 +0x260（608），把 transfer manager（来自 buffer vtable +0x58）作为第一个实参传入，同时传入 raw buffer、descriptor future 和 on_done callback；该调用调度经 DCN 的 cross-host DMA。（在 raw-buffer-acquisition 失败路径中，wrapper 改为通过 a3+0x10 处的 callback 对象直接调用 on_done，且 bool = false。）实际 wire transport，即 ICI link、megascale routing、DMA descriptor 格式，归 ../ici/dma-descriptor.md 和 ../megascale/overview.md 所有；本页止于 C-ABI 和 buffer-side pin。

GOTCHA — CopyToRemoteDevice 只通过 on_done 报告错误，绝不通过返回值报告（C wrapper 在 send 被调度后返回 NULL，即使 send 之后失败）。重新实现者如果检查返回的 PJRT_Error* 来判断 transfer 成功，会漏掉所有运行时失败。on_done 的第二个 bool 参数区分 completed send 与 aborted send；absl::Status 携带失败。请把完成逻辑接到 callback，而不是调用返回值。
QUIRK — descriptor 以 PjRtFuture<std::string> 承载，而不是结构体。接收方产生的序列化 recv-descriptor 字节对 libtpu 的 C-ABI 是不透明的；它们只是发送方转发的字符串。这是有意的：它允许 descriptor 格式演进（它编码 ICI/megascale endpoint 地址）而不改变 C API。不要试图在 PJRT 层解析它。

函数映射

Function	Addr	Role
`pjrt::PJRT_Transfers_PJRT_Buffer_CopyToRemoteDevice`	`0xf85ce20`	C wrapper；descriptor promise + on_done bridge
`tsl::internal::PromiseMaker<std::string>::Make`	(inlined)	构建 serialized-descriptor future
`xla::CommonPjRtBufferImpl::CopyToRemoteDevice`	`0xf91c8c0`	buffer vtable `+0xc8`；pin + 调度 DCN send
`xla::CommonPjRtBuffer::AcquireScopedRawBuffer`	(called)	为 transfer 固定 source device 字节
`xla::MegaScalePjRtBuffer::CopyToRemoteDevice`	`0xe6eb100`	cross-pod buffer override
`xla::TfPjRtBuffer::CopyToRemoteDevice`	`0x10852700`	TF-wrapper override（non-TPU client）
`std::function<void(absl::Status,bool)>::__call_func<...$_2>`	`0xf85ec60`	on_done C-callback policy thunk

点对点发送 / 接收（扩展 +0x28 / +0x30）

Client_CrossHostReceiveBuffers（0xf85bba0）和 Client_CrossHostSendBuffers（0xf85c2a0）是描述符握手之下的更底层 primitive。它们在显式命名的 peer 之间移动 buffer，不经过 MakeCrossHostReceiveBuffers 描述符往返；peer set 由 args 提供（在 backing symbol xla::PjRtClient::CrossHostSendBuffers(Span<const PjRtBuffer*>, Span<const GlobalDeviceId>, ...) @ 0xe6edac0 中可见，是 GlobalDeviceId 的 span）。它们是描述符形式组合出的构建块；JAX 的 pipeline-parallel 路径使用描述符形式，而 collective-style bulk movement 可以直接使用这些接口。

这两个函数除了 C wrapper 和 backing symbol 外未逐字节追踪；args-struct 字段布局未确认（symbol/backing identity 标为 HIGH，精确偏移标为 LOW）。重新实现者应对照 upstream pjrt_c_api.h 的 PJRT_Transfers_PJRT_Client_CrossHost{Send,Receive}Buffers_Args 交叉检查 args 结构。

Function	Addr	Backing
`pjrt::PJRT_Transfers_PJRT_Client_CrossHostReceiveBuffers`	`0xf85bba0`	`xla::PjRtClient::CrossHostReceiveBuffers` @ `0xe6edb00`
`pjrt::PJRT_Transfers_PJRT_Client_CrossHostSendBuffers`	`0xf85c2a0`	`xla::PjRtClient::CrossHostSendBuffers` @ `0xe6edac0`

两个接口的关系

DmaMap 和 cross-host receive 是独立 API，解决相邻问题；重新实现者应保持二者区分：

DmaMap 固定 host 内存，使本地 TPU DMA 引擎可直接访问。它是 client-wide、无 buffer 的操作；工作单元是 host 虚拟地址范围。它支撑 pinned_host memory space（见 buffer-and-memory.md）和零拷贝 host ingest/egress。
Cross-host receive 经 DCN 在不同 host 之间移动 device 字节。工作单元是 PJRT_Buffer，传输是 ICI/megascale，而不是本地 DMA-mapper ioctl。

二者可以组合：cross-host receive 最终把字节落入 device buffer，而 pinned host 区域可以作为此类 transfer 的 host 侧 staging endpoint；但二者互不依赖，并且经过完全不同的后端栈（KernelDmaMapper vs. transfer manager + ICI/megascale）。

	DmaMap / DmaUnmap	CrossHostTransfers (type 12)
API 形式	main-table 槽位（始终存在）	chain extension（可选，type id 12）
句柄	`PJRT_Client` + raw `(void*, size_t)`	`PJRT_Client` / `PJRT_Buffer`
范围	本地 host ↔ 本地 TPU DMA engine	host ↔ remote host over DCN
单元	host VA range	device buffer + recv descriptor
Async	同步（返回 mapped status）	descriptor future + `on_done` callback
后端	`TpuClient::DmaMap` → `KernelDmaMapper`	inner `CopyToRemoteDevice` / `MakeCrossHostReceiveBuffers` → ICI/megascale

交叉引用

PJRT API 概览 — 140 槽 PJRT_Api 和这些接口接入的 extension-chain 模型
API Vtable 重建 — 双数字 struct_size versioning 以及如何恢复 slot/vtable 偏移
Buffer ABI 与内存布局 — recv buffer 复用的 272 字节 PJRT_Buffer，以及 DmaMap 支撑的 pinned_host memory space
RawBuffer 扩展（type 8） — 非类型化 raw-byte 接口和 CopyToRemoteDevice 依赖的 AcquireScopedRawBuffer pin
Extension Chain — consumer 如何发现 type-12 CrossHostTransfers node
Remaining Extensions — flat-creator 模式和 type-12 chain entry 摘要
Client 与 Device — DmaMap 和 MakeCrossHostReceiveBuffers 解引用的 PJRT_Client wrapper
Collectives 与 Communicator — 进程内 collectives 接口；与跨 host 点对点移动对比
Host Callbacks — execute-time transfer 的 host-side send/recv rendezvous
ICI 概览 — 承载 cross-host 字节的 device-side inter-chip transport
ICI DMA Descriptor — CopyToRemoteDevice 之下的 on-wire descriptor 格式
Megascale 概览 — 支撑 MegaScale client/buffer override 的 cross-pod DCN runtime

DMA 映射与跨主机接收 ​

摘要 ​

DMA Map（main-table 槽位） ​

目的 ​

入口点 ​

TpuClient::DmaMap 扇出 ​

C 到 C++ notifier 适配器 ​

函数映射 ​

Buffer_CopyToRemoteDevice（扩展 +0x20） ​

目的 ​

算法 ​

函数映射 ​

点对点发送 / 接收（扩展 +0x28 / +0x30） ​

两个接口的关系 ​

交叉引用 ​

DMA 映射与跨主机接收

摘要

DMA Map（main-table 槽位）

目的

入口点

TpuClient::DmaMap 扇出

C 到 C++ notifier 适配器

函数映射

Buffer_CopyToRemoteDevice（扩展 +0x20）

目的

算法

函数映射

点对点发送 / 接收（扩展 +0x28 / +0x30）

两个接口的关系

交叉引用