Sharding Propagation

本页中的所有地址都适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会不同；请把每个 VA 都视为绑定到该版本。

摘要

xla::ShardingPropagation 是 GSPMD 的推断 pass：给定一个 HLO module，其中用户（或前端）已经给部分指令标注了 HloSharding，它会沿数据流图传播这些已知 sharding，为其他每条指令补全 sharding，直到不再变化。它是向 partitioner 提供输入的三个 sharding 生产者之一，负责手动/部分标注场景；另外两个是 auto-sharding ILP 搜索和 Shardy import，二者记录在 Auto-Sharding and SPMD Partitioner。这个 pass 不会通过优化来选择 sharding；它会像类型推断 pass 从部分标注推导类型一样，推导已有标注所蕴含的唯一或最具体 sharding。

熟悉 XLA-on-GPU 的读者已经具备所需框架：这就是上游发布的同一套 sharding_propagation.cc 机制，只是由 TPU pipeline 注入的 custom-call sharding helper (TpuCustomCallShardingHelper) 做了重定向。核心是运行到固定点的 worklist 数据流分析。每次迭代都会对指令图做两次扫描：一次前向扫描，根据操作数推断一条指令的 sharding（post-order，操作数先于用户），一次后向扫描，根据用户推断它（reverse post-order，用户先于操作数）。每个候选都通过同一个门控提交，即 MaybeImproveInstructionSharding；该门控只有在候选比指令已有 sharding 严格更具体，或能与其兼容合并时才接受。由于门控是单调的（sharding 只会变得更具体，绝不会变得更不具体），固定点保证会终止。aggressiveness 级别（0..3）会跨迭代逐步放宽逐 op 规则，使保守推断先稳定，推测性规则稍后才启用。

本页负责说明 propagation pass 本身：pre-pass 标注扫描、前向和后向推断数据流、逐 op 推断规则（elementwise、dot、conv、reshape、gather、reduce-window）、MaybeImproveInstructionSharding 提交门控、shard-group 与跨 computation 传播、domain 规范化，以及迭代到固定点的循环。它不负责 ILP strategy 搜索或将 collectives 物化的 SPMD partitioner；这些内容在相邻页面。结构如下：(1) pass driver 与 pre-pass 扫描，(2) 提交门控，(3) 前向传播 + 逐 op 规则，(4) 后向传播，(5) shard-group / cross-computation / domain 传播，(6) 固定点循环，(7) custom-call helper 分发。

对于重新实现，契约是：

推断数据流。 每次迭代对指令图做两次有向扫描：从操作数前向传播、从用户后向传播；访问顺序为 post-order / reverse-post-order，使单次扫描尽可能收敛。
提交门控。 MaybeImproveInstructionSharding(new, instr, may_combine_partial, replace_existing) 是决定候选 sharding 是否替换、细化或被拒绝的单调规则，也是单调性保证终止的原因。
逐 op 规则。 每个 opcode 家族如何把操作数/用户 sharding 映射到指令 sharding（以及反向映射）：elementwise 透传、dot/conv 的 contracting-vs-batch 维度映射、reshape 维度拆分/合并、gather/scatter 索引处理、reduce-window 的 dilation guard。
种子与约束。 Sharding / SPMDFullToShardShape / SPMDShardToFullShape custom call、shard-group（shard_as / shard_like）标注，以及 ShardBarrier 标记如何为推断播种并约束传播范围。


