MSA AllocateSegment

地址、结构体偏移和位掩码值适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。

摘要

MsaAlgorithm::AllocateSegment (0x1dc73ca0) 是内存空间分配的内层循环体。驱动部分，即缓冲区区间排序、同址分组以及 autotuner "IOR" 重放表面，位于 msa-overview.md；本页记录驱动按排序顺序为每个 (AllocationValue, use) 对调用一次的逐段放置尝试。一个 "segment" 是某个值从其定义（或前一次分配的结束）到单个 use 之间的 live range；AllocateSegment 决定该 segment 的缓冲区物理上位于何处，即在 alternate memory (VMEM) 中（带或不带 copy），或退回 default memory (HBM)，并发出实现该决定的 Allocation 对象（以及任何 async CopyStart/CopyDone 对）。

熟悉上游 XLA MsaAlgorithm::AllocateSegment 的读者会认出其形状：先协调 required assignments，然后运行一条固定的放置原语级联，每个原语返回一个 AllocationResult 位掩码，并在第一个 kSuccess 处短路。这个构建以字节精度确定的是：(1) 六阶段级联顺序及其短路/取消提交语义，(2) 级联读取的约 ~0x210 字节 AllocationRequest 结构体，该结构体从写入者 CreateAllocationRequest (0x1dc72820) 恢复，并与函数体中读取的每个字段交叉核对，(3) AllocationResult 失败位掩码，即 10 个命名失败位，范围 0x001..0x200，从 SingleFailureResultToString (0x1dc7bc80) 恢复。这里的贪心性质是具体的：AllocateSegment 内部没有回溯；失败的 prefetch 或 evict 若留下 pending chunks，会 OR 入 FailRequiresUncommit (0x40)，使驱动回滚，而该 segment 只会在稍后驱动重试并放宽窗口后再次尝试。

本页记录重实现者必须复现的三项内容：级联主体及其阶段门控、alternate-memory fit check（AllocateInAlternateMemoryNoCopy 和 prefetch chunk-candidate 搜索），以及 prefetch/eviction 插入（Prefetch/Evict 和为它们提供前驱 allocation 的 find_if）。prefetch 咨询的 interval-picker sweep、async-copy resource model，以及逐版本数值默认值在此概述，并在相邻页面详述。

对于重实现，契约是：

级联。 Required-assignment reconcile → direct pin → free coloring-reserved + alternate no-copy → force-min-time（当 +0x208 == 1）→ prefetch → evict → default fallback；每一步返回一个 AllocationResult，并在 kSuccess 处短路。
请求。 级联读取的 AllocationRequest 字段布局；尤其是 +0x208..+0x20b 处用于门控各阶段是否运行的四个尾部状态字节。
结果。 失败位掩码（10 个命名位，0x001..0x200）、字节精确的 bit→name 映射，以及把回滚交给驱动的 |0x40 (FailRequiresUncommit) 规则。
fit check + insertion。 no-copy 路径和 prefetch 路径各自如何向 best-fit heap 查询 chunk，以及 prefetch/evict 如何寻找其前驱 allocation。


方法	`xla::memory_space_assignment::MsaAlgorithm::AllocateSegment(AllocationRequest&)`
入口	`0x1dc73ca0`（约 5.5 KB；1301 行 decompile）
请求写入者	`MsaAlgorithm::CreateAllocationRequest` @ `0x1dc72820`
结果字符串化器	`MsaAlgorithm::SingleFailureResultToString` @ `0x1dc7bc80`; `ResultToString` @ `0x1dc68c80`
No-copy fit	`AllocateInAlternateMemoryNoCopy` @ `0x1dc7ca80`
Prefetch / evict	`Prefetch` @ `0x1dc7e4a0` · `Evict` @ `0x1dc7d660`
前驱 `find_if`	`AllocateSegment::$_2` @ `0x1dc7e420`
IR 层级	XLA HLO (`HloValue` / `HloUse`)、heap-simulator 时间线（逻辑时间）
内存空间	`kDefault` = HBM (`0`), `kAlternate` = VMEM (`1`)；VMEM 字节预算和 word size 随版本而定，见 msa-reservation-hbm-policy.md
置信度	CONFIRMED（字节锚定），除非某行或标注另有说明

