自动分片和 SPMD 分区器

此页面上的所有地址适用于 libtpu-0.0.40-cp314 轮子中的 libtpu.so(构建 ID 89edbbe81c5b328a958fe628a9f2207d)。其他版本会有所不同。

摘要

两个耦合引擎将部分注释的 XLA 程序转换为每分区 SPMD 程序。第一个，自动分片(xla::TpuAutoSharding 封装了开源 xla::AutoSharding)，为每条缺少分片的指令选择一个分片：它为每个操作枚举一组离散的候选策略，为每个操作附加一个成本向量，在指令 DAG 上构建一个策略/成本图，并将整个过程交给使用 OR-Tools 解决的真正的混合整数程序operations_research::MPSolver(CP-SAT整数后端，GLOP用于LP松弛)。第二个，xla::jellyfish::TpuSpmdPartitioner，具体化决策：它逐条指令遍历 HLO，并将每个操作重写到它计算的每个分区片中，插入集合(AllReduce / AllGather / AllToAll / CollectivePermute)和光环交换，使重写在数字上正确。

了解 XLA-on-GPU 的读者已经拥有大部分框架：这与上游船舶相同的 GSPMD 机器，使用 TPU 成本模型(ICI 网格上的 alpha-beta 链路模型)和瘦 TPU 分区子类进行重新定位。自动分片 ILP 是阿尔帕公式 — 每个(节点、策略)一个二进制策略变量、每个(边、策略对)一个二进制重新分片变量、峰值内存约束以及线性目标求和每节点计算/通信成本和每边重新分片成本。分区器是开源 SpmdPartitioner 驱动程序，具有 SpmdPartitioningVisitor 调度约 53 个 Handle* 方法； TPU 子类仅覆盖其中的四个，并且不提供特殊的集体创建者 - 它使用 GetDefaultCollectiveOpsCreator，并且所有特定于 TPU 的集体整形发生在 SPMD 后集体重写的“下游”(请参阅集合通信)。

该页面拥有解算器、分区重写和集体物化。从显式注释“推断”分片的传播过程——其前向/后向推理规则、定点循环、自定义调用分片助手——记录在分片传播上；本页链接它并且不重述其规则。该页面的结构如下：(1) 自动分片策略/成本/求解管道，(2) MIP 公式，(3) TPU 分区驱动程序，(4) 每操作访问者和分区重写，(5) 集体插入和光环交换，(6) TPU 特定扩展(windowed-einsum、scaled-dot、SparseCore 两级分区)和 (7) 标志目录。

重新实现，合约为：

策略枚举模型。 每操作候选生成(DotHandler/ConvHandler 加上内置枚举器)，ShardingStrategy =(分片，成本)元组，以及每条指令保存一个候选集的 StrategyGroup。
成本向量和 ILP。 alpha-beta 通信成本模型、节点与边缘成本以及由 FormulateAndSolveMIPFromProblem 求解的整数程序(变量、约束、目标)。
分区重写。 SpmdPartitioner::RunImpl 预处理→分区→最终确定脊柱和生成每个分区 HLO 的 SpmdPartitioningVisitor 每个操作调度。
集体插入。 5 回调 SPMDCollectiveOpsCreator，决定发出哪个集体的每个操作码规则，以及用于窗口操作的 ExchangeHalo 系列。


自动分片入口	`xla::TpuAutoSharding::RunImpl` @ `0x1118d3a0` → `xla::AutoSharding::RunImpl` @ `0x1280a180`
自动分片核心	`xla::AutoShardingImplementation::RunAutoSharding` @ `0x128055a0` (12.4 KB)
MIP 求解器条目	`xla::spmd::FormulateAndSolveMIPFromProblem` @ `0x128407c0` (22.5 KB)
解算器后端	OR-工具 `operations_research::MPSolver` — CP-SAT (`SatInterface`) 整数 + GLOP (`GLOPInterface`) LP
分区程序条目	`xla::jellyfish::TpuSpmdPartitioner::RunImpl` @ `0x127a2a80` (4.8 KB) → 底座 `xla::spmd::SpmdPartitioner::RunImpl` @ `0x1c7fe400`
每操作调度	`xla::spmd::SpmdPartitioningVisitor` — ~53 `Handle*` 方法
集体创作者	`xla::spmd::GetDefaultCollectiveOpsCreator` @ `0x1c7fc120`(5 个 `std::function` 字段；无 ReduceScatter 回调)
通管道主机	`xla::jellyfish::AddTpuPartitioningPasses` @ `0x1278a440`
通行证名称	`"auto-sharding-automatic-partition"`、`"tpu-spmd-partitioning"`、`"tpu-partition-assignment"`
自动入口门	标志 `xla_tpu_spmd_auto_partitioning`(默认 false)
信心	HIGH(符号字节锚定；密钥体反编译验证)，除非行/标注另有说明

分区管道

用途

分片和分区不是一次传递，而是在 RunHloPasses 内部添加的嵌套管道。 AddTpuPartitioningPasses(0x1278a440，~3.7 KB)是其主机。它的工作是采用一个可能带有“一些”分片注释的 HLO 模块，并生成一个模块，其中每条指令都被注释，然后重写为其每个分区的形式。

入口点

反编译的 AddTpuPartitioningPasses (0x1278a440) 按名称引用下面的每个通道。 ShardingPropagation 出现三次 - 它在创建新传播机会的转换之间重新运行。两个上游通道 (TpuAutoSharding / ShardyXLA) 由同级阶段 AddAutoShardingAndRelatedPasses 添加；此页面涵盖自动路径和分区器，并链接传播/Shardy。

