-sdy-close-shardings

关闭张量分片并舍弃复制的轴。

-sdy-constant-or-scalar-merger

合并具有匹配分片的相同常量和标量扩展。

对具有相同分片的常量执行轻量级 CSE。

导入流水线会拆分并复制常量和标量扩展,以便在常量子计算的不同用途之间不传播分片。如果常量在传播后具有相同的分片,此传递会合并它们以节省编译时间。如需了解详情,请参阅 -sdy-constant-or-scalar-splitter。

-sdy-convert-global-to-local

将 SDY 计划从全球形状转换为本地形状。

通过根据分片属性对逻辑维度进行分区,将 SDY 程序从全局形状转换为本地形状。

此 pass 利用类型转换器将 RankedTensorType 从全局逻辑形状映射到设备本地物理形状。

选项

-per-dim-all-gather                     : Keep per-dimension all-gather without combining them into a single all-gather.
-combine-multi-dimension-reduce-scatter : Combine multi-dimension reduce-scatter into a single reduce-scatter.
-enable-rgv3                            : Use StableHLO ReplicaGroupV3 (mesh-axes based) for collectives.
-replica-count                          : Number of replicas (data parallelism).
-partition-count                        : Number of partitions (model parallelism).

-sdy-drop-sharding-and-mesh

从程序中移除网格操作和分片表示法。

-sdy-drop-sharding-rules

从所有已注册的 op 中舍弃 OpShardingRuleAttr

-sdy-export-named-computations

概述来自 NamedComputationOp 的通话。

创建一个传递,该传递通过名为 NamedComputationOp 的新私有函数将 NamedComputationOp 转换为 CallOpname。新的 FuncOpCallOp 与原始 NamedComputationOp 的操作数/结果具有相同的分片。

如果模块中存在与 NamedComputationOp 同名的函数,MLIR 符号表会将其更改为 {name}_#

-sdy-insert-explicit-reshards

插入显式重新分片,使所有操作都具有兼容的分片。

兼容分片本质上是指,相应操作可以接受分片的运算对象并生成分片的结果,而无需任何重新分片通信(请注意,相应操作可能仍需要通信,例如全归约或光晕交换)。

传播后,某些操作可能仍具有不兼容的分片。

请注意,当一个轴(或子轴)用于跨多个张量对不对应的维度(例如 matmul 中的非收缩维度)进行分片时,或者当一个轴对一个张量中的维度进行分片,但不对另一个张量中的相应维度进行分片时,我们称该操作存在分片冲突。因此,在此传递之后,操作将不再冲突。

此传递会显式注入重新分片操作,以便对于每个操作,相应维度在所有操作数和结果中以相同方式分片,并且每个轴(或子轴)只能用于对单个维度类型进行分片。

示例:

输入:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

输出:

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

在上面的示例中,lhsrhs 都在其非收缩维度上沿“x”轴分片,这是不兼容的。该 pass 在点运算之前在 rhs 上插入显式 reshard,以便点运算具有兼容的分片。

选项

-enable-full-version : Enable full version.

-sdy-insert-func-call-reshards

插入了针对函数和调用分片冲突的重新分片。

为函数插入重新分片,并针对结果调用分片冲突。

-sdy-optimize-collectives

优化集体通信操作。

通过分解拆分维度并融合到组合的全到全,消除了所有到全链之前冗余的集体排列操作。

-sdy-pad-for-divisibility

将具有不可分分片的张量填充为可分形状。

-sdy-per-instruction-partitioning

_选择性地提取目标指令,对其运行分区器流水线以生成设备本地代码,并将其封装在 sdy.manual 计算中。

出于调试和二分目的,使用独立的 Shardy 分区器流水线(resolve-permutation-factors、reshard-to-collectives、pad-for-divisibility、convert-global-to-local)选择性地对各个指令进行分区,并将生成的设备本地代码封装在 sdy.manual_computation 中。这样可以保持其余全局程序的完整性,以便 XLA SPMD 仅对未封装的全局指令进行分区。

选项

-filter               : Filter string for selective partitioning. Can be empty (all sharded ops), comma-separated op name substrings (e.g. 'dot, pad'), or key-values like 'selectLow=0, selectHigh=10'.
-enable-halo-exchange : Implement halo exchange logic for windowed operations inside the sub-pipeline.
-replica-count        : Number of replicas (data parallelism).
-partition-count      : Number of partitions (model parallelism).

-sdy-propagate-to-func-results

将分片从函数终止符传播到函数结果。

将函数终止符值的分片复制到相应的 func.func 结果,同时保留主函数上的现有分片。

-sdy-remove-all-gather-reduce-scatter-for-cmv1

_移除了 CMV1 的 sdy.all_gather 和 sdy.reducescatter。

