-sdy-close-shardings

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

-sdy-drop-sharding-rules

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

-sdy-insert-explicit-reshards

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

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

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

请注意,如果某个轴(或子轴)用于跨多个张量分片不对应的维度(例如 matmul 中的非收缩维度),或者某个轴分片一个张量中的维度,但不分片另一个张量中的对应维度,则表示该操作存在分片冲突。因此,在此传递之后,操作将不会发生冲突。

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

一个说明性示例:

输入:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"y"},{"x"}\]>}
%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 张量在点积运算之前会被显式重新分片,以便仅按其第一个维度和轴“x”进行分片。这样,点运算就变得兼容了。

-sdy-remove-sharding-groups

在传播后移除 ShardingGroupOps。

-sdy-sharding-constraint-to-reshard

将 ShardingConstraintOp 转换为 ReshardOp。

-sdy-sink-data-flow-edges

将所有 DataFlowEdgeOp 传送到其输入。

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

TODO(tomnatan):考虑将分片移至所有可以附加分片的目标。

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

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

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