LHS post_layout 变体

本页中的所有地址都适用于来自 libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。由未剥离、完整符号 ELF 的静态分析重建（nm -C 可解析每个方法）；其他版本会有所不同。

摘要

XLA LatencyHidingScheduler (LHS) 的 post_layout 调用，是 v0.0.40 在生产中实际运行的调度器 pass。它位于 RunHloScheduler (0x1096fac0) 中，也就是 DeepseaCompilerBase::RunHloPasses 的最终调度阶段，该阶段在布局分配和主融合之后执行。由于布局分配已经运行，它生成的 schedule 能看到真实张量布局、内存空间放置，以及关键的 SparseCore offload assignment 结果。这就是 wiki 另外两个变体页面用相对本页的差异来描述的 LHS：lhs-post-layout-pre-fusion（同一主体接线在死掉的 "Pre main fusion" 槽位）和 lhs-ilp-variant（同一函数内部由 flag 控制的异步分类器切换）。

重新实现者首先会意外发现，RunHloScheduler 构建并顺序运行两个 xla::HloPassPipeline 对象，而不是一个嵌套流水线。第一个名为 final_scheduler，用 HloMemorySchedulerWithBrkgaFallback 放下一个基础 schedule，即 LHS 作为前置条件要求的、最小化内存压力的顺序（module->has_schedule()）。第二个名为 async_scheduling，包含 LHS 本身以及围绕它的异步重叠重写。LHS 在资源模型下将异步 start/done op 对与 compute 重叠，行为与上游 XLA LatencyHidingScheduler 完全一样，但由 TPU 特定的 SchedulerConfig、TPU LatencyEstimator 和了解 ICI link、HBM bank、megacore queue、host-DMA tap 以及 SparseCore offload queue 的 TPU AsyncTracker 驱动。

最重要的结构性事实是，LHS RunImpl 主体在所有变体之间共享：只有一个已编译的 RunImpl (0x136321a0)，即未修改的上游 latency_hiding_scheduler.cc 主体。变体之间只在传给 AddPass<LatencyHidingScheduler> 的两个输入上不同：SchedulerConfig POD 和 SchedulingContext（估算器 + tracker）。因此本页记录 post_layout pass 的主体（共享 RunImpl 及其内存压力重试循环）和 post_layout 的配置差异（GetSchedulerConfig 与 GetTpuAsyncTracker 在这里注入、而融合前变体不会注入的内容）。

对于重新实现，约定如下：

RunHloScheduler 的双流水线结构：先运行 final_scheduler（基础内存 schedule），再运行 async_scheduling（LHS），顺序执行，并且 LHS 受两个 IsPassDisabled 检查控制。
由 GetSchedulerConfig 构建的 SchedulerConfig POD：按 kind 的重叠限制、memory_limit 三态，以及区分首次运行与重跑的 v45 post-SC-assignment rerun gate。
SchedulingContext 接线：选择哪个 LatencyEstimator（approximate、real cost model 或 PGLE），以及 TpuAsyncTracker 资源模型在此阶段额外加入什么（reserved-SparseCore 集合）。
共享 RunImpl 主体：has_schedule() 前置条件、候选列表调度器调用，以及带 relax-by-factor next_memory_limit() 的、感知碎片的内存压力重试循环。


驱动器	`xla::jellyfish::(anonymous)::RunHloScheduler` — `0x1096fac0`（2075 行反编译）
阶段	第 7 阶段 — 最终调度，在 LayoutAssignment 和主融合之后
流水线	`final_scheduler`（基础 sched）→ `async_scheduling`（LHS）— 两个局部 `HloPassPipeline` 对象
AddPass 位置	`0x1096fac0` 第 1137 行（常规）和第 1411 行（SparseCore-queue 路径）
Pass 主体	`xla::LatencyHidingScheduler::RunImpl` — `0x136321a0`（共享，1217 行反编译）
配置构建器	`xla::jellyfish::(anonymous)::GetSchedulerConfig` — `0x10974aa0`（137 行反编译，137 字节 POD）
估算器构建器	`xla::jellyfish::(anonymous)::GetLatencyEstimator` — `0x10974e00`
Tracker 构建器	`xla::jellyfish::(anonymous)::GetTpuAsyncTracker` — `0x10975520` → `TpuAsyncTracker::Create` `0x10ffb3e0`
Scheduler core	`xla::DefaultSchedulerCore::DefaultSchedulerCore` — `0x10976ce0`
在 0.0.40 中已接线？	是 — 这是规范 LHS；融合前槽位是死的