text

RunHloPasses
  AddAutoShardingAndRelatedPasses        ── choose shardings
    ShardingPropagation                  ── manual / partial-annotation flow  (see sharding-propagation.md)
    TpuAutoSharding                       ── auto flow ("auto-sharding-automatic-partition")
    ShardyXLA                             ── Shardy (JAX-native) flow  (see sharding-propagation.md)
  AddTpuPartitioningPasses  0x1278a440   ── materialize per-partition HLO
    SpmdPrepare              0x12dfc300
    ConvOperandSwapper
    TpuSpmdConcatRewriter    0x1278f900   ── "tpu-spmd-concat-rewriter"
    HloConstantSplitter
    TpuPartitionAssignment   0x1278f040   ── "tpu-partition-assignment"  (gated, default off)
    TpuSpmdPartitioner       0x127a2a80   ── "tpu-spmd-partitioning"  ← the rewrite
    RecognizeReduceWindow
    CollectivePermuteCSE
    WholeGraphManualPass
```text

> **注意 —** 有三个分片*生产者*和一个*消费者*。生产者——`ShardingPropagation`(GSPMD 推理)、`TpuAutoSharding`(ILP 搜索)和 `ShardyXLA`(Shardy 导入)——每个模块都是互斥的，并且都集中在同一个产品上：每个 `HloInstruction` 都带有一个 `HloSharding`。消费者 `TpuSpmdPartitioner` 并不关心哪个生产者运行。重新实现者可以单独针对分片注释构建分区器。

---

## 自动分片 — 策略搜索

### 用途

`TpuAutoSharding`(“auto-sharding-automatic-partition”，vtable `0x21822a80`)是自动路径：当用户不提供分片时，它通过将选择制定为离散优化来*导出*最佳分片。它是围绕开源 `xla::AutoSharding` 的薄 TPU 包装；包装器添加 `ApplyShardingConfig`(`0x1118c100`，加载隐藏的配置)和 `ExtractShardingConfig`(`0x1118c820`，序列化所选分片以进行重播)，然后 `RunImpl`(`0x1118d3a0`)分派到基础中。仅当设置 `xla_tpu_spmd_auto_partitioning` 时输入(默认 false)。

### 入口点

```text
TpuAutoSharding::RunImpl                 0x1118d3a0 (0.85 KB)  ── apply config, dispatch
  AutoSharding::RunImpl                  0x1280a180 (11.0 KB)
    AutoShardingImplementation::RunAutoSharding  0x128055a0 (12.4 KB)  ── the multi-phase core

算法

RunAutoSharding (0x128055a0) 携带 10 个源位置 ::site 数据对象($_2..$_11，位于 0x2230c9c0..0x2230ca98，相距 24 个字节) —标记每个阶段的错误路径的状态注释站点。从这些站点及其周围的辅助符号中恢复的相序：

function RunAutoSharding(module, exec_threads):              // 0x128055a0
    SaveAndRemoveShardingAnnotation(module)                  // 0x128020a0 — stash user shardings, clear them
    CanonicalizeLayouts(module)                              // 0x12803860
    cost = BuildHloCostAnalysis(module, tpu_shape_size_fn)   // custom shape-size lambda
    meshes = ComputeMeshShapeCandidates(option)              // try_multiple_mesh_shapes
    for mesh in meshes:                                       // single mesh unless search enabled
        aliases = BuildAliasSet(module)                      // 0x12e174a0 — params/outputs forced equal
        strat_graph = BuildStrategyAndCost(module, cost, mesh)  // per-op StrategyGroup + edge costs
            // DotHandler/ConvHandler register matmul/conv strategies
        ComputeAliasCompatibility(strat_graph)               // 0x12e188e0 — trim alias-incompatible
        TrimOrGenerateStrategiesBasedOnExistingSharding(...)  // honor surviving user shardings
        problem = LowerToIopddl(strat_graph, mesh, budget)    // iopddl::Problem (the MIP IR)
        FindShavedStrategies(problem)                         // 0x12851a60 — drop infeasible strategies
        output = CreateAutoShardingSolverRequestAndCallSolver(problem)  // 0x127f24a0 → MIP solve
        if output.feasible: break                            // else try next mesh
    if !output: error "could not find a solution for any of the mesh shapes tried"
    GenerateReduceScatter(strat_graph, output)               // 0x127fefc0 — RS opportunities
    InsertReshardReshapes(module, output)                    // 0x127f7ee0
    SetHloSharding(module, output)                           // 0x127f7540 — commit chosen shardings
