MRB 链分配器

本页中的所有地址都适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，未 strip；下文所有符号都是已 demangle 的 C++ 名称）。段映射：.text/.rodata VMA == 文件偏移；.data.rel.ro VMA - 0x200000 == 文件偏移；.data VMA - 0x400000 == 文件偏移。其他版本会有所不同。

摘要

MrbChainAllocator 是 Stage-2 pass，用来决定每条 matmul 累加链占用哪个 matrix-result buffer (MRB) FIFO 项，以及占用多久。TPU matmul 不会把结果写入寄存器；它把操作数推入 systolic array，许多周期之后，部分和落入一小组 result-FIFO buffer（MRB）中的一个。当多个 matmul 累加到同一个运行中的和时，也就是 K-tiled reduction，它们形成一条累加链，必须从该链第一个 matmul 开始一直持有同一个 MRB 项，直到结果被消费。分配器的工作是在程序顺序时间线上把每条链放到有限的 MRB 池中，并在每个项的结果 retire 的瞬间释放它，使后续链可以复用。

自然的参照框架是单一物理寄存器类上的 linear-scan register allocation，但有一个变化：这些“寄存器”是具有固定 retire latency 的 FIFO，因此驱逐按时间作为键，而不是按 liveness interval。分配器按 matmul issue 顺序遍历链；每一步都会推进单调时钟（AdvanceTimeTo），驱逐所有结果在“现在”之前已经可用的链，并把这些链的 MRB 项回收到按 chunk 划分的空闲池（ReleaseMrbReservation）；随后为当前链保留一个项：要么是新的 insert，要么如果该链已经因为较早的 matmul 持有一个项，则执行 replace_data，延长该项不可驱逐的截止时间。若某条链必须持有其项超过可用窗口，就在当前时间切开（SplitAccumulationChain），并将其尾部推迟到之后再重新保留。这是 MXU assignment bin-packer 的 MRB 对应物；不同之处在于，MXU packer 在 N 个可互换 lane 上执行贪心 min-makespan 搜索，而 MRB 分配器是一次确定性的程序顺序扫描，带有按时间索引的驱逐池，因为 result FIFO 是有序、固定 retire latency 的资源，而不是 N 条并行 matmul lane。

本页记录时间线算法和链分配逻辑：核心 chains_unevictable_until_ boost bimap（一侧按 ProgramOrder 排序，另一侧按 chain pointer 哈希）、AssignMrbEntriesToChains driver loop、AdvanceTimeTo 驱逐引擎、ReleaseMrbReservation 按 chunk 的回收池，以及 SplitAccumulationChain 的 defer-tail 机制。推进 reservation end-time 的每步成本，也就是 ExtendMrbReservation 的两条并行推进 push_time + min(throughput, cap) 和 push_time + min(latency, cap)，是这条时间线消费的成本单元；throughput 与 edge 项由 CycleTable（CycleTableInstruction）和 LatencyBetween 定价，各自被 option cap（+0x20/+0x28）从上方截断，并写入不同的 per-MXU timeline array，而不是相加。本页后半部分补齐公开 edge-latency 表面：可选 LatencyBetween jitter（由 flag 控制的均匀随机扰动，默认关闭）以及 trace-instruction edge clamp。

重新实现时，契约如下：

chains_unevictable_until_ bimap relation type，以及两条 reservation 分支（新链用 insert，已有链用 replace_data 延长）。
driver loop 纪律：按 matmul program order 对 accumulations 做 introsort，然后对每个 accumulation 执行 Split -> AdvanceTimeTo -> FIFO-push rounding -> Reserve -> ExtendMrbReservation。
驱逐规则：推进时钟会释放所有 matres_program_order 严格小于新时间的链，并把每个释放的 MrbEntry 回收到按 chunk 划分的空闲池。
split 规则：链不能在过去被 split；切分保留 head 中已消费的 accumulations，并把 tail 放入以 program order 为键的 btree_map 延后处理。
LatencyBetween wrapper：base edge -> optional jitter（仅当安装了 BitGen*）-> trace-edge clamps。


Driver	`(anon)::AssignMrbEntriesToChains` `@0x10f4ac60` (`Target&, CycleTable&, LatencyTable&, Span<const AccumulationChain>, int, MrbAccumulationOptions`)
每步成本单元	`MrbChainAllocator::ExtendMrbReservation` `@0x10f58800`（两条并行的 `push + min(throughput, cap)` / `push + min(latency, cap)` timeline 推进；由 `CycleTable`/`LatencyBetween` 定价）
驱逐引擎	`MrbChainAllocator::AdvanceTimeTo(StrongInt<ProgramOrder>)` `@0x10f5e9e0`
回收池	`MrbChainAllocator::ReleaseMrbReservation(MrbEntry)` `@0x10f5f9e0`
Split / defer-tail	`MrbChainAllocator::SplitAccumulationChain` `@0x10f598e0`
核心结构	`chains_unevictable_until_` — boost `bimap<set_of<StrongInt<ProgramOrder>>, unordered_set_of<AccumulationChainAfterSplit*>, with_info<AccumulationWithOriginalChain>>`
Jitter wrapper	`LatencyTable::LatencyBetween(LloValue,LloValue)` `@0x1c89f820`；base ctor `@0x1c89f800`
源文件	`platforms/xla/service/jellyfish/mxu_accumulation.cc`（RetChecks 位于 lines 1017, 1332, 1371）
置信度	CONFIRMED（byte-anchored），除非某行或 callout 另有说明