流水线位置

目的

RunHloScheduler 是 DeepseaCompilerBase::RunHloPasses (0x1093a420) 的第七个也是最后一个阶段。第 1-6 阶段是 HLO→HLO 重写流水线（PreOptimization、sharding、SPMD、layout-through-optimize、LayoutAssignment、PostOptimization）。当控制到达 RunHloScheduler 时，每个张量都有具体 layout 和内存空间，fusion 已经定稿，并且在此函数内部，SparseCore offload op 将完成 assignment。这里的任务是生成附加到 module 的指令 schedule（module->set_schedule），先作为最小化内存的基础顺序，再由 LHS 精修，使异步传输与 compute 重叠。

入口点

text

DeepseaCompilerBase::RunHloPasses (0x1093a420)        [Phase 7 / final]
  └─ RunHloScheduler (0x1096fac0)
        Pipeline A — "final_scheduler"  (&v318, inline 16B SSO at line 373)
          MaybeAddInvariantCheckers                       (line 404)
          [opt] LatencyHidingLayerScheduler               (line 504, default OFF)
          HloMemorySchedulerWithBrkgaFallback             (line 567)  ← base schedule
          ConstantDeferring / BarnaCoreFeedScheduler      (lines 611/615)
          Run(Pipeline A)                                 (line 620)
        Pipeline B — "async_scheduling"  (&v341, inline 16B SSO at line 679)
          AsyncOpScheduler / AsyncCollectiveMerger / …    (lines 746-754)
          Run(Pipeline B, first wave)                     (line 765)
          if !IsPassDisabled("latency-hiding-scheduler"):
            GetSchedulerConfig(memory_limit, env, target, is_lem)   (line 802)
            GetLatencyEstimator(...)                       (line 828)
            GetTpuAsyncTracker(..., &reserved_sc)          (line 987 / 1349)
            make_unique<DefaultSchedulerCore>(ctx, cfg)    (line 1046 / 1400)
            AddPass<LatencyHidingScheduler>(ctx, core)     (line 1137 / 1411)  ← THIS PASS
          Run(Pipeline B, LHS wave)
        TpuCopyInsertion / HloSchedule::Update / set_schedule (lines 1722-1971)
