优化屏障

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。

摘要

XLA 优化屏障 — HloOpcode::kOptimizationBarrier，opcode 78（0x4E），助记符 opt-barrier — 是 TPU 编译器的调度与融合栅栏。它是单操作数的 identity op，结果形状与其操作数完全相同（可能是 tuple），唯一职责是作为一个不透明节点留在数据流图上，使任何 pass 都不知道如何跨过它进行融合、公共子表达式消除、重排或提升。屏障生产者侧的 op 与消费者侧的 op 由真实的 producer → barrier → consumer 边隔开，而由于每个 XLA pass 本来就遵守数据依赖，仅这一项结构事实就足以把所有内容固定在原位 — 不需要该 op 携带任何显式的 “do-not-cross” 标志。

熟悉 LLVM 的读者可以把它看成 llvm.assume 风格优化栅栏和调度钉之间的混合体：它没有副作用，不发出指令，在成本模型中消耗零成本，但它只要出现在依赖图上，就会阻止融合 pass、CSE/GVN、延迟隐藏调度器和 LICM 将计算移动穿过它。在 TPU 后端，屏障纯粹是 HLO 层级的提示：它会在整个 HLO 优化和调度流水线中被遵守，然后由一个后期 pass 擦除 — 从不降低为硬件同步 — 这个 pass 的 name() 很直白地是 "cse_barrier_expander"。该 pass 以及总是跟在它后面的 TupleSimplifier 之后，到达 LLO 的程序中完全不含 opt-barrier。

本页记录屏障在 TPU 后端生命周期的三个部分：(1) 它的语义 — 结构不变量，以及将其作为不透明透传节点遵守的十二个 Handle… handler；(2) 它的插入点 — 屏障从何处进入 TPU HloModule（前端、StableHLO→HLO 转换器，以及上游生产者 pass，包括 rematerialization / host-offload / collective-pipeliner 家族，MSA 中的内存压力机制依赖这些 pass）；(3) 它的降低 — OptimizationBarrierExpander::RunImpl 按字节精确恢复的擦除算法、添加它的三个流水线位置，以及 TensorCore async-op 调度器阶段在屏障完成其调度用途后丢弃它所处的 scheduler-mode gate。本页不会重新推导延迟隐藏调度器、融合 pass 本身，或无关的 collective/megascale/bundle-packer “barriers” — 这些只为划清边界而在此命名。

对于重新实现，契约是：

该 op 是单操作数、同形状 identity，没有内存效应且零成本；它不引入新的 HloValue，并在 points-to 层级别名其操作数的缓冲区。
它通过结构阻塞优化，而非通过标志。 Fusion 不能穿过不透明 op 融合；CSE/GVN 不会跨具有不同用户的不同指令合并；默认 visitor 将它路由到 DefaultAction，因此不覆盖它的 pass 会把它当作通用不透明 op。
十二个类提供显式 handler；除 AlgebraicSimplifier 外全部都是透传，后者允许穿过屏障剪除死 tuple 元素（但绝不跨它重排或合并计算）。
它在后期擦除，而不是降低。 OptimizationBarrierExpander（名称 "cse_barrier_expander"）将每个用户重新连接到 operand(0)，把控制依赖迁移到操作数上，重写 scheduled sequence 以丢弃屏障，并删除该指令。不发出硬件同步。


