-mpmd-copy-topology-from-main
トポロジをメイン関数から mpmd.call で参照される関数にコピーします。
トポロジ属性をメイン関数から mpmd.call で参照される関数にコピーします。また、mpmd.call の呼び出し先を非公開に設定して、エントリ ポイント関数と間違われないようにします。
-mpmd-enforce-input-output-equisharding
MPMD 関数の均等シャーディング制約を適用します。
必要に応じて TransferOps を導入することで、MPMD 関数の入出力均等シャーディング制約を適用します。
オプション
-constraints : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.
-mpmd-generate-sdy-meshes-from-topology
MPMD トポロジに基づいてシャーディングされたメッシュを生成します。
このパスでは、既存のシャーディングされたメッシュ オペレーションを削除し、MPMD トポロジに基づくオペレーションに置き換えます。また、新しいメッシュ オペレーションを参照するようにテンソル シャーディングを更新します。
-mpmd-infer-mesh-assign-mesh-func-leaves
__use_set 分析と srcset 分析を使用して、未使用の計算、関数の出力、関数の入力にメッシュを割り当てます。
このパスでは、use_set 情報と src_set 情報を使用して、AssignOps を作成するか、型を変更することで、関数本体のリーフ（未使用の計算の結果、未使用の関数引数、関数の出力）にメッシュを割り当てます。
また、分析のために、特定の中間値をリーフとして扱います。つまり、mpmd.reduce と mpmd.broadcast のオペランドはリーフとして扱われ、それらに割り当て解除ペアが作成されます。
この割り当てにより、リーフ以外のすべてのオペレーションの use_set がクリアされます。これは、以前にアノテーションが付けられた使用は、reduce オペレーションの推論によって一部の値の use_set が変更されるため、古くなるためです。最初の use_set 伝播は reduce オペレーションを認識しませんが、reduce オペレーションを推論したため、伝播は異なります。
リーフ オペレーションの use セットと src セットが正しく入力されていない場合、このパスは失敗し、エラーが生成されます。
前提条件: すべてのオペレーションに空でない src セットがあるか、転送を推論します。
オプション
-infer-transfers : Whether to create transfers when needed, instead of erroring.
-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.
-mpmd-infer-mesh-assign-using-input-output-constraints
入出力割り当て制約に従って、入力と出力にメッシュを割り当てます。
このパスでは、入出力均等割り当て制約を使用して、入力と出力の両方を同じメッシュに割り当てます。
入力が後で他のメッシュに転送されるかどうかに関係なく、入力が同じメッシュ上にあることが保証されます。ただし、このパスの後に populate-src-set を実行しないでください。
__要件:
- 均等割り当て制約の一部である可能性があるエントリ ポイント関数の入力
iの場合:iに MeshTensorType があるか、明確に定義された use セットがあります。 - 均等割り当て制約の一部である可能性があるエントリ ポイント関数の出力
oの場合:oの型が MeshTensorType であるか、明確に定義された src セットと use セットがあります。
値の明確に定義された use セットには、mpmd.assign オペレーションを介して値が（推移的に）割り当てられるすべてのメッシュが含まれ、他のメッシュは含まれません。明確に定義された src セットには、テンソルが存在できるすべてのメッシュが含まれ、他のメッシュは含まれません。
これはエントリ ポイント関数でのみ実行されますが、既存の関数に対するすべてのパスが完了してから実行する必要があるため、モジュール オペレーション パスにします。たとえば、これを EntryPointFunctionPass にすると、パス マネージャーは、エントリ ポイント以外の関数で検証が完了する前にこのパスを実行する可能性があります。
オプション
-verbose-logging : Whether to enable verbose logging
-constraints : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.
-mpmd-infer-mesh-convert-reduce-ops
アノテーション付きの reduce オペレーションを mpmd.reduce オペレーションに変換し、reduce オペレーションのチェーンをフラット化します。
アノテーション付きの reduce オペレーションを mpmd.reduce オペレーションに変換し、これらの reduce オペレーションのチェーンをフラット化します。
記号の場合:
x = add(w0, w1) {mpmd.reduce = #mpmd.reduce
オプション
-infer-cross-mesh-reductions : Whether to infer cross-mesh reductions. Will be enabled by default once stable.
-mpmd-infer-mesh-finalize
パターン メッシュ推論の後に最終的なクリーンアップを適用します。
オプション
-infer-transfers : Whether to create transfers when needed, instead of erroring.
-mpmd-infer-mesh-populate-src-set
__UnassignOps と関数引数の src_set を初期化し、srcset を伝播します。
このパスでは、src_set を初期化して伝播し、グラフに src_set 情報を入力します。
前提条件: 関数引数に src_set を設定するには、use_set を入力する必要があります。
初期化: UnassignOp の src_set は、割り当てるメッシュに設定されます。関数引数の src_set は、その use_set に設定されます。
伝播: src_set は、オペランドからオペレーション自体に順方向に伝播し、オペランドの交差を取ります。詳細については、PropagateSrcSet をご覧ください。
-mpmd-infer-mesh-populate-use-set
__AssignOps の use_set を初期化し、useset を伝播します。
このパスでは、use_set を初期化して逆方向に伝播し、グラフに use_set 情報を入力します。
初期化: AssignOp の use_set は、割り当てるメッシュに設定されます。
伝播: use_set は、ユーザーからオペレーション自体に逆方向に伝播し、ユーザーの和集合を取ります。use_set は推移的な使用のセットであるため、オペレーションの use_set は定義上、そのユーザーの use_set の和集合です。
-mpmd-infer-mesh-rewrite-using-analysis
__useset に従ってオペレーションを書き換えます。
このパスでは、use_set 分析と src_set 分析を使用して、メッシュレス オペレーションをフラグメントでラップして割り当てます。
また、クリーンアップの一環として use_set 属性と src_set 属性を削除します。これは、この後、分析が不要になるためです。
前提条件: すべてのオペレーションに use_set があること（分析が完了していること）。 前提条件: エントリ ポイント以外の関数のすべての引数は、少なくとも 1 つのターミネータ以外のオペレーションで使用されていること。
TODO: jupvfranco - 分析にあまり依存しなくなったため、このパスの名前の変更を検討してください。
オプション
-max-clones : How many copies of a meshless operation we allow. Setting it to 1 means we never clone the op.
-mpmd-infer-mesh-validate-no-additional-transfers-needed
メッシュ割り当てに追加の転送が必要ないことを検証します。
このパスでは、追加の転送を導入せずに、すべてのメッシュレス オペレーションにメッシュ割り当てが可能であることを検証します。
関数オペレーションではないメッシュレス オペレーションの場合、次の場合はエラーになります。
- use_set が特定のオペレーションの src_set に含まれていない場合（転送が必要な場合）。
関数オペレーションの場合、関数はメッシュレス オペレーションまたはブロック引数を返すため、関数引数に対して上記の条件を確認するだけで十分です。
オプション
-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.
-mpmd-infer-mesh-validate-src-set-not-empty
__すべてのメッシュレス オペレーションに空でない srcset があることを検証します。
このパスでは、すべてのメッシュレス オペレーションを検証し、オペレーションをどこかに割り当てられることを確認します。つまり、関数オペレーションではないメッシュレス オペレーションの場合、オペレーションで src_set が空の場合、またはクロスメッシュ削減と推論されたが変換されていない場合はエラーになります。これは、関数リーフ割り当ての前提条件です。
関数オペレーションの場合、関数はメッシュレス オペレーションまたはブロック引数を返すため、関数引数に対して上記の条件を確認するだけで十分です。
呼び出し先でエラーが発生した場合は呼び出し元を出力する必要があるため、これはモジュールレベルのパスである必要があります。
前提条件: このパスを実行する前に、クロスメッシュ削減を reduce オペレーションに変換する必要があります。
オプション
-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.
-mpmd-inline-nested-user-exposed-ops
__名前付き計算にネストされているユーザー公開の mpmd オペレーションをインライン化します。__
名前付き計算にネストされている名前付き計算、名前付きテンソル、ブロードキャスト、reduce オペレーションをインライン化し、そのメッシュ割り当て（定義されている場合）が親のメッシュ割り当てと一致することを確認します。
オプション
-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.
-mpmd-insert-nameless-clone-of-negligible-ops
名前付き計算の外部で無視できるオペレーションを複製します。
無視できるオペレーション（単一の結果、ゼロ オペランド オペレーション）は、計算の return オペレーションで使用されるたびに、名前付き計算の外部で複製され、名前付き計算の結果がクローンに置き換えられます。このような結果が異なるメッシュに割り当てられた名前付き計算で使用されると、メッシュ推論の競合が発生する可能性があるため、これが必要です。このパスを適用することで、メッシュ推論でこれらの無視できるオペレーションを複製できます。
このパスでは、名前付き計算はまったく変更されません。
-mpmd-introduce-transfers
ユーザー メッシュ割り当てに基づいてデータ転送を作成します。
ユーザー メッシュ割り当てに基づいて転送オペレーションを導入するパスを作成します。これには以下が該当します。
- UnassignOp の結果が後で呼び出し先で割り当てられる場合は、mpmd 呼び出しに UnassignOp をプッシュします。
- UnassignOp の AssignOp を TransferOp に置き換えます。
- フラグメント間にメッシュレスの追加がある場合は、追加を消費メッシュに割り当てて転送を導入します。
-mpmd-map-input-output-to-mesh
関数入力と出力にメッシュを割り当てます。
ユーザー定義のメッシュ割り当てに基づいて、関数入力/出力をメッシュにマッピングするパスを作成します。
入力引数の場合、このパスは次のようになります。
- メッシュに配置する入力テンソルをメッシュテンソルにキャストします。
- 関数シグネチャを更新します。
- テンソルが使用される前に mpmd.unassign を追加します。
出力引数の場合、このパスでは、テンソルが返される前に mpmd.assign を追加し、関数シグネチャを更新します。
要件: 各入出力インデックスが有効であり、マッピングされた各メッシュがトポロジ内の有効なメッシュであること。
オプション
-input-assignment : Mapping between function input indices and assigned mesh names.E.g., '0@m0,1@m1' defines that input with index 0 will be assigned to mesh m0 and input with index 1 will be assigned to mesh m1.
-output-assignment : Mapping between function output indices and assigned mesh names.E.g., '0@m0,1@m1' defines that output with index 0 will be assigned to mesh m0 and output with index 1 will be assigned to mesh m1.
-mpmd-map-named-ops-to-mpmd-ops
ユーザー定義のオペレーションにメッシュを割り当てます。
必要に応じて mpmd.named_tensor を Assign(Unassign(%v)) に割り当て（assignment にエントリがあるかどうかによって異なります）、名前付き計算とメッシュ名のユーザー定義マッピングを使用して、各名前付き計算をメッシュにマッピングするパスを作成します。つまり、各名前付き計算をフラグメントに置き換え、これらのフラグメントのオペランドに AssignOps を作成し、結果に UnassignOps を作成します。新しく導入されたパターン Assign(Unassign(%v)) は Transfer(%v) に書き換えられます。
このパスの後に名前付き計算/名前付きテンソル オペレーションは存在しません。
要件: すべての名前付き計算と名前付きテンソルが関数の最上位レベルに存在すること。
オプション
-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.
-mpmd-simplify-named-computations
名前付き計算オペレーションの入力と出力を簡略化します。
名前付き計算ごとに個別に簡略化します。具体的には次のとおりです。
- 結果と対応する戻り値を重複排除します。
- オペランドと対応するブロック引数を重複排除します。
- 対応する戻りオペランドがオペレーションのブロック引数である結果を削除します。
- 対応するブロック引数が使用されなくなった（または最初から使用されていなかった）オペランドを削除します。
- 未使用の結果を削除します。
- 名前付き計算内のパターン
arg -> stablehlo.optimization_barrier -> returnをパターンarg -> returnに置き換えて、 さらに簡略化できるようにします。
-mpmd-validate-named-ops-in-mpmd-func
名前付きオペレーションが mpmd 関数にのみネストされていることを検証します。
NamedComputationOp と NamedTensorOp が mpmd 関数（トポロジ属性を持つ関数）にのみネストされていることを検証します。