```text

> **明白了 —** 自动分片 **不** 支持 `shard_as`/`shard_like` 组注释(字符串 `"Auto-sharding currently does not support shard_as/shard_like sharding annotations"`)。这些是仅传播的功能 ([分片传播](sharding-propagation.md))。连接自动路径的重新实现者必须拒绝分片组输入，而不是默默地忽略它们。

### 策略枚举

策略是一个操作的候选分片及其成本；一个操作的候选集合是 `StrategyGroup`(元组递归)。候选者来自每个 op-family 的生成器。 `DotHandler`(matmul，在 `DotHandler::RegisterStrategies` `0x12825140` 确认)枚举了一种分片点批量/收缩/非收缩维度的策略，通过 `GenerateNameForDotSharding(lhs_spec, rhs_spec)` 命名每个策略：

| 点策略 | 分片模式 | 分区时隐含的集体 |
|---|---|---|
| `Replicated` | 所有暗淡均已复制 | 无 |
| `Tile[batch]` | 批量并行 | 无(已针对每个分区) |
| `Tile[contracting] + AllReduce` | 威震天张量并行 | AllReduce 输出 |
| `Tile[non-contracting] + AllGather` | 仅输出分片 | AllGather 在分片操作数上 |
| `Tile[contracting] + ReduceScatter` | FSDP型 | ReduceScatter(具体化下游) |
| `ag_windowed_einsum_o%dt%dm%d` | windowed-einsum (AllGather) | 每分区 Dot + 循环内 CollectivePermute |
| `rs_windowed_einsum_t%dm%d` | 窗口化 einsum (ReduceScatter) | 内循环 AllReduce → DynamicSlice |

两个窗口 einsum 附加器被确认为它们自己的 TU 私有函数：`DotHandler::AppendAllGatherWindowedEinsumStrategyForOperand` (`0x12826780`) 和 `DotHandler::AppendReduceScatterWindowedEinsumStrategy` (`0x128270a0`)。 `ConvHandler` 共享相同的 `HandlerBase` 底座。其他操作系列使用内置枚举器：

| 发电机 | 地址 | 生成用于 |
|---|---|---|
| `EnumerateAllPartition` | `0x127ec800` | 数组操作的 n 维平铺 |
| `EnumerateAll1DPartition` | `0x127eaa60` | 单轴平铺机 |
| `AddReplicatedStrategy` | `0x127e8940` | 始终存在的复制后备 |
| `FollowReduceStrategy` | `0x127e4640` | 减少操作数分片后的操作 |
| `CreateReshapeStrategies` | `0x127f1ce0` | 重塑布局 |
| `GenerateOutfeedStrategy` | `0x127e7780` | 出料 |
| `HandlePartialReduce` | `0x127e1e00` | TPU `partial_reduce_handler::kPartialReduce` 定制调用 |
| `MaybeFollowInsStrategyGroup` | `0x127e4340` | tuple / get-tuple-element(继承后面的操作集) |

### 成本模型 — ICI 网格上的 Alpha-Beta

`xla::spmd::ClusterEnvironment`(ctor `0x12808800`)拥有成本模型。它包含物理一维 `DeviceMesh`、重塑的 N 维逻辑网格以及每个网格尺寸的 `device_mesh_alpha`(固定延迟，秒)和 `device_mesh_beta`(每字节成本，秒/字节)向量。每个集体成本都是标准的 alpha-beta 环成本公式，其中 `n[d]` = 沿着暗淡 `d` 的网格大小，`B` = 移动的字节数：

```text
AllReduceCost(B, d)     = α[d] + 2·(n[d]-1)/n[d]  · B · β[d]    // 0x12de7980
AllGatherCost(B, d)     = α[d] +   (n[d]-1)/n[d]  · B · β[d]    // 0x12de77a0
ReduceScatterCost(B, d) = α[d] +   (n[d]-1)/n[d]  · B · β[d]    // 0x12de7b80
AllToAllCost(B, d)      = α[d] +   (n[d]-1)/n[d]² · B · β[d]    // 0x12de7d60

注意 — 四个成本方法地址是字节确认的；封闭式表达式是根据方法名称和操作数形状重建的规范 Alpa/XLA 环成本形式，不是读取反汇编的算术(对精确系数的中等置信度)。 ReshardingCost (0x12de9860) 将它们与 CollectivePermuteCost (0x12de8f00) 结合起来，对任意源→目标分片更改进行定价； GetMeshDimPermutationOrderInShardingSpec(0x12e2ed00)决定重新分片是否经过AllToAll或CollectivePermute。 TPU 网格的默认 alpha/beta 值是特定于版本的(Jellyfish v3 与 Pufferfish v5 …)，并且未固定在二进制文件中(低)；字符串 "If not sure how to set device_mesh_alpha and device_mesh_beta, please leave them empty and default values will be used." 显示默认值存在。

这些每个策略的成本变成两个成本数组：节点成本(计算+策略对操作本身施加的通信)和边缘成本(当生产者的策略和消费者的策略不一致时的重新分片成本)，由 ComputeCommunicationCost、CommunicationReshardingCostVector 和 MemoryReshardingCostVector 填充。

功能图

功能	地址	角色
`TpuAutoSharding::RunImpl`	`0x1118d3a0`	TPU 包装，发送至基地
`TpuAutoSharding::ApplyShardingConfig`	`0x1118c100`	应用隐藏配置
`TpuAutoSharding::ExtractShardingConfig`	`0x1118c820`	序列化所选分片
`AutoSharding::RunImpl`	`0x1280a180`	基本条目
`AutoShardingImplementation::RunAutoSharding`	`0x128055a0`	多相铁芯
`SaveAndRemoveShardingAnnotation`	`0x128020a0`	存储 + 清除用户分片
`BuildAliasSet`	`0x12e174a0`	强制别名参数/输出相等
`DotHandler::RegisterStrategies`	`0x12825140`	matmul策略枚举
`DotHandler::AppendAllGatherWindowedEinsumStrategyForOperand`	`0x12826780`	AG 窗口化 einsum 候选者
`DotHandler::AppendReduceScatterWindowedEinsumStrategy`	`0x128270a0`	RS 窗口 einsum 候选者
`ClusterEnvironment::ClusterEnvironment`	`0x12808800`	成本模型状态
`StrategyShaverForProblem::FindShavedStrategies`	`0x12851a60`	下降不可行的策略
`SetHloSharding`	`0x127f7540`	提交