Pass 名称	`"sharding-propagation"` (`name()` @ `0x1c8615a0`)
Pass 入口	`xla::ShardingPropagation::RunImpl(HloModule*, exec_threads)` @ `0x213aa140`（33.6 KB — 文件中最大的 `RunImpl`）
固定点循环	`ShardingPropagation::RunToFixPoint` @ `0x1c85ae60`（22.4 KB）
前向扫描	`InferShardingFromOperands` @ `0x1c856780`（9.9 KB）
后向扫描	`InferShardingFromUsers` @ `0x1c859fa0`（2.2 KB）+ `GetShardingFromUser` @ `0x1c8531e0`（7.3 KB）
提交门控	`(anonymous)::MaybeImproveInstructionSharding` @ `0x1c84e980`; sub-shape variant `MaybeImproveInstructionSubSharding` @ `0x1c850820`
源路径	`third_party/tensorflow/compiler/xla/service/sharding_propagation.cc`（rodata）
Custom-call helper	`xla::jellyfish::TpuCustomCallShardingHelper`（`InferShardingFromOperands` @ `0x1278bf80`）覆盖 base `xla::CustomCallShardingHelper` @ `0x1c864120`
相关 flag	`xla_tpu_sharding_metadata` — 在 `tpu_compilation_environment.cc` 中注册的命令行 flag；控制构造函数的 `propagate_metadata`（随 sharding 携带 op-name provenance）。默认值未在反编译中固定（LOW）。
置信度	HIGH（符号逐字节锚定；关键函数体经反编译验证），除非某行/标注另有说明

Pass Driver

目的

ShardingPropagation::RunImpl (0x213aa140) 是唯一入口。它每次调用运行一次，但内部会迭代到固定点。它的任务是接收一个带有稀疏 HloSharding 标注集合的 HLO module，并将其稠密化：返回时，pass 能到达的每条指令都会带有推断出的 sharding，布尔返回值报告是否有任何变化（从而让外层 HloPassPipeline 知道是否需要重新运行依赖 pass）。反编译函数体锚定到 sharding_propagation.cc，并发出完成诊断 "Sharding propagation completed after N iterations"。

TPU pipeline 会在 partitioning pipeline 内部运行这个 pass 三次（见 Auto-Sharding and SPMD Partitioner；AddTpuPartitioningPasses 三次引用 ShardingPropagation），因为中间 rewrite（例如 concat decomposition）会创建新指令，而这些指令的 sharding 必须重新推断。重新实现者必须让该 pass 具备幂等性：在已收敛 module 上重新运行时必须返回 false（无变化）。

构造函数

该 pass 携带塑造其规则的配置。它只有一个构造函数（xla::ShardingPropagation::ShardingPropagation @ 0x1094a2c0）；调用点是 0x1094a180 / 0x1093a300 处的 AddPass<ShardingPropagation> 模板实例化（后者以 absl::InlinedVector<bool, 1> 形式传递 span）：

// xla::ShardingPropagation::ShardingPropagation  @ 0x1094a2c0
ShardingPropagation(
    bool is_spmd,                  // SPMD (true) vs replicated/MPMD mode
    bool propagate_metadata,       // carry sharding "metadata" (op-name provenance) forward
    absl::Span<const bool> allow_spmd_sharding_propagation_to_output,      // per-output: may infer onto root?
    absl::Span<const bool> allow_spmd_sharding_propagation_to_parameters,  // per-param:  may infer onto param?
    bool cse_prevention_only,      // restricted mode: only annotate to block bad CSE
    std::unique_ptr<CustomCallShardingHelper> sharding_helper);  // TPU passes TpuCustomCallShardingHelper
