-sdy-close-shardings

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

-sdy-constant-merger

合并具有匹配分片的完全相同的常量。

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

导入流水线会拆分和复制常量,以便在常量子计算的不同用法之间不会传播分片。如果常量在传播后具有相同的分片,此传递会将它们合并以节省编译时间。

-sdy-drop-sharding-rules

从所有已注册的操作中移除 OpShardingRuleAttr

-sdy-insert-explicit-reshards

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

兼容的分片本质上意味着操作可以接受分片运算数并生成分片结果,而无需任何重新分片通信(请注意,操作可能仍需要通信,例如全局求和或 Halo 交换)。

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

请注意,如果某个轴(或子轴)用于跨多个张量分片不对应的维度(例如 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 上插入显式分片,以便点运算具有兼容的分片。

-sdy-remove-sharding-groups

在传播后移除 ShardingGroupOps。

-sdy-reshard-to-collectives

将 ReshardOp 转换为各种 Shardy 集合操作。

匹配分片操作并将其重写为各种 Shardy 集合操作。此传递后,模块中将不再有任何分片操作。此传递假定已插入显式分片 (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"}。第二个维度保持不变。

-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-temp-explicit-reshards-for-optimizations

为特定优化插入显式分片。

在我们能够默认启用 sdy-insert-explicit-reshards 卡券之前,此卡券只是一种临时解决方案。

这样,我们就可以在特定操作中插入显式分片,以进行优化。

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

使 FuncOp 输入/输出均匀分片,从而消除了因不可分割的分片而导致的任何填充需求。

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