-sdy-close-shardings

关闭张量分片并删除复制的轴。

-sdy-constant-or-scalar-merger

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

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

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

-sdy-convert-global-to-local

将 SDY 程序从全局形状转换为本地形状。

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

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

-sdy-drop-sharding-rules

从所有已注册的运算中删除 OpShardingRuleAttr

-sdy-insert-explicit-reshards

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

兼容的分片本质上是指运算可以接受分片的操作数并生成分片的结果,而无需任何重新分片通信(请注意,运算可能仍然需要通信,例如 all-reduce 或 halo-swaps)。

传播后,某些运算可能仍然具有不兼容的分片。

请注意,当一个轴(或子轴)用于跨多个张量对不对应的维度(例如 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”进行分片,这是不兼容的。该传递会在点运算之前在 rhs 上插入显式重新分片,以便点运算具有兼容的分片。

选项

-enable-full-version                  : Enable full version.
-avoid-reshards-on-named-computations : Avoid explicit reshards/collectives on named computations.

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

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

移除模式 all-gather + dot 中的 all-gather。移除模式 dot + reduce-scatter 中的 reduce-scatter。此传递用于与集体 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 集体运算。

匹配重新分片运算并将其重写为各种 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"} ,因此我们推断出我们已 all-gather {"y", "z"}。第二个维度未更改。

选项

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

-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-update-non-divisible-input-output-shardings

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

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