Skip to content

LatencyHidingScheduler 核心

地址适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(build-id 89edbbe81c5b328a958fe628a9f2207d)。其他版本会不同。所有 .text.rodata 地址都是虚拟地址;对这个二进制来说,.text VMA == 文件偏移 0xe63c000.rodata VMA == 文件偏移 0x84a0000

摘要

LatencyHidingScheduler(LHS)是一个列表调度器,它会重新排序已经做过内存最小化的 HLO 指令序列,使长延迟的 async 工作(collective starts、host/ICI DMA、async copies)尽早发出,并把它的延迟隐藏在独立计算之下。它不会从零构建调度:基础调度已经存在(由 HloMemorySchedulerWithBrkgaFallback 产生,也就是 LHS 的 has_schedule() 前置条件),LHS 会用一个重叠 async 边的调度替换它。具体工作委托给每个 computation 一个 DefaultSchedulerCore;TPU 构建接入了 TpuAsyncTracker,用于建模每个 async op 消耗的物理资源。

熟悉 LLVM 的读者应记住一个类比和一个差异。类比是:DefaultSchedulerCore::ScheduleComputation 是经典的关键路径列表调度器,也就是一个 ready-set drain,每一步中每个 DAG-ready 节点都是候选者,并由多键比较器选出一个。差异是:这个调度器反向构建序列(从 roots 开始,用一个随着向后遍历而递减current_time 时钟自底向上进行),优先级主要由 async critical-path depth and height 而不是普通 dependency height 主导,并且比较器把逐候选者的内存峰值预算门与延迟隐藏键交织在一起,所以同一个隐藏延迟的比较器在预算紧张时可以被驱动去用重叠换取更低的 HBM 峰值。

本页是第 VIII 部分中核心算法的锚点:ScheduleComputation drain loop 及其 SchedulingStep、22 键 ReadySetLt 候选比较器(以 async-depth / async-height 优先级和 max(0, current_time − ready_time) 延迟隐藏项为中心)、计算优先级所读取的逐节点 depth/height 的 InitializeGraphAnalysis pass,以及 TpuAsyncTracker 资源模型如何决定哪些 async op 可以重叠。pipeline-placement 变体(post_layoutpost_layout_pre_fusion、ILP swap)有各自页面,并在下方链接;比较器消费的逐节点代价来自感知 bundle 的代价模型