核心数据结构

目的

分配器必须在每个 loop step 以 O(1)/O(log n) 回答两个问题：(a) 哪条链下一个 retire（这样时钟可以驱逐它），以及 (b) 这个 chain pointer 是否已经持有 reservation（这样 reserve 会变成 extend，而不是新的 entry）？单个 Boost bimap 同时回答二者：一侧按 retire time 排序，另一侧按 chain pointer 哈希。

`chains_unevictable_until_`

relation type 是 byte-exact 的，从 final_erase_ 实例化 @0x10f61d60 demangle 得到，并在 SplitAccumulationChain @0x10f598e0 中内联重新确认：

// boost::bimaps::relation::mutant_relation<...> — chains_unevictable_until_
bimap<
    set_of<            StrongInt<ProgramOrder>,         std::less<...> >,   // LEFT  — ORDERED by retire time
    unordered_set_of<  AccumulationChainAfterSplit*,    boost::hash<...> >, // RIGHT — HASHED by chain pointer
    with_info<         AccumulationWithOriginalChain >                      // INFO  — next_accumulation
>;
```text

- **LEFT key** = 链的 `matres_program_order`：链的 matrix result 可用时的程序顺序时间，也就是其 MRB 项*在此之前*不可被驱逐的时间。LEFT index 是有序（`set_of`）树，因此其 front node 始终是下一条 retire 的链。
- **RIGHT key** = `AccumulationChainAfterSplit*` chain pointer，位于由 `boost::hash` 哈希的 `unordered_set_of` 中。reserve 会在这里查找链，以决定 insert-vs-extend。
- **INFO** = `AccumulationWithOriginalChain`（`next_accumulation`），即当前正在保留的 accumulation 之后的那个 accumulation；它随 relation 携带，而不是作为 key。

> **NOTE —** bimap 是*驱逐*索引；它不同于 `AdvanceTimeTo` post-check 的 per-MXU `chains_next_accumulating_at_[mxu_id]` array（RetCheck string `"chains_next_accumulating_at_[mxu_id].left.begin()->first >= to"`，mxu_accumulation.cc:1371），也不同于 per-MXU `mrb_entries_[mxu_id][curr_level]` size-class block map（insert string `"mrb_entries_[mxu_id][curr_level].insert(MrbBlock{.offset = 2 * mrb_block.offset + 1}).second"`）。一个 allocator 中三种结构共存：按时间排序的 eviction bimap、per-MXU next-accumulation index，以及 per-MXU level-keyed block map。只有第一个驱动驱逐；本页记录它。

### 分配器对象布局

由四个 traced method 的 field access 重建。Field role 为 CONFIRMED；option cap 与 chunk array 的 struct-name binding 为 INFERRED（无 struct-layout string）。

| Offset | Field | 含义 |
|---|---|---|
| `+0x00` | `Target*` | JF `Target`（vtable gate `[+0x390]` MRB-support，`[+0x5e0]` FIFO granule） |
| `+0x08` | `CycleTable*` | `GetCyclesForThroughput` 来源，即 `ExtendMrbReservation` 的 throughput 轴 |
| `+0x10` | `LatencyTable*` | `LatencyBetween` 来源，即 edge 轴 |
| `+0x20` | `s64` throughput cap | `min(throughput, cap)` 单元，从 `MrbAccumulationOptions` 复制 |
| `+0x28` | `s64` latency cap | `min(latency, cap)` 单元，从 `MrbAccumulationOptions` 复制 |
| `+0x70` | `s64 latest_matmul_` | 单调时钟；`AdvanceTimeTo`/`SplitAccumulationChain` RetCheck guard |
| `+0xd8`..`+0xe8` | `chains_unevictable_until_` | 驱逐 bimap（LEFT-ordered head at `+0xd8`，RIGHT-hashed at `+0xe0`） |
| `+0x7a0` | per-chunk free pool | 回收项的 `flat_hash_set<MrbEntry>`，按 chunk_id 分桶（`ReleaseMrbReservation`） |
| — | `btree_map<long, unique_ptr<AccumulationChainAfterSplit>>` | 推迟的 split-off chain tail，以 program order 为键 |

`+0x70` 时钟带有规范名称 `latest_matmul_`：`AdvanceTimeTo` 与 `SplitAccumulationChain` 都读取 `*(this+0x70)`，并在它们的 RetCheck format string 中称为 `latest_matmul`（`@0x10f5ea08`，`@0x10f59920`）。

---

## `AssignMrbEntriesToChains` — Driver

### 目的

`AssignMrbEntriesToChains` `@0x10f4ac60` 是入口点。它以 `absl::Span<const AccumulationChain>`（demangled signature，line 1）接收某个 sequence group 的 matmul accumulation chain，将它们复制成可变工作形态，在栈上构造 `MrbChainAllocator`，并驱动 per-accumulation loop。整个函数就是分配器的 `main`。

### 入口点

```text
AssignMrbEntriesToChains  @0x10f4ac60   ── (Target&, CycleTable&, LatencyTable&, Span<AccumulationChain>, int, MrbAccumulationOptions)
  ├─ copy Span<AccumulationChain> → working vector<AccumulationChainAfterSplit>
  ├─ introsort accumulations by matmul program order        @0x10f4b092  (gtl::OrderBy<Accumulation::program_order, Less>)
  ├─ construct MrbChainAllocator (stack)                    @0x10f4b113
  │     ├─ copy MrbAccumulationOptions+0x4..+0x14 → allocator+0x1c..+0x2c   (16-byte vmovups)
  │     └─ __size_returning_new per-MXU chunk arrays         @0x10f4b2cf  (count clamped ≥ 8)
  └─ per accumulation (in program order):
        SplitAccumulationChain    @0x10f4c984
        AdvanceTimeTo(now)        @0x10f4c9c3
        FIFO-push rounding        @0x10f4ca06
        Reserve (insert | replace) → ExtendMrbReservation  @0x10f4d9a8 / @0x10f4dfca