```text

> **NOTE —** `cse_prevention_only` 是一种受限模式，其中 pass 不会尝试完整推断；它只附加那些唯一目的在于阻止 CSE pass 合并两个必须在每个 partition 中保持不同的 op 的 sharding。CSE-prevention 标注会带上 rodata 后缀 `"_sharding_propagation_cse_prevention"`（见于 `RunImpl`）。跳过此模式的重新实现仍会得到*正确* sharding，但后续 CSE 可能会折叠 partition-distinct computation。

TPU pipeline 总是用 `std::make_unique<TpuCustomCallShardingHelper>()` 构造该 pass。一个四模板参数 `AddPass` 实例化（`0x14bbc2e0`, `<ShardingPropagation, bool, bool, Span<bool>, Span<bool>>`）只传递前四个构造函数参数；`cse_prevention_only` 和 helper 使用默认值，并且两个 span 参数都使用普通 `absl::Span<const bool>`，而不是 `InlinedVector` 形式。

### 入口点

```text
ShardingPropagation::RunImpl                 0x213aa140 (33.6 KB)  ── driver: scan + fixpoint
  ├─ (pre-pass linear scan, inline)                                ── gather annotations, domains, groups
  └─ RunToFixPoint                           0x1c85ae60 (22.4 KB)  ── iterate forward+backward to convergence
       ├─ InferShardingFromShardGroup        0x1c856420 (0.83 KB)  ── shard_as / shard_like group seeding
       ├─ InferShardingFromOperands          0x1c856780 (9.9 KB)   ── FORWARD: operand → instr
       │     └─ MaybeImproveInstructionSharding  0x1c84e980        ── commit gate (forward)
       ├─ InferShardingFromUsers             0x1c859fa0 (2.2 KB)   ── BACKWARD driver
       │     └─ GetShardingFromUser          0x1c8531e0 (7.3 KB)   ── BACKWARD: user → operand sharding
       ├─ MaybeComputationPropagation        0x1c85a860 (0.2 KB)   ── cross-computation (call/while/cond)
       ├─ GetRelatedInstructions             0x1c861140 (0.84 KB)  ── shard-group reachability
       └─ NormalizeDomain                    0x1c852dc0 (1.0 KB)   ── domain-boundary conflict resolution

算法 — pre-pass 扫描

在固定点循环之前，RunImpl 会遍历 module 一次，以收集种子和约束。反编译函数体引用了下面的每个 marker string。

function RunImpl(module, execution_threads):                  // 0x213aa140
    // 1. Seed: collect explicit annotations and the boundary markers.
    for instr in module.instructions (post-order):
        if instr.has_sharding():
            provided_shardings.insert(instr)                  // frozen — never overwritten
        if instr.IsCustomCall("Sharding"):                    // consumer-side sharding hint
            register_domain_bracket(instr)                    // re-attached later by HloDomainRemover
        if instr.IsCustomCall("SPMDFullToShardShape") or       // manual-region entry
           instr.IsCustomCall("SPMDShardToFullShape"):        // manual-region exit
            mark_manual_boundary(instr)                       // inference stops at the boundary
        if instr.IsCustomCall("ShardBarrierFrom") or
           instr.IsCustomCall("ShardBarrierTo"):              // user-frozen sharding fence
            mark_shard_barrier(instr)                         // propagation may not cross
        scan_shard_group(instr, &instruction_to_shard_group_id,   // shard_as / shard_like
                         &shard_group_id_to_shard_as_group,
                         &shard_group_id_to_shard_like_group)

    // 2. Validate per-output / per-parameter propagation permission.
    CHECK(allow_spmd_sharding_propagation_to_output_.size() matches root tuple arity)

    // 3. Run to fixed point.
    iterations = 0
    changed = RunToFixPoint(aggressiveness=0..3, ..., module, ..., &iterations)
    LOG("Sharding propagation completed after " << iterations << " iterations")
    return changed
```text

> **陷阱 —** `provided_shardings` 中的指令（即用户/前端已标注的指令）是**冻结的**。前向和后向推断可以*读取*它们，但提交门控拒绝覆盖它们；它们是数据流问题的边界条件。把用户标注当作又一个候选的重新实现，会让推断覆盖用户意图。该集合会作为 `const absl::flat_hash_set<const HloInstruction*>& provided_shardings` 传递给每个 helper。

---

## 提交门控 — `MaybeImproveInstructionSharding`

### 目的

前向或后向推断产生的每个候选 sharding 都会汇入同一个函数 `(anonymous)::MaybeImproveInstructionSharding` (`0x1c84e980`)。它是单调性保证：只有当提交候选会使指令的 sharding *严格更具体*，或能与现有 sharding *兼容合并*时，它才会提交。由于 sharding 只沿 lattice 上升（replicated → partially-tiled → fully-tiled）而不会下降，worklist 必然达到固定点。

### 算法

反编译签名为 `MaybeImproveInstructionSharding(HloSharding sharding, HloInstruction* instr, bool may_combine_partial_sharding, int replace_existing, ...)`。重建逻辑：

```c
// (anonymous)::MaybeImproveInstructionSharding  @ 0x1c84e980
function MaybeImproveInstructionSharding(new_sharding, instr,
                                         may_combine_partial_sharding,
                                         replace_existing):
    if not instr.has_sharding():
        instr.set_sharding(new_sharding)                  // first assignment — always commit
        return true
    old = instr.sharding()
    if old == new_sharding:
        return false                                       // no progress
    if may_combine_partial_sharding and
       both old and new are partially-tiled and compatible:
        merged = MergeSharding(old, new_sharding)          // union of tiled dims — strictly tighter
        instr.set_sharding(merged)
        return true                                        // "Refined partial sharding"
    if IsStrictlyMoreSpecific(new_sharding, old) or replace_existing:
        instr.set_sharding(new_sharding)
        return true
    return false                                           // candidate rejected — keep old