Opcode	`HloOpcode::kOptimizationBarrier` == `78` (`0x4E`); mnemonic `opt-barrier`
Opcode-string 表	`HloOpcodeString` @ `0x1e5ef000` — `off_21D30500[opcode + 0x80]` → `"opt-barrier"`
前端 builder	`xla::XlaBuilder::OptimizationBarrier(XlaOp)` @ `0x1e409da0` (AddInstruction opcode 78, 1 operand)
自由函数 wrapper	`xla::OptimizationBarrier(XlaOp)` @ `0x1e421e40`
StableHLO→mhlo importer	`StablehloToHloOpConverter<OptimizationBarrierOp>::matchAndRewrite` @ `0x16b5c540`
擦除 pass	`xla::OptimizationBarrierExpander::RunImpl` @ `0x164b3ba0`; `name()` @ `0x164b4380` → `"cse_barrier_expander"`
Pass-add trampoline	`HloPassPipeline::AddPass<OptimizationBarrierExpander>` @ `0x10972e60` (3 callers)
擦除 gate (TPU)	`(DWORD)(GetTpuCompEnv(...) + 5348) == 2` in `jellyfish::RunHloScheduler` @ `0x1096fac0`, line 1052
遵守它的 handlers	12 个不同的 `HandleOptimizationBarrier` / value-set updater（见下表）
降低为硬件同步？	否 — 在 LLO 前擦除；到达 LLO 的程序没有 opt-barrier
Opt-barrier 专用 flag	`FLAGS_xla_tpu_aggressive_opt_barrier_removal`（唯一一个；decision site 未隔离）
置信度	CONFIRMED（byte-anchored），除非某行或 callout 另有说明

Op 本身：Opcode 78

二进制中以三种独立方式确认了 opcode 值，并且所有十二个遵守它的 handler 都以同一个字节为 gate。

每个 handler 都测试 *(BYTE*)(inst + 12) == 78 — opcode 字节位于 HloInstruction 的 offset +12。这是例如擦除循环（if ( *(_BYTE *)(v21 + 12) == 78 )）和值集合 updater 的 CHECK 中的字面测试。

XlaBuilder::OptimizationBarrier (0x1e409da0) 通过 vtable AddInstruction 槽构造 instruction proto，使用 opcode = 78 和 operand_count = 1：