算法

function AssignMrbEntriesToChains(target, cycles, latency, chains, num_mxus, opts):  // @0x10f4ac60
    // 1. Materialize a mutable working list of accumulations.
    work = copy(chains)                          // 0x650-byte AccumulationChain → 0x70-byte AccumulationChainAfterSplit

    // 2. Process chains in matmul *issue* order, so the clock only ever moves forward.
    introsort(work.accumulations,                // @0x10f4b092
              OrderBy(Accumulation::matmul_program_order, Less))   // ascending

    // 3. Build the allocator; copy the two ExtendMrbReservation caps out of the options.
    alloc = MrbChainAllocator(target, cycles, latency, opts, num_mxus)   // @0x10f4b113
    //   alloc.throughput_cap (+0x20) and alloc.latency_cap (+0x28) copied from MrbAccumulationOptions
    //   alloc.chunks[max(num_mxus, 8)] = __size_returning_new(...)      // @0x10f4b2cf

    // 4. Main loop — one iteration per accumulation, already in program order.
    for accum in work.accumulations:
        chain = accum.chain
        now   = accum.matmul_program_order

        // 4a. If a held chain cannot keep its entry to its result time, cut it now and defer the tail.
        if chain_must_yield_entry(chain, accum):
            SplitAccumulationChain(chain, accum)      // @0x10f4c984

        // 4b. Advance the monotone clock; this evicts + recycles every retired chain.
        AdvanceTimeTo(now)                            // @0x10f4c9c3

        // 4c. Size the result-FIFO push and pick the MXU instance.
        push = LloInstructionPushesToResultFifo(accum.matmul)            // @0x1d4f3600
        if target.vtable[+0x390]() > 0:               // MRB-support gate  (line 1871)
            granule = target.vtable[+0x5e0]()         // FIFO-push granule (byte) (line 1874)
            push    = ceil(push / granule) * granule  // round up to a granule multiple (idiv)
        size_class = bsr(push)                        // log2 size class   (_BitScanReverse, line 1887)
        mxu        = accum.matmul.unit_id() & 3        // MXU-instance index from matmul unit_id (line 1864/1716)

        // 4d. Reserve — new entry or extend existing — then price the step.
        it = alloc.chains_unevictable_until_.right.find(&chain)
        if it == end:                                  // first matmul of this chain
            alloc.chains_unevictable_until_.right.insert(
                { &chain, accum.matres_program_order, accum.next_accumulation })
            ExtendMrbReservation(chain, accum.next_accumulation)         // @0x10f4d9a8
        else:                                          // chain absorbs another matmul
            alloc.chains_unevictable_until_.right.replace_data(
                it, accum.matres_program_order)
            ExtendMrbReservation(chain, accum.next_accumulation)         // @0x10f4dfca