固定点循环在成功时发出的两个日志字符串 "Add sharding (forward-pass): " / "Refined partial sharding (forward-pass): "（以及 (backward-pass) / (shard group) 变体）区分了首次赋值路径和partial merge 路径。MaybeImproveInstructionSubSharding 变体 (0x1c850820) 对 tuple-shaped 指令的某个元素执行同样操作，通过 ShapeIndex 索引。

怪癖 — may_combine_partial_sharding 允许从两个不同方向推断出的两个不同 partial sharding 通过交叉其 tiled-dimension 集合而共存，而不是让一个拒绝另一个。没有它，一个 producer 推断为按行 sharded、consumer 推断为按列 sharded 的张量会震荡（每次扫描都拒绝对方选择），而不是收敛到二者共同蕴含的 row-and-column（2-D tiled）sharding。merge 是前向和后向扫描协作而不是互相对抗的原因。

前向传播 — 从操作数推断

目的

前向扫描 InferShardingFromOperands (0x1c856780，最大的 helper，10.1 KB) 根据其操作数上已有的 sharding 为一条指令计算候选 sharding，然后把它交给提交门控。它以 post-order 访问指令（每个操作数先于其用户），因此在单次扫描中，新推断出的操作数 sharding 会立即供其用户使用。

逐 op 推断规则

前向规则依赖 opcode。规则空间的形状由反编译函数体中的 opcode-specific 诊断确认（"Not applying sharding to reduce window because dilatation isn't supported yet"、"sort" 处理、SPMDShardToFullShape 边界检查），本质上是对 HloOpcode 家族的 switch：

Opcode 家族	前向规则（操作数 sharding → 指令 sharding）
Elementwise (`add`, `mul`, `select`, unary math, …)	透传：指令采用其操作数的（合并后）sharding；所有操作数和结果维度一一对应
`broadcast`	将操作数 sharding 提升到 broadcast 结果上，把操作数维度映射到其 broadcast 位置；新维度 replicated
`reshape`	将操作数 tiled 维度通过 reshape 映射：split dim 传播到 split factors，merged dim 只有在两个输入一致时才传播；否则不做推断（留给后向）
`transpose` / `reverse`	根据维度映射对操作数 tile assignment 做置换/镜像
`dot`	LHS/RHS batch dims → output batch dims；sharded contracting dim 表示输出是未规约（partial-sum），并会在之后强制 AllReduce；非 contracting dim 映射到输出维度
`convolution`	与 dot 相同的 dimension-numbers 机制：batch/feature/spatial dims 从 LHS 映射到 output；sharded spatial dims 传播到 output spatial dims
`reduce`	从操作数 sharding 中丢弃被规约维度；sharded reduced dim 会产生 unreduced result
`reduce-window`	类似 `reduce`，但有门控：如果 window 有 dilation，该规则退出（`"… because dilatation isn't supported yet"`）并让 op 保持未 sharded
`gather`	operand-data sharding 通过 gather dim-numbers 映射到 gathered output；index sharding 单独处理
`scatter`	映射 updates/operand shardings；index operand 约束哪些 output dims 可以 sharded
`pad` / `slice` / `dynamic-slice`	未受影响维度上的 tiling 透传；padded/sliced dims 只有在边界与 tile 对齐时传播
`sort`	sorted dim 必须 replicated（不能切分 comparison axis）；其他维度透传
`kCustomCall`	委托给 `CustomCallShardingHelper` vtable（见下文）

