缓冲区捐赠与别名

本页中的所有地址、结构体偏移、源文件/行号引用和符号名均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d）。该镜像未被 strip；反混淆后的 C++ 符号名和嵌入的 CHECK 字符串源码路径均按原文引用。其他版本会有所不同。

摘要

一个会就地更新张量的 TPU 计算，例如优化器步骤后的权重、decode token 后的 KV cache，不能在每次调用时都为新的 HBM 分配和拷贝付费。XLA 的机制是输入/输出别名：在编译时，程序声明输出 o 将复用输入 p 的 HBM；在运行时，runtime 不为 o 分配新缓冲区，而是把输入 p 的缓冲区交给可执行程序并允许程序覆盖它。为保证安全，调用方必须承诺调用后不再读取输入 p，这个承诺就是缓冲区捐赠，通过每次执行的 ExecuteOptions::non_donatable_input_indices 表达（索引不在该集合中即表示捐赠）。本页记录这个三段式契约：编译时的 xla::HloInputOutputAliasConfig（哪个输出复用哪个输入，以及复用是强制还是可选）、PJRT 缓冲区上的运行时捐赠握手（ScopedHold → AcquireDonation → ConfirmDonation），以及运行时兑现点 tfrt::tpu::AllocateOutputBuffersWithInputReuse (0xf7ba9a0)，它遍历输出 tuple，并对每个叶子要么复用已捐赠输入的 tpu::TpuBuffer，要么重新分配。

熟悉 GPU 上 XLA 的读者可以先保留一个类比，然后立刻把它复杂化。在 GPU 上，捐赠同样进入 HloInputOutputAliasConfig；TPU 的差异在于复用的是什么。被复用的对象是位于 TpuSharedMemoryLocation 的 tpu::TpuBuffer，也就是 allocator 分配出的 HBM 偏移，内部持有 tiled on-device layout 中的字节。因此，复用意味着输出的带 padding 的设备 shape（§buffer layout §2）必须与被捐赠输入按字节兼容，并且被捐赠的缓冲区不能是 compactor 可以自由重定位的对象。alias config 是编译后程序载荷中 ComputationLayout 和 CompilerMetadata 的同级数据（serialization），因此重新加载的程序会在无需重新编译的情况下重放同一套别名关系。

本页负责说明捐赠机制 + input_output_alias 配置 + 运行时别名兑现。它不复述支撑被捐赠和新分配缓冲区的 HBM free-list allocator（见 hbm-allocator.md）、两种 shape 必须一致的 tiled 字节布局（见 tpu-buffer-layout.md）、承载该配置的可执行序列化容器（见 ../compiler/tpu-program-serialization.md），或位于 ScopedHold 之上的 PJRT_Buffer 外部引用计数 ABI（见 ../pjrt/buffer-and-memory.md）。

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

配置 — xla::HloInputOutputAliasConfig：以输出 shape index 为键的 Alias{parameter_number, parameter_index, AliasKind} 映射；AliasKind ∈ {kMayAlias, kMustAlias}（proto MAY_ALIAS/MUST_ALIAS）。
编译时构建 — OptimizeInputOutputBufferAlias::Build (0x164d9b00) 根据参数/结果 shape 兼容性合成配置和 HloBufferDonorConfig；前端通过 XlaBuilder::SetUpAlias 声明显式别名。
捐赠标志 — 每次执行的 ExecuteOptions::non_donatable_input_indices；当且仅当输入不在该集合中时，它被捐赠。不在捐赠集合中的 kMustAlias 输入会触发运行时 fatal。
运行时握手 — ScopedHold(kDonation) → AcquireDonation → ConfirmDonation：输入的 TrackedTpuDeviceBuffer 所有权从 donor PjRtBuffer 转移到 execute，并将 donor 标记为 invalidated。
运行时兑现 — tfrt::tpu::AllocateOutputBuffersWithInputReuse (0xf7ba9a0)：对每个输出叶子执行 GetAliasedParameter(output_index) → 若存在别名，则复用已捐赠输入的 TpuBuffer；否则执行 AllocateTpuBufferWithRetry。