```text

### 两条 Reserve 分支

reserve 是内联的 `ReserveMrbEntry`（一个双分支 `absl::Overload` visitor，lambda `@0x222f5f10`/`@0x222f5f40`）。分支由 chain pointer 是否已经位于 `chains_unevictable_until_.right` 中选择：

| 分支 | 条件 | Bimap op | `ExtendMrbReservation` 位置 | 源字符串 |
|---|---|---|---|---|
| **new reservation** | chain absent（第一个 matmul） | `right.insert({&chain, matres_program_order, next_accum})` | `@0x10f4d9a8` | 由 `RetCheck("!chain->mrb_entry.has_value()")` cc:1542 保护；该链必须尚未持有 entry |
| **extend reservation** | chain present（吸收一个 matmul） | `right.replace_data(it, matres_program_order)` | `@0x10f4dfca` | `"chains_unevictable_until_.right.replace_data(it, curr_accumulation.accumulation.matres_program_order)"` `@0xa104fb1` |

> **QUIRK —** 两条分支的差异只在于它们*如何*触碰 bimap，而不在于如何定价。二者都会调用 `ExtendMrbReservation`，用 per-step cost 推进链的 reservation end-time。`replace_data` 分支会覆写 LEFT key（`matres_program_order`），使驱逐顺序跟踪 multi-matmul chain 的*最新* matmul。重新实现若在 extend 路径上重新 `insert` 而不是 `replace_data`，会在有序索引中留下陈旧副本，并过早释放 entry。

### FIFO-Push 舍入

在 driver 中 byte-exact 确认（lines 1869-1887）。matmul 推送的 result-FIFO slot 数是 `LloInstructionPushesToResultFifo(matmul)` `@0x1d4f3600`；当 target 支持 MRB（gate `vtable[+0x390]() > 0`，line 1871）时，push 会通过整数除法向上舍入为 per-gen granule（`vtable[+0x5e0]()`，一个 byte，line 1874）的倍数；随后 `bsr(push)` 产生 log2 size class（cc:875 的 `RetCheck(absl::has_single_bit(entries_needed))` 保护其为 power-of-two，并检查 `target_level < kLevelsInMrbAllocator`）。MXU instance 是 `accumulation.matmul->unit_id() & 3`，即 matmul unit id 的低 2 位，**不是**其 program order（line 1864；unit id 在 line 1716 读取，并由 cc:1518 的 `"accumulation.matmul->unit_id().has_value()"` 进行 `RetCheck` 保护）。

> **NOTE —** 这里的 `Target` vtable 层**不是** `JellyfishTarget` base `@0x21cc6bc0`；`[+0x390]`（MRB-support count）和 `[+0x5e0]`（FIFO-push granule）的具名 accessor 未解析。slot offset、`>0` gate 以及 round-up arithmetic 是 byte-exact；accessor *names* 为 INFERRED。

---

## `AdvanceTimeTo` — 驱逐引擎

### 目的

`AdvanceTimeTo(StrongInt<ProgramOrder> new_time)` `@0x10f5e9e0` 让 allocator 成为一条*时间线*。将时钟推进到当前 matmul 的 program order，会释放所有结果已经 retired 的 MRB entry，也就是所有 `matres_program_order` 严格小于 `new_time` 的链，并将每个释放项回收到 per-chunk pool，同时从 eviction bimap 中擦除。

### 算法

```c
function AdvanceTimeTo(new_time):                        // @0x10f5e9e0
    // 1. The clock is monotone: you may never advance backward.
    if this.latest_matmul_ /*+0x70*/ > new_time:         // line 128
        FATAL "MrbEntries must be reserved in program order "
              "(latest_matmul=%v, current_matmul=%v)"    // mxu_accumulation.cc:1332  @0xa0f5d91

    // 2. Evict every chain whose result became available before now.
    while front = chains_unevictable_until_.left.begin():    // ordered-by-ProgramOrder head  @+0xd8
        if front.matres_program_order >= new_time:           // line 175 (cmp [node-0x88], to)
            break                                            // nothing left retires before now
        chain = front.chain
        if front.complete /* relation-node byte *(node-8) != 1 → chain is done */:
            ReleaseMrbReservation(chain.mrb_entry)            // recycle  (lines 178-205)
            // chain+0x2c is RetChecked nonzero here (BUG() if 0) before reading mrb_entry @chain+0x20
        else: // *(node-8) == 1: chain is now "evictable, next Accumulation is ..." (deferred, not released)
        chains_unevictable_until_.final_erase_(front)         // remove from bimap  @0x10f61d60

    // 3. Post-condition: each per-MXU next-accumulation index head is now >= new_time.
    for mxu in 0 .. num_mxus:                                 // lines 706-739
        RetCheck(chains_next_accumulating_at_[mxu].left.begin()->first >= new_time)  // cc:1371

    this.latest_matmul_ = new_time                            // commit the clock  (line 754)