// 0x1e409da0, decompiled (vtable AddInstruction at +280)
(*(...)( *(_QWORD *)a1 + 280LL))(&v19, a1, v15, /*opcode=*/78, v21, /*nops=*/1, v6);
```text

HloOpcodeString (0x1e5ef000) 是一个扁平表查找 off_21D30500[(opcode + 0x80)]；slot 78 处的 .rodata 助记符是字符串 "opt-barrier"。

结构不变量

verifier 和分析会强制执行一个最小契约：

恰好一个操作数。 ShapeVerifier::HandleOptimizationBarrier (0x1e441ce0) 调用 CheckOperandCount(inst, 1)，失败时附加源位置 hlo_verifier.cc:342。HloDimensionInfoPropagation (0x10a07d80) 用锚定在 hlo_dimension_analysis.cc:371 的人类可读消息 "Optimization barrier must have exactly one operand." 重复同一检查。
结果形状 == operand(0) 形状（可能是 tuple）。ShapeVerifier 在操作数数量检查后调用 CheckShape(inst, operand(0)->shape())。
透传值身份。 该 op 不引入新的 HloValue。它原样转发操作数的 value set、points-to set、evaluated literal 和 value-semantics（见 handler 表）。因此 Buffer assignment 和 MSA 会把屏障输出缓冲区视为与其操作数相同的缓冲区家族。

NOTE — +12 opcode-byte offset 和 +88 shape-pointer offset（ShapeVerifier 中的 *(const xla::Shape **)(v3 + 88)）是此 build 的 HloInstruction 布局事实；它们在本页 handler 中稳定，但并非跨 libtpu 版本的 ABI 承诺。

MLIR 方言表面

程序到达 XLA HLO 之前，可以在四种 MLIR 方言中携带该 op，它们全都是带 PairwiseSameOperandAndResultType + InferTypeOpInterface 的可变参数 identity op。形状/类型推断是 identity 映射：mlir::hlo::inferOptimizationBarrierOp (0x18145ae0) 将每个结果类型设为对应操作数类型。

方言 / op	助记符	备注（已验证）
`mlir::TF::XlaOptimizationBarrierOp`	`tf.XlaOptimizationBarrier`	Traits 包括 `ConditionallySpeculatable`, `AlwaysSpeculatableImplTrait`, `MemoryEffectOpInterface`, `DerivedAttributeOpInterface`; `getEffects()` @ `0x1041f600` 为空（无声明的内存效应）
`mlir::mhlo::OptimizationBarrierOp`	`mhlo.optimization_barrier`	build/create/parse/print @ `0x17da03c0`/`04e0`/`0580`/`0760`
`mlir::stablehlo::OptimizationBarrierOp`	`stablehlo.optimization_barrier`	reference interpreter @ `0x12645ee0` 复制 operands → results（identity）
`mlir::vhlo::OptimizationBarrierOpV1`	`vhlo.optimization_barrier_v1`	versioned serialization form

TF op 的 foldHook (0x1012fac0) 是通用 op-fold trampoline；它不会把屏障折叠掉变成其操作数（没有屏障的 const-fold）。因为该 op 携带 AlwaysSpeculatableImpl 且 getEffects 为空，其 MLIR 层语义纯粹来自它作为结构性数据流节点（operand → result 依赖），而不是副作用或 memory-fence 属性。栅栏是这条边，而不是一个标志。

为什么屏障会阻塞优化（机制）

这是重新实现者的关键点：并不存在每个 pass 都查询的显式 “do-not-cross” bit。屏障通过三个每个 XLA pass 已经遵守的属性来阻塞 fusion / CSE / reorder / LICM。

1. 它是数据流节点，不是 no-op

链 A → barrier → B 在 A 与 B 之间放入一条真实的 producer→consumer 边。TPU fusion passes（TpuInstructionFusion / TpuMultiOutputFusion）遍历生产者-消费者对；屏障是它们不知道如何穿过融合的不透明 op，因此 A 和 B 永远不能合并为一个 fused computation。延迟隐藏调度器遵守同样的边：它既不能把 B 提升到屏障之上，也不能把 A 下沉到屏障之下，因为那会违反依赖。

2. 它别名其操作数的缓冲区，但又是不同的值持有者

TuplePointsToAnalysis::HandleOptimizationBarrier (0x1e387940) 取 operand(0) 并调用 CreateCopiedPointsToSet(barrier, operand(0))：

// 0x1e387940, decompiled
v4 = (const xla::HloInstruction *)xla::HloInstruction::operand(a2, 0);
xla::TuplePointsToAnalysis::CreateCopiedPointsToSet(this, a2, v4, a3, a4);

屏障的 points-to set 是操作数 points-to set 的副本，因此 buffer assignment / MSA 会把屏障输出视为与操作数相同的缓冲区家族 — 屏障本身不会强制一次物化拷贝。同时，HloDataflowAnalysis::UpdateOptimizationBarrierValueSet (0x1e4d18c0) 将 operand(0) 的 value set 复制进屏障，并先断言 opcode：

// 0x1e4d18c0, decompiled CHECK
CHECK("barrier->opcode() == HloOpcode::kOptimizationBarrier");  // hlo_dataflow_analysis.cc:674
```text

即使缓冲区别名，CSE/GVN 比较的是*值编号*，而屏障是带有自己用户的独立指令。因此两个相同的子图如果分别流入两个独立屏障，就**不会跨屏障合并**。擦除 pass 名为 `cse_barrier_expander`，正是因为 CSE 是屏障一直抑制到后期的对象。

### 3. 默认 visitor 将它路由到 `DefaultAction`

`DfsHloVisitorWithDefaultBase<…>::HandleOptimizationBarrier`（`HloInstruction*` 实例化为 `0x10946240`，const-pointer 实例化为 `0x10c5b560`）转发到 offset `+1120` 的 vtable `DefaultAction` 槽。任何**没有**覆盖 `HandleOptimizationBarrier` 的 pass — 其中包括大多数 pass，如 LayoutAssignment、MSA 和 BufferAssignment — 因而会把屏障当作通用不透明 op，遵守依赖并把 layout / buffers 穿过它串起来，而不会尝试将它简化掉。

成本模型同意该 op 免费：`HloCostAnalysis::HandleOptimizationBarrier` (`0x1e4812e0`) 返回 OK 且不添加成本。屏障不会在融合排序中惩罚周围 op；它只是不能被穿过融合。

---

## 遵守屏障的 Handler