MIP 配方

用途

策略选择是作为真正的混合整数程序来解决的，而不是贪婪的启发式程序。反编译的FormulateAndSolveMIPFromProblem(0x128407c0，约22.5 KB)直接在OR-Tools上构建 - 其局部变量为operations_research::MPSolver*、MPObjective*、MPConstraint*和MPVariable*，并且它构造行约束(MakeRowConstraint)和一个目标。两个求解器后端链接：operations_research::SatInterface(CP-SAT、Solve 0x1285dce0)用于整数阶段(默认)，operations_research::GLOPInterface(LP 单纯形、Solve 0x12ef4760)用于松弛。第三方 LP 求解器(Gurobi、HiGHS 等)都没有作为代码链接，仅作为描述符原型链接，因此实时后端是 GLOP + CP-SAT。

算法 — 整数规划

该模型是 Alpa 公式：每个节点一个单热策略向量、每个边一个单热重新分片向量、峰值内存约束和线性目标。中间表示为 iopddl::Problem(节点、边、策略)，被确认为求解器条目的第一个形式参数 - 其完整签名为 FormulateAndSolveMIPFromProblem(const iopddl::Problem&, const xla::spmd::AutoShardingSolverParams&)。

text

DECISION VARIABLES
  s[v][k] ∈ {0,1}        per node v, per candidate strategy k
  e[(u,v)][i,j] ∈ {0,1}  per edge (u,v), per (u-strategy i, v-strategy j) pair

CONSTRAINTS
  Σ_k s[v][k] = 1                      ∀ v                    // exactly one strategy
  Σ_{i,j} e[(u,v)][i,j] = 1            ∀ (u,v)                 // one resharding per edge
  e[(u,v)][i,j] ≤ s[u][i]                                     // edge–node consistency
  e[(u,v)][i,j] ≤ s[v][j]
  Σ_{v live at t} mem(s[v][·]) ≤ memory_budget_per_device     // peak-memory, per time step t
  alias-follow: aliased params/outputs must take matching strategies   // BuildAliasSet
  group: shard_as/shard_like members forced identical (propagation path only)

OBJECTIVE  (minimize)
  Σ_v Σ_k  node_cost[v][k]·s[v][k]
+ Σ_(u,v) Σ_(i,j)  edge_cost[(u,v)][i,j]·e[(u,v)][i,j]
```text

两个与 TPU 相关的预处理步骤在求解之前缩小了程序。 `CheckDominance`(`0x12850cc0`，~2 KB)删除了每个成本轴上占主导地位的策略； `StrategyShaverForProblem::FindShavedStrategies` (`0x12851a60`) — Google 内部扩展 — 删除了无法出现在任何可行解决方案中的策略； `ReduceMemoryTerms` 将峰值内存约束折叠为更少的变量。求解后，`Evaluate`(`0x128473c0`，~7.7 KB)对所选分配进行重新定价以进行报告。

进入求解器的请求是 `AutoShardingSolverRequest` 原型(解析表 `AutoShardingSolverRequest::_table_` @ `0x218f17c0`、vtable `0x218f13e8`)：嵌套 `Nodes`(每节点策略列表)、`Edges`(端点对)、 `Costs`/`Coeff`(成本向量和系数矩阵)、`Pair`((i，j)条目)、`Group`(分片组约束)、`SolverTimeout`和`Names`(调试)。 `solver_type`选择后端(默认为`SOLVER_TYPE_CP_SAT`)，`solver_specific_parameters`携带文本格式的`SatParameters`或`GlopParameters`原型。

### 功能图

| 功能 | 地址 | 角色 |
|---|---|---|
| `FormulateAndSolveMIPFromProblem` | `0x128407c0` | 通过 OR-Tools `MPSolver` 构建 + 求解 MIP |
| `CreateAutoShardingSolverRequestAndCallSolver` | `0x127f24a0` | 构建请求、调度 |
| `(anon)::SolveAndExtractSolution` | `0x2139b540` | 调用`MPSolver::Solve`，读取变量 |
| `CheckDominance` | `0x12850cc0` | 剪枝主导策略 |
| `Evaluate` | `0x128473c0` | 解决后重新定价 |
| `operations_research::SatInterface::Solve` | `0x1285dce0` | CP-SAT 整数后端 |
| `operations_research::GLOPInterface::Solve` | `0x12ef4760` | GLOP LP 放松 |

> **QUIRK —** 请求携带 `SolverTimeout` 并且代码发出三个不同的失败字符串：“在给定时间限制内找不到有效的解决方案。请将其报告为错误！”、“只能在给定时间限制内找到非最佳解决方案。”以及无网格可行消息。重新实现者必须将 CP-SAT 求解视为易错并在超时时回退(当节点成本无限时，`use_sharding_propagation_for_default_shardings` 选项通过 [传播](sharding-propagation.md) 计算默认值)。

---

## 分区器 — `TpuSpmdPartitioner`

### 用途

`xla::jellyfish::TpuSpmdPartitioner`(“tpu-spmd-partitioning”，ctor `0x127a1e40`)是消费者：给定一个每条指令都有分片的模块，它将全局程序重写为单个分区运行的 SPMD 程序。它是开源 `xla::spmd::SpmdPartitioner` 的瘦子类；反编译的 `RunImpl` (`0x127a2a80`) 调用基础 `SpmdPartitioner::RunImpl` (`0x1c7fe400`) 并仅添加 TPU 布局/精度修复。已确认源路径：`platforms/xla/service/jellyfish/spmd/tpu_spmd_partitioner.cc`。

### 入口点

```text
TpuSpmdPartitioner::RunImpl   0x127a2a80 (4.8 KB)
  SpmdPartitioner::RunImpl    0x1c7fe400 (7.7 KB)            ── base driver
    PreprocessSharding        0x1c804ae0 (2.9 KB)            ── fill Replicated, validate
    PreprocessCallSites       0x1c800260 (5.1 KB)            ── flatten call/while/conditional
    PreprocessHlos            0x1c8016c0 (9.5 KB)            ── canonicalize for emission
    PartitionComputation      0x1c7fda60                     ── per-computation visitor walk
      SpmdPartitioningVisitor  ── ~53 Handle* methods (the rewrite)
    ConvertUnreducedSharding  0x1c803d00 (3.5 KB)            ── unreduced → AllReduce/ReduceScatter
    RecordInputsOutputsSharding  0x1c7fe0a0                  ── annotate alias config

