型推論

StableHLO はもともと MHLO 言語からブートストラップされ、型推論の MHLO 実装を継承しました。実装の進行状況は status.md で追跡されます。

以下のガイドライン案は、StableHLO 演算に高品質の検証ツールとシェイプ関数を確実に実装することを目的としています。

提案

これらの提案は、既存の実装の見直しと、包括的な対応までの新しい運用の実現の両方に適用されます。

（P1）信頼できる情報源として StableHLO の仕様を使用する

spec は、StableHLO 演算のすべての検証ツールとシェイプ関数の信頼できる情報源です。すべての op の既存の検証ツールとシェイプ関数を見直し、仕様に完全に適合させる必要があります。仕様ドキュメントは進化し続けています。演算の仕様が利用できない場合は、代わりに XLA 実装（xla/service/shape_inference.cc や xla/service/hlo_verifier.cc など）を信頼できる情報源として使用する必要があります。XLA 実装は無限ダイナミズムに対応しないため、制限なしダイナミズムの場合は、ダイナミズム RFC が利用可能になるまで常識が適用されます。

（P2）ODS を最大限に活用する

ODS ファイル（StablehloOps.td など）は、各オペランド/属性/結果の特性と型でオペレーションを定義し、検証を行います。したがって、ODS ですでに保証されているプロパティ用の検証ツールやシェイプ関数では、確認コードは必要ありません。確認コードが ODS で重複している場合はトリガーされないため、削除します。

ODS の制約に対するテストを追加する必要があるか？テストのガイドラインを確立するをご覧ください。

（P3）検証ツールとシェイプ関数で確認コードを保持する

両方:

  • verifiers: Op::verify() によって実装されます。
  • シェイプ関数: Op::inferReturnTypes() や Op::inferReturnTypeComponents などの InferTypeOpInterface によって実装されます。

オペランド、属性、結果を確認するための確認コードが含まれることがある。最初の分割では、検証ツールにオペランド/属性をチェックさせ、次にシェイプ関数で推測される結果の型のみを計算し、実際の結果型に対する互換性をチェックします。ただし、実際には、この分割にはいくつかの問題があります。

  • シェイプ関数は、最初に検証ツールを呼び出さなくても、自動生成された build() 関数によって呼び出すことができます。関連する入力もシェイプ関数で検証する必要があります
  • 重複コード: たとえば、検証ツールでは、オペランドに対してなんらかの処理を行ってから、中間結果を検証します。そして、シェイプ関数では、これらの中間結果が最終結果を推測するために役立ちます。この中間結果は 2 回計算する必要があります。
  • メンテナンスの負担: 演算の検証は 2 つの異なる方法に含まれます。

解決策は次のとおりです。

  1. リージョンを使用しないほとんどの演算（PadOp など）: すべての確認コードをシェイプ関数に含め、検証ツールを完全に破棄します。

  2. リージョンを含む演算の場合（ReduceOp/IfOp など。完全なリストはこちら）: 自動生成されたビルダーはリージョンをパラメータとして受け取りません。そのため、これらのビルダーに型推論が含まれている場合、シェイプ関数は空のリージョンで呼び出されます（こちらの例を参照）。

    1. 領域が型推論に不要な場合（ReduceOp など）は、領域関連の検証ロジックを形状関数ではなく検証ツールに配置します。避けられない場合は、コードを複製します。

    2. 型推論（IfOp/CaseOp/MapOp）に領域が必要な場合、ODS によって Op の定義に領域が存在することがすでに保証されている場合でも、シェイプ関数は領域が明示的に空でないことを検証する必要があります。

（P4）テストのガイドラインを確立する

ODS の対象となる検証用のテストを追加または維持する必要はありますか？

いいえ。テストでは、検証ツールとシェイプ関数に焦点を当てる必要がありますが、ODS の変更では、この演算の見直しが必要です。

ただし、欠落している部分には注意してください。たとえば、演算に形質のみをチェックし、要素タイプはチェックしない SameOperandsAndResultShape トレイトが含まれている場合、オペランド/結果の要素タイプの検証にはテストが必要です。

検証機能と型推論のテストはどこに配置されますか？

ops_stablehlo.mlir には、演算の陽性ケースと、（少なくとも）検証エラーごとに（少なくとも）1 件の陰性テストが含まれています。また、推定される戻り値の型が実際の結果の型と互換性がある（同じではない）こともチェックできます。

infer_stablehlo.mlir は、hlo_test_infer.get_return_type_components"(%x):... を使用して演算のシェイプ関数の存在を行ごとに検証し、推測された型が期待どおりに一致することを確認します。一般的に、オペレーションごとに 1 回の陽性テスト。

手順

演算の検証器やシェイプ関数を実装または再使用する場合:

  1. すべての正のケースを ops_stablehlo.mlir に置きます。

  2. インターフェースをテストするために、infer_stablehlo.mlir に陽性テストを 1 つ追加します。

  3. （省略可）演算が複雑で多数のテストを含む可能性がある場合は、verify_<op_name>.mlir または verify_<your_topic>.mlir という名前の別のテストファイルを同じフォルダに追加することを検討してください。