配置对象	`xla::HloInputOutputAliasConfig`（`Alias` 的输出 shape tree）；`GetAliasedParameter` @ `0x1e580200`, `GetAliasedOutput` @ `0x1e5800a0`, `SetUpAlias` @ `0x1e57e140`
别名种类	`AliasKind::kMayAlias` / `kMustAlias`; proto tokens `MAY_ALIAS` @ `0xc18c23a`, `MUST_ALIAS` @ `0xc18c249`
前端声明	`xla::XlaBuilder::SetUpAlias` @ `0xfb21220`; `PopulateInputOutputAliasAndBufferDonor`（发出 `HloModuleProto`）
编译时 pass	`xla::OptimizeInputOutputBufferAlias::Build` @ `0x164d9b00`（构建配置和 `HloBufferDonorConfig`）
必须捐赠集合	`xla::ComputeParametersThatMustBeDonated(HloInputOutputAliasConfig&, int, bool)` @ `0x1d7f4700` (`xla/pjrt/utils.cc:824`)
捐赠标志	`ExecuteOptions::non_donatable_input_indices`（diagnostic @ `0x87a48c6`; proto field-name @ `0xbf85686`）
捐赠 hold	`CommonPjRtBuffer::ScopedHold::AcquireDonation` @ `0xf93d600`; `ConfirmDonation` @ `0xf93dca0`; `DropDonationHold` @ `0xf93d900`
tracked-buffer 确认	`xla::TrackedTpuDeviceBuffer::ConfirmDonation` @ `0xf840660` (`tracked_tpu_device_buffer.cc:88`)
运行时兑现（TPU）	`tfrt::tpu::AllocateOutputBuffersWithInputReuse` @ `0xf7ba9a0`
运行时兑现（generic）	`xla::CommonPjRtClient::AllocateOutputBuffersWithInputReuse` @ `0xf91ec20`
dispatch 规划器	`xla::InferDispatchInfo(CommonPjRtClient*, ComputationLayout&, HloInputOutputAliasConfig&, …)` @ `0xf90cb40`
配置恢复	`HloInputOutputAliasConfig::CreateFromProto(Shape, HloInputOutputAliasProto&)`
置信度	CONFIRMED（byte-anchored），除非某行或 callout 另有说明

1. `HloInputOutputAliasConfig`

目的

整个子系统都围绕一个对象：xla::HloInputOutputAliasConfig。它回答 runtime 需要的唯一查询：“对这个输出叶子，它别名到哪个输入参数（如果有），以及该别名是否强制？”，查询接口是 GetAliasedParameter。重新实现者必须复现这个对象，因为它是编译后程序载荷中携带、并在每次 execute 时逐字消费的值；runtime 从不重新计算别名关系，而是重放冻结的配置。

数据模型

该配置包装程序的结果 shape，并为每个带别名的叶子打上一个 Alias 记录。内部结构是输出 shape tree 上的 xla::internal::IndexTable（已确认：GetAliasedParameter 访问 xla::internal::IndexTable::GetEntry((char*)a2 + 80, …)，并在 hlo_input_output_alias_config.cc:175 CHECK ShapeUtil::IndexIsValid(alias_.shape(), output_index)）。每个条目是：

struct xla::HloInputOutputAliasConfig::Alias {   // 每个带别名的输出叶子一个
    int64_t   parameter_number;   // 哪个 execute 输入（扁平参数索引）
    ShapeIndex parameter_index;   // 该输入 tuple tree 内的叶子
    AliasKind kind;               // kMayAlias | kMustAlias
};