```text

> **怪癖 —** 两个流水线名称都不会作为可用作 pass 参数的、以 null 结尾的 `.rodata` C 字符串保留下来。二者都是在栈上以内联 16 字节 SSO `std::string` 构建的：`final_scheduler` 通过 `strcpy(ptr, "final_scheduler")` 构建，并设置 SSO 长度字节 `ptr[23] = 15`（第 373 行），随后在第 379-380 行由 `vmovdqu` 将 buffer 复制到字段方向；`async_scheduling` 通过长度字节 `ptr[23] = 16`（第 679 行）以及 `vmovups xmm0, cs:xmmword_865DCAB`（第 682 行）构建。字节锚点验证：`.rodata` 中 VA `0x865dcab` 保存 16 个原始字节 `async_scheduling`（xmmword immediate），而在二进制上 `rg -a -c final_scheduler` 返回 0，说明 `final_scheduler` 只作为内联到函数体中的 `strcpy` 源存在。重新实现者若在二进制中 grep `final_scheduler`，不会把它作为离散诊断字符串找到；`async_scheduling` 恰好出现一次，即该 immediate。

### 两条流水线，不是一个嵌套流水线

`final_scheduler` (`&v318`) 和 `async_scheduling` (`&v341`) 是两个独立的局部 `HloPassPipeline` 对象，各自由自己的 `HloPassInterface::Run` 驱动；它们不是一个嵌套流水线。`final_scheduler` 先运行并完成（第 620 行的 `Run`），随后才构造 `async_scheduling` (`&v341`)，两者只共享 module。由 `HloMemorySchedulerWithBrkgaFallback`（第 567 行）产生的基础 schedule 满足了 LHS 前置条件 `module->has_schedule()`。这正是为什么 LHS 可以在*这里*运行，却无法在第 6 阶段的 "Pre main fusion" 槽位运行：那里没有先前的内存 schedule。

---

## Gate

### 目的

post_layout LHS 受两个 `IsPassDisabled` 检查和一组 async-enable 谓词控制。如果任一 gate 失败，schedule 就只是基础内存 schedule。

### 算法

```c
function RunHloScheduler_lhs_gate(module, env, target, config):   // 0x1096fac0
    // Gate 1: is the async_scheduling pipeline built at all? (line 666)
    if IsPassDisabled("async-op-scheduler", config, env):         // 0x12fd8340
        return base_schedule_only
    if !(ShouldUseAsyncAllGather(target, ...)                     // line 674
         || env[184] == 2                                         // line 676
         || (env[230] >= 2 && env[184] == 0)
         || GetTpuCompEnv(env)[3664]):                            // line 677
        return base_schedule_only

    // ... AsyncOpScheduler etc. run as the first wave (line 765) ...

    // Gate 2: is the LHS itself enabled? (line 798)
    if IsPassDisabled("latency-hiding-scheduler", config, env):
        return async_first_wave_only
    // build config + context + core, then AddPass<LatencyHidingScheduler>

IsPassDisabled (0x12fd8340) 会将 pass 名称与以下内容比对：(a) 反向 enabled-passes 列表，(b) --xla_disable_hlo_passes disabled 列表，以及 (c) 由 debug-options version 字段（env+472）控制的、二分查找的静态 kDisabledPasses 表。

注 — EnableLatencyHidingLayerScheduler (0x1d6b8960, env+2144) 默认 FALSE，因此会跳过 Pipeline A 的 LatencyHidingLayerScheduler 分支（第 434-545 行）。运行的是 async_scheduling 中的规范 LatencyHidingScheduler。

开关（效果）	来源	默认值
`async-op-scheduler` pass disable	`IsPassDisabled` 名称	enabled
`latency-hiding-scheduler` pass disable	`IsPassDisabled` 名称	enabled
`EnableLatencyHidingLayerScheduler`	`0x1d6b8960` / env+2144	false
`xla_latency_hiding_scheduler_rerun`	rerun count	1

SchedulerConfig POD（配置差异）

目的

GetSchedulerConfig (0x10974aa0) 构建 137 字节 SchedulerConfig POD，DefaultSchedulerCore 会复制它，RunImpl 会读取它。在 post_layout 路径中，第 802 行以 GetSchedulerConfig(out, memory_limit_bytes=a9, env, target, is_lem) 调用，其中 is_lem = ModuleContainsLEMSparseCoreInstruction(module)。该构建器与融合前变体相同；差异在于调用点馈入的值，主要是 memory_limit_bytes（最终预算 vs 融合前预算）和 v45 post-SC-assignment rerun 字节。

算法