十二个类为 opcode 78 提供显式 handler。除 `AlgebraicSimplifier` 外，每一个都是透传 / 不透明的，后者可以剪除死 tuple 分支（但不能重排或合并计算）。每个 handler 都以命名地址处的反编译函数作字节锚定。

| # | 类 / handler | 地址 | 行为（已恢复） | 锚点 |
|---|---|---|---|---|
| 1 | `ShapeVerifier` | `0x1e441ce0` | `CheckOperandCount(inst, 1)`；然后 `CheckShape(inst, operand(0)->shape())` | `hlo_verifier.cc:342` |
| 2 | `HloDataflowAnalysis::UpdateOptimizationBarrierValueSet` | `0x1e4d18c0` | barrier value set := `operand(0)` value set；`CHECK` opcode == `kOptimizationBarrier` | `hlo_dataflow_analysis.cc:674` |
| 3 | `TuplePointsToAnalysis` | `0x1e387940` | `CreateCopiedPointsToSet(barrier, operand(0))` — 别名 operand buffers | — |
| 4 | `HloCostAnalysis` | `0x1e4812e0` | 返回 OK；零成本 | — |
| 5 | `HloEvaluator` | `0x1ddb0040` | `SetEvaluatedLiteralFor(barrier, Clone(GetEvaluatedLiteralFor(operand(0))))`（const-fold 透传） | — |
| 6 | `AlgebraicSimplifierVisitor` | `0x1dd27360` | dead-tuple-element pruning：遍历 GTE（opcode 64）用户，重建更小的 tuple operand + 重新索引 GTE；`CHECK use->opcode() == kGetTupleElement` | `algebraic_simplifier.cc:5299,5302` |
| 7 | `SpmdPartitioningVisitor` | `0x1c7f3540` | thunk 到 `HandleElementwise`（sharding 原样流过） | — |
| 8 | `HloValueSemanticsPropagation` | `0x111bece0` | `DeepCopyHloValueSemantics(barrier, GetInstructionSemantics(operand(0)))` | — |
| 9 | `HloDimensionInfoPropagation` | `0x10a07d80` | 复制 `operand(0)` DimensionInfo 通过；`CHECK operand_count == 1` | `hlo_dimension_analysis.cc:371,378` |
| 10 | `jellyfish::XPrecisionRewriteVisitor` (TPU) | `0x1115be20` | precision rewrite 时，将 barrier output shape 设为 `operand(0)` new shape | `xprecision_rewriter.cc:4465` |
| 11 | `jellyfish::semantics_guided_sharding::DimLabelPropagation` (TPU) | `0x11197a80` | 向前转发 dim labels | — |
| 12 | `DfsHloVisitorWithDefaultBase` (ptr / const-ptr) | `0x10946240` / `0x10c5b560` | 路由到 `DefaultAction`（vtable `+1120`）— 通用不透明处理 | — |

### 唯一的重写 handler：AlgebraicSimplifier 死 tuple 剪枝

Handler #6 是唯一会在屏障*周围*重写的 pass，而且它从不穿越屏障。当屏障包裹一个 tuple 时，simplifier 会执行 XLA 标准的“把屏障输入 tuple 缩小到只包含活元素”变换。反编译显示它在用户为 GTE 时遍历屏障的用户（`while ( *((_BYTE *)*v9 + 12) == 64 )`，opcode 64 = `kGetTupleElement`），当 GTE 有多个用户、被索引操作数是模块入口 root、或操作数本身有副作用时，将索引标为 live。死分支被丢弃，创建一个新的更小 `kTuple`（使用 `HloInstruction::CreateGetTupleElement` 重新索引保留下来的 GTE），并把屏障的 operand-0 替换为更小 tuple（`ReplaceOperandWithDifferentShape`）。防御性 `CHECK use->opcode() == HloOpcode::kGetTupleElement`（`algebraic_simplifier.cc:5302`）保护“每个用户都是 GTE”的假设：