enum AliasKind {                   // proto HloInputOutputAliasProto.Kind
    kMayAlias,   // MAY_ALIAS  — 若已捐赠则复用；否则分配 + 拷贝
    kMustAlias,  // MUST_ALIAS — 输入必须被捐赠；未捐赠即 fatal
};
```text

该配置有两种 `const` 查询：

| 查询 | 符号 | 地址 | 返回 |
|---|---|---|---|
| 输出 → 输入 | `GetAliasedParameter(output_index)` | `0x1e580200` | `optional<Alias>`（空 ⇒ 输出新分配） |
| 输入 → 输出 | `GetAliasedOutput(param_number, param_index)` | `0x1e5800a0` | `optional<ShapeIndex>` |
| 添加别名 | `SetUpAlias(output_index, param_number, param_index, kind)` | `0x1e57e140` | 修改表 |

> **NOTE —** 两个 `AliasKind` 值是安全模型的主干。`kMayAlias` 是*性能提示*：如果调用方捐赠该输入，就复用它；否则 runtime 分配新的输出并拷贝。`kMustAlias` 是*要求*：程序在 lower 时假定了就地更新（没有 copy-in fallback），因此如果匹配输入未被捐赠，runtime 会 fatal。字面 proto enum 字符串 `MAY_ALIAS` (`0xc18c23a`) 和 `MUST_ALIAS` (`0xc18c249`) 在 `.rodata` 中相邻（位于序列化的 `HloInputOutputAliasProto` enum descriptor 中），diagnostic `print-must-aliases` (`0x855e010`) 只转储强制集合。

### 配置从何而来

有两条路径会填充同一个配置对象：

1. **显式、由前端声明。** 前端（JAX `donate_argnums`、TF `XlaCallModule` aliasing）调用 `xla::XlaBuilder::SetUpAlias(output_index, param, param_index, kind)` (`0xfb21220`)。构建 HLO module 时，`XlaBuilder::PopulateInputOutputAliasAndBufferDonor` 将这些信息序列化进 `HloModuleProto`（以及并行的 `HloBufferDonorConfig`，用于 donate-but-don't-alias 参数）。MLIR ingestion 通过 `xla::ImportInputOutputAlias` / `ConvertInputOutputAlias` 镜像这一过程。
2. **编译器合成。** `xla::OptimizeInputOutputBufferAlias::Build(Span<const Shape> param_shapes, Shape& result_shape, HloInputOutputAliasConfig* out, HloBufferDonorConfig* donors)` (`0x164d9b00`) 通过匹配可捐赠参数叶子与字节大小相同的结果叶子，贪心发现别名机会（反编译显示 `DonorEntry`/`DoneeEntry` 向量经 `__stable_sort` 排序后匹配）。这个 pass 将“此参数可捐赠”转化为“此输出复用它”。

> **GOTCHA —** 配置按**输出** tuple 做 shape-indexed，而不是按输入。`GetAliasedParameter(output_index)` 是运行时热查询（兑现循环对每个输出叶子调用一次）。反向的 `GetAliasedOutput` 用于捐赠规划侧（给定一个已捐赠输入，是否有东西复用它？）。如果重新实现者用输入参数作为表的 key，会让兑现循环从每输出 O(1) 退化为每输出 O(inputs)。

### 函数映射

| 函数 | 地址 | 作用 |
|---|---|---|
| `HloInputOutputAliasConfig::GetAliasedParameter` | `0x1e580200` | 输出叶子 → `optional<Alias>`（运行时查询） |
| `HloInputOutputAliasConfig::GetAliasedOutput` | `0x1e5800a0` | 输入叶子 → `optional<ShapeIndex>` |
| `HloInputOutputAliasConfig::SetUpAlias` | `0x1e57e140` | 记录一个别名条目 |
| `HloInputOutputAliasConfig::ForEachAliasWithStatus` | （由 `0x1d7f4700` 使用） | 遍历所有条目，可失败 |
| `HloInputOutputAliasConfig::CreateFromProto` | (`xla` core) | 从 `HloInputOutputAliasProto` 重新水合 |
| `XlaBuilder::SetUpAlias` | `0xfb21220` | 前端别名声明 |
| `XlaBuilder::PopulateInputOutputAliasAndBufferDonor` | (`xla` core) | 将 aliases + donors 序列化进 `HloModuleProto` |
| `OptimizeInputOutputBufferAlias::Build` | `0x164d9b00` | 编译器合成 alias/donor 发现 |
| `ImportInputOutputAlias` / `ConvertInputOutputAlias` | (`xla` core) | MLIR ↔ 配置转换 |

---

## 2. 捐赠标志和必须捐赠集合

### 目的

别名是编译时的*计划*；捐赠是运行时的*许可*。计划说“输出 *o* 可以复用输入 *p*”；许可说“在这次具体调用中，调用方之后不会再碰输入 *p*，所以你可以覆盖它”。本节说明许可如何表达，以及如何与计划对账；不匹配变成 fatal error 的地方就在这个对账过程。

### 捐赠标志 — `non_donatable_input_indices`

PJRT 在 execute 调用中不携带每缓冲区的 “donate me” 位。它携带的是*反向*集合：`xla::ExecuteOptions::non_donatable_input_indices`，即调用方想要**保留**的输入索引集合。契约是：

> 当且仅当输入索引*不在* `ExecuteOptions::non_donatable_input_indices` 中时，该输入被**捐赠**。

证据：`0x87a48c6` 处的运行时 diagnostic — *"pinned_host buffers do not support donation denial at runtime via `ExecuteOptions::non_donatable_input_indices`"* — 以及 `0xbf85686` 处的 proto field-name 字符串 `non_donatable_input_indices`。默认构造的 `ExecuteOptions` 具有空集合，因此默认情况下，配置标为 aliasable 的每个输入都会被捐赠。

> **GOTCHA —** 极性是 opt-*out*，不是 opt-*in*。调用方若想保留一个它意外标成可捐赠的输入，必须把该索引加入 `non_donatable_input_indices`；什么都不做就会捐赠它。这是最常见的就地更新 bug 类别来源：调用方依赖默认空 deny-list，导致某个缓冲区被 execute 静默消耗（invalidated）。donor `PjRtBuffer` 在调用后变得不可用（其 tracked device buffer 被移走，见 §4）。
>
> **QUIRK —** `kPinnedHbm` 缓冲区（运行时锁定层，见 [overview §2](overview.md#2-the-memoryspace-enum)）**不能**在运行时被 deny-list（字符串 `0x87a48c6`）。配置标为 aliasable 的 pinned-host 输入会无条件被捐赠；deny-list 机制对该 memory space 静默无效。重新实现者必须要么拒绝 pinned 缓冲区上的 `non_donatable_input_indices`，要么明确记录该标志在那里不起作用。

### 必须捐赠集合 — `ComputeParametersThatMustBeDonated`

dispatch 前，runtime 会计算程序*要求*捐赠的参数索引集合，即 `kMustAlias` 参数。`xla::ComputeParametersThatMustBeDonated(const HloInputOutputAliasConfig&, int param_count, bool)` (`0x1d7f4700`, source `xla/pjrt/utils.cc:824`) 遍历配置并返回排序后的 `vector<int>`：

```c
// xla::ComputeParametersThatMustBeDonated(config, param_count, kind_filter)  0x1d7f4700
function ComputeParametersThatMustBeDonated(config, param_count, filter):
    donated = vector<int>()                 // reserve param_count
    status = config.ForEachAliasWithStatus(
        [&](output_index, Alias alias) -> Status {
            if (alias.kind matches filter)   // e.g. kMustAlias only
                donated.push_back(alias.parameter_number)
            return Ok
        })
    CHECK(status.ok())                       // else StatusBuilder @ utils.cc:824
    std::sort(donated.begin(), donated.end())
    return donated                           // sorted, possibly with duplicates