`AllocationRequest`

CreateAllocationRequest (0x1dc72820) 为每个 (value, use) 构建一个 AllocationRequest，并按引用传给 AllocateSegment。该结构体约为 ~0x210 字节；下面的布局从写入者的字段存储区域 (0x1dc730d8..0x1dc731e5) 恢复，并与级联主体中每个 [rbx+N] / [a2+N] 读取交叉核对。decompile 直接确认了门控字节偏移：Evict 以 *((unsigned __int8 *)a2 + 520)（即 +0x208）作为其 force 参数被调用，而 Prefetch 以 *((unsigned __int8 *)a2 + 521)（即 +0x209）作为其 window 参数被调用。

Offset	Field	Meaning	Confidence
`+0x00`	`inclusive_start_time`	`max(use_time, earliest_prefetch)` 或 preferred-override time	CONFIRMED
`+0x08`	`end_time`	use 的 latest time bound	CONFIRMED
`+0x10`	`latest_prefetch_time`	修正后的 use time bound	CONFIRMED
`+0x18`	`require_start_time`	required def time	CONFIRMED
`+0x20`	`size`	`BufferInterval` size，字节	CONFIRMED
`+0x28`	`allow_no_copy_alternate_mem`	bool	CONFIRMED
`+0x29`	`require_no_copy_alternate_mem`	bool	CONFIRMED
`+0x2a`	`require_copy_allocation`	bool (arg)	CONFIRMED
`+0x2b`	`no_copy_chunk_required`	bool	CONFIRMED
`+0x2c`	`allow_prefetch`	bool	CONFIRMED
`+0x30`	`earliest_prefetch_time`	packed: low byte flag, time `>>8`	HIGH
`+0x38`	`has_preferred_prefetch_time`	bool	CONFIRMED
`+0x40`	`preferred_prefetch_time`	packed (low-byte flag \| `time<<8`)	HIGH
`+0x48`	`prefer_no_copy`	bool	CONFIRMED
`+0x50`	`preferred_offset`	`AliasedOffset*` (arg)	CONFIRMED
`+0x58`	`use`	`const HloUse*`，正在放置的 use	CONFIRMED
`+0x60`	`allocation_value`	`AllocationValue*`	CONFIRMED
`+0x68`	`all_use_times`	16-B span head/span	HIGH
`+0x88`	`required_copy_done_time`	int64	HIGH
`+0x90`	`allow_sync_conversion`	bool (arg)	CONFIRMED
`+0xc0`	`split_shape`	`Shape`，仅当 `+0x200 == 1` 时构造	HIGH
`+0x200`	`has_split_shape`	bool	HIGH
`+0x208`	force-evict flag	作为 `Evict` 的 `bool force` 传入	CONFIRMED
`+0x209`	allow-prefetch-window flag	作为 `Prefetch` 的 `bool window` 传入	CONFIRMED
`+0x20a`, `+0x20b`	uncommit / retry state bits	用于门控 retry/evict/bail	HIGH

+0x208..+0x20b 处的四个尾部字节是 "already-attempted / requires-uncommit" 状态。+0x208（cmp 1）门控 ForceAlternateMemoryAllocationForMinTime 路径，并且是传给 Evict 的 force 参数；+0x209 是传给 Prefetch 的 window 参数；+0x20a|+0x20b OR 在一起后门控 colored-reserved free 路径。CreateAllocationRequest 触达的 sync-copy candidate 测试（经由 IsAsyncConversionSliceCandidate → IsAsyncConversion{Copy,Slice}Candidate）为 copy 匹配 HLO opcode 44 (kCopy)，为 slice 匹配 119/54 (kSlice/kDynamicSlice)，使 MSA 能把已经调度的同步 kCopy/kSlice 转换为由 MSA 管理的 async copy（见 Async-copy resource and sync replacement）。