function GetSchedulerConfig(out, mem_limit_bytes, env, target, is_lem):  // 0x10974aa0
    v45 = env[4097]              // xla_tpu_rerun_latency_hiding_scheduler_post_sc_assignment
    sc_queue = EnableSparseCoreOffloadQueuingInLhs(env, target)          // line 30
    multislice = target.GetMultiSliceTopology()

    out.collective_broadcast_overlap_limit = 1                          // off 0
    out.collective_permute_overlap_limit   = env[3964]                  // off 8
    out.all_to_all_overlap_limit           = 1                          // off 16
    out.all_gather_overlap_limit           = 1                          // off 24

    if sc_queue:                                                        // lines 37-42
        out.all_reduce_overlap_limit      = INT32_MAX                   // off 32
        out.send_recv_overlap_limit       = INT64_MAX                   // off 40
        out.host_send_recv_overlap_limit  = INT64_MAX                   // off 48
    else:
        out.all_reduce_overlap_limit      = env[3960]
        out.send_recv_overlap_limit       = AutoOr<long>(env+1120, dflt=1)
        out.host_send_recv_overlap_limit  = AutoOr<long>(env+1128, dflt=1)

    out.parallel_collective_overlap_limit = env[3644]                   // off 56
    out.max_concurrent_host_send_recv     = MaxConcurrentHostSendRecv(env,...) // off 64
    out.max_async_call_overlap_limit      = env[5000]                   // off 72

    // memory_limit (off 80): -1 (unlimited) UNLESS a rerun or pressure-tracking
    out.memory_limit = -1                                              // lines 66-71
    if v45:                          out.memory_limit = mem_limit_bytes
    if EnableSchedulerMemoryPressureTracking(env, multislice, target):
                                     out.memory_limit = mem_limit_bytes

    out.aggressive_scheduling_policies = env[5384]                     // off 88
    out.force_send_recv_to_use_resource = env[4912]                    // off 96
    out.selective_resource_min_limit   = 1                             // off 104
    out.target_scheduling_rule_priority = 2                            // off 112

    out.schedule_send_recvs = v45 | ScheduleSendRecvs(env,...)         // off 120, line 76
    out.flush_partial_schedule_on_busy_resource = 0                    // off 121
    // off 122-125: packed bool quad via vpinsrb (lines 81-88):
    //   schedule_send_recv_host_on_same_resource, use_real_cost_model,
    //   use_aggressive_scheduling  ← use_real_cost_model lands at off 123
    out.enable_release_start_policy        = env[4583]                 // off 126
    out.enable_send_recv_aggregation       = env[4789] | v45           // off 127, line 90
    out.flush_loop_collective_overlap      = env[4946]                 // off 128
    out.treat_loop_collective_as_non_collective = env[4871]           // off 129
    // off 130-133: packed via vpinsrw/vpshufb (lines 105-117):
    //   release_start_policy_with_resource = AutoOr<bool>(env+1928)
    //   release_one_start_policy_per_iter  = AutoOr<bool>(env+1936)
    //   prioritize_non_collective_over_idle_cycles = AutoOr<bool>(env+2432)
    out.disable_loop_consume_resource = (~AutoOr<bool>(env+2104) & 0x101)==0  // off 134
    out.enable_idle_count_extension   = (~AutoOr<bool>(env+2496) & 0x101)==0  // off 135
    out[136] = 0
    out[144..159] = 0   // vmovdqu xmm0 zero — resource arrays
    out[160] = 0        // scheduler-rerun-count slot