```c
// 0x1dd27360, decompiled fragments
while ( *((_BYTE *)*v9 + 12) == 64 )            // walk GTE users (kGetTupleElement)
  ...
TupleElement = xla::HloInstruction::CreateGetTupleElement(v25, v21, v20);   // re-index
...
CHECK("use->opcode() == HloOpcode::kGetTupleElement");   // algebraic_simplifier.cc:5302

这证明屏障对 fusion/CSE 是不透明的，但仍允许对未使用 tuple 分支做死代码风格清理。它不允许跨栅栏重排或合并计算。

NOTE — MemorySpaceAssignment / LayoutAssignment / BufferAssignment 中缺少 HandleOptimizationBarrier override 这一点本身就是证据：这些 pass 使用 DefaultAction 不透明路径（handler #12）。屏障别名其操作数缓冲区（handler #3），因此 MSA 会把两者视为一个缓冲区家族 — 但 MSA 是否拒绝跨屏障 prefetch 是从不透明 op 路径推断出来的，而不是直接 trace 到的。请把跨屏障 prefetch 策略视为低置信度，直到有直接 MSA trace；见 MSA overview。

插入：屏障从何处进入 TPU Module

屏障从两个表面到达 TPU HloModule — 前端 / IR import，以及上游生产者 pass — 二者都会调用（可能内联的）HloInstruction::CreateOptimizationBarrier，opcode 参数为 78。

前端 / IR 表面（完整恢复）

text

jax.lax.optimization_barrier / XLA client
  → xla::XlaBuilder::OptimizationBarrier(XlaOp)   0x1e409da0     [opcode 78, 1 operand]

OR via MLIR:
  stablehlo.optimization_barrier
    → StablehloToHloOpConverter<OptimizationBarrierOp>::matchAndRewrite   0x16b5c540
      → mlir::mhlo::OptimizationBarrierOp::create
    → mhlo → HLO importer → HLO kOptimizationBarrier
```text

自由函数 `xla::OptimizationBarrier(XlaOp)` (`0x1e421e40`) 是一个薄 wrapper，会转发到 builder 方法。

### 生产者 pass（部分 — create 调用被内联）

`CreateOptimizationBarrier` factory 被内联进若干上游 pass 的 helper 方法，因此无法作为独立符号隔离。候选生产者按行为和 opcode-78 构造点的存在识别；在此 build 中每个都作为反编译函数存在：

| 生产者 pass | 入口 | 为什么用屏障包裹值（已恢复行为） |
|---|---|---|
| `HloRematerialization::RunImpl` | `0x1d6c8300` | 内存压力下的 rematerialization — 屏障固定重新计算的值，使其不会被重新 fusion/CSE 回原值，迫使重新计算在选定点物化 |
| `HostOffloader::RunImpl` | `0x1107d5e0` | host-offload buffer wrapping — host buffers 被包裹，以强制跨 offload 边界物化 |
| `CollectivePipeliner::RunPipeliner` | `0x12ffb100` | collective pipelining ordering — 屏障固定 pipelined collective 的 start/done 边界，使其不能被重排进 steady state |
| `WhileLoopUnroller::RunImpl` | `0x12eeb500` | while-loop unrolling — 屏障分隔不得被 CSE 重新合并的已展开迭代体 |

> **WARNING —** 内联的 `CreateOptimizationBarrier` create-site **无法**在约 12 次反编译 grep 中隔离为独立符号。上面的生产者列表是通过行为和这些 pass 在二进制中的存在恢复的，而不是通过逐生产者干净 trace opcode-78 构造得到的。精确生产者 trace 是后续工作；请将逐生产者*理由*按标注视为 HIGH/MEDIUM，而非字节精确。

与内存压力的关联是该 op 对编译器后端重要的原因：rematerialization（经典的“重新计算而非 spill”权衡）正是 [MSA](msa-overview.md) 内存压力信号的消费者，并且它使用 opt-barriers 让重新计算的值*保持*重新计算，而不是被重新优化回一起。屏障就是让有意复制的计算保持复制状态的契约。

---

## 降低：屏障被擦除，而不是同步