GOTCHA — guard 在 latest_matmul_ > new_time（line 128）时触发，也就是时间必须非递减（current > new 时 FATAL，current <= new 时成功）。因为 driver 以 introsorted program order 喂入 accumulations，new_time 在正常操作中不会回退，guard 不会触发；它是结构性 invariant check，不是运行时分支。相比之下，loop 中的 eviction comparison 是严格 <：结果恰好落在 new_time 的链尚未 retired，并继续保留其 entry。
NOTE — eviction loop 根据 relation-node 的 "complete" byte *(node-8) 分支：当它 != 1 时，链已完成，因此其 MRB entry 通过 ReleaseMrbReservation 回收（VLOG "... releasing MRB entry, as it is complete"，cc:1355）；当它 == 1 时，链仍有 pending next-accumulation，因此记录日志 "Chain holding ... is now evictable, next Accumulation is ..."（mxu_accumulation.cc:1347），并在不 release 的情况下擦除。在 release 路径上，读取 chain+0x20/chain+0x28 的 MrbEntry 前，会断言链自身的 *(chain+0x2c) validity byte 非零（若为 0 则 BUG()）。

`ReleaseMrbReservation` — 回收池

目的

ReleaseMrbReservation(MrbEntry) @0x10f5f9e0 将释放的 result-FIFO entry 返回到 per-chunk 可用 entry 池，使后续 reservation 可以复用它，而不是分配新的 chunk slot。MRB pool 是真正的 FIFO，具有固定 retire latency；回收保留该顺序。

`MrbEntry` 与池

MrbEntry 是打包的 { short chunk_id, int fifo_index, byte format }。this+0x7a0 处的 pool 是 Abseil flat_hash_set<MrbEntry>，其 bucket 由 chunk_id 寻址，每个 bucket 持有一个 absl::linked_hash_set<MrbEntry>（按插入顺序，因此 recycle list 本身是 FIFO）。

function ReleaseMrbReservation(entry):                   // @0x10f5f9e0
    // 1. Address the per-chunk bucket: chunk_id selects a 0x40-byte sub-table.
    bucket = pool_base(this+0x7a0) + (entry.chunk_id << 6)    // line 33: (int16)entry << 6

    // 2. Hash the whole {chunk_id, fifo_index, format} with the absl CRC32 mixer.
    h = MixingHashState(kSeed @0x22042400)               // _mm_crc32_u64 over the 3 fields (lines 42-43)

    // 3. Insert a recycle node; on first use, grow the SOO table.
    node = operator new(0x20)                            // line 111
    node.entry  = entry                                  // {chunk_id, fifo_index, format}
    node.format = entry.format
    push_front(bucket.free_list /* head @+0x28 */, node) // lines 114-119
    ++bucket.count                                       // free-entry count @+0x38 (line 118)
    return node
    // VLOG "Freeing MRB entry " @0xa1b5c8d
```text

> **QUIRK —** recycle node 是 `operator new(0x20)`，并被链接到每个 chunk 的*双向链表* `linked_hash_set` 中，而不是压入普通 free stack。linked hash set 既去重（entry 不能被两次释放到同一 chunk），又保留 recycle order。使用裸 LIFO free list 的重新实现会以错误顺序复用 entry，并偏离二进制的确定性 FIFO 复用。

---

## `SplitAccumulationChain` — 切分与延后

### 目的

`SplitAccumulationChain(AccumulationChainAfterSplit&, AccumulationWithOriginalChain const&)` `@0x10f598e0` 处理 contention 情况：无法将其 MRB reservation 保持到当前 matmul 的链，会在 split point 被切开。head 保留目前已经消费的 accumulations；tail 变为一条新链，延后到之后的 program order 中重新 reserve。

### 算法

```c
function SplitAccumulationChain(chain, split_point):     // @0x10f598e0
    // 1. You cannot split in the past — the cut must be at or after the clock.
    if split_point.matmul_program_order /*accum+0x8*/ < this.latest_matmul_ /*+0x70*/:   // line 63
        FATAL "Cannot split an AccumulationChain in the past"   // mxu_accumulation.cc:1017  @0x84dcedc
        // RetCheck expr: "split_point.accumulation.matmul_program_order >= latest_matmul_"  @0x87b7fa6

    // 2. If the chain still occupies an MRB entry, free it — the split forces a wait.
    if chain.holds_mrb_entry /* *(chain+0x2c) == 1 */:                   // line 74
        chains_unevictable_until_.final_erase_(chain)                    // line 251
        ReleaseMrbReservation(chain.mrb_entry)                           // line 287
        // VLOG "Freeing MRB entry ..., since the occupying chain has been split at ..."  @0xa1daa56 (cc:1033)

    // 3. Cut the accumulation span at the split point.
    consumed   = split_point.matmul - chain.begin /*chain+0x8*/         // line 120
    CHECK(consumed <= chain.len /*chain+0x18*/)                         // "pos > size()" guard
    tail_begin = split_point.matmul                                     // line 119 (a4[13])
    tail_span  = chain.span_ptr /*chain+0x10*/ + 0x60 * consumed        // 0x60-byte stride
    tail_len   = chain.len - consumed
    chain.len  = consumed        /* the head keeps the already-consumed accumulations */   // line 128

    // 4. Defer the tail: store it as a heap-owned chain keyed by program order.
    tail = operator new(0x30) { begin=tail_begin, span=tail_span, len=tail_len, holds_mrb=false }  // line 195
    deferred_tails /* btree_map<long, unique_ptr<AccumulationChainAfterSplit>> */
        .insert({ split_point.matmul_program_order, tail })            // @0x10f5d2e0 (line 204)