```text

### `v45` post-SC-assignment rerun gate

`v45 = env[4097]` 是 `xla_tpu_rerun_latency_hiding_scheduler_post_sc_assignment`。它是让 post_layout **重跑**区别于首次 post_layout 运行的*唯一*字段。置位时，它强制开启三个本来有条件的字段：

| 字段（偏移） | 首次运行 | 重跑（`v45` 置位） |
|---|---|---|
| `memory_limit` (80) | `-1`，除非压力跟踪 | `mem_limit_bytes`（始终真实） |
| `schedule_send_recvs` (120) | `ScheduleSendRecvs(env)` | 强制**开启**（`v45 \| …`） |
| `enable_send_recv_aggregation` (127) | `env[4789]` | 强制**开启**（`env[4789] \| v45`） |

理由是：SparseCore assignment 之后，host/ICI send/recv op 必须围绕已经放置的 SparseCore 工作紧密排序，并且 schedule 要基于真实内存预算重新推导，而不是首次 pass 的无限预算。

> **陷阱 —** SparseCore-offload-queue override (`sc_queue`) 与 `v45` rerun gate 是相互独立的。`EnableSparseCoreOffloadQueuingInLhs` 分支将 collective/send-recv *重叠限制*强制为 INT32/INT64_MAX，使 collective 可与 SparseCore offload stream 自由重叠；`v45` 则强制开启 *send/recv 调度和聚合*。重新实现若混淆二者，会在非重跑 SC 构建或重跑非 SC 构建中错误设置重叠包络。

### 与融合前变体的差异

构建器相同，137 字节偏移 schema 相同。行为差异完全在调用点值中：融合前槽位会传入融合前内存预算和 `v45 == 0`；post_layout 调用传入最终预算，并且在 SC-assignment 重跑时可能携带非零 `v45`。由于融合前 LHS 在 0.0.40 中**没有接线**，该构建器实际为可运行 LHS 产生的唯一 `SchedulerConfig` 是 post_layout 配置（第 802 行）。

---

## SchedulingContext（估算器 + tracker）

### 目的

`AddPass<LatencyHidingScheduler>` 接受两个 ctor 参数：`shared_ptr<const SchedulingContext>` 和 `unique_ptr<DefaultSchedulerCore>`。`SchedulingContext` 包装 `LatencyEstimator`（op 延迟模型）和 `AsyncTracker`（资源模型）。在此阶段，估算器在 LHS 与 SparseCoreQueueAssignment pass 之间共享（一个 `shared_ptr`，第 849 行构建），tracker 携带 reserved-SparseCore 集合。

### LatencyEstimator 选择

`GetLatencyEstimator` (`0x10974e00`) 根据 `SchedulerConfig` 偏移 123 (`use_real_cost_model`) 分支，该字节由 `SchedulerUsingRealCostModel(env, multislice, target, is_lem)` 设置：

```c
function GetLatencyEstimator(target, config, env):   // 0x10974e00
    if config.use_real_cost_model != 1:              // off 123, line 55
        // lightweight ApproximateLatencyEstimator wrapping DefaultGetCanonicalAsyncOp
        if PGLE_profile_path is set:
            FATAL "Attempted to use profile-guided latency estimator without "
                  "the real cost model. Either enable "
                  "--xla_tpu_scheduler_using_real_cost_model or unset "
                  "--xla_tpu_impure_latency_hiding_scheduler_profile_path."  // line 70
        return ApproximateLatencyEstimator()
    else:
        est = CostModelLatencyEstimator(target, config, env,           // 0x10ff8a60, line 91
                  cost_model_opts = (env[5070] == 1))   // EmitterLearnedCostModel toggle
        if PGLE_profile_path non-empty:
            proto = tsl::ReadBinaryProto(path)         // line 99; InvalidArgument on failure
            est = ProfileGuidedLatencyEstimator(est, proto)            // line 104
        return est

在 post_layout 路径中，估算器构建两次：第 828 行用于 LHS context，第 1195 行用于 SparseCoreQueueAssignment，但二者通过 context 的 refcount block 共享同一个实例。有关 CostModelLatencyEstimator 内部，见 learned-cost-model-client；有关成本模型族，见 cost/overview。

AsyncTracker / 资源模型

GetTpuAsyncTracker (0x10975520) → TpuAsyncTracker::Create (0x10ffb3e0)。基础 xla::AsyncTracker 资源类别来自二进制中已验证存在的三个 .rodata 常量：XLA_LATENCY_HIDING_SCHEDULER_RESOURCE_SHARING、..._RESOURCE_SERIALIZING、..._ENABLE_SELECTIVE_RESOURCES，由对应的 xla_latency_hiding_scheduler_resource_{sharing,serializing} / ..._enable_selective_resources flag 驱动。TPU 子类为 ICI link、HBM bank、per-megacore queue、host-DMA tap，以及（当 EnableSparseCoreOffloadQueuingInLhs 时）SparseCore offload queue 添加资源类别。完整资源分类记录在 scheduler-resourcetype-model。