NOTE — ~0x210 并非精确值。 结构体大小由观察到的最高 store 限定下界（≥0x210）；+0x20b 之后的 padding 未列举。偏移 +0x68/+0xc0/+0x200 和两个 packed-time 字段是根据 store 宽度以及 0x20cd6660 处的条件 Shape ctor 解码的，因此为 HIGH 而非 CONFIRMED。

`AllocationResult` 位掩码

返回类型不是 enum，而是一个失败 bitset（10 个命名位，0x001..0x200），其中 0 表示成功，每个置位 bit 表示一个失败原因。SingleFailureResultToString (0x1dc7bc80) 先用 0x7f 做范围检查（cmp 0x7f），将 ≤0x7f 的值路由到 0xb564d94 处 65 项 jump table（按单 bit 值 0..0x40 索引），并将三个高位（0x80/0x100/0x200）路由到显式比较；每个分支存储其结果字符串。Success、FailOutOfMemory 和 UnknownResult 在 decompile 中是 inline strcpy 字符串字面量；其余名称从 rodata 加载（因此不会以 inline string args 形式出现），这与 jump-table-plus-rodata 结构一致。

Bit	Name
`0x000`	`Success`（未置位）
`0x001`	`FailOutOfMemory`
`0x002`	`FailPrevAllocationNotInAlternateMem`
`0x004`	`FailLiveRangeTooLong`
`0x008`	`FailLiveRangeTooShort`
`0x010`	`FailOutOfAsyncCopies`
`0x020`	`FailViolatesAsyncCopyResource`
`0x040`	`FailRequiresUncommit`
`0x080`	`AllSlicesHaveTheSameStartTime`
`0x100`	`FailConflictingPreferredOffsets`
`0x200`	`FailSyncDataMoveReplacement`
else	`UnknownResult`

设置了多个 bit 的结果会由 ResultToString (0x1dc68c80) 逐名分解并连接。两个谓词 helper 读取该掩码：result_requires_uncommit() 测试 bit 0x40；result_failed_because_of_async_copy() 测试 0x10|0x20。对级联重要的语义是：

FailRequiresUncommit (0x40) 不是由某个原语产生的；AllocateSegment 在 evict 失败之后（以及 residual-fail fallback 路径上）将其 OR 入，用来告诉驱动在下一次尝试前必须回滚该值的 pending chunks。decompile 在三个位置显示了这个 OR（v69 | 0x40, v240 | 0x40, v83 | 0x40）。
FailOutOfAsyncCopies (0x10) 与 FailViolatesAsyncCopyResource (0x20) 区分“达到已配置 outstanding-copy cap”和“copy 无法放入任何空闲 time-bucket window”；二者都来自 async-copy resource model，而不是简单计数器。

级联

AllocateSegment 运行一条固定的放置尝试序列，每次返回一个 AllocationResult，并在第一个 kSuccess (== 0) 处返回。各阶段按函数体顺序如下：