// InferShardingFromOperands  @ 0x1c856780  (post-order visit)
function InferShardingFromOperands(instr, computation_map, is_spmd,
                                   aggressiveness, sharding_helper,
                                   call_graph):
    if instr in provided_shardings: return false           // frozen
    if instr.IsCustomCall("SPMDShardToFullShape"): return  // boundary: do not look inside

    switch opcode_family(instr):
        ELEMENTWISE:   cand = MergeOperandShardings(instr)
        DOT, CONV:     cand = MapByDimensionNumbers(instr)  // batch/contracting/spatial map
        REDUCE:        cand = DropReducedDims(operand_sharding, instr.dimensions())
        REDUCE_WINDOW: if HasDilation(instr.window()): return false   // guard
                       cand = DropReducedDims(...)
        RESHAPE:       cand = ReshapePropagate(operand_sharding, instr)  // may be no-op
        GATHER/SCATTER:cand = IndexAndDataPropagate(instr)
        SORT:          cand = ReplicateSortedDim(operand_sharding, instr)
        CUSTOM_CALL:   cand = sharding_helper.InferShardingFromOperands(instr)
        ...
    if cand.has_value():
        may_combine = (aggressiveness >= 1)                 // partial merge enabled past iter 0
        return MaybeImproveInstructionSharding(cand, instr, may_combine, /*replace=*/false)
    return false
```text

> **NOTE —** 从 `RunToFixPoint` 传入的 `aggressiveness` (0..3) 参数会跨迭代单调放宽规则。级别 0 只启用明确、无需 merge 的推断（因此最确定的 sharding 先稳定）；更高级别会启用 partial-sharding merge 和推测性传播（例如传播到操作数仅*部分*一致的 op 上）。精确的逐 opcode 阈值表编码在 switch 分支中，未被穷尽反编译；level-gating 为 HIGH 置信度，精确逐 opcode cutoffs 为 LOW。

---

## 后向传播 — 从用户推断

### 目的

后向扫描是对偶过程：`InferShardingFromUsers` (`0x1c859fa0`) 负责驱动，`GetShardingFromUser` (`0x1c8531e0`) 执行具体工作，计算一个*用户*对其某个操作数所蕴含的 sharding，然后交给同一个提交门控。它以 **reverse post-order** 访问（每个用户先于其操作数），因此用户的 sharding 可在一次扫描中流回输入。这使得标注在 root 上的输出 sharding 能一路回传到参数。

### 算法

`GetShardingFromUser` 是前向逐 op 映射的反向形式：给定 `user` 以及 `instr` 占据的操作数索引，它把 `user.sharding()` 反投影到该操作数的 shape 上。反编译函数体中的诊断 `"update_index <= operand_count"` 确认了 dynamic-update-slice / scatter 反投影路径（update operand 的 sharding 从 output sharding 派生）。

```c
// InferShardingFromUsers  @ 0x1c859fa0  (reverse-post-order)
function InferShardingFromUsers(instr, ..., aggressiveness, is_spmd, sharding_helper):
    if instr in provided_shardings: return false
    for user in instr.users():
        if not user.has_sharding(): continue
        cand = GetShardingFromUser(*instr, *user, aggressiveness,   // 0x1c8531e0
                                   is_spmd, call_graph, sharding_helper)
        if cand.has_value():
            changed |= MaybeImproveInstructionSharding(
                           cand, instr, /*may_combine=*/aggressiveness >= 1, false)
    return changed