TPU 优化屏障**永远不会到达 LLO**。它由后期 HLO pass `xla::OptimizationBarrierExpander` 移除，其 `name()` (`0x164b4380`) 返回 `"cse_barrier_expander"` — 反编译中逐字确认：

```c
// 0x164b4380
absl::string_view name() const { return "cse_barrier_expander"; }

擦除算法（RunImpl @ 0x164b3ba0）

RunImpl（1981 bytes，99 basic blocks）收集 opcode-78 指令，重写 scheduled sequence 以丢弃它们，然后重新连接并删除每个屏障。反编译确认了每一步：MakeNonfusionComputations（line 85）、opcode == 78 测试（line 153）、HloSchedule::set_sequence（lines 373/384），以及 CopyAllControlDepsTo → ReplaceAllUsesWith → DropAllControlDeps → RemoveInstruction 尾部（lines 427/433/439/445）。

cpp

// xla::OptimizationBarrierExpander::RunImpl   @ 0x164b3ba0
// name() == "cse_barrier_expander"   (optimization_barrier_expander.cc)
absl::StatusOr<bool> RunImpl(
    HloModule* m,
    const absl::flat_hash_set<absl::string_view>& execution_threads) {
  std::vector<HloInstruction*> barriers;
  // NB: MakeNonfusionComputations — fusion bodies are skipped on purpose.
  for (HloComputation* comp : m->MakeNonfusionComputations(execution_threads)) {
    bool comp_has_barrier = false;
    for (HloInstruction* inst : comp->instructions())
      if (inst->opcode() == HloOpcode::kOptimizationBarrier) {     // *(BYTE*)(inst+12) == 78
        barriers.push_back(inst);
        comp_has_barrier = true;
      }
    if (comp_has_barrier && m->has_schedule()) {                   // gated on module schedule flag
      // Rebuild the computation's scheduled sequence WITHOUT the barriers.
      std::vector<HloInstruction*> seq;
      for (HloInstruction* i : m->schedule().sequence(comp).instructions())
        if (i->opcode() != HloOpcode::kOptimizationBarrier) seq.push_back(i);
      m->mutable_schedule()->set_sequence(comp, seq);              // 0x164b4165 / 0x164b412a
    }
  }
  for (HloInstruction* b : barriers) {
    HloInstruction* op = b->mutable_operand(0);
    TF_RETURN_IF_ERROR(b->CopyAllControlDepsTo(op, op));           // preserve ordering edges
    TF_RETURN_IF_ERROR(b->ReplaceAllUsesWith(op));                // users read the operand directly
    TF_RETURN_IF_ERROR(b->DropAllControlDeps());
    TF_RETURN_IF_ERROR(b->parent()->RemoveInstruction(b));        // delete the barrier
  }
  return !barriers.empty();
}
```text

因此在降低时屏障会被**擦除**：每个用户重新连接到屏障的操作数，控制依赖迁移到操作数上，使相对顺序得以保留，随后删除指令。它纯粹是 HLO 层级的调度 / 优化提示，**不发出硬件同步** — 这确认了它完成调度用途后就消失的假设。（与下面的 collective 和 bundle-packer barriers 对比，后者*会*发出真实硬件同步。）

> **NOTE —** `MakeNonfusionComputations` 遍历意味着最终落入 fusion body 的屏障不会在这里被触碰。实践中 fusion 永远不会吸收屏障（机制点 #1），因此这不是缺口；只有某个其他 pass 把屏障埋进 fusion body 时才会相关，但未观察到这种情况。

### 擦除如何保留 async-window pinning

async-copy / collective window 的 pinning 是结构性数据依赖加控制依赖的结果，而不是专门机制 — 擦除 pass 会在屏障消失后小心保留它：

- 调度期间，async op 是 start/done 对（`kAsyncStart` opcode 17 / `kAsyncDone`，或 collective-start / collective-done）。如果屏障位于 async 结果的消费者与后续 op 之间，延迟隐藏调度器不能将该 op 移到屏障之上，因此 async window 不能被非法折叠。
- 在屏障存在时产生的 schedule 已经由 `set_sequence` 冻结（屏障会在原地从 sequence 中过滤掉，周围顺序保持不变）。
- `CopyAllControlDepsTo(operand)` 将屏障携带的任何控制边迁移到其操作数上，因此即使屏障指令被删除，相对顺序仍然保留。