// MsaAlgorithm::AllocateSegment(AllocationRequest& req)            // 0x1dc73ca0
function AllocateSegment(req):
    // ---- STAGE 0: REQUIRED-ASSIGNMENT RECONCILIATION ----
    def_req   = RequiredMemoryAssignmentAt(value, def_time)        // 0x1dc48840
    use_req   = RequiredMemoryAssignmentAt(value, use_time)        // 0x1dc48840
    alias_req = AliasedRequiredAssignmentForUse(use)               // 0x1dc769c0
    if conflict(def_req, use_req):                                 // (memory_space, offset, padding)
        FATAL(def_req.ToString(), use_req.ToString(), use)         // 0x1dc47680
    required = reconcile(def_req, use_req, alias_req)

    // ---- STAGE 1: DIRECT PIN (colored / required alternate-memory) ----
    if required == kAlternate:
        a = new PinnedAllocation(position, kAlternate, chunk,      // 0x1dcdc400
                                 def_time, use_time)
        a->AddUse(use)                                             // 0x1dcda700
        value.allocation_sequence.push_back(a)                     // 0x1dc55740
        CreateOrAddToAliasedOffset(*a, req.preferred_offset)       // 0x1dc79160
        return kSuccess

    result = kSuccess                                              // sentinel cleared

    // ---- STAGE 2: free coloring-reserved, then dual-live + no-copy fit ----
    FreeAlternateMemoryColoringReservedAllocations(req)            // 0x1dc7c640 (line 717)
    if alternate_mem_allowed(req):            // gate over +0x29/+0x2a/+0x2b/+0x20a/+0x20b
        CheckAndUpdateForDualLiveAllocationValues(prev, req)       // 0x1dc7bf20 (line 722)

        // ---- STAGE 3: ALTERNATE NO-COPY (the fit check) ----
        result = AllocateInAlternateMemoryNoCopy(req)              // 0x1dc7ca80 (line 723)
        if result == kSuccess: return kSuccess

    // ---- STAGE 3b: FORCE-MIN-TIME (only when +0x208 == 1) ----
    if req[+0x208] == 1:                                           // line 738
        if ForceAlternateMemoryAllocationForMinTime(req):          // 0x1dc7d460 (line 749)
            return result                  // non-zero -> propagate

    // ---- STAGE 4: PREFETCH (insert async copy) ----
    prev_alloc = *find_if(value.allocation_sequence, $_2)          // 0x1dc7e420
    if req[+0x209] or req.allow_prefetch:      // +0x2c
        result |= Prefetch(req, *prev_alloc, /*window=*/req[+0x209]) // 0x1dc7e4a0
        if result == kSuccess: return kSuccess

    // ---- STAGE 5: EVICT ----
    if has_prior_alternate_alloc:
        result = Evict(req, /*force=*/req[+0x208])                 // 0x1dc7d660
        if result == kSuccess: return kSuccess
        result |= 0x40   // FailRequiresUncommit -> driver rolls back pending
        return result

    // ---- STAGE 6: DEFAULT-MEMORY FALLBACK ----
    d = new PinnedAllocation(position, kDefault/*0*/, none, def, use) // 0x1dcdc400
    d->AddUse(use)                                                 // 0x1dcda700
    value.allocation_sequence.push_back(d)                         // 0x1dc55740
    return result   // may carry residual fail bits (also OR'd 0x40 on a fail path)