注 — post_layout 的显著行为是 reserved-SparseCore 集合。RunHloScheduler 第 869-983 行从 GetReservedSparseCores（第 873 行）构建 flat_hash_set<int64_t>，并将其作为最后一个 Create 参数传入。这些是从重叠池中保留出来的 SparseCore ID，这些信息只在 SparseCore assignment 之后存在，因此融合前槽位（更早运行）会传入空集合。这就是两个变体之间的 AsyncTracker 差异。
陷阱（低）— GetTpuAsyncTracker (0x10975520) 从 env 读取并转发到 TpuAsyncTracker::Create 的数字参数，与具体物理资源的对应关系，是通过偏移恢复的，但没有逐一命名到 ICI、HBM、megacore queue。观测到的读取包括 env+4568 / +5320 的两个 _QWORD 资源限制、env+4867 / +5086 的字节 flag、env+5412 的 _DWORD，以及五个 AutoOr<long>(env+{2368, 2376, 2384, 2392, 2400}) 按 kind 限制（加上 env+2696）。重新实现者必须把每个限制所约束的具体资源视为推断结果。

Pass 主体 — 共享 RunImpl

目的

AddPass<LatencyHidingScheduler>（0x10975c40 处的模板化 add）在两个互斥位置调用，二者都馈入同一个 RunImpl：

位置	行	何时采用	Core 构建于
1	1137	常规路径	第 1046 行的 `DefaultSchedulerCore`
2	1411	SparseCore-queue 路径 — 新的 `GetTpuAsyncTracker` (1349) 用 reserved-SC 集合重新给 context 建 key，随后新的 `make_unique<DefaultSchedulerCore>` (1400) 替换 core	第 1400 行的 `DefaultSchedulerCore`

两者都调用同一个已编译 RunImpl，地址为 0x136321a0，即未修改的上游 xla::LatencyHidingScheduler::RunImpl。二进制中只有一个 RunImpl；变体只在它收到的配置和上下文上不同。

DefaultSchedulerCore 接线

core ctor (0x10976ce0) 复制 137 字节配置 POD（从 [rcx]…[rcx+0x69] 到 [core+0x10]…[core+0x79] 的 vmovups 块），并安装四个 std::function 槽：

Core 偏移	槽	post_layout 默认值
`+0xB8`	`target_overlap_limit` `fn(Shape const&)→long`	`RunTensorCoreAsyncOpScheduler::$_0`（第 1040 行）
`+0xD8`	`candidate_compare_` `fn(ScheduleCandidate&,…)→optional<CandidateResult>`	空（第 1000 行）— 内置列表调度器启发式
`+0xF8`	`post_step_mutator_` `fn(SchedulingState&)→void`	空（第 1002 行）
`+0x118`	`should_schedule_` `fn(HloInstruction const*)→bool`	`$_1`（ILP）或 `$_2`（常规）（第 1086/1097 行）

should_schedule_ 分类器切换是 ILP 变体（见 lhs-ilp-variant）：EnableIlpLatencyHidingScheduler(env)（第 1084 行）在更宽的 $_1 异步集合与更窄的 $_2 之间选择。该选择发生在 AddPass 之前，并且只改变 core 将哪些指令视为异步候选，RunImpl 主体不变。

怪癖 — candidate_compare_ 在 post_layout（以及融合前）变体中是空的。空 compare 函数意味着 RunImpl 回退到内置列表调度优先级：先 DAG-ready，再 memory-pressure-reducing，再 longest-critical-path-remaining，最后稳定 instruction index。重新实现者若期待这里存在自定义比较器，将找不到它；TPU 后端依赖默认启发式加资源模型，而不是专用优先级函数。

算法 — RunImpl 主体

function LatencyHidingScheduler_RunImpl(module, ctx, core):   // 0x136321a0
    VLOG(5) << module->ToString()                             // source 4106
    CHECK(module->has_schedule(),                             // line 280, source 4110
          "LatencyHidingScheduler expects a base schedule "
          "that minimizes memory pressure.")

    computations = MakeComputationPostOrder(module)           // line 320
        .filter(c -> c.caller_instructions() async-relevant)  // line 332
    VLOG(2) << "Computations to schedule " << computations    // line 578

    for each computation:
        VLOG << "] Statistics before scheduling:"             // line 673
              << LatencyHidingStatistics(...).ToString()      // 0x13630b60 / 0x13631a40
        ScheduleComputation(core, computation)   // candidate-list list scheduler
        // memory-pressure retry loop below
        VLOG << "] Statistics after scheduling:" << ...       // line 1146
    set_schedule(module, new_order)