这就是它与 collective / async-tree barrier（见下文）的精确区别：优化屏障在*编译时的 HLO schedule 中*固定顺序，然后消失；collective barrier 会发出实际的运行时同步。

---

## Expander 在何处运行（流水线位置）

`HloPassPipeline::AddPass<OptimizationBarrierExpander>` (`0x10972e60`) 有三个 caller；三个目标函数都作为反编译符号存在。

| Caller | 入口 | 角色 | 与 TPU 相关？ |
|---|---|---|---|
| `xla::cpu::CpuCompiler::RunHloPassesThroughLayoutAssn` | `0x14bb99e0` | CPU backend | 否 — TPU 可忽略 |
| `xla::jellyfish::(anon)::RunHloScheduler` | `0x1096fac0` | TPU scheduler stage；`RunTensorCoreAsyncOpScheduler` 内 AddPass，带 gate，随后是 `TupleSimplifier` | **是** |
| `xla::jellyfish::(anon)::PostOptimizationPipeline` | `0x1093fd40` | TPU late-cleanup stage；另一个擦除位置 | **是** |

### 调度器阶段 gate（已验证）

在 `RunHloScheduler` 内部，只有当 TPU 编译环境的 scheduler-mode 字段等于 2 时才会添加 expander，并且其后紧跟 `TupleSimplifier`。反编译显示 gate 位于 line 1052，两个 `AddPass` 调用位于 1054–1055，正好在 1058–1059 的 `AsyncCollectiveMerger` 之前：

```c
// 0x1096fac0 — jellyfish::RunHloScheduler, lines 1052..1059
if ( *(_DWORD *)(xla::jellyfish::GetTpuCompEnv(v389[0], ptr) + 5348) == 2 ) {
  xla::HloPassPipeline::AddPass<xla::OptimizationBarrierExpander>(&v341);   // line 1054
  xla::HloPassPipeline::AddPass<xla::TupleSimplifier>(&v341);               // line 1055
}
xla::jellyfish::AsyncCollectiveMerger::PreLatencyHidingSchedulerConfig(...); // line 1058
xla::HloPassPipeline::AddPass<xla::jellyfish::AsyncCollectiveMerger,...>(...); // line 1059

因此 expander 会在 LatencyHidingScheduler / AsyncOpScheduler 已经把屏障用作调度钉之后运行，并且在不再需要它的 async-collective merge 步骤之前运行。expander 后面的 TupleSimplifier 很重要：由于屏障通常包裹 tuple，消费者通过 GTE 访问其分支，擦除屏障会留下 GetTupleElement(Tuple(...)) 链，随后由 TupleSimplifier 折叠。

TPU 路径的顺序直觉

text

... LayoutAssignment → fusion (see fusion-patterns) → scheduling
    (LatencyHidingScheduler / AsyncOpScheduler honour the barrier edges to
     pin async-copy / collective boundaries)
  → cse_barrier_expander + TupleSimplifier   (barrier erased; GTE/tuple
     chains it left behind are cleaned up)
  → AsyncCollectiveMerger
  → LLO lowering — sees NO opt-barrier
```text

### 唯一 opt-barrier 专用旋钮

`FLAGS_xla_tpu_aggressive_opt_barrier_removal` 是二进制中唯一的 optimization-barrier-specific flag。它控制优化屏障被剥离的激进程度（推测允许比默认 scheduler-gated 位置更早或更多地移除）。

> **WARNING —** `xla_tpu_aggressive_opt_barrier_removal` 的决策点（读取该 flag 的分支）**没有**被隔离，并且 `+5348` 处的 `GetTpuCompEnv` 字段是一个 scheduler-mode enum，其其他取值和字段名尚未恢复。请将 `== 2` gate 视为*观察到的*启用值，而不是有文档的 enum — 对 enum 完整含义为**低置信度**。