算法

驱动程序是一个固定的三阶段脊柱：规范化图，通过访问者遍历每个计算，然后最终确定。

function TpuSpmdPartitioner::RunImpl(module, exec_threads):   // 0x127a2a80
    status = SpmdPartitioner::RunImpl(module, exec_threads)   // 0x1c7fe400 — base does the work
        PreprocessSharding(module)        // every op gets a sharding; Replicated default
        PreprocessCallSites(module)       // push shardings through kCall/kWhile/kConditional
        PreprocessHlos(module)            // unfold tuple Sharding annotations, etc.
        for comp in module.computations():
            visitor = CreateVisitor(comp, num_partitions, collective_ops_creator)  // 0x127a2520
            PartitionComputation(comp, visitor)   // dispatch each op to Handle<Op>
        ConvertUnreducedSharding(module)  // 0x1c803d00 — IsUnreduced outputs → collective
        RecordInputsOutputsSharding(module)
    return status                         // tpu_spmd_partitioner.cc source-loc on the error paths
```text

### TPU 覆盖

TPU 子类恰好重写了四个基本方法；其他一切，包括访问者，都是开源基础——TPU 特定的行为是通过集体创建者和下游重写而不是通过子类来传递的。

| 覆盖 | 地址 | 改变了什么 |
|---|---|---|
| `AllGatherShards` | `0x127a2880` | 37 字节 (`0x25`) 存根 → 转发到基础 `AllGatherShardsInternal`，传递尾随的 per-dim-communication bool |
| `AllReduceAlongShardingDims` | `0x127a28c0` (0.43 KB) | 通过 `MayIncreaseBF16AllReduceAccumulationAccuracy` (`0x127a22c0`) 进行 F32-vs-BF16 累积选择 |
| `CreateVisitor` | `0x127a2520` | 构造基础 `SpmdPartitioningVisitor`(无 TPU 访问者子类) |
| `UpdateLayout` | `0x127a4100` (0.42 KB) | 将每个分区形状重写为 TPU 规范(子通道/通道)布局 |

> **注意 —** `AllReduceAlongShardingDims` 使用签名 `(SpmdBuilder*, HloInstruction*, const HloSharding&, int64_t*, absl::Span<const int64_t>, SPMDCollectiveOpsCreator)` 进行字节确认。其精度门`MayIncreaseBF16AllReduceAccumulationAccuracy`也得到了确认，采用`ObjectView<TpuCompilationEnvironment>`和创建者；它查询 `xla_tpu_spmd_f32_accum_for_bf16_ar` 和 `_min_subgroup_size` 伴随标志来决定是否将 BF16 约简向上转换为 F32 累加。决策是纯粹由阈值驱动还是由配置文件驱动尚未解决(低)。

---

## 每操作重写 — `SpmdPartitioningVisitor`

### 用途

访问者是全局→每个分区重写发生的地方，一次一个 HLO。每个 `Handle<Op>` 都采用全局指令及其操作数的 `PartitionedHlo` 包装器，计算每个分区的替换，并插入使替换正确的任何集合。约 53 名处理人员被找到。访问者没有 TPU 子类——基类处理所有操作码，TPU 专门化存在于集体回调和 SPMD 后重写中。

### 算法 — 重写轴

处理程序不是 53 行操作码→处理程序，而是聚集成少量重写形状。与重新实现相关的轴是*处理程序针对操作维度上的分片分歧所做的操作*：

| 重写形状 | 代表处理人 | 它发出什么 |
|---|---|---|
| **传递**(分片已最终) | `HandleElementwise`、`HandleBroadcast`、`HandleTranspose`、`HandleOptimizationBarrier`、`HandleCollectivePermute` | 每分区操作，无集合 |
| **收缩调光减少** | `HandleDotHelper`、`HandleConvolution`、`HandleReduce` | 每分区计算 + **AllReduce** 输出 |
| **计算前重新分片** | `HandleDotWithoutConflicts`、`HandleReshape`、`HandleSlice`/`HandleDynamicSlice` | reshard操作数(AllGather / SliceValidData)然后op |
| **索引数据分割** | `HandleGather`、`HandleScatter`、`HandleSort` | 收集索引 (AllGather) / 每个分区操作 + AllReduce / AllToAll |
| **空间光环** | `HandleConvolution`(空间)、`HandleReduceWindow`、`HandleSelectAndScatter` | **光环交换** (CollectivePermute) + 每分区窗口操作 |
| **递归** | `HandleConditional`、`HandleWhile`、`HandleCall`、`HandleTuple` | 分区每个子区域/元素 |
| **自定义呼叫调度** | `HandleCustomCall`(`0x1c716540`，5.1 KB) | 通过 `custom_call_target` 调度表 → `_SPMDInternal_*` / TopK / 部分归约处理程序 |

点/转换处理程序共享一个内核 `PartitionDot<...>`，由函子参数化：

| 函子实例化 | 地址 | 用于 |
|---|---|---|
| `HandleDotHelper<CreateShardedDotFunctor>` | `0x1c7191c0` | 通用 matmul |
| `HandleDotHelper<CreateShardedConvolutionFunctor>` | `0x1c7200e0` | 卷积(已验证：`HandleConvolution`调用此) |
| `HandleDotHelper<CreateShardedScaledDotFunctor>` | `0x1c71c420` | 缩放点(NVFP4 / 缩放-FP8；`PartitionedHloMX`) |

> **QUIRK —** `HandleConvolution` (`0x1c703120`) 本身并不实现卷积分区 - 反编译体显示它调用 `HandleDotHelper<CreateShardedConvolutionFunctor>`。卷积被划分为*作为点*；只有空间光环交换部分(`PartitionConv` `0x1c76bea0`)是特定于转换的。编写单独的转换分区器的重新实现是复制点内核。

### PartitionedHlo 抽象

进入处理程序的每个操作数都是 `PartitionedHlo`：每个分区的 HLO 值加上其当前分片。需要不同分片的处理程序调用重新分片，它通过 `ClusterEnvironment::ReshardingCost` 对更改进行定价并发出相应的集合。执行重新分片的访问者内部助手：`ShuffleDataWithAllToAll`(`0x1c791340`，通过每个等级 AllToAll 进行完整复制)、`GetAllToAllSharding` (`0x1c7d8400`)、`ClampGatherIndices` (`0x1c793b20`) 和 `GetPerGroupCollectiveOpsCreator` (`0x1c824140`，推送子创建者以进行分层重新分片)。

---

## 集体插入

### 用途

分区器从不直接构造 `HloInstruction` 集合体 - 它通过 `xla::spmd::SPMDCollectiveOpsCreator` 回调表进行调用，因此相同的分区器服务于每个后端。 TPU 编译器安装*默认*创建者；它不添加自定义回调。所有特定于 TPU 的集体整形(组合、异步、F32 累积、跨片 MegaScale)都发生在 [集体重写](../collectives/overview.md) 的下游，而不是在这里。

### 创建者结构体

反编译的 `GetDefaultCollectiveOpsCreator` (`0x1c7fc120`) 构建了一个由 **五个** `std::function` 字段组成的结构，每个字段都是 TU 私有的 `$_N` 闭包，其中捕获了 `num_replicas`/`num_partitions`。字段布局和签名是字节确认的：

```c
struct SPMDCollectiveOpsCreator {                 // built at 0x1c7fc120
  // $_0  CollectivePermute — source-target pair list
  fn<HloInstruction*(SpmdBuilder*, HloInstruction* operand,
                     vector<pair<int64,int64>>& source_target_pairs,
                     int64 next_channel_id)>            create_cross_partition_collective_permute;
  // $_1  PartitionId — emits the kPartitionId constant for the visitor
  fn<HloInstruction*(SpmdBuilder*)>                     create_partition_id;
  // $_2  AllReduce — pure reduction, no group split
  fn<HloInstruction*(SpmdBuilder*, HloInstruction* operand,
                     HloComputation* reduction,
                     const CollectiveDeviceListBase& device_list,
                     int64 next_channel_id)>            create_cross_partition_all_reduce;
  // $_3  AllGather — concat per-partition slices along all_gather_dim
  fn<HloInstruction*(SpmdBuilder*, Span<HloInstruction* const> operands,
                     const CollectiveDeviceListBase& device_list,
                     int64 all_gather_dim,
                     optional<int64> next_channel_id)>   create_cross_partition_all_gather;
  // $_4  AllToAll — operand split / output gather
  fn<HloInstruction*(SpmdBuilder*, HloInstruction* operand,
                     const Shape& output_shape,
                     const CollectiveDeviceListBase& device_list,
                     int64 split_dim, int64 concat_dim)> create_cross_partition_all_to_all;
};

