OpenXLA Project

このページは Cloud Translation API によって翻訳されました。

`-sdy-close-shardings`

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

`-sdy-constant-merger`

一致するシャーディングを持つ同じ定数を統合する。

同じシャーディングを持つ定数に対して軽量 CSE を実行します。

インポートパイプラインは、定数の分割と複製を行い、定数のサブ計算の異なる使用間でシャーディングが伝播されないようにします。伝播後に定数に同じシャーディングがある場合、このパスは定数を統合してコンパイル時間を節約します。

`-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, \[{"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>

上記の例では、lhs と rhs はどちらも、収縮しないディメンションの軸「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"}, {}\] としてシャーディングされています。次に、\[{"x"}, {}\] として再シャーディングする reshard オペレーションがあります。最初の軸では、シャード再構成後に接尾辞 {"y", "z"} が削除されているため、{"y", "z"} がすべて収集されたと推測されます。2 つ目のディメンションは変更されません。

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