// GetShardingFromUser  @ 0x1c8531e0 — inverse per-op map
function GetShardingFromUser(operand, user, ...):
    switch user.opcode():
        ELEMENTWISE: return user.sharding()                 // same shape, same tiling
        TRANSPOSE:   return InversePermute(user.sharding(), user.dimensions())
        RESHAPE:     return ReshapeBackProject(user.sharding(), operand.shape())
        DOT/CONV:    return BackProjectByDimNumbers(user, operand_index)
        DYNAMIC_UPDATE_SLICE / SCATTER:
                     return user.sharding()  // operand inherits output tiling; index handled separately
        GATHER:      return DataOperandSharding(user)        // only for the data operand
        ...
        default:     return nullopt                          // cannot back-infer

陷阱 — 后向推断有意比前向推断更弱：许多 opcode 会反投影到 nullopt（例如无法安全地从 reduced result 推断 reduce operand 的 sharding）。前向推断是主力；后向推断主要用于把输出 sharding（root 标注或 SPMDShardToFullShape 蕴含）推向参数，并解析前向扫描留空的 shape-changing op。让后向与前向同样激进的重新实现，会引入虚假 sharding，随后前向扫描又不得不与其对抗。

Group、Cross-Computation 与 Domain 传播

RunToFixPoint 内部有三个辅助机制与两次主扫描并行运行。

Shard groups — `InferShardingFromShardGroup`

shard_as 和 shard_like 标注声明一组指令必须共享 sharding，而不管数据流如何。InferShardingFromShardGroup (0x1c856420) 会传播：对每个至少有一个成员带显式 sharding 的 shard-group ID，它会把该 sharding 复制到所有其他成员。两个 flavor 的 shape 前置条件不同（从循环中的 "Aligning shard group: "、"Shard-As group "、"Shard-Like group " 字符串恢复）：

shard_as 要求成员具有相同 shape（ShapeUtil::SameDimensions 加 element-type match）。
shard_like 只要求相同维度（element type 可以不同）。

GetRelatedInstructions (0x1c861140) 会沿 group 边收集传递闭包，因此一个成员变化会在下一次迭代中重新排队整个 group。提交通过同一个门控，并记录为 "Add sharding (shard group): " / "Refined partial sharding (shard group): "。

怪癖 — shard group 受 ShardBarrierFrom / ShardBarrierTo custom call 限定。barrier 标记用户已冻结的 sharding（通常由 Shardy import 插入以保留用户约束）；前向、后向或 group 传播都不能越过它。pre-pass 扫描记录 barrier 指令，扫描时跳过它们，因此重新实现必须把 barrier 视为推断防火墙，而不是 no-op。

Cross-computation — `MaybeComputationPropagation`

当一个 kCall、kWhile 或 kConditional 指令获得 sharding 时，必须把它对被调用 computation 的参数和 root 所蕴含的 sharding 传播到该 computation 内，反之亦然。MaybeComputationPropagation (0x1c85a860) 执行此操作，并把受影响 computation 重新加入 worklist（它记录 "Consider computation: "）。这就是 sharding 穿过控制流边界的机制；partitioner 的 PreprocessCallSites 之后依赖这些 sharding 保持一致。

Domain normalization — `NormalizeDomain`

Sharding custom call 会创建 kDomain 括号（成对的 enter/exit 指令，界定某个 sharding 的区域）。推断完成后，NormalizeDomain (0x1c852dc0) 会解决 domain 内部推断出的 sharding 与 domain 边界上的显式 sharding 之间的冲突，选择最紧的一个（replicated < partial < tiled）。随后 HloDomainRemover("sharding", …) 会移除这些括号。该 pass 与 hlo_sharding_util::CanonicalizeLayoutAfterShardingPropagation (0x1e3d55a0) 协作；后者在 propagation 收敛后修正张量布局，使其与选定 tiling 一致。

固定点循环 — `RunToFixPoint`

目的

RunToFixPoint (0x1c85ae60) 是重复扫描直到收敛的外层 driver。其重建签名（从 demangled 符号和周围 by-reference 参数恢复）是一长串参数，携带 worklist 状态、shard-group 映射，以及 int64_t& iterations out-parameter。

算法