GOTCHA — step 3 保留短的 head，并将 tail 作为新的 deferred chain 发出；head 的 len 被设为 consumed（*(chain+0x18) = consumed，line 128），而不是 len - consumed；堆上的 tail node 得到 len - consumed。split point 是边界：它之前的一切留在原处，它及之后的一切被延后，并在时钟之后到达 split_point.matmul_program_order 时重新 reserve。
NOTE — tail 由以 program order 为键的 btree_map<long, unique_ptr<AccumulationChainAfterSplit>> 拥有，因此 deferred remainder 会按顺序重新进入时间线。这是 allocator 唯一在 heap 上分配 chain 的地方；原始 chain 位于 driver 的 working vector 中。

总体策略与 MXU 对比

分配器是一次单向 forward sweep，每一步有三个按时间索引的效果：split-if-contended、advance-and-evict、reserve-or-extend。这里没有 makespan search，也没有 backtracking。per-step advance 是 ExtendMrbReservation 的两条并行写入：push_time + min(CycleTableInstruction, cap@+0x20) 写入 per-MXU throughput timeline，push_time + min(LatencyBetween, cap@+0x28) 写入 per-MXU result-entry timeline；throughput 轴来自 CycleTable，edge 轴来自 LatencyBetween，二者都被 option cap（+0x20/+0x28）从上方截断。两者不相加。

方面	MRB Chain Allocator（本页）	MXU Assignment Bin-Packer
资源形态	一个有序 FIFO，具有固定 retire latency	N 个可互换的 MXU pass
纪律	确定性的程序顺序扫描	贪心 min-makespan 搜索
驱逐	按时间索引（`AdvanceTimeTo` evict-by-`matres_program_order`）	n/a — 只放置，从不驱逐
评分	无；placement 由时间强制决定	每个 candidate 的 `LatchLatencyChangeAfterAdding`
contention 处理	split + defer tail（`SplitAccumulationChain`）	选择不同的 MXU pass
共享成本输入	`CycleTable` matmul throughput + `LatencyBetween` edge	相同

结构差异是关键：result FIFO 是有序资源，其 entry 按时钟 retire，因此最优 placement 由时间强制决定，不需要搜索。MXU lane 是可互换的，因此对 assignment 的搜索是有收益的。二者共享成本单元，但不共享分配算法。

`LatencyBetween` — 公开 Edge Wrapper

目的

LatencyTable::LatencyBetween(LloValue from, LloValue to) @0x1c89f820 是 read-after-write edge latency 的公开 dispatcher，每个 scheduler stage 都以它定价，包括上文 ExtendMrbReservation 的 edge 轴。它用两个 gen-invariant correction 包装 per-gen LatencyBetweenInternal（base model，记录于 bundle-aware-cost）：可选随机 jitter 与 trace-instruction edge clamp。本页负责这两个 correction。

算法

function LatencyBetween(from, to):                       // @0x1c89f820
    raw = this->vtable[+0x18](from, to)                  // base: per-gen LatencyBetweenInternal

    // --- OPTIONAL JITTER ---  (line 27)
    bitgen = this->jitter_bitgen  /* this+0x10 */
    if bitgen != null:                                   // null by default → skip
        raw += Uniform(bitgen, 0, 101)                   // bounds (0,101) passed to absl UniformDistributionWrapper<int>
                                                         // via DistributionCaller<BitGen> @0xfa7c9a0 (Randen PRNG)

    // --- TRACE-EDGE CLAMPS ---  (lines 33-52)
    if from.opcode == 0x84 /*trace-arg*/ and to.opcode == 0x84:
        cfg = AutoOr<int>::FromProtoOrDie(AutoProto)     // @0x10979760  ([from+0x10]→[+0x38]→deref→[+0xc78])
        edge = cfg.is_set ? cfg.value : 16               // default 16
        raw  = max(raw, edge)
    else if from.opcode == 0x82 /*set-tracemark*/ and (to.opcode - 0x82) <= 2:   // to ∈ {0x82,0x83,0x84}
        raw = max(raw, 2)                                // min-2 floor on tracemark → trace* edges

    return raw                                            // VLOG latency_table.cc:83 @0x878d9cb