```text

decompile 直接确认了调用顺序：`FreeAlternateMemoryColoringReservedAllocations` 在 line 717，`AllocateInAlternateMemoryNoCopy` 在 723，`ForceAlternateMemoryAllocationForMinTime` 在 749，`Evict(..., a2[+0x208])` 在 829（紧随其后 line 832 有 `| 0x40`），以及 `Prefetch(..., a2[+0x209])` 在 988。`PinnedAllocation` ctor (`0x1dcdc400`) 在函数体中出现两次，一次用于 Stage 1 的 direct alternate pin，一次用于 Stage 6 default fallback；而 `AddUse` (`0x1dcda700`) 在该 segment 的发出 allocation 和附加 use 路径中共出现四次。

> **NOTE — 函数内没有回溯。** Stages 3–6 不会撤销彼此的副作用。一旦某个原语把 chunk 提交给 heap，只有驱动的 uncommit 路径（由 `FailRequiresUncommit` 触发）会回收它们。这就是“带重试的贪心”（MSA 的实际行为）与真正回溯搜索之间的区别。

### Stage 0 — 协调 required assignments

首先执行三次 required-assignment 查找。`RequiredMemoryAssignmentAt(value, t)` (`0x1dc48840`) 返回该值在逻辑时间 `t` 携带的任何硬性 memory-space pin，一次用于 def，一次用于 use。`AliasedRequiredAssignmentForUse(use)` (`0x1dc769c0`) 合入通过 colocation/alias group 继承的任何 pin（例如 while-loop carried value，或 layout pass 着色的 buffer）。如果 def 和 use 在 `(memory_space, offset, padding)` 上不一致，pass 处于不可满足状态，并会发出 `LogMessageFatal` 打印三个 `ToString()`；真正的冲突是编译器 bug，而不是 fallback。否则，一致的 assignment 成为 `required`。

### Stage 1 — direct pin

如果 `required == kAlternate`，该值*必须*驻留在 VMEM 中（它被着色或携带 required alternate pin），因此无事可决策：在 `kAlternate` 中构建覆盖 `[def_time, use_time]` 的 `PinnedAllocation`，附加该 use，将其推入该值的 allocation sequence，并通过 `CreateOrAddToAliasedOffset` 穿接 `AliasedOffset`，使 colocation group 的每个成员共享一个 VMEM offset。无条件返回 `kSuccess`；chunk reservation 已在 coloring 期间发生。

### Stage 2 — 释放 reserved，然后 dual-live + no-copy

`FreeAlternateMemoryColoringReservedAllocations` 在该区域顶部无条件运行（line 717），释放 coloring 预留但该 segment 现在可能声明使用的 VMEM。随后，根据尾部状态字节（`+0x20a | +0x20b | ...`）和 `require_no_copy_alternate_mem` flag (`+0x29`) 共同计算出的门控决定是否尝试 alternate placement；若尝试，`CheckAndUpdateForDualLiveAllocationValues` 处理该值同时在两个位置 live 的情况，然后运行 Stage 3 的 `AllocateInAlternateMemoryNoCopy`，如果成功放置 buffer，就短路到 `kSuccess`。

`ForceAlternateMemoryAllocationForMinTime` 路径是**独立且更晚**的（decompile line 749，位于 723 的 no-copy 尝试之后）：仅当 `+0x208 == 1` 时到达，并强制 buffer 至少在其最小 live interval 中进入 VMEM；如果它返回非零，该结果会直接传播出去（不再进入后续阶段）。

---

## fit check — `AllocateInAlternateMemoryNoCopy`

Stage 3 是成本最低的放置：使 buffer 在 `[def, use]` 全程保留在 VMEM 中，且**完全不 copy**；这只有在请求大小的空闲 VMEM chunk 在整个 interval 内持续存在时才可能。`AllocateInAlternateMemoryNoCopy` (`0x1dc7ca80`) 通过一次 `FindBestChunkCandidate` (`0x1dc7fb00`) 调用向 best-fit heap 查询该 chunk，而后者封装了 `GlobalDecreasingSizeBestFitHeap::FindChunkCandidate` (`0x1e48f960`)。该 heap 是 MSA 在 VMEM 字节预算上维护的 `GlobalDecreasingSizeBestFitHeap`，扣除了任何 scoped reservation，并按 VMEM word 对齐 chunk offsets/sizes；逐版本预算和 word size 记录在 [msa-reservation-hbm-policy.md](msa-reservation-hbm-policy.md)，allocator 细节见 [../memory/vmem-allocator.md](../memory/vmem-allocator.md)。

```c
// AllocateInAlternateMemoryNoCopy(req)                            // 0x1dc7ca80
result = kSuccess
chunk  = FindBestChunkCandidate(req, /*interval=*/[def, use])     // 0x1dc7fb00 -> best-fit
                                                                  //   over VMEM byte budget
if not chunk:
    return FailOutOfMemory            // 0x01 (no persistent free chunk)
heap.Commit(chunk)
a = new PinnedAllocation(position, kAlternate, chunk, def, use)
a->AddUse(use)
value.allocation_sequence.push_back(a)
return kSuccess