```text

> **陷阱 —** `has_schedule()` CHECK 是 `LogMessageFatal`（第 280 行）。若没有先前的内存 schedule 就调用 LHS，它会中止编译，而不是优雅地 no-op。这是 `final_scheduler` 必须先运行的*原因*，也是融合前槽位无法在 0.0.40 中运行可用 LHS 的*原因*。

### 内存压力重试循环

从 `RunImpl` 反编译恢复（第 873-919 行，source 4182）。当生成的 schedule 超过内存限制时，调度器会按乘法因子（`0.9`，缩小 10%；见字节锚定的 [core 页面](latency-hiding-scheduler-core.md)）**收紧**限制并重新调度，从而让比较器的内存压力键更强地触发：

```c
// only if config.estimate_fragmentation_size (a SchedulingContext config byte)
int64_t frag = estimate_fragmentation_size
             ? EstimateFragmentationSize(module, alias_analysis, alias_info)  // TPU-private
             : 0;
int64_t mem_limit   = core.memory_limit();        // vtable +80
int     max_retries = core.memory_limit_retries();// vtable +88
for (int attempt = 0; attempt < max_retries; ++attempt) {
    int64_t used = frag + core.computed_use();     // vtable +64
    if (used <= mem_limit) break;
    LOG(INFO) << "LatencyHidingScheduler current memory usage: " << used
              << " bytes, does not fit in initial limit: " << mem_limit
              << ". Setting the new limit to " << core.next_memory_limit();  // line 881
    if (core.SchedulerInitialize(module) != ok) break;
    mem_limit = core.next_memory_limit();          // vtable +80, relax-by-factor
    core.set_memory_limit(mem_limit);
    reschedule_all_computations();
    if (estimate_fragmentation_size) frag = EstimateFragmentationSize(...);
}

next_memory_limit() 是 AVX shrink-by-factor：vmulsd xmm0, xmm0, [qword_A2DFD10]（缩放常量 0.9），随后是 vsubsd/vcvttsd2si clamp 和 sign-extend，即 new_limit = clamp(floor(old_limit * 0.9))。在 post_layout 变体中，只要 EnableSchedulerMemoryPressureTracking 为 true 或 v45 rerun 字节置位，该循环就是可运行的（预算真实）；否则 memory_limit == -1，循环惰性不动。

陷阱（低）— EstimateFragmentationSize 是 TPU 私有的（与 OSS 同名，不同 TU）；除 core 页面上命名的 GlobalDecreasingSizeBestFitHeap HeapSimulator 运行之外，它的 peak-live-with-fragmentation 模型和 allocator-alignment 假设未被进一步追踪。应将碎片模型视为推断结果。

post_layout 与 pre-fusion 概览

方面	pre-fusion	post_layout（本页）
流水线 / 阶段	"Pre main fusion"（第 6 阶段）	`RunHloScheduler` 中的 `async_scheduling`（第 7 阶段）
在 v0.0.40 中已接线？	否（没有到达的 AddPass 位置）	是（第 1137 / 1411 行）
`RunImpl`	`0x136321a0`（共享）	`0x136321a0`（共享，主体相同）
基础 schedule 来源	（需要先前的内存 sched）	`HloMemorySchedulerWithBrkgaFallback`（第 567 行）
`SchedulerConfig` 构建器	`GetSchedulerConfig`（相同）	`GetSchedulerConfig`，最终预算 + `v45` rerun gate
LatencyEstimator	`GetLatencyEstimator`（相同 helper）	`GetLatencyEstimator`（相同 helper）
AsyncTracker	`GetTpuAsyncTracker`，无 reserved-SC	`GetTpuAsyncTracker` + reserved-SparseCore 集合
`candidate_compare_`	空（启发式）	空（启发式）
异步分类器	`$_2`（或 ILP gate 下的 `$_1`）	`$_2`（或 ILP gate 下的 `$_1`）
`memory_limit` (off 80)	`-1`，除非压力跟踪	`-1`，除非压力跟踪或 `v45` 重跑