---

## 不是这个 Op：另外三种 “Barriers”

libtpu 中 “barrier” 一词命名了四种无关机制。只有 #1 是本页主题；其他只为划清边界而列出，并且与优化屏障不同，它们会发出真实硬件同步。

| 变体 | 含义 | 代表性符号 / 地址 | 发出真实同步？ |
|---|---|---|---|
| **1. HLO optimization barrier**（本页） | opcode 78；编译期 CSE/fusion/reorder 阻塞器；后期擦除 | `OptimizationBarrierExpander` `0x164b3ba0`（name `cse_barrier_expander`）；所有 `HandleOptimizationBarrier` handlers | **否** — LLO 前擦除 |
| **2. Collective / async-tree barrier** | `CustomCall("BarrierStart")` / `"BarrierDone"`；固定 collective start/done；megascale & ICI | `async_barrier_util::GetBarrierStartFromAsyncStart` `0x11007c80`; `AllToAllEmitterBase::EmitBarrierStartImpl` `0x10f07240` / `EmitBarrierDoneImpl` `0x10f07d40`; `BarrierStartEmitter` `0x10e8ece0` | 是 — collective sync |
| **3. Megascale / barrier-assignment config** | 带 `barrier_type()` 的 `BarrierConfig` proto；per-core sflag/DMA barrier | `InferBarrierConfig` `0x1376c240`; `BarrierAssignment::BarrierAssignment` `0x109c6e40`; `TensorCoreBarrierAssignment::ForEachCollective` `0x109c7060`; `SparseCoreBarrierAssignment::AssignBarriersForKernels` `0x109c6080` | 是 — runtime sync (sflag/DMA) |
| **4. LLO bundle-packer scheduling barrier** | packer 不能跨过移动 op 的 VLIW-bundle “point of no return” | `BundlePacker::BarrierHere` `0x14021de0` (bundle_packer.h:207); `BundlePacker::EstablishBarrier` `0x14021d00`; `SlotTracker::UpdatePointOfNoReturn` | N/A — packer-internal fence |

二进制中的区分性 opcode 测试：

- opt-barrier: `inst.opcode() == 78`
- async-start（collective barrier 上下文）：`inst.opcode() == 17`
- barrier-start tuple wrapping：operand `opcode() == 0x81` (`kTuple`), `IsCustomCall(…, "BarrierStart")`

二进制中有 25 个 barrier-family flags；**只有一个**（`xla_tpu_aggressive_opt_barrier_removal`）涉及优化屏障。其他 24 个（`xla_tpu_enable_megascale_barrier`、`xla_tpu_force_global_barriers`、`xla_tpu_use_custom_tree_barrier`、`xla_tpu_use_dma_for_custom_barrier` 等）都配置 collective/megascale/hardware-sync barriers — 也就是变体 2–4，而非此 op。

---

## 交叉引用

- [编译器概览](overview.md) — 屏障经过的五阶段主干和 IR 层栈。
- [编译阶段](compile-phases.md) — HLO 优化、调度和 LLO 降低所在位置；屏障会在 LLO crossing 前擦除。
- [融合模式](fusion-patterns.md) — 屏障对其不透明的 `TpuInstructionFusion` / `TpuMultiOutputFusion` pass。
- [HLO 预 pass](hlo-pre-passes.md) — 调度前运行的早期 HLO 层变换。
- [MSA 概览](msa-overview.md) — 驱动屏障插入的内存压力 / rematerialization 消费者，以及屏障 points-to 副本保留的缓冲区家族别名。

优化屏障 ​

摘要 ​

Op 本身：Opcode 78 ​

结构不变量 ​

MLIR 方言表面 ​

为什么屏障会阻塞优化（机制） ​

1. 它是数据流节点，不是 no-op ​

2. 它别名其操作数的缓冲区，但又是不同的值持有者 ​

插入：屏障从何处进入 TPU Module ​

前端 / IR 表面（完整恢复） ​

擦除算法（RunImpl @ 0x164b3ba0） ​

TPU 路径的顺序直觉 ​