反编译函数体确认了逐迭代结构：它记录 "Sharding propagation iteration N" 和运行计数器 "\n instructions already sharded: "、"\n shardings inferred from operands: "、"\n shardings inferred from users: "、"\n shardings inferred from shard group: "、"\n total instructions: "，以及当前 "\n aggressiveness: "。

// ShardingPropagation::RunToFixPoint  @ 0x1c85ae60
function RunToFixPoint(aggressiveness_max, propagate_shard_group,
                       computation_map, provided_shardings, call_graph,
                       module, exec_threads, unspecified_dims,
                       shard_group maps..., &iterations):
    any_changed = false
    for aggressiveness in 0 .. aggressiveness_max:        // 0..3 outer schedule
        do:                                               // inner fixpoint at this level
            changed = false
            if propagate_shard_group:
                changed |= InferShardingFromShardGroup(...)         // 0x1c856420
            for instr in post_order(module):                        // FORWARD sweep
                changed |= InferShardingFromOperands(instr, ..., aggressiveness)
            for instr in reverse_post_order(module):                // BACKWARD sweep
                changed |= InferShardingFromUsers(instr, ..., aggressiveness)
            for comp_instr in control_flow_instrs:
                changed |= MaybeComputationPropagation(...)          // 0x1c85a860
            iterations += 1
            LOG("Sharding propagation iteration " << iterations)
            any_changed |= changed
        while changed                                     // until no instruction's sharding moved
    NormalizeDomain(...)                                  // 0x1c852dc0 — boundary cleanup
    return any_changed
