-sdy-close-shardings

テンソル シャーディングを閉じて、レプリケートされた軸を削除します。

-sdy-drop-sharding-rules

登録されているすべてのオペレーションから OpShardingRuleAttr を削除します。

-sdy-insert-explicit-reshards

明示的なリシャードを挿入して、すべてのオペレーションに互換性のあるシャーディングを適用します。

互換性のあるシャーディングとは、基本的に、オペレーションがシャーディングされたオペランドを受け入れて、シャーディングされた結果を生成できることを意味します（オペレーションで、all-reduce や halo-swap などの通信が必要な場合もあります）。

伝播後も、一部のオペレーションには互換性のないシャーディングが残っている場合があります。

軸（またはサブ軸）を使用して、対応しないディメンション（matmul の非収縮ディメンションなど）を複数のテンソルにシャーディングする場合、または軸が 1 つのテンソルのディメンションをシャーディングするが、他のテンソルの対応するディメンションをシャーディングしない場合、オペレーションにシャーディング競合があるといいます。したがって、このパスの後、オペレーションは競合がなくなります。

このパスは、リシャーディング オペレーションを明示的に挿入します。これにより、オペレーションごとに、対応するディメンションがすべてのオペランドと結果で同じ方法でシャーディングされ、すべての軸（またはサブ軸）は 1 つのディメンション タイプをシャーディングするためにのみ使用できます。

わかりやすい例:

入力:

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>

上記の例では、lhs テンソルと rhs テンソルの両方が、収縮しないディメンションの軸「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 のユーザーは、テンソルのパディングを必要とせずに、関数の入力と出力を均等に分割またはシャーディングできることを期待しています。伝播により、入力と出力のシャーディングが分割できない場合があるため、このパスでは、均等にシャーディングされた元のシャーディングで最大のディメンション シャーディング接頭辞に更新します。