明白了 — 5 字段结构中 没有 ReduceScatter 回调(已验证 — 仅写入 $_0..$_4)。 “ReduceScatter”不会在访问者时发出。它稍后会具体化为“AllReduce then DynamicSlice”或通过专用下游 TpuAllReduceScatterFusion 通道(请参阅 ReduceScatter)。添加第六个回调的重新实现者与二进制文件不同； TPU 分区器依靠融合通道从 AR+DS 恢复 RS。

每个操作码插入规则

操作需要哪个集合的决定不在创建者中——而是在策略生成器(用于自动路径)和访问者处理程序(用于重写)中。恢复后的规则集：

分片 HLO 模式	插入集体	发出者
`dot` / `conv` 带分片收缩暗淡	AllReduce 输出	`PartitionDot` + `HandleAllReduce`
`dot` 分片批次、复制签约	无(已针对每个分区)	`HandleDotWithoutConflicts`
`dot` 仅按输出分片	分片操作数上的 AllGather	`HandleDotHelper`
`dot` 窗口式 einsum (AG)	每分区 Dot + 循环内 CollectivePermute	`AppendAllGatherWindowedEinsumStrategyForOperand`
`dot` 窗口式 einsum (RS)	内循环 AllReduce → DynamicSlice	`AppendReduceScatterWindowedEinsumStrategy`
`reduce` 沿分片暗淡	AllReduce	`HandleReduce`
`convolution` 跨空间维度分片	光环交换(CollectivePermute)+每部分转换	`PadEachPartitionWithHaloExchange`
`reduce-window` 分片空间	halo交换+每分区缩减窗口	`HandleReduceWindow` + `ExchangeHaloAndGetValidData`
`slice` / `dynamic-slice` 跨分区	`SliceValidData` + AllGather	`HandleSlice`，`HandleDynamicSlice`
`gather` 具有分片索引	AllGather索引，然后按分区收集	`PartitionGather`
`scatter` 具有分片索引	每分区分散 + AllReduce	`PartitionScatter`
`concat` 沿分片暗淡	复制+连接(TPU重写) 分片暗淡上的	`TpuSpmdConcatRewriter`
`sort`	按分区排序 + AllToAll	`HandleSort`
`unreduced` 输出	AllReduce或ReduceScatter(取决于消费者)	`ConvertUnreducedSharding`
`partial-reduce`(TPU定制调用)	每分区部分归约 + AllReduce	部分减少访客
跨分区 FFT	每分区 FFT + CollectivePermute	`GetFinalFftUsingCollectivePermute`