移除模式 all-gather + 点中的 all-gather。移除了模式点 + reduce-scatter 中的 reduce-scatter。此 pass 用于与集合 matmul V1 (CMV1) 保持兼容性。这是针对 b/432019089 的临时解决方案。

-sdy-remove-propagation-debug-info

在导出期间移除传播调试信息(传播边和来源分片)。

-sdy-remove-sharding-groups

在传播后移除 ShardingGroupOps。

-sdy-remove-sub-axes-in-input-output-shardings

移除了输入/输出分片中的子轴。

部分 Shardy 用户希望函数输入/输出具有不含子轴的分片。此传递会从输入/输出开放维度分片中移除子轴及其尾随轴。此传递通常在 sdy-update-non-divisible-input-output-shardings 之后进行,以确保移除子轴不会引入任何不可分割的分片。

-sdy-reshard-to-collectives

将 ReshardOp 转换为各种 Shardy 集体运算。

匹配 reshard 操作并将其重写为各种 Shardy 集体操作。在此传递之后,模块中不再有任何重新分片操作。

(可选)如果 keepRedundantReshards 为 true,则剩余的唯一重新分片操作是冗余操作。默认情况下,它会假定已插入显式重新分片 (sdy-insert-explicit-reshards),并且不会保留冗余的重新分片。如果可能尚未插入显式重新分片,则应保留冗余重新分片。

示例:

输入:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

输出:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.all_gather \[{"y", "z"}, {}\] %arg0 out_sharding=<@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

在上面的示例中,张量 %0 : tensor<16x2xf32> 被分片为 \[{"x", "y", "z"}, {}\]。然后,有一个 reshard 操作将其重新分片为 \[{"x"}, {}\]。在第一个轴上,由于在重新分片后移除了后缀 {"y", "z"},因此我们推断已全收集 {"y", "z"}。第二个维度未更改。

选项

-keep-redundant-reshards : Whether it keeps redundant reshards or removes.

-sdy-resolve-permutation-factors

解决了映射到 kPermutation 因素的维度上的分片问题。

具有 kPermutation 个因子的分片维度可能需要跨设备通信(例如,窗口的 Halo 交换或反向的集体排列)。

如果 enableHaloExchange 为 true,则该 pass 会使用可用的优化通信逻辑来解析置换因子。否则,该 pass 只会插入 sdy.reshard 操作来复制这些维度。enableHaloExchange 的默认值为 true。

选项

-enable-halo-exchange : Implement halo exchange logic for windowed operations.
-replica-count        : Number of replicas (data parallelism).
-partition-count      : Number of partitions (model parallelism).

-sdy-resolve-single-device-sharding

将受目标设备 ID 保护的单设备操作降级为 stablehlo。

选项

-replica-count   : Number of replicas per partition.
-partition-count : Number of partitions per replica.

-sdy-sharding-constraint-to-reshard

将 ShardingConstraintOp 转换为 ReshardOp。

-sdy-sink-data-flow-edges

将所有 DataFlowEdgeOp 汇集到其输入中。

将每个 DataFlowEdgeOp 的分片移至其输入(边缘的根目标),并将该操作替换为其输入。

选项

-sink-debug-sharding-origins          : Whether to sink the debug sharding origins info. See `debug-sharding-origins` option in propagation for more info.
-sink-debug-propagation-edge-sharding : Whether to sink the debug propagation edge sharding info. See `debug-propagation-edge-sharding` option in propagation for more info.

-sdy-sink-func-data-flow-edges

将所有 FuncDataFlowEdgeOp 汇集到其输入中。

将每个 FuncDataFlowEdgeOp 的分片移至其输入,并将相应操作替换为其输入。

-sdy-unflatten-call-graph

取消扁平化调用图。

取消扁平化图表。它会根据附加到函数的“original_func_name”属性描述的相同输入/输出分片相同来源,对函数进行去重。

选项

-dedup-functions-fully : If true, regardless of the input and output shardings of functions, it keeps one callee function for each caller function. The default is false, meaning it will deduplicate only if the input and output shardings are the same.

-sdy-update-non-divisible-input-output-shardings

使 FuncOp 输入/输出均匀分片,从而无需因分片不可分而进行填充。

Shardy 的用户希望函数输入/输出能够均匀地划分/分片,以避免需要填充张量。传播可能会导致输入/输出具有不可整除的分片,因此此传递会将它们更新为原始分片中可均匀分片的最大维度分片前缀。

-sdy-verify-unreduced-axes

验证未缩减轴使用情况的一致性。

验证对于每个运算,如果其操作数具有未缩减的轴,则该运算要么显式缩减这些轴(例如,通过 sdy.reshard),要么将它们传递给其结果(要么是 func.call 等边界运算)。