runtime 会把该集合与有效捐赠集合（不在 non_donatable_input_indices 中的输入）交叉检查。一个未被捐赠的 kMustAlias 参数会触发 fatal：

"An input was configured to be must-alias at compile time but not donated at runtime: %s" — strings 0x8588a34（format）和 0xa27831b（prefix）。
NOTE — TestBufferDonationClashes (0x1d7f4be0) 是配套 validator：它检测两个参数是否解析到同一个底层 device buffer（例如同一个 PjRtBuffer 作为两个参数传入）且其中一个被捐赠；捐赠共享缓冲区会破坏另一个别名。重新实现者必须在兑现捐赠前运行这个检查，否则两个别名输出可能踩同一个缓冲区。

函数映射

函数	地址	作用
`ComputeParametersThatMustBeDonated(config, int, bool)`	`0x1d7f4700`	must-donate params 的排序 `vector<int>`
`ComputeParametersThatMustBeDonated(HloModule, bool)`	`0x1d7f4580`	同上，从 module 提取配置
`TestBufferDonationClashes`	`0x1d7f4be0`	检测两个 args 解析到一个 donated buffer
`InferDispatchInfo(client, ComputationLayout&, config&, …)`	`0xf90cb40`	按 launch 规划 dispatch（layouts + aliasing）
(deny-list field) `ExecuteOptions::non_donatable_input_indices`	string `0xbf85686`	每次执行的 opt-out 捐赠标志

3. 运行时兑现 — `AllocateOutputBuffersWithInputReuse`

目的

这里是编译时计划和运行时许可转化为物理决策的地方：对每个输出叶子，是分配新的 HBM 缓冲区，还是复用已捐赠输入的缓冲区。该函数是整页内容的运行时半边；之前的一切都是 setup。它有两个实例：generic 的 xla::CommonPjRtClient::AllocateOutputBuffersWithInputReuse (0xf91ec20，基于 ScopedHold) 和 TPU 具体的 tfrt::tpu::AllocateOutputBuffersWithInputReuse (0xf7ba9a0，基于 tpu::TpuBuffer/TpuCoreLocation)。这里记录 TPU 版本，因为它是面向硅片的路径。