对重新实现而言,契约是:

  • drain loop 是 while (ready_set || pending) { SchedulingStep; }SchedulingStep 通过 FindAndExtractBestNodeAvailable 取出单个最佳 ready 节点,追加它,并用该节点的 bundle cost 推进 current_time。序列在最后反转。
  • 优先级是一个 22 键有序比较器(ReadySetLt:force-flags 优先,然后是内存峰值预算门,再是 async-overlap 键(schedule-async-start/done、resource-conflict、async-depth、async-height、unlock-start/done),然后是 ready-node fan-out,最后是稳定的原始顺序 tie-break。第一个区分两个候选者的键获胜。
  • critical-path depth/height 由 InitializeGraphAnalysis 预计算:每个节点的 [+0x38] async depth 是对 predecessors 的 max,其 [+0x40] async height 是对 successors 的 max/+。比较器从不重新计算它们;只读取字段。
  • 延迟隐藏项是 remaining = max(0.0, current_time − node.ready_time):内联计算(状态 +0x138current_time 减去节点 +0x28vsubsd,再由 vmaxsd clamp 到 0)。operand 尚未“retired”的节点(它的 ready_time 仍在时钟前方)会被惩罚;这是调度时浮现的 dependency stall。
  • Async overlap 由 TpuAsyncTracker 按物理资源门控:一个 0..46 的 TpuResourceType enum(基础 OSS collective 类 + TPU ICI 方向 + host-DMA + SparseCore + VMEM + DCN),带有逐资源 hazard class;两个 async op 只有在不争用同一个 hazardous resource 时才会重叠。
逐 computation 入口(薄封装)DefaultSchedulerCore::ScheduleComputation(comp) @ 0x1362e920 — 通过 vtable +24 构建状态,force-delay RE2 gate,分派 vtable +40
Drain-loop workerDefaultSchedulerCore::ScheduleComputation(comp, state) @ 0x1362eb60(LHS src 3633–3783)
单步DefaultSchedulerCore::SchedulingStep(state) @ 0x1362d480(LHS src 3466–3471)
候选选择 / 比较器DefaultSchedulerCore::FindAndExtractBestNodeAvailable @ 0x13618880(2458 行反编译;22 个 ReadySetLt 键)
Depth/height 预计算HloScheduleGraph::InitializeGraphAnalysis @ 0x1362a860(1264 行)
Async 资源模型TpuAsyncTracker::GetResourceName @ 0x10fff420; GetResourceHazardType @ 0x110015e0
逐节点代价来源CostModelLatencyEstimator::GetLatencyBetween @ 0x10ff8f00 → bundle MaxResourceCycles(见 Bundle-Aware Cost
Mem-budget retry 驱动器LatencyHidingScheduler::RunImpl @ 0x136321a0(LHS src 4182);relax factor qword_A2DFD10 = 0.9
状态偏移ready_set @ state+208(qw 26);pending @ state+504(qw 63);current_time(double)@ state+0x138
源文件third_party/tensorflow/compiler/xla/service/latency_hiding_scheduler.cc.rodata 中的行号引用)
置信度CONFIRMED(byte-anchored),除非某行另有说明

Drain Loop — ScheduleComputation

目的

ScheduleComputation 调度单个 HloComputation。公开的逐 computation 入口(0x1362e920)是一个薄封装,它构建逐 computation 的 SchedulingState,应用 force_delay / scheduling-group frontend-attribute gate,并转发到运行实际 drain loop 的 worker(0x1362eb60)。worker 是调度器的核心:它 drain 一个 ready-set,直到每条指令都被放置,然后翻转反向序列,并(可选地)捕获 ScheduleProto

入口点

text
DefaultSchedulerCore::ScheduleComputation(comp)        0x1362e920   ── thin wrapper
  ├─ vtable+24  MakeSchedulingState(comp)              0x1362e620   ── per-comp state
  │     └─ RE2::RE2 / RE2::PartialMatchN               (src 76-79)  ── force_delay / scheduling-group attr match
  └─ vtable+40  ScheduleComputation(comp, state)       0x1362eb60   ── worker, the drain loop
        ├─ HloScheduleGraph::InitializeGraphAnalysis    0x1362a860   ── precompute depth/height
        ├─ MemoryPressureTracker::Reset                              ── seed pressure tracker
        ├─ FindTopRoots / FindBottomRoots                            ── seed ready_set (roots first)
        ├─ [loop] SchedulingStep(state)                 0x1362d480
        └─ ComputationScheduleToProto                   0x13630780   ── optional capture (src 3783)
```text

这个薄封装通过 `DefaultSchedulerCore` vtable 分派,因此 SparseCore 变体(`SparseCoreSchedulerCore`、`SparseCoreCostModelSchedulerCore`)可以替换自己的 `MakeSchedulingState` / worker,同时复用同一个 `SchedulingStep` 机制。vtable slot 在 `0x1362e920` 处被逐字读取:`(*(...)(*(_QWORD*)core + 24))(...)` 构建状态,`(*(...)(*(_QWORD*)core + 40))(...)` 运行 worker。

### 算法

```c
// DefaultSchedulerCore::ScheduleComputation(comp, state)  @ 0x1362eb60  (LHS src 3633-3783)
StatusOr<HloInstructionSequence>
ScheduleComputation(const HloComputation *comp, shared_ptr<SchedulingState> st) {
    CHECK(st != nullptr);                                    // src 3649 "sched_state != nullptr"
    HloScheduleGraph *g = &st->sched_graph;                  // state+8
    g->InitializeGraphAnalysis();                            // 0x1362a860 — fills depth/height (+0x38/+0x40)
    g->ResetScheduling();
    MemoryPressureTracker::Reset(st->pressure_tracker /*qw51*/, comp,
                                 ModulePressureState::GetPressureStateForComputation(...));

    // seed the ready set with the DAG roots; bottom-up by default, top-down if
    // st->schedule_top_down (byte +440) == 1.
    if (st->schedule_top_down /*byte+440*/) g->FindTopRoots(&roots);     // src ~3700
    else                                    g->FindBottomRoots(&roots);
    for (HloGraphNode *r : roots) {
        r->ready_time = 0;                                   // node[+0x28] = 0   (root)
        st->ready_set.push_back(r);                          // unless force-delayed → nodes_pending
    }
    VLOG(5) << "Initial memory pressure for " << comp << ": " << peak;   // src 3717

    // ---- THE DRAIN LOOP ----  (src 3719)
    while (st->ready_set_size /*qw26 @ +208*/ || st->pending_size /*qw63 @ +504*/) {
        VLOG(10) << "Current ready time: " << st->current_time;          // src 3721 (state+0x138)
        VLOG(2)  << "Current ready queue:";                             // src 3722
        // first try the non-annotated ops; if none can be placed, retry an
        // annotation group with minimum resources (TryScheduleOneAnnotationGroup).
        Status s = TryScheduleOneAnnotationGroup(comp, st, /*use_max=*/1);// src ~3732
        if (!s.ok()) {
            VLOG(3) << "Failed to schedule any non-annotated ops, trying again "
                       "with minimum resources for annotation groups";   // src 3745
            s = TryScheduleOneAnnotationGroup(comp, st, /*use_max=*/0);   // src ~3755
            TF_RETURN_IF_ERROR(s);
        }
    }

    if (st->schedule_top_down /*byte+440*/ == 0)
        reverse(st->new_sequence_reversed);                  // built reversed → flip to program order
    CHECK(st->new_sequence_reversed.size() ==
          g->GetOriginalInstrList().size());                 // src 3779 "Not all instructions ... scheduled"

    if (st->capture_schedule_proto /*byte+376*/)
        *captured += ComputationScheduleToProto(comp, *st, *latency_estimator);  // src 3783
    return st->new_sequence;
}

TryScheduleOneAnnotationGroup 是逐迭代主体,最终调用 SchedulingStep;它的两次尝试结构(annotation groups 使用 max-resources,然后使用 min-resources)是该循环唯一的 retry。drain 条件在 0x1362eb60 第 539 行按字节精确读取:while ( *((_QWORD *)state + 26) || *((_QWORD *)state + 63) ),qword index 26 是 state+208(ready-set size),index 63 是 state+504(annotation-pending count)。src 3779 的终止 CHECK 会将 new_sequence_reversed.size() 与原始指令数比较;如果不同,则 FATAL "Not all instructions have been scheduled <n> vs <m>"

QUIRK — 调度是反向构建的。 默认情况下,ready-set 从 FindBottomRoots(computation 的输出)播种并向后遍历;current_time 在 roots 处从 0 开始,并随着调度器向 inputs 移动而增加。发出的 new_sequence_reversed 随后在结尾翻转为程序顺序。从 inputs 播种并向前遍历的重新实现会产生镜像优先级,并在每条 async edge 的错误一侧隐藏延迟。方向由 byte+440 标志(schedule_top_down)决定;设置后(annotation-group 模式)改用 FindTopRoots,并跳过反转。

函数映射

函数地址作用
ScheduleComputation(comp)0x1362e920薄封装:构建状态 + force-delay RE2 gate + vtable 分派
ScheduleComputation(comp, state)0x1362eb60worker:drain loop(src 3633–3783)
MakeSchedulingState0x1362e620逐 computation SchedulingState ctor(vtable +24
InitializeScheduler0x1362c180一次性 scheduler init
TryScheduleOneAnnotationGroup0x1362e360逐迭代主体;max-resource 后 min-resource retry
ComputationScheduleToProto0x13630780可选的 schedule-dump 填充器(src 3783)

单步 — SchedulingStep

目的

SchedulingStep 正好执行一次放置:选择最佳 ready 节点,追加它,推进时钟。它是一个间接点,在这里比较器(FindAndExtractBestNodeAvailable)和节点放置逻辑(ScheduleNode)通过 core vtable 被调用,因此子类可以覆盖其中任一半。

算法

c
// DefaultSchedulerCore::SchedulingStep(state)  @ 0x1362d480  (LHS src 3466-3471)
Status SchedulingStep(SchedulingState *st) {
    HloGraphNode *node = nullptr;
    // vtable+128 : FindAndExtractBestNodeAvailable(state, pred = empty)
    //   picks the single best ready node by the ReadySetLt comparator and POPS
    //   it out of ready_set (swap-with-last + shrink).
    Status s = core->vtable[+128](&result, core, st, /*pred=*/{});       // 0x13618880
    node = result.node;
    TF_RETURN_IF_ERROR_WITH_SOURCE(s, src 3466);
    CHECK(node != nullptr);                                              // src 3467 "node != nullptr"

    // vtable+112 : ScheduleNode(node, state)
    //   appends node to new_sequence_reversed, advances current_time by the
    //   node's bundle NodeCost, updates the MemoryPressureTracker, releases the
    //   node's resources, and pushes newly-ready successors into ready_set.
    //   The new clock value comes back in xmm0 (var_38).
    double new_time = core->vtable[+112](&done, core, node, st);         // ScheduleNode
    TF_RETURN_IF_ERROR_WITH_SOURCE(done, src 3469);
    st->current_time = new_time;                                        // state+0x138  (src 3470)

    VLOG(2) << "Scheduled: " << node->name();                           // src 3470
    VLOG(5) << node->ToString();                                        // src 3471
    return OkStatus();
}
```text

两个 vtable 间接调用在 `0x1362d480` 处逐字读取:`(*(...)(*(_QWORD*)core + 128))(...)` 是候选选择器,`(*(...)(*(_QWORD*)core + 112))(...)` 是 `ScheduleNode`。时钟写回是第 60 行的 `vmovsd qword ptr [r14+138h], xmm0`,`current_time` 位于 `state+0x138`,也是比较器下一步读取的唯一可变标量。

> **NOTE — `ScheduleNode`(vtable `+112`)在这里没有逐行反编译。** 它的效果从比较器读取和 `current_time` 写回推断而来:它按节点的 bundle cost([Bundle-Aware Cost](../cost/bundle-aware-cost.md))推进时钟,修改 `MemoryPressureTracker`,并用最后一个 predecessor 刚 retired 的 successors 刷新 `ready_set`。它驱动的 resource-occupier bookeeping 位于 `AddOccupierToResource`(`0x1361dc00`)/ `DeleteOccupierFromResource`(`0x1361d9c0`),这是比较器后续步骤读取的 resource-conflict 和 release 键的动态状态。将 clock-advance 和 ready-set refresh 视为 CONFIRMED(写回有 byte-anchored 证据);逐资源 occupier push/pop 为 HIGH(由比较器读取推断)。

---

## Critical-Path Depth 和 Height — `InitializeGraphAnalysis`

### 目的

比较器最强的非 force、非 memory 键是 **async depth****async height**,即节点从 async leaves 到自身以及从自身到 async roots 的关键路径距离,以 cost units 计量。它们*不会*在每次比较时重新计算;在 drain loop 开始前,它们由 `InitializeGraphAnalysis`(`0x1362a860`)一次性写入每个 `HloGraphNode`,比较器只读取字段。

### 算法

```c
// HloScheduleGraph::InitializeGraphAnalysis()  @ 0x1362a860  (1264-line decompile)
//   Two reverse-topological sweeps over the schedule graph.
void InitializeGraphAnalysis() {
    // (1) ASYNC DEPTH  — node[+0x38]
    //   for each node in forward topo order:
    //       depth = max over predecessors p of  p.depth  (raised by async edge cost)
    //   implemented as the vmaxsd chain over [pred+0x38] writing [node+0x38]
    for (node in topo_order) {
        double d = 0.0;
        for (pred in node.predecessors)
            d = max(d, pred.depth);                          // vmaxsd [r12+38h], [rsi+38h] (lines 539-597)
        node.depth = d;                                      // vmovsd [r12+38h], xmm0
    }
    // (2) ASYNC HEIGHT — node[+0x40]
    //   for each node in reverse topo order:
    //       height = (max / sum) over successors s of  s.height + edge_latency
    //   the async edge contributes its GetLatencyBetween cost; the chain at
    //   lines 946-1000 mixes vmaxsd (critical path) with vaddsd (async accumulation).
    for (node in reverse_topo_order) {
        double h = 0.0;
        for (succ in node.successors)
            h = max(h, succ.height /* + edge cost */);       // vmaxsd / vaddsd [rax+40h] (lines 946-1000)
        node.height = h;                                     // vmovsd [rax+40h], xmm0
    }
}

depth sweep 是第 481–597 行的 vmaxsd xmm0, xmm0, [pred+38h] / vmovsd [node+38h], xmm0 序列;height sweep 是第 946–1000 行围绕 [node+40h]vmaxsd/vaddsd 簇。两者都是纯 max-propagation(height 额外累加 async-edge latency),所以 [+0x38] 是从 leaves 进入该节点的最长代价路径,[+0x40] 是从该节点出去到 roots 的最长 async-cost 路径。馈入 height propagation 的 edge cost 是来自 GetLatencyBetween 的 bundle/latency cost(Bundle-Aware Cost)。

QUIRK — 这里的 “depth” 和 “height” 是 async critical-path costs,不是图层级。 [+0x38]/[+0x40] 是通过 async edges 累加的 double 周期代价,不是整数 DAG depths。比较器偏好更深的 async depth(键 15)和更高的 async height(键 16),正是因为位于长 async critical path 上的节点应当被尽早发出,好让其 dependent async op 的延迟隐藏在它下面的一切工作之下。使用整数拓扑层级的重新实现,只有在所有 edge cost 都相同的时候才会正确排序候选者,而现实中并非如此。


候选比较器 — ReadySetLt(22 键)

目的

FindAndExtractBestNodeAvailable0x13618880,2458 行反编译)迭代 ready-set,维护一个当前“best”候选者,并对每个竞争者调用从 OSS latency_hiding_scheduler.cc 内联的 ReadySetLt 比较器。比较器是固定的有序 tie-break 键列表:每个键自上而下求值,第一个区分两个候选者的键决定赢家,并记录一个 CandidateResult 原因字符串(kXxx enum)。当无法选出候选者时,函数会发出两种诊断之一并失败。

这里没有单独的用户比较器:这个构建中的 candidate_compare_post_step_mutator_ SchedulerConfig hook 都为空,因此本文记录的内置 ReadySetLt 链就是实际运行的内容。ILP 变体(ILP-LHS)只替换 async classifier,从不替换这条链。

算法

下面是按求值顺序排列的键,包括被比较的 HloGraphNode 字段、CandidateResult 原因字符串(22 个都已在反编译中验证存在)以及获胜方向:

#原因字符串比较对象(HloGraphNode 字段)偏好
1kForceEarly[+16] force-early flagflag-set 优先
2kForceDelay[+18] force-delay flagflag-NOT-set 优先(delay later)
3kForceDelayAfterTarget[+18] force-delay flag(post-target 变体)flag-NOT-set 优先(delay later)
4kForceDelayPriority[+12] force-delay priority(int)LOWER priority 优先
5kPreference[+0x14] preference(float)HIGHER 优先
6kMemoryPeakOverLimit(live_peak + ΔMemoryPressure) vs memory_limit保持 UNDER limit 的那个
7kOnlyDecreaseMemoryOverLimitover limit 时 only-this-reduces-pressure会降低的那个
8kDecreaseMemoryOverLimit二者都 ΔMemoryPressure < 0,比较幅度更大的 reduction
9kScheduleAsyncStart[+24] & 9 == 9(async-start kind)async-start 优先(当 “schedule ASAP”)
10kScheduleDoneReadySetLt::ShouldScheduleAsyncDone(node)应触发的 async-done
11kScheduleStart[+17] == 1 且 operand-shape matchmatching start
12kReleaseNonextendable[+0x28] ready_time vs current_time,occupier count ≥ 2释放 non-extendable resource 的那个
13kLessSerialResourceConflictGetNumConflictingSerialResources(node)FEWER conflicts 优先
14kNotValuableForSelectiveOverlap[+24] >> 7 selective bit 和 ready-min ≤ cfg[+88]valuable 的那个(由 config +130 门控)
15kFreeBackedupResource[+40] frees-backed-up + UpdateCandidateResourceConstrainedfreeing 的那个
16kAsyncDepth[+0x38] async depthDEEPER 优先(当 schedule-ASAP)
17kAsyncHeight[+0x40] async height(到 async 的 critical path)HIGHER 优先
18kUnlockStart[+24] & 0x20(unlock-start bit)unlocking-start 优先
19kUnlockDone[+24] & 0x10(unlock-done bit)unlocking-done 优先
20kCreatesMoreReadyNodes[+100] #successors-made-ready(uint)MORE 优先
21kDecreaseMemoryΔMemoryPressure sign(一个 <0,一个 ≥0降低 pressure 的那个
22kOriginalOrder[+28] original sequence index(int)LOWER index 优先(稳定;schedule-ASAP 时翻转)

键 6–8 由 config byte +129(prioritize-pressure)结合 over-limit 状态门控;键 14 由 config byte +130 门控。schedule-ASAP 标志(反编译中的 v36)会翻转若干键(async-start、async-depth、original-order)的方向,使得在 ASAP regime 中,async work 和长 critical-path 节点会被尽可能早地拉出。

延迟隐藏项

延迟隐藏启发式的核心是一个出现在键 12、16 和 17 中的单个 clamp 后减法:

c
// remaining latency of a candidate node  (lines 988-991, 1025-1028, 1463-1493)
double remaining = max(0.0, st->current_time - node->ready_time);
//   vmovsd  xmm1, [r14+138h]        ; current_time      (state+0x138)
//   vsubsd  xmm1, xmm1, [r12+28h]   ; - node.ready_time (node+0x28)
//   vmaxsd  xmm1, xmm2, xmm1        ; clamp at 0
```text

`ready_time`(`node[+0x28]`)是考虑其 operands 后节点*能够*被调度的最早 cycle;`current_time`(`state+0x138`)是当前 step 的时钟。当 `current_time` 已经过了节点的 `ready_time`,`remaining` 为 0,也就是节点的 inputs 已 retired,可以无 stall 发出。当 `ready_time` 仍在时钟前方时,该节点会 stall,而比较器用这个项(连同 async height 键 17)来偏好现在发出一个*不同*节点,并稍后再发出这个会 stall 的节点。这正是来自 `GetLatencyBetween`([Bundle-Aware Cost](../cost/bundle-aware-cost.md))的 dependency stall 在调度时重新浮现。当 async-height(`[+0x40]`)和 async-depth(`[+0x38]`)键都无法区分这一对候选者时,比较器落到一个 clock-distance tie-break(第 12071218 行):它为每个候选者形成 `first_root.ready_time − current_time − candidate[+0x30]`,通过在第 1212 行 against `qword_A2DF238`(`0x7FFFFFFFFFFFFFFF`,IEEE abs-value mask,byte-confirmed)的 `vandpd` 取绝对值,并偏好更小的幅度,也就是 `[+0x30]` 估计更接近当前时钟的候选者。

> **GOTCHA — memory-peak gate(键 6)可以覆盖所有后续键。** `kMemoryPeakOverLimit` 会把候选者的 `MemoryPressureTracker::MemoryPressureDifference` 加到 live peak 上,并将结果与 `memory_limit`(`state→config + 672` 处的 `SchedulerConfig` 副本)比较。如果只有一个候选者能保持 peak `≤ limit`,它就会获胜,*不管 async depth、height 或它下面任何 latency-hiding 键如何。* 把 memory 键放在 async 键后面的重新实现会以 OOM 为代价隐藏延迟。这个 gate 有意位于位置 6,高于所有 overlap heuristics。

### 失败诊断

当 ready-set 非空但没有节点通过逐指令 `should_schedule_` predicate 或 annotation-group overlap limits 时,候选者会被收集到一个 `InlinedVector<pair<HloGraphNode*, SkipNodeReason>, 2>`,函数会用两个消息之一失败,二者都锚定在反编译中:

```text
"There is a scheduling group which exceeds the overlap limits. Annotation id: <n>.   (line 1742, LHS src 2092)
 It needs <x> resources, but the limit is <y>."
"FindAndExtractBestNodeAvailable failed to find a node to schedule, skipped nodes: <list>"   (line 1939)

比较器内部的 VLOG 诊断会叙述选择:"Choosing from ready (<name>) Reason: First Candidate"(第 534 行)以及 "... mem pressure chosen <m> mem pressure other <n>"(第 840 行)。

函数映射

函数地址作用
FindAndExtractBestNodeAvailable0x13618880ready-set 迭代 + 22 键 ReadySetLt 链 + pop
FindAndExtractBestAnnotatedNode0x1361e2c0annotation-group-restricted 变体
ReadySetLt::ShouldScheduleAsyncDone(inlined)键 10 predicate;在 pair 上调用两次(lines 969–977)
ReadySetLt::UpdateCandidateResourceConstrained(inlined)将 resource occupancy 折叠进键 13/15(lines 581–583);body 未 dump
GetNumConflictingSerialResources(inlined)键 13 conflict count;恢复为 call target

Async 资源模型 — TpuAsyncTracker

目的

两个 async op 是否可以重叠由每个 op 消耗的物理资源及该资源的 hazard class 决定。TpuAsyncTracker 用 TPU-specific resources 扩展 OSS AsyncTracker:六个方向性 ICI links、host-DMA taps、SparseCore queues、VMEM 和 DCN bandwidth。比较器的 resource-conflict 键(12 和 14)以及 ScheduleNode occupier bookkeeping 会读取这个模型来门控 overlap。完整的 0..46 TpuResourceType enum 和 hazard taxonomy 编目在 ResourceType Taxonomy 页面;本节固定命名和分类资源的 dispatch。

算法

c
// TpuAsyncTracker::GetResourceName(resource)  @ 0x10fff420  (src 1152)
char *GetResourceName(long r) {
    CHECK(r <= 46);                                          // FATAL "resource_type < ...kTpuResourceTypeEnd"
    if (r < 13)  return AsyncTracker::GetResourceName(r);    // 0..12 base OSS collective/copy/sendrecv
    if ((0x17FFF >> (r - 13)) & ((r - 13) < 0x11))           // 13..29 with the bitmask gap at idx 28
        return off_2181E148[r - 13];                         // TPU names (ICI / host-DMA / SparseCore / VMEM)
    if ((r - 30) < 0x10) return "kCustomCollective";         // 30..45
    return &nptr;                                            // 46 sentinel
}

// TpuAsyncTracker::GetResourceHazardType(resource)  @ 0x110015e0  (src 1848)
long GetResourceHazardType(long r) {
    CHECK(r <= 46);
    if (r >= 13) {
        if ((0x109FF >> (r - 13)) & ((r - 13) < 0x11))
            return dword_AC0B2C0[r - 13];                    // [0,1,1,1,1,1,1,0,0,0,0,2,0,0,0,0,2]
        return 3 * (unsigned)((r - 30) >= 0x10) + 1;         // kCustomCollective → 1
    }
    // base resources 0..12:
    if (this->track_sync_op_resource /*byte+202*/ != 1)
        return AsyncTracker::GetResourceHazardType(r);       // default: 4*(r != 5)
    // TPU collective-serialization override:
    if (r == 3 /*kAllReduce*/ || r == 6 /*kReduceScatter*/ ||
        (r == 2 /*kAllGather*/ && this->byte_314))
        return 3;                                            // kSerial — collective engine single-occupancy
    return AsyncTracker::GetResourceHazardType(r);
}
```text

分割是逐字读取的:`r > 46` 在源代码第 1152 行(name)/ 1848 行(hazard)FATAL;`r < 13` 委托给基础 `AsyncTracker`;13..29 索引 `off_2181E148`(names)和 `dword_AC0B2C0`(hazards);30..45 全部是 `kCustomCollective`。基础 hazard 是 `4 * (r != 5)`,也就是说每个基础 class 都是**可共享的**(hazard 4,overlap 到逐 kind limit),*除了* `kCopy`(`r == 5`,hazard 0 = unsharable,async copies 在 copy engine 上串行化)。

hazard codes 是:`0` = unsharable(single-issue),`1` = serial(one in flight,FIFO),`2` = nonextendable(不能被延后到其窗口之外,即 `kVmem`),`3` = serial-TPU-collective-override,`4` = shareable。TPU resource hazard table 是 `dword_AC0B2C0 = [0,1,1,1,1,1,1,0,0,0,0,2,0,0,0,0,2]`,对应 indices 13..29:DCN bw unsharable,六个 ICI directions、两个 host-DMA taps 和 `kSparseCore` serial,SparseCore sub-queues unsharable,`kVmem` nonextendable。

> **QUIRK — collectives 按 ICI *direction* overlap,而不是按单个 “collective” counter。** `MayAddIciLinks`(`0x10fffb20`)解码 collective 的 opcode,用 `CostModel::GetCycles` 为 transfer 定价(第 187 行),并把代价存入*六个方向性 ICI counters*(resources 14..19)之一,方向由 tracker byte `+208` 处的逐 collective host-to-device flag 选择(第 130 行:`kHostToDevice ? r+1 : 2-r`)。不同 ICI axes 上的两个 collectives 可以自由 overlap;同一 axis 上的两个会串行化(hazard 1)。把 “one collective resource” 作为模型的重新实现会错误地把 X+ 上的 all-reduce 与 Y+ 上的 all-gather 串行化。
>
> **NOTE — TPU collective-serialization override 由配置门控。** `byte+202` “track-sync-op-resource” 标志和 `byte+314` all-gather 标志决定 `kAllReduce`/`kReduceScatter`/`kAllGather` 是否从 shareable(4)提升为 serial(3)。当标志未设置时,基础 OSS hazard(shareable)适用,同 kind collectives 可以 overlap 到其逐 kind limit。逐 kind limits 本身来自 `TpuAsyncTracker` ctor args(`a13..a20`,通过 `AutoOr<long>` 默认到 `INT64_MAX`);只有 DCN(`a5 = env+4568`)和 host/ICI pool(`a6 = env+5320`)limits 锚定到具体 env offsets。将 override gating 视为 CONFIRMED,将精确逐 kind default limits 视为 PARTIAL。

### 函数映射

| 函数 | 地址 | 作用 |
|---|---|---|
| `TpuAsyncTracker::GetResourceName` | `0x10fff420` | resource id → name;0..46 dispatch |
| `TpuAsyncTracker::GetResourceHazardType` | `0x110015e0` | resource id → hazard class;collective override |
| `TpuAsyncTracker::GetNumAvailableResources` | `0x10fff600` | resource id → 逐 kind overlap limit |
| `TpuAsyncTracker::MayAddIciLinks` | `0x10fffb20` | 将 collective cost 存入方向性 ICI counter |
| `TpuAsyncTracker::MayAddHostTransfers` | `0x11000280` | host send/recv → host-DMA counters |
| `TpuAsyncTracker::ctor` | `0x10ffba60` | 赋值逐 kind limits(`a5/a6` env-anchored) |

---

## Memory-Budget 驱动器 — `RunImpl`

上面的 drain loop 只调度*一次*尝试。`LatencyHidingScheduler::RunImpl`(`0x136321a0`,LHS src 4182)把它包装在一个 memory-pressure retry loop 中。每次调度尝试后,它检查结果的 HBM footprint 峰值是否符合预算;如果不符合,就收紧预算并重新运行,迫使比较器的 memory 键(68)在下一轮更积极地触发。

```c
// LatencyHidingScheduler::RunImpl memory retry  @ 0x136321a0  (lines 882-926)
//   peak footprint = EstimateFragmentationSize() + computed_use()
//   while ( frag + computed_use() > memory_limit  &&  retries-- > 0 ):
//       memory_limit = next_memory_limit();    // TIGHTENS by 10%
//       re-run ScheduleComputation
double next_memory_limit(uint64 current) {
    double d = (double)current;                  // exact u64->f64 (2^52/2^84 biases)
    d *= 0.9;                                     // qword_A2DFD10 = 0x3FECCCCCCCCCCCCD   <<< RELAX FACTOR
    return (int64)floor(d);                       // vcvttsd2si, with 2^63 (qword_A2DF388) signed-cast fixup
}

NOTE — next_memory_limit 会缩小预算;它并不会放宽预算。 OSS 名称读起来像是向上放宽预算,但 byte-exact factor 是 qword_A2DFD10 = 0.9(在 RunImpl 第 892/918 行读取),所以每次 retry 都会把预算缩小 10%。失败的尝试会要求列表调度器适配一个小 10% 的预算,从而推动它偏好降低 memory-pressure 的候选键(6–8),胜过 latency-hiding 键,也就是用 overlap 换取更低的 HBM 峰值。除非 EnableSchedulerMemoryPressureTracking 或 post-SC-assignment rerun byte 让 limit 有限,否则该循环是惰性的(memory_limit == -1)。

peak-footprint estimate 的 fragmentation 项来自 EstimateFragmentationSize0x13633c00,LHS src 334/336),它不是闭式 alignment 公式,而是完整的 heap simulation:它构建一个 GlobalDecreasingSizeBestFitHeap<HloValue>(largest-buffer-first best-fit,alignment 1)并运行 HeapSimulator::Run,返回 max(0, heap.fragmentation_size),即 allocator gaps 损失的字节数。vmaxsd/vmulsd 常量(qword_A2DFD10qword_A2DF388)和 heap 构建在反编译中 byte-confirmed。


Worked Example — Overlap vs Stall

6acc60406(TensorCore clock f MHz)上的 4-op 片段,基础调度已经由 memory scheduler 放置:

text
%ar   = all-reduce-start(...)      ; collective on ICI X+  (resource 16, hazard 1)
%ar.d = all-reduce-done(%ar)
%mm   = dot(%a, %b)                ; bf16 matmul, 212-cycle bundle
%add  = add(%ar.d, %mm)
```text

matmul 的 `NodeCost` 是 `212 / (f·1e6)` s;all-reduce 的 transfer cost 是存入 ICI X+ counter(resource 16)的 `CostModel::GetCycles`。在 `%ar` 于 `t0` 发出后走比较器:

- `ready_set = {%mm}`(`%ar.d` 要到 ICI latency 流逝后才 ready;`%add` 等待二者)。`%mm` 是唯一候选者,所以它在 `t0` 被调度,`current_time` 推进 212 cycles。在这个窗口中,`%ar` 的 ICI transfer 正在一个*不同*资源(ICI X+,不是 MXU slot)上 in flight → **OVERLAP**:matmul 隐藏 all-reduce latency。async-height 键(17)此前已经偏置调度器去尽早发出 `%ar`(它的 `[+0x40]` 到达 `%add`),而 matmul 填补了 bubble。
- 如果 ICI latency `< 212` cycles,`%ar.d` 会在 `%mm` 完成前 ready;`kScheduleDone`(键 10)接着触发它,然后是 `%add`。总计 ≈ 212 cycles,也就是 collective 完全隐藏,zero stall。
- 如果 ICI latency `> 212` cycles(cross-host DCN),在 `%mm` 之后没有其他 ready non-async node,所以 `current_time` 跳到 `%ar.d` 的 `ready_time`。**STALL** = `(latency − 212)` cycles,计为 wasted all-reduce cycles。第二个独立 matmul 会让键 17/20(`kAsyncHeight`/`kCreatesMoreReadyNodes`)塞入 bubble 并减少 stall。
- 如果 `memory_limit` 紧张,`kMemoryPeakOverLimit`(键 6)可以强制 `%ar.d`(释放 all-reduce buffer)排在 `%mm` 前面,即使这会丢失 overlap,也就是用 latency-hiding 换取更低的 HBM 峰值。0.9 retry shrink 是跨 reruns 把调度器推入这一 regime 的驱动力。

---

## 置信度汇总

| 论断 | 证据 |
|---|---|
| Drain loop 是 `while (ready_set @ +208 \|\| pending @ +504)` | `ScheduleComputation(comp,state)` @ `0x1362eb60` 第 539 行 |
| `SchedulingStep` = vtable+128 select,vtable+112 `ScheduleNode`,写入 `current_time` @ `+0x138` | `0x1362d480` 第 32/53/60 行 |
| 调度反向构建,结尾翻转;`byte+440` 选择方向 | `FindTopRoots`/`FindBottomRoots` + reverse @ `0x1362eb60` |
| 比较器是 22 键 `ReadySetLt` 链(所有 reason strings 存在) | `0x13618880`,第 534–1553 行(22 个 `kXxx` strings) |
| 延迟隐藏项 `max(0, current_time − ready_time)` | `vsubsd [r14+138h]−[r12+28h]` + `vmaxsd` @ 第 988-991 行 |
| Async depth/height 在 `[+0x38]`/`[+0x40]` 读取,由 max-propagation 预计算 | `InitializeGraphAnalysis` @ `0x1362a860` 第 481-1000 行 |
| Memory-peak gate(键 6)先于所有 overlap keys | `kMemoryPeakOverLimit` branch @ 第 771 行,位于 async keys 前 |
| `TpuAsyncTracker` 0..46 分割;base hazard `4*(r!=5)`;TPU table `[0,1,1,…,2]` | `GetResourceName`/`GetResourceHazardType` @ `0x10fff420`/`0x110015e0` |
| Collectives 按 ICI direction overlap(byte `+208`),通过 `GetCycles` 定价 | `MayAddIciLinks` @ `0x10fffb20` 第 130/187 行 |
| `next_memory_limit` 每次 retry 将预算缩小 0.9 | `qword_A2DFD10` @ RunImpl `0x136321a0` 第 918 行 |
| Fragmentation = `GlobalDecreasingSizeBestFitHeap` HeapSimulator run | `EstimateFragmentationSize` @ `0x13633c00` 第 39/52/73 行 |
| `ScheduleNode` 内部(clock-advance/resource-release)为推断,未 dump | 比较器读取 + `current_time` 写回 |
| 逐 kind async overlap limits(`a13..a20`)默认 `INT64_MAX` | `TpuAsyncTracker` ctor @ `0x10ffba60`;只有 `a5/a6` env-anchored |

---

## 交叉引用

- [Scheduler Overview](overview.md) — LHS 在 TPU scheduling pipeline 中的位置以及两个 `AddPass` callsites。
- [LHS post_layout](lhs-post-layout.md) — 驱动此核心的 post-layout LHS pipeline placement。
- [LHS post_layout_pre_fusion](lhs-post-layout-pre-fusion.md) — pre-fusion 变体的 `SchedulerConfig` 和 hook table。
- [LHS ILP variant](lhs-ilp-variant.md) — 在此比较器前替换 async classifier 的 `EnableIlpLatencyHidingScheduler` gate(比较器链不变)。
- [ResourceType Taxonomy](scheduler-resourcetype-model.md) — 本页 `TpuAsyncTracker` dispatch 消费的完整 0..46 `TpuResourceType` enum、hazard classes 和逐资源 availability counts。
- [Cost Model Overview](../cost/overview.md) — 每个 cycle number 背后的三类 cost-class families 和 `Target` clock wiring。
- [Bundle-Aware Cost](../cost/bundle-aware-cost.md) — 馈入此比较器的 `NodeCost` 和 async-height 键的 `MaxResourceCycles` bundle cost 与 `LatencyBetween` dependency latency。
- **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 — [back to index](../index.md)