光环兑换

窗口操作(卷积、缩减窗口、选择和分散)需要每个分区查看相邻分区边界数据的“光环”。 ExchangeHalo 系列通过滑动窗口 CollectivePermutes 实现了这一点。 ExchangeHaloAndGetValidData (0x1c825660) 是条目：其字节确认签名采用操作数、其 Shape、两个 OffsetCalculation(窗口边界)、四个 int64 光环大小、HloSharding、填充值和SPMDCollectiveOpsCreator。家人：

助手	地址	角色
`ExchangeHaloAndGetValidData`	`0x1c825660`	全光环交换+有效数据掩码(条目)
`ExchangeHalo`	`0x1c822340`	核心基于CollectivePermute的边界交换
`ExchangeHaloCompact`	`0x1c81d3e0`	压缩变体(更少的排列)
`PadEachPartitionWithHaloExchange`	`0x1c790640`	用相邻分区的边缘填充每个分区
`TileToPartialReplicateHaloExchange`	`0x1c81ccc0`	过渡图块→部分复制时的光环
`GetFinalFftUsingCollectivePermute`	`0x1c791980`	FFT专用滑动窗口

注意 — 光环宽度是根据窗口的扩张/步幅/填充计算的，而不是单独根据分片计算的 — ExchangeHaloAndGetValidData 精确地采用 OffsetCalculation 参数，以便在交换后可以屏蔽每个分区的有效区域。交换固定宽度光环的重新实现器将在数组边缘读取垃圾；有效数据掩码是强制性的。

TPU 特定扩展

除了标准 GSPMD 原语之外，TPU 分区器还添加了多个分片感知构造。这些都没有引入新的集体回调——它们决定了标准集体的走向。

Windowed-einsum. 一个操作数太大而无法预先进行 AllGather 的 matmul 被划分为 while 循环，该循环与计算和通信重叠。 AG 变体在每次迭代中点一个每个分区的 LHS 切片，然后 CollectivePermute 将切片移动一个分区；输出最终完全复制。 RS 变体通过每次迭代的部分 AllReduce 累积每个分区的输出切片。两者均在策略时选择(ag_/rs_windowed_einsum_* 策略)，并由 xla_tpu_enable_windowed_einsum_for_all_gather / _for_reduce_scatter 进行门控，其中 xla_tpu_spmd_unroll_windowed_einsum、_bidirectional_windowed_einsum 和 xla_jf_spmd_threshold_for_windowed_einsum_mib 大小阈值控制循环形状。 WindowedEinsumLoopConfig 记录所选的配置。
Scaled-dot (CreateShardedScaledDotFunctor). 对于 NVFP4/scaled-FP8 matmuls，分区器通过 PartitionedHloMX 将 (operand, scale) 对视为标记元组，将尺度张量与其操作数共同分片，以便重新分片将两者一起移动。
MultiPad / MultiSlice / MultiRotate / RotateRight. TPU xla.spmd_internal.* 自定义调用，用于批量按分区操作，每个调用都有自己的 HandleCustomCallSPMDInternal_* (0x1c70e7e0..0x1c715b60)。 RotateRight 是跨分区的右循环移位，通过 CollectivePermute 实现。
partial_reduce_handler::kPartialReduce. 唯一具有非平凡分片传播的 TPU 自定义调用；它注册自己的 SpmdPartitioningVisitor 并发出由 kReductionDimKey、kLog2ReductionKey、kRecallTargetKey 控制的 Sort+TopK 骨架。
分层 SparseCore 分区。 对于嵌入密集型模型，条目计算分为两个粒度 - TensorCore 和 SparseCore。 SparseCoreHierarchicalSpmdPartitioner (RunImpl 0x13c7ee20, ~10.6 KB) 填充 SC 输入 (PadSparseCoreProgramInputs)，取消填充输出，并显式分区 SC 条目计算；内部SparseCoreSpmdPartitioner(ctor 0x13c818a0)和SparseCorePartitioningVisitor覆盖HandleSort/HandleScatter/HandleAllToAll并添加PartitionSharedMemoryParallelScatter。来源：platforms/xla/sparse_core/hlo/sparse_core_spmd_partitioning.cc。
分片屏障和自定义调用帮助程序。 ShardBarrierFromPartitioner / ShardBarrierToPartitioner、TpuLogCustomCallPartitioner(_xla_log 调试)和超大规模 MetadataCustomCallPartitioner 是 CustomCallShardingHelper 子类；它们冻结或通过特定自定义调用目标的分片(在传播期间咨询 - 请参阅定制呼叫降低)。