```text

反编译控制流显示在 aggressiveness schedule 上有嵌套的 `do { … } while (changed)`（循环体顶部的最外层 `do` 包住每个级别的收敛过程）。终止有两个保证：**单调提交门控**（没有 sharding 会变松，且 lattice 高度有限，受张量维数约束）以及一个内部迭代上限。

> **NOTE —** 迭代上限以 `long&` (`iterations`) 形式传入/观察，但与之比较的*常量*未在反编译中固定；循环主要依赖 no-change 条件，上限只是安全网。把上限存在视为 HIGH，精确值视为 LOW；重新实现应防御性地限制迭代次数（lattice-height bound 是有原则的界限），而不是信任无界循环。

---

## Custom-Call Sharding Helper 分发

### 目的

`kCustomCall` 指令对通用逐 op 规则而言是不透明的，因此两次扫描都会委托给穿过 pass 的 `CustomCallShardingHelper` 虚表。base `xla::CustomCallShardingHelper::InferShardingFromOperands` (`0x1c864120`) 返回**无 sharding**（`std::nullopt`；其函数体把结果的 presence byte 设为 0 后返回），把该指令留给数据流的其他部分处理。TPU pipeline 注入 `xla::jellyfish::TpuCustomCallShardingHelper`（`InferShardingFromOperands` @ `0x1278bf80`），它为识别出的 TPU custom call 覆盖规则。

### TPU helper 处理的目标

反编译的 `TpuCustomCallShardingHelper::InferShardingFromOperands` (`0x1278bf80`) 会按下面固定顺序对一组 target 测试 `IsCustomCall`；任何未匹配的 target 都返回 `std::nullopt`（no-sharding sentinel，与 base 相同）。匹配分支如下：

| Custom-call target | Sharding 规则 |
|---|---|
| `xla-sdc-checker-get-checksums` | 返回 `std::nullopt` — SDC-checker debug custom call 保持未 sharded（落到与未匹配 target 相同的出口） |
| `QrDecompositionBlock` | 非平凡：当 operand 0 携带 sharding 时，构造一个**双元素 tuple sharding**，即操作数 sharding 加派生 block sharding（`(anonymous)::DeriveQrBlockShardingFromOtherSharding`）；否则返回 `nullopt` |
| `MoveToHost` | 跟随操作数：结果采用 operand 0 的 sharding，**除非**指令自身 sharding 已经是 replicated，此时返回 `nullopt` |
| `PartialReduce` | 非平凡：当 operand 0 已 sharded 时，从 backend config 读取 `reduction_dim`（带 `"PartialReduce backend config cannot be null."` guard）；被规约维度丢弃，其余维度透传 |

该函数还匹配一个通过带 float 参数的 `IsCustomCall` overload 与长度 7 比较的 target，其分支跟随操作数；反编译中没有恢复该字面量（比较读取 `.data` 指针而非 `.rodata` 字符串），因此这里省略。

> **NOTE —** `PartialReduce` *不是*唯一非平凡规则：`QrDecompositionBlock` 也会改变维度/shape（它合成 tuple sharding）。此函数中**没有 `MoveToDevice`** 分支，并且 `xla-sdc-checker-get-checksums` 返回 `nullopt`，而不是跟随其操作数。

兄弟 helper `ShardBarrierFromPartitioner` / `ShardBarrierToPartitioner`（`InferShardingFromOperands` @ `0x1c863740` / `0x1c863800`）和 `InspectShardingCallPartitioner`（`jax::…` @ `0xe8b8cc0`）实现了上文引用的 barrier-fence 和 debug-inspect 行为。

---

## 相关组件

| Component | Relationship |
|---|---|
| `TpuAutoSharding` / `AutoSharding` | 另一种 sharding *生产者*（ILP 搜索）— 每个 module 与此 pass 互斥；见 [Auto-Sharding and SPMD Partitioner](auto-sharding-spmd.md) |
| `ShardyXLA` (`"shardy-xla"`) | 另一种生产者（Shardy/JAX-native import）；当 Shardy 和 `xla.sharding` attrs 共存时回退到 GSPMD propagation |
| `TpuSpmdPartitioner` (`"tpu-spmd-partitioning"`) | *消费者*：把推断出的 sharding 物化为 per-partition HLO + collectives；见相邻页面 |
| `HloDomainRemover` | 在 `NormalizeDomain` 解决冲突后，移除 `Sharding` custom call 创建的 `kDomain` 括号 |
| `CanonicalizeLayoutAfterShardingPropagation` (`0x1e3d55a0`) | 在 propagation 收敛后修正张量布局，使其与选定 tiling 一致 |

## 交叉引用

- [Auto-Sharding and SPMD Partitioner](auto-sharding-spmd.md) — ILP strategy *搜索*与 SPMD *partitioner*（collective 物化）；此 pass 输入的另外两个 sharding 生产者和消费者
- [The TPU Compiler](overview.md) — Part V 导览；sharding/SPMD 在五阶段下降中的位置
- [Compile Phases](compile-phases.md) — 承载 `RunHloPasses` 的 PjRt partial-program phase 主干
- [HLO Pre-Passes](hlo-pre-passes.md) — sharding 推断之前运行的 HLO-level passes
- [HLO Pass Registry](hlo-pass-registry.md) — `ShardingPropagation` 如何添加到 pipeline（`AddPass<>` 实例化）
- [Dot / Conv → MXU Lowering](dot-conv-mxu-lowering.md) — dot/conv op 的下游 lowering；本 pass 会推断这些 op 的 contracting/batch-dim sharding
- [Collectives Overview](../collectives/overview.md) — sharding 推断后 partitioner 发出的 post-SPMD collectives

Sharding Propagation ​

摘要 ​

Pass Driver ​

目的 ​

构造函数 ​

算法 — pre-pass 扫描 ​

前向传播 — 从操作数推断 ​

目的 ​

逐 op 推断规则 ​

Group、Cross-Computation 与 Domain 传播 ​

Shard groups — InferShardingFromShardGroup ​

Cross-computation — MaybeComputationPropagation ​

Domain normalization — NormalizeDomain ​

固定点循环 — RunToFixPoint ​

目的 ​

算法 ​

Sharding Propagation

摘要

Pass Driver

目的

构造函数

算法 — pre-pass 扫描

前向传播 — 从操作数推断

目的

逐 op 推断规则

Group、Cross-Computation 与 Domain 传播

Shard groups — `InferShardingFromShardGroup`

Cross-computation — `MaybeComputationPropagation`

Domain normalization — `NormalizeDomain`

固定点循环 — `RunToFixPoint`

目的

算法