```text

### Jitter 开关

jitter `BitGen*` 位于 `LatencyTable+0x10`，并在 base constructor `@0x1c89f800` 中**初始化为 0**（`*(this+0x10) = 0`，decompiled），且各 per-gen derived constructor 都没有设置它，因此 jitter **默认关闭**。仅当命令行 flag `FLAGS_xla_jf_random_latency`（flag name `"xla_jf_random_latency"` `@0x84bcebc`；flag global `@0x223b47c8`）启用时才会安装。

| 开关 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| `xla_jf_random_latency` | BOOL/BitGen install | **off (`+0x10` == null)** | 设置后在 `LatencyTable+0x10` 安装 Randen-AES `BitGen*`；此后每条 edge 都会增加一个以 bounds `(0, 101)` 抽取的均匀随机扰动 |

> **QUIRK —** jitter 是*鲁棒性/fuzzing*开关，不是硬件 latency 项。打开后，scheduler 会在 latency variation 下运行，从而 stress-test 其决策对 edge-latency noise 的敏感性；它不建模任何真实硅片延迟。为实际 TPU cycles 定价的重新实现必须保持 `+0x10` 为 null；确定性的 `LatencyBetweenInternal` 才是硬件模型。
>
> **NOTE —** 将该 flag 连接到 `BitGen*` install 的位置没有 byte-walk；base ctor 中的 null default 和 flag 的 name/global 为 CONFIRMED，但 install 发生在*何时*（哪个 factory path 读取 `GetFlag(xla_jf_random_latency)`）未追踪（LOW）。`0x84/0x84` trace-arg edge 的 `AutoOr<int>` proto field 是结构性解码的（`@0x10979760`）；其默认值 16 为 byte-exact，但 proto field *name* 未固定（LOW）。

### Trace-Edge Clamp

clamp 使 profiling-marker instruction edge 不会在 schedule 中坍缩到一起。三个 opcode（从 `opcode_string` table `@0x21cd0d60` 经 `.rela.dyn` RELATIVE-addend decode 解析）为 `0x82` = set-tracemark，`0x83` = trace，`0x84` = trace-arg。

| `from` opcode | `to` opcode | 应用的 floor | 来源 |
|---|---|---|---|
| `0x84` trace-arg | `0x84` trace-arg | `max(raw, AutoOr<int> default 16)` | line 36-46 |
| `0x82` set-tracemark | `0x82` / `0x83` / `0x84` | `max(raw, 2)` | line 49-51 |

trace-arg -> trace-arg edge 采用一个可配置 floor（默认 16 cycles）；set-tracemark -> any-trace edge 采用固定 min-2 floor。二者都确保插入的 trace marker 保持最小间距，避免它们堆入同一个 bundle，从而破坏发出的 trace stream。

---

## 函数映射

| Function | Address | 角色 |
|---|---|---|
| `(anon)::AssignMrbEntriesToChains` | `0x10f4ac60` | driver：copy -> introsort -> allocator ctor -> main loop |
| `MrbChainAllocator::ExtendMrbReservation` | `0x10f58800` | per-step cost cell（throughput 和 edge，各自经 `min` cap，写入分离的 per-MXU timeline） |
| `MrbChainAllocator::AdvanceTimeTo` | `0x10f5e9e0` | monotone clock + evict-by-`matres_program_order` |
| `MrbChainAllocator::ReleaseMrbReservation` | `0x10f5f9e0` | 将释放的 `MrbEntry` 回收到 per-chunk pool |
| `MrbChainAllocator::SplitAccumulationChain` | `0x10f598e0` | 在时钟处切分 chain；延后 tail |
| `bimap final_erase_` | `0x10f61d60` | `chains_unevictable_until_` relation type（demangled） |
| `btree_map<long, unique_ptr<…>>::insert` | `0x10f5d2e0` | 存储 deferred split-off chain tail |
| `LloInstructionPushesToResultFifo` | `0x1d4f3600` | FIFO-push count，舍入到 per-gen granule |
| `Target::vtable[+0x390]` | — | MRB-support count gate（`> 0`） |
| `Target::vtable[+0x5e0]` | — | FIFO-push rounding granule（byte） |
| `LatencyTable::LatencyBetween` | `0x1c89f820` | base edge -> optional jitter -> trace clamps |
| `LatencyTable::LatencyTable(TpuVersion)` | `0x1c89f800` | base ctor：`+0x10 = 0`（jitter off） |
| `DistributionCaller<BitGen>::Impl<UniformDistributionWrapper<int>>` | `0xfa7c9a0` | `Uniform(0,101)` jitter draw（Randen PRNG） |
| `AutoOr<int>::FromProtoOrDie` | `0x10979760` | trace-arg `0x84/0x84` edge config（default 16） |
| `LloOpcodeString` | `0x1d631360` | opcode -> name table `@0x21cd0d60`（0x82/0x83/0x84） |

---

## 注意事项

- **Program-order 是唯一排序。** 每个 guard（`AdvanceTimeTo`、`SplitAccumulationChain`）都以 `matmul_program_order` 和 `latest_matmul_` 时钟为键。`@0x10f4b092` 的 introsort 是正确性的前置条件，不是优化：如果 accumulations 不是按 matmul order 升序排列，monotone-clock RetCheck（cc:1332）会触发。
- **驱逐是严格 `<`，时钟 guard 是 `>`。** `matres_program_order == new_time` 的链在这一步*不会*被驱逐（其结果尚未 retire）；但时钟可以推进*到*该时间。这两个比较有意使用不同关系。
- **bimap 和 per-MXU index 必须保持一致。** `AdvanceTimeTo` 在驱逐后会 post-check 每个 `chains_next_accumulating_at_[mxu]` front（cc:1371）。只维护 eviction bimap 的重新实现会通过 eviction loop，但会违反该 invariant。
- **`MrbChunkState` / `AccumulationChainAfterSplit` layout 是功能性重建。** 0x70-byte `AccumulationChainAfterSplit` body（`+0x8` begin、`+0x10` span ptr、`+0x18` len、`+0x20`/`+0x28` 为 `MrbEntry`、`+0x2c` 为 "has mrb entry" byte）以及 per-MXU `MrbChunkState` element 是由 access pattern 重建的，而不是来自 struct-layout string（exact byte layout 为 LOW；field *roles* 从使用中 CONFIRMED）。

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `mxu-assignment-binpacker` | MXU sibling：在 N 条 lane 上贪心 min-makespan，对比本页的确定性 FIFO sweep |
| `mxu-sequence-struct` | `MxuSequence`/`SequenceInfo` record，其 accumulation chain 输入此 allocator |
| `mrb-fifo-msr-placement` | 下一层：把保留的 `(chunk_id, fifo_index)` `MrbEntry` 转换成物理 FIFO/MSR address |
| `bundle-aware-cost` | 拥有 base `LatencyBetweenInternal` edge 和此 allocator 定价所依赖的 `MaxResourceCycles` throughput |
| `jf-cycletable` | `GetCyclesForThroughput` 单元，也就是 `ExtendMrbReservation` 的 throughput 轴 |

---

## 交叉引用

- [TPU 调度流水线](overview.md) — 此 allocator 位于 HLO scheduler 与 bundle packer 之间的 Stage 2 placement。
- [MXU 分配装箱器](mxu-assignment-binpacker.md) — MXU sibling allocator；这个对比解释了为什么 MRB 不需要 makespan search。
- [MxuSequence / SequenceInfo](mxu-sequence-struct.md) — 保存此 pass 消费的 accumulation chain 的 per-sequence record。
- [MRB FIFO / MSR 放置](mrb-fifo-msr-placement.md) — 下游 pass，将保留的 `MrbEntry` 物化为 result-FIFO 和 MSR address。
- [Latch 分配与 Overrun](latch-assignment-overrun.md) — bundle packer 作为 slot-legality input 读取的 latch-index commit。
- [Bundle-Aware Cost](../cost/bundle-aware-cost.md) — 此 allocator 的 cost cell 使用的 base `LatencyBetween`/`LatencyBetweenInternal` edge axis 和 `MaxResourceCycles` throughput。
- [JF CycleTable](../cost/jf-cycletable.md) — 驱动 `ExtendMrbReservation` recurrence 的 per-gen matmul throughput cycles。
- [MXU Latency Overview](../cost/mxu-latency-overview.md) — 与此 result-buffer allocator 互补的 MXU occupancy reservation model。
- [MXU Slot](../isa/slot-mxu.md) — 此 allocator 放置其 accumulation chain 的 LLO MXU instructions（matmul/matpush/matres）。
- [LLO Opcode Enum](../isa/llo-opcode-enum.md) — `0x82`/`0x83`/`0x84` trace-instruction identity 背后的 opcode numbering。
- **Binary:** `extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so`（build-id `89edbbe81c5b328a958fe628a9f2207d`）
- **Index entry:** Part VIII — Instruction Scheduling & Bundle Packing — [返回索引](../index.md)

MRB 链分配器 ​

摘要 ​

核心数据结构 ​

目的 ​

chains_unevictable_until_ ​

算法 ​

ReleaseMrbReservation — 回收池 ​

目的 ​

MrbEntry 与池 ​

总体策略与 MXU 对比 ​

LatencyBetween — 公开 Edge Wrapper ​

目的 ​

算法 ​

MRB 链分配器

摘要

核心数据结构

目的

`chains_unevictable_until_`

算法

`ReleaseMrbReservation` — 回收池

目的

`MrbEntry` 与池

总体策略与 MXU 对比

`LatencyBetween` — 公开 Edge Wrapper

目的

算法