旗帜目录

分片/SPMD/集体标志通过 TpuCompilationEnvironment 读取(从 AbslFlag*Gen* 符号系列恢复；默认从 AbslFlagDefaultGenFor*)。入口门和分区器相关子集：

标志	型	默认	控制
`xla_tpu_spmd_auto_partitioning`	布尔值	假	进入自动分片(ILP)路径
`xla_tpu_spmd_auto_partitioning_search_mesh_shapes`	布尔值	假	`try_multiple_mesh_shapes`
`xla_tpu_spmd_run_partition_assignment`	布尔值	假	运行`TpuPartitionAssignment`
`xla_tpu_spmd_skip_partitioning`	布尔值	假	完全跳过SPMD(调试)
`xla_tpu_spmd_decompose_sharded_concats`	布尔值	正确	`TpuSpmdConcatRewriter`
`xla_tpu_spmd_f32_accum_for_bf16_ar`	布尔值	—	F32 BF16 AllReduce 累积
`xla_tpu_spmd_f32_accum_for_bf16_ar_min_subgroup_size`	int64	—	上述最小子组
`xla_tpu_enable_windowed_einsum_for_all_gather`	布尔值	—	允许 AG windowed-einsum
`xla_tpu_enable_windowed_einsum_for_reduce_scatter`	布尔值	假	允许 RS 窗口化 einsum
`xla_tpu_spmd_unroll_windowed_einsum`	布尔值	假	展开 WE 循环
`xla_tpu_spmd_bidirectional_windowed_einsum`	布尔值	假	前进+后退班次表
`xla_jf_spmd_threshold_for_windowed_einsum_mib`	int64	—	启用 WE 的大小阈值 (MiB)
`xla_tpu_auto_spmd_partitioning_memory_budget_gb`	int64	—	`memory_budget_per_device`
`xla_tpu_auto_spmd_partitioning_memory_budget_ratio`	双	—	`memory_budget_ratio`
`xla_use_shardy`	布尔值	假	使用Shardy生产者而不是GSPMD

注意 — 破折号标记默认值未从此二进制文件中的 AbslFlagDefaultGenFor* 固定；门控布尔值默认关闭，因此开箱即用的编译使用传播，而不是 ILP。重新实现者应该将自动分片视为选择加入的路径。

未解决的问题

确切的 AutoShardingOption C++ 字段偏移(名称已恢复；布局需要在 0x12e0ce00 处行走 CheckAndSetup)。低的。
构造求解器所用的确切 MPSolver::OptimizationProblemType 枚举值 — 从 SatInterface 链接和 auto_sharding_cpsat_for_problem.cc 源路径推断出 CP-SAT，而不是读取构造函数参数。中等的。
每个 TpuVersion 的默认 device_mesh_alpha / device_mesh_beta 值。低的。
MayIncreaseBF16AllReduceAccumulationAccuracy 是仅阈值驱动还是轮廓驱动。低的。
TpuExp0PartitioningAlgorithm(0x1278eea0、0xd1字节中唯一注册的PartitioningAlgorithm、Run)委托给$_0 lambda；没有追踪到它实施了什么实验启发法。默认情况下它是关闭的。低的。

组件	关系
`ShardingPropagation`	GSPMD 生产者；推断分区器消耗的分片
`ShardyXLA`	Shardy(JAX 原生)制作人； GSPMD 的替代方案
`TpuPartitionAssignment`	门控预通道，可以在分区器之前选择分区算法
`TpuSpmdConcatRewriter`	TPU 预通道分解分片串联
`TpuAllReduceScatterFusion`(以及集体重写)	下游；恢复ReduceScatter并塑造所有特定于TPU的集体

交叉引用

分片传播 — 推理规则、定点循环和自定义调用分片助手；此页面的兄弟制作人
TPU编译器 — 分区管道位于 RunHloPasses 中
编译阶段 — 承载 AddTpuPartitioningPasses 的五相主干
定制呼叫降低 — partial_reduce_handler、Mosaic 和 shard-barrier 自定义调用表面
RaggedDot 和卷积几何降低 — 分区器发出的每分区点/转换如何进一步降低
集体概述 - 后 SPMD 集体重写，为 TPU 网格塑造 AllReduce/AllGather/AllToAll
ReduceScatter —其中具体化了ReduceScatter，因为分区器不直接发出任何信号

自动分片和 SPMD 分区器 ​

摘要 ​

分区管道 ​

用途 ​

入口点 ​

算法 ​

功能图 ​

MIP 配方 ​

用途 ​

算法 — 整数规划 ​

算法 ​

每个操作码插入规则 ​

光环兑换 ​

TPU 特定扩展 ​

旗帜目录 ​

未解决的问题 ​

相关组件 ​

交叉引用 ​