-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 の非縮約ディメンションなど）をシャーディングする場合や、軸が 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 に明示的なリシャーディングを挿入し、ドット オペレーションが互換性のあるシャーディングを持つようにします。

オプション

-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 を削除します。このパスは、Collective 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"} が削除されるため、 {"y", "z"} が all-gather されたと推測されます。2 番目のディメンションは変更されません。

オプション

-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 のユーザーは、テンソルのパディングを回避するために、関数の入力/出力が均等に分割/シャーディング可能であることを想定しています。伝播により、入力/出力に分割できないシャーディングが含まれる可能性があるため、このパスは、均等にシャーディングされた元のシャーディングの最大ディメンション シャーディング プレフィックスに更新します。