算法 — 每叶子复用或分配循环

tfrt::tpu::AllocateOutputBuffersWithInputReuse (0xf7ba9a0) 接收 alias config、输出 shapes（ArrayRef<AsyncValueRef<Shape>>）、TpuCoreLocation、已捐赠输入 TpuBuffers 和一个 System。它迭代输出 tuple 的叶子（*(_QWORD*)(out_shape + 16) 是 tuple element count；如果结果不是 tuple，会在 xla/shape.h:803 CHECK Expected a tuple shape），并对每个叶子执行：

// tfrt::tpu::AllocateOutputBuffersWithInputReuse  0xf7ba9a0  (decompiled control flow)
function AllocateOutputBuffersWithInputReuse(config, out_shapes, core_loc,
                                             donated_inputs, out_buffers, system, rm):
    shared_mem = TpuCoreLocation::LocalSharedMemory(core_loc,
                     TpuChipParts::PreferredSharedMemoryType(chip_parts))
    result_shape = config.shape()                       // the program result shape
    CHECK(result_shape.IsTuple())                       // xla/shape.h:803
    for i in [0 .. result_shape.tuple_shapes_size()):
        output_index = {i}                              // top-level leaf index
        alias = config.GetAliasedParameter(output_index)   // 0x1e580200  -> optional<Alias>
        if alias.has_value():                           // v77 == 1 branch
            // ---- REUSE: hand the donated input's buffer straight through ----
            buf = donated_inputs[ alias.parameter_number ]   // $_0::operator()
            out_buffers.push_back(buf)                       // no allocation, no copy
        else:
            // ---- ALLOCATE: size the padded device shape, alloc fresh HBM ----
            n_bytes = TransferSizeUtil::ShapeSizeCompact(topology, out_shapes[i])  // see buffer-layout §4
            buf = AllocateTpuBufferWithRetry(system, shared_mem, n_bytes, …)       // OOM -> defrag -> retry
            out_buffers.push_back(buf)
    return
```text

反编译分支非常明确：由 `GetAliasedParameter` 结果设置的每叶子标志（`v77`）在*复用*分支（lambda `$_0::operator()`，把已捐赠 `TpuBuffer` 移入输出向量）和*分配*分支（`ShapeSizeCompact` 计算叶子大小，然后 `AllocateTpuBufferWithRetry`）之间选择。复用分支不执行任何 HBM 分配，也不执行任何字节拷贝；输出就是输入的缓冲区，程序会通过它写入。

> **GOTCHA —** 复用分支信任被捐赠输入的缓冲区与输出叶子的*带 padding 设备 shape* 按字节兼容。编译时 `OptimizeInputOutputBufferAlias::Build` 强制相等字节大小，而 [layout assignment](../compiler/layout-assignment.md) 确保两个叶子获得相同的 `(SublaneCount, LaneCount)` tile，因此被捐赠字节已经处于输出期望的布局中（[tpu-buffer-layout §3](tpu-buffer-layout.md#3-the-tile-and-the-padding-rules)）。若重新实现把不同 tiling 或 `memory_space` color 的两个叶子别名起来，就会复用一个程序会错误寻址的缓冲区；该 bug 是静默的（没有分配错误）并会破坏输出。
>
> **NOTE —** 分配分支调用 `AllocateTpuBufferWithRetry`，它用 OOM→defragment→retry 循环包装 `tpu::AllocateBuffer`。关键是，*捐赠*（复用）的缓冲区在这次分配期间**不**受 defrag 重定位影响；在兑现循环运行前，它已被捐赠 hold 固定（§4），因此 [compactor](on-device-compaction.md) 不能把它从程序脚下移走。这就是与 `kPinnedHbm` 的联系：被捐赠缓冲区在 execute 期间表现得像 pinned 缓冲区。

### 函数映射

| 函数 | 地址 | 作用 |
|---|---|---|
| `tfrt::tpu::AllocateOutputBuffersWithInputReuse` | `0xf7ba9a0` | 对 donated `TpuBuffer`s 逐叶子复用或分配 |
| (reuse lambda) `…::$_0::operator()` | `0xf7bb1c0` | 将已捐赠输入缓冲区移入输出槽 |
| `CommonPjRtClient::AllocateOutputBuffersWithInputReuse` | `0xf91ec20` | generic（基于 ScopedHold）的 sibling |
| `tfrt::tpu::AllocateTpuBufferWithRetry` | （在 `0xf7ba9a0` 中） | 新 HBM 分配，带 OOM→defrag→retry |
| `TransferSizeUtil::ShapeSizeCompact` | `0x1d6ae8a0` | 计算输出叶子的带 padding 设备 shape 大小 |
| `HloInputOutputAliasConfig::GetAliasedParameter` | `0x1e580200` | 每叶子 alias 查询 |

---

## 4. PJRT 缓冲区上的捐赠握手

### 目的

在兑现循环能复用已捐赠输入之前，runtime 必须从调用方的 `PjRtBuffer` 中原子且仅一次地取得该输入 device buffer 的所有权，并将 `PjRtBuffer` 标记为已消费，使后续任何使用都报错，而不是与就地写入竞争。这就是 `ScopedHold` 捐赠握手。重新实现者必须完全正确地实现该状态机；binary 会对每个非法转换执行 `CHECK` fatal。

### `ScopedHold` 状态机

`xla::CommonPjRtBuffer::ScopedHold` 是作用于 `PjRtBuffer` 的 RAII guard，意图有多种；捐赠意图是 `kDonation`。execute-prepare 阶段（`ExecutePrepare` → `PrepareArguments`，见 [the adapter execute path](../pjrt/stream-executor-pjrt-adapter.md)）将每个输入固定为 `ScopedHold`，并对已捐赠输入执行：

```c
// the donation handshake (decompiled, abstract_tracked_device_buffer.cc + tracked_tpu_device_buffer.cc)