因此本页记录的是该构建中唯一可运行的 LHS。融合前槽位存在于代码中，但在 0.0.40 中是死的；ILP 变体是同一函数内部由 flag 控制的分类器切换。

函数映射

函数	地址	作用
`RunHloScheduler`	`0x1096fac0`	第 7 阶段驱动器；构建两条流水线，控制并添加 LHS
`LatencyHidingScheduler::RunImpl`	`0x136321a0`	共享 pass 主体；has_schedule CHECK、列表调度器、重试循环
`GetSchedulerConfig`	`0x10974aa0`	137 字节 `SchedulerConfig` POD；`v45` rerun gate
`GetLatencyEstimator`	`0x10974e00`	Approximate / CostModel / PGLE 估算器选择
`GetTpuAsyncTracker`	`0x10975520`	包装 `TpuAsyncTracker::Create`
`TpuAsyncTracker::Create`	`0x10ffb3e0`	TPU 资源模型 + reserved-SC 集合
`CostModelLatencyEstimator` ctor	`0x10ff8a60`	真实成本模型（对象大小 `0x3D0`）
`DefaultSchedulerCore` ctor	`0x10976ce0`	复制配置；安装四个 `std::function` hook
`AddPass<LatencyHidingScheduler, …>`	`0x10975c40`	模板化 pass 添加器（位置 1137/1411）
`IsPassDisabled`	`0x12fd8340`	Pass 启用/禁用解析
`EnableLatencyHidingLayerScheduler`	`0x1d6b8960`	Layer-scheduler gate（默认 false）
`EnableSchedulerMemoryPressureTracking`	`0x1d6b66e0`	`memory_limit` 三态
`SchedulerUsingRealCostModel`	`0x1d6b6580`	设置 `use_real_cost_model`（off 123）
`EnableIlpLatencyHidingScheduler`	`0x1d6b7e00`	`$_1`/`$_2` 分类器 gate
`HloMemorySchedulerWithBrkgaFallback::RunImpl`	`0x10abe6e0`	基础内存 schedule

交叉引用

LatencyHidingScheduler 核心 — 所有变体共享的 ScheduleComputation 候选列表循环和 async tracker
LHS post_layout_pre_fusion 变体 — 相同 RunImpl，接线在死掉的第 6 阶段 "Pre main fusion" 槽位
LHS ILP 变体 — 同一函数内部的 $_1/$_2 异步分类器切换（EnableIlpLatencyHidingScheduler）
Scheduler ResourceType 模型 — AsyncTracker 强制执行的 ICI/HBM/megacore/host-DMA/SparseCore 资源分类
调度概览 — 调度在 lowering 与 encoding 之间的位置
成本模型概览 — 供给 GetLatencyEstimator 的延迟/成本模型族
Learned Cost Model Client — CostModelLatencyEstimator 真实成本模型路径
第 VIII 部分索引 — Instruction Scheduling & Bundle Packing

LHS post_layout 变体 ​

摘要 ​

流水线位置 ​

目的 ​

入口点 ​

相关开关 ​

SchedulerConfig POD（配置差异） ​

目的 ​

算法 ​

AsyncTracker / 资源模型 ​

Pass 主体 — 共享 RunImpl ​

目的 ​

DefaultSchedulerCore 接线 ​

算法 — RunImpl 主体 ​

post_layout 与 pre-fusion 概览 ​

函数映射 ​

交叉引用 ​

LHS post_layout 变体

摘要

流水线位置

目的

入口点

相关开关

SchedulerConfig POD（配置差异）

目的

算法

AsyncTracker / 资源模型

Pass 主体 — 共享 RunImpl

目的

DefaultSchedulerCore 接线

算法 — RunImpl 主体

post_layout 与 pre-fusion 概览

函数映射

交叉引用