该 fit check 是时间性的，而不仅是空间性的：chunk 必须在 heap-simulator 时间线上的整个 [def, use] interval 内空闲，因为 no-copy allocation 会连续占用 VMEM。如果没有单个 chunk 跨越整个 interval，结果就是 FailOutOfMemory (0x01)，级联会落入 prefetch。

Prefetch 和 eviction 插入

如果 buffer 无法在整个 interval 中免费驻留 VMEM，MSA 会尝试及时把它 copy 进 VMEM（prefetch），或者如果该值已有更旧的 VMEM allocation，则把它 copy 回 HBM（evict）。两个原语都需要该值的前驱 allocation，即其 interval 与当前 use 相邻的 allocation；AllocateSegment 用 lambda 谓词 $_2 (0x1dc7e420) 对 value.allocation_sequence 执行 find_if 来定位它。

Prefetch (Stage 4)

Prefetch (0x1dc7e4a0) 决定在 use 之前将 buffer 从 HBM copy 到 VMEM 是否有收益且可行。下面勾勒的是 prefetch 路径，不是 122 行 Prefetch 函数的逐字转录：alternate-memory-benefit 权重在 AllocateSegment 自身中计算（它调用 GetAlternateMemoryBenefit 的 HloPosition 和 HloUse 两个 overload），picker sweep 和 fit/resource 循环位于 Prefetch + CheckPrefetchFit (0x1dc820e0)，sliced emission 位于 AddAsyncSlicesForPrefetch (0x1dc7ea40)。

// prefetch path: AllocateSegment + Prefetch (0x1dc7e4a0) + CheckPrefetchFit (0x1dc820e0)
benefit = CostAnalysis::GetAlternateMemoryBenefit(use)            // 0x1dcec480 (in AllocateSegment)
        // weighed against OperandBytesAccessed / OutputBytesAccessed
        //   (0x1dceb340 / 0x1dceb360)
picker.Begin(use, earliest_prefetch, latest_prefetch, preferred)  // 0x1dcd7a00
while not picker.Done():
    t_start = picker.Next()                                       // 0x1dcd7e20
    chunk   = FindBestChunkCandidate(req, [t_start, use_time])    // 0x1dc7fb00
            //   -> GlobalDecreasingSizeBestFitHeap::FindChunkCandidate (0x1e48f960)
    if not chunk: continue
    if not AsyncCopyResource.HasEnoughResource(t_start, use_time, cost):
        result |= FailViolatesAsyncCopyResource  // 0x20
        continue
    emit CopyAllocation(kCopyStart @ t_start, kCopyDone @ use_time)
    //   or SlicedCopyAllocation (make_unique @ 0x1dc7f7e0, via AddAsyncSlicesForPrefetch
    //   @ 0x1dc7ea40) on the sliced path
    return kSuccess
return result   // accumulated fail bits
```text

interval picker 以两个 cursor 从 preferred time 向外扫描候选 copy-start times（一个 increasing cursor 从 `earliest_prefetch_time` 初始化，一个 decreasing cursor 从 `latest_prefetch_time` 初始化），并使用方向 flag，同时积分 bandwidth-idle time，使 copy 放置在最能隐藏在 compute 后面的位置。对于每个候选 start time，`FindBestChunkCandidate` (`0x1dc7fb00`) 会向 VMEM heap 请求覆盖 `[t_start, use_time]` 的 chunk（chunk 只需从 copy 的*到达*时间开始空闲，而不必从 def 开始空闲，这就是 prefetch 能在 no-copy 失败时成功的原因），async-copy resource model 则验证该 copy 是否适合 DMA-bandwidth time budget。picker、它的三个 overlap ratios，以及 `/100.0` bandwidth-idle 积分步骤在 [msa-overview.md](msa-overview.md) 中详述；逐版本 ratio 默认值见 [msa-per-version-defaults.md](msa-per-version-defaults.md)。

> **CONFIDENCE — picker integration constant。** picker `Next` 中的 `/100.0`（`0x84a2fb8` 处的 float）积分粒度是在 disassembly 层面跟踪到的；C decompile 对 float 算术的呈现不透明（该常量没有在 pseudocode 中以字面量出现），因此*常量值*为 HIGH，而 picker 函数的存在和签名由 demangled symbols 确认为 CONFIRMED。

### Evict (Stage 5)

Eviction 是反向操作：已有 VMEM (`kAlternate`) allocation 的值可能需要被 copy 回 HBM 以腾出空间。`Evict` (`0x1dc7d660`) 仅在 `has_prior_alternate_alloc` 成立时到达，并从 `AllocationRequest +0x208` 取得 `force` flag：

```c
// Evict(req, force = req[+0x208])                                 // 0x1dc7d660
prev = the prior kAlternate allocation for this value
emit CopyAllocation(kCopyStart in VMEM @ prev.end,
                    kCopyDone in HBM)         // copy buffer OUT to default memory