// 1. Acquire: move the tracked device buffer OUT of the PjRtBuffer into the hold.
ScopedHold::AcquireDonation()                       // 0xf93d600
    REQUIRE(state == kUsable)                        // CHECK "!ok()" @ abstract_tracked_device_buffer.cc:180
    buffer_ptr_ = take(buffer_)                       // ownership transfer
    REQUIRE(buffer_ptr_ != nullptr)                   // CHECK @ :192
    state = kDonated

// 2. On successful dispatch: confirm — the donor PjRtBuffer is now permanently empty.
ScopedHold::ConfirmDonation()                       // 0xf93dca0
    CommonPjRtBuffer::ConfirmDonation(buffer_ptr_)    // 0xf93dd40
        TrackedTpuDeviceBuffer::ConfirmDonation()     // 0xf840660
            REQUIRE(in_use_)                          // CHECK "in_use_" @ tracked_tpu_device_buffer.cc:88
            in_use_ = false                           // release the in-use marker
            release all AsyncValueRef usage/definition events  // refcount decrements

// 3. On failure / not-actually-donated: give the buffer back to the PjRtBuffer.
ScopedHold::DropDonationHold(unique_ptr<AbstractTrackedDeviceBuffer>)   // 0xf93d900
    return ownership to buffer_  (PjRtBuffer becomes usable again)

AcquireDonation (0xf93d600) 的反编译显示了所有权移动（*(_QWORD*)(a1 + 32) = take(*a2)）、abstract_tracked_device_buffer.cc:192 处的 post-move CHECK(buffer_ptr_ != nullptr)，以及 :180 处禁止从已处于 error/used 状态的 hold 获取捐赠的 guard CHECK(!ok())。TrackedTpuDeviceBuffer::ConfirmDonation (0xf840660) 在 tracked_tpu_device_buffer.cc:88 断言 in_use_，清除它，并在缓冲区的 absl::Mutex 下遍历缓冲区的 usage/definition AsyncValue 向量递减引用计数。

GOTCHA — 三状态生命周期（Acquire → 要么 Confirm 要么 Drop，绝不二者皆有，也绝不二者皆无）由 fatal 强制，而不是优雅错误。AcquireDonation 在 dispatch 保证成功之前就把缓冲区移走；如果随后 dispatch 失败，runtime 必须调用 DropDonationHold 归还缓冲区，否则 donor PjRtBuffer 会变成 leaked-empty（看似可用但没有后备存储）。忘记 drop path 的重新实现会把一次临时 dispatch 失败变成永久死亡的缓冲区。
NOTE — PjRtBuffer::DonateWithControlDependency(Future<void>) (0xe6eb260, TPU impl CommonPjRtBufferImpl::DonateWithControlDependency @ 0xf92a740, PJRT C shim PJRT_Buffer_DonateWithControlDependency @ 0xf86f2e0) 是独立捐赠 API：它允许调用方捐赠一个受 future 门控的缓冲区，独立于 execute 调用（例如在传输完成后捐赠）。一旦控制依赖 resolve，它会汇入同一套 ConfirmDonation 机制。

配置和握手如何在 dispatch 处相遇

xla::InferDispatchInfo (0xf90cb40) 是把两者绑定起来的规划器。它接收 ComputationLayout（输入/输出设备 shape，与可执行程序一起 serialized）和 HloInputOutputAliasConfig，再加上 device assignment，并生成描述哪些输入要捐赠以及输出如何映射的 DispatchInfo。execute 路径随后：

text

PJRT_LoadedExecutable_Execute
  └─ CommonPjRtLoadedExecutable::Execute
       └─ ExecutePrepare → PrepareArguments
            ├─ pin each input as ScopedHold
            ├─ for inputs absent from ExecuteOptions::non_donatable_input_indices
            │      and aliased by the config:  AcquireDonation               (§4)
            ├─ AllocateOutputBuffersWithInputReuse(config, out_shapes, …)    (§3)
            │      → reuse donated TpuBuffer  |  AllocateTpuBufferWithRetry
            └─ verify ComputeParametersThatMustBeDonated ⊆ donated set       (§2)
                 (else FATAL "must-alias at compile time but not donated at runtime")
       └─ ExecuteLaunch → … → tpu::System::Execute   (program overwrites the reused buffer in place)
       └─ on success: ConfirmDonation for each donated input   (donor PjRtBuffer permanently consumed)
```text

### 函数映射

| 函数 | 地址 | 源锚点 | 作用 |
|---|---|---|---|
| `ScopedHold::AcquireDonation` | `0xf93d600` | `abstract_tracked_device_buffer.cc:180/192` | 将 tracked buffer 移出 PjRtBuffer |
| `ScopedHold::ConfirmDonation` | `0xf93dca0` | — | finalize：donor 永久为空 |
| `CommonPjRtBuffer::ConfirmDonation` | `0xf93dd40` | — | 分派到 tracked-buffer confirm |
| `TrackedTpuDeviceBuffer::ConfirmDonation` | `0xf840660` | `tracked_tpu_device_buffer.cc:88` | 清除 `in_use_`，释放 events |
| `TrackedCpuDeviceBuffer::ConfirmDonation` | `0xf916b20` | — | CPU sibling |
| `CommonPjRtBuffer::DropDonationHold` | `0xf93d900` | — | 失败 / 未实际捐赠时归还缓冲区 |
| `CommonPjRtBuffer::GetBufferForDonationHoldLocked` | `0xf93d3a0` | — | 在锁下为 hold 取出缓冲区 |
| `PjRtBuffer::DonateWithControlDependency` | `0xe6eb260` | — | 独立的 future-gated donation |
| `CommonPjRtBufferImpl::DonateWithControlDependency` | `0xf92a740` | — | 上述 API 的 TPU 实现 |
| `PJRT_Buffer_DonateWithControlDependency` | `0xf86f2e0` | — | PJRT C-ABI shim |
| `InferDispatchInfo(…, ComputationLayout&, config&, …)` | `0xf90cb40` | — | 规划 donation + output mapping |

---

## 5. 配置如何随可执行程序携带

### 目的

一个带捐赠、带别名的程序会在多次 execute 间复用，也可能从编译缓存加载而不重新编译。为此，alias config 必须随编译后程序载荷一起传输。重新实现者必须把它序列化为程序的一等同级数据，而不是重新生成它。

### 被携带的配置

编译后程序载荷是一个构造一次并在程序整个生命周期携带的 tuple。`HloInputOutputAliasConfig` 是每种载荷变体的成员，位于 `ComputationLayout`（或 `vector<Shape>`）和 `xdb::CompilerMetadata` 旁边：

| 载荷类型 | 构造函数 | 成员（相关） |
|---|---|---|
| `xla::jellyfish::TpuJitResult::Program` | `0xf8b7720` / `0xf8b7500` | `unique_ptr<TpuCoreProgram>`, `ComputationLayout`, **`HloInputOutputAliasConfig`**, `CompilerMetadata`, `HostTransferProto[]`, `HostExecutionProto[]` |
| `tfrt::tpu::TpuCompilationCacheEntry::Program` | `0xf7bc500` / `0xf7bbd60` | `int`, `AsyncValueRef<TpuCoreProgram>`, `vector<Shape>×2`, **`HloInputOutputAliasConfig`**, `CompilerMetadata`, `HostTransferProto[]` |
| `tfrt::tpu::TpuJitResultTF::Program` | `0xf7c8d60` | `AsyncValueRef<TpuCoreProgram>`, **`const HloInputOutputAliasConfig`**, `CompilerMetadata`, `HostTransferProto[]` |

磁盘形式是 `xla::HloInputOutputAliasProto`（带嵌套的 `HloInputOutputAliasProto_AliasEntryProto`；`CopyFrom`/`MergeImpl`/`InternalSwap`/`Clear` 均存在），`HloInputOutputAliasConfig::CreateFromProto(result_shape, proto)` 会在加载时重新水合内存中的配置。这就是为什么 cache-hit reload 的程序能精确保留别名关系：配置是数据，不会被重新计算。

> **NOTE —** alias config 和 `ComputationLayout` 是互补的：`ComputationLayout` 固定输入和输出的*带 padding 设备 shape*（[tpu-buffer-layout §2](tpu-buffer-layout.md#2-host-shape--device-shape)），alias config 固定*哪个输出设备 shape 复用哪个输入设备 shape*。兑现循环（§3）依赖的字节兼容性是两者的结合：相同 layout（来自 `ComputationLayout`）加一条 alias edge（来自配置）。完整载荷容器见 [../compiler/tpu-program-serialization.md](../compiler/tpu-program-serialization.md)。

---

## 相关组件

| 组件 | 关系 |
|---|---|
| `xla::HloInputOutputAliasConfig` | 本页围绕的 output-indexed alias map |
| `xla::CommonPjRtBuffer::ScopedHold` | 在 execute 时转移缓冲区所有权的 RAII donation hold |
| `tfrt::tpu::AllocateOutputBuffersWithInputReuse` | 按输出叶子复用或分配的运行时兑现循环 |
| `xla::OptimizeInputOutputBufferAlias` | 合成 alias + donor config 的编译时 pass |
| `xla::InferDispatchInfo` | 按 launch 对账 config + `ComputationLayout` + donation flag 的规划器 |
| `tpu::TpuBuffer` / `TpuSharedMemoryLocation` | 被复用的 HBM 对象；由 [hbm-allocator.md](hbm-allocator.md) 分配 |

## 交叉引用

- [tpu-buffer-layout.md](tpu-buffer-layout.md) — 带 padding 设备 shape，以及别名输入/输出叶子为保证复用按字节正确而必须一致的 `(SublaneCount, LaneCount)` tiling
- [hbm-allocator.md](hbm-allocator.md) — 位于复用（捐赠）缓冲区和新 `AllocateTpuBufferWithRetry` 分配背后的 `BestFitAllocator`
- [overview.md](overview.md) — `memory_space` 层级；`kPinnedHbm` 以及为什么 donated buffers 在 execute 期间表现为 pinned
- [on-device-compaction.md](on-device-compaction.md) — 可重定位空闲 HBM、但不能移动被 donation `ScopedHold` 持有缓冲区的 compactor
- [../compiler/tpu-program-serialization.md](../compiler/tpu-program-serialization.md) — 在 `ComputationLayout` 和 `CompilerMetadata` 旁携带 `HloInputOutputAliasConfig` 的编译后程序载荷
- [../compiler/layout-assignment.md](../compiler/layout-assignment.md) — 固定每叶子 tile，使两个 alias 叶子 layout-compatible
- [../pjrt/stream-executor-pjrt-adapter.md](../pjrt/stream-executor-pjrt-adapter.md) — 驱动捐赠握手的 `ExecutePrepare` → `PrepareArguments` → `ExecuteLaunch` 路径
- [../pjrt/buffer-and-memory.md](../pjrt/buffer-and-memory.md) — 位于 `ScopedHold` 之上的 `PJRT_Buffer` ABI 和外部引用计数
- [../pjrt/executable-execution.md](../pjrt/executable-execution.md) — 发起 execute 调用的 executable-execution surface
- [back to index](../index.md) — Part X — On-Chip Memory & DMA

缓冲区捐赠与别名 ​

摘要 ​

1. HloInputOutputAliasConfig ​

目的 ​

数据模型 ​

函数映射 ​

3. 运行时兑现 — AllocateOutputBuffersWithInputReuse ​