if not AsyncCopyResource.HasEnoughResource(...):
    return FailViolatesAsyncCopyResource      // 0x20 (or FailOutOfAsyncCopies 0x10)
return kSuccess

evict 失败时，级联不会静默落入后续阶段；它把 FailRequiresUncommit (0x40) OR 入结果并返回，通知驱动回滚该 segment 已 pending 的任何 chunks。这是唯一一个原语失败会升级为强制驱动 uncommit，而不是落入下一阶段的位置。

Stage 6 — default-memory fallback

如果该值从未有过先前的 VMEM allocation（因此 evict 不适用）且 prefetch 失败，buffer 就简单地留在 HBM：在 kDefault (0) 中为该 use 构建 PinnedAllocation，附加 use，并 push。返回的 result 仍可能携带 prefetch 尝试留下的 residual fail bits（并且在 fail path 上 OR 入 0x40），由驱动解释；这里的非成功结果意味着 MSA 无法为该 segment 改善 HBM 放置，而不是编译失败。

Async-copy resource 和 sync replacement {#async-copy-resource-and-sync-replacement}

prefetch 和 evict 都通过 AsynchronousCopyResource 对其 copy 做门控；这是一个按逻辑时间划分的 free-resource bucket model（不是简单计数器）。HasEnoughResource(start, end, resource) (0x1dc78d40) 遍历 [start, end) 上的 buckets，检查 copy 的 cost 是否能从可用 DMA-bandwidth windows 中扣除；ConsumeResource (0x1dc77540) 进行扣账；AddCopy/RemoveCopy (0x1dc78580 / 0x1dc787c0) 负责 commit/undo。无法适配任何 bucket window 的 copy 产生 FailViolatesAsyncCopyResource (0x20)；超过配置的 outstanding-copy cap 产生 FailOutOfAsyncCopies (0x10)。

从 CreateAllocationRequest 到达的 kCopy opcode (44, 0x2c) 匹配供给 sync-copy-replacement 路径：当 alternate-memory benefit 足够时，一个已经调度的同步 kCopy（或 kSlice/kDynamicSlice，opcodes 119/54）可以转换成 MSA 管理的 async copy，并由 xla_msa_enable_sync_copy_replacement / _sync_slice_replacement 门控、由 copies_limit_for_sync_mem_op_conversion 限制。失败的转换表现为 FailSyncDataMoveReplacement (0x200)。resource-model 循环体和 sync-replacement 状态机尚未在 basic-block 层级完全解码（见 Limits）。

worked example — 一个 while-body activation buffer

跟踪一个 while-body activation %act（大小 2 MiB，由逻辑 t=100 的 matmul 定义，在 t=160 被消费）经过级联的过程（VMEM 字节预算和 word size 随版本而定，见 msa-reservation-hbm-policy.md）：

Driver sort。 buffer-interval comparator 按 memory_boundedness_score × size 对 %act 排名；产生它的 matmul 受 bandwidth 限制，因此 score 很高，%act 会较早被访问（见 msa-overview.md）。
Request。 CreateAllocationRequest 构建 AllocationRequest：+0x00 start = earliest_prefetch，+0x08 end = 160，+0x20 size = 2 MiB，+0x2c allow_prefetch = 1，+0x58 use = %act@160。
Stage 3 fit check。 AllocateInAlternateMemoryNoCopy 向 heap 请求一个在 [100, 160] 上空闲的 2-MiB chunk。如果存在，则放置一个 kAlternate PinnedAllocation → kSuccess，结束。
Stage 4 prefetch。 否则执行 picker.Begin(use=%act@160, earliest, latest)；Next() 扫描候选 copy-start times；对每个候选，FindBestChunkCandidate 请求覆盖 [t_start, 160] 的 2-MiB chunk（只需要从 copy 到达时开始空闲），AsyncCopyResource.HasEnoughResource(t_start, 160, cost) 验证该 copy 的 DMA window。成功时发出 CopyAllocation（kCopyStart @ t_start, kCopyDone @ 160）→ kSuccess。两个 chunk 都按 VMEM word 对齐。
Stage 5 evict / Stage 6 fallback。 如果 prefetch 失败且存在更旧的 kAlternate allocation，Evict 会把更旧的 buffer copy 出去；若失败，result |= 0x40，驱动回滚 pending chunks。如果不存在先前 VMEM allocation，%act 留在 HBM（kDefault PinnedAllocation）。
Handoff。 选定的 Chunk{offset, size} 随 Allocation 传递；经过驱动的 Process() 后，它成为静态 program-memory offset，由 on-device allocator pin 住且永不重定位，见 ../memory/on-device-compaction.md。

限制 {#limits}

本页没有以字节精度确定的内容：

数值 ratio/cap 默认值。 每个 TPU 版本的字面量 min/preferred/max_overlap_to_async_copy_ratio 和 max_outstanding_prefetches/evictions 存在于逐版本 TpuCompilationEnvironment proto field-default overlays 中，而不在代码中；这些 flag 都已接线（见 msa-per-version-defaults.md）。inefficient_use_to_copy_ratio / memory_budget_ratio 也有同样缺口。（数值为 LOW；flag 名称为 CONFIRMED。）
ConsumeResource spill/delay loop。 已解码到调用层级；将 copy 分摊到后续 buckets 并记录诱导 delay 的内层循环未列举。（HIGH。）
Sync-replacement state machine。 AsyncConversionResult 值布局以及 copies_limit_for_sync_mem_op_conversion 的精确执行点尚未完全解码。（HIGH。）
Picker integration constant，见 Prefetch 中的 CONFIDENCE 标注。（100.0 值为 HIGH。）

交叉引用

msa-overview.md — MSA 驱动：buffer-interval sort、colocation grouping、autotuner IOR replay surface，以及完整的 prefetch interval picker / async-copy resource model。
msa-per-version-defaults.md — 每个 TPU generation 的 overlap ratios 和 outstanding-copy caps；5-variant per-family flag override scheme。
msa-reservation-hbm-policy.md — VMEM byte budget、scoped reservations，以及该级联供给的 HBM fallback policy。
layout-assignment.md — 写入 memory_space 并为 buffers 着色的 pass，产生 Stage 0 协调的 required-assignment pins。
compile-phases.md — MSA 在编译流水线中的位置。
../memory/vmem-allocator.md — no-copy 和 prefetch fit checks 查询的 VMEM best-fit allocator。
../memory/on-device-compaction.md — MSA 分配的静态 chunk offsets 如何在运行时被 pin 住。

MSA AllocateSegment ​

摘要 ​

AllocationRequest ​

AllocationResult 位掩码 ​

级联 ​

Prefetch 和 eviction 插入 ​

Prefetch (Stage 4) ​

Stage 6 — default-memory fallback ​

Async-copy resource 和 sync replacement {#async-copy-resource-and-sync-replacement} ​

worked example — 一个 while-body activation buffer ​

限制 {#limits} ​

交叉引用 ​