型推論

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

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

提案

これらの提案は、既存の実装の再検討と、包括的なカバレッジまでの新しいオペレーションの実現の両方に適用されます。

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

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

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

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

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

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

両方:

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

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

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

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

  1. リージョンを使用しないほとんどのオペレーション（PadOp など）: すべての確認コードをシェイプ関数に入れて、ベリファイアを完全に破棄するようにします。戻り値の型を推測できないため（ReshapeOp や BroadcastInDimOp などで）これが不可能な場合は、必要な検証ロジックを含む検証ツールを作成します。通常、AddOp などの推定可能な演算では、追加の検証を行うために検証ツールが必要となることがあります。これは、指定された戻り値の型の制約を検証しているためで、型/形状推論メソッドではアクセスできないためです。

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

    1. 領域が型推論に不要な場合（ReduceOp など）は、領域関連の検証ロジックをシェイプ関数ではなく検証ツールに配置します。避けられない場合は、一部のコードを重複させる。

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

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

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

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

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

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

ops_stablehlo.mlir には、ops の陽性ケースと、（少なくとも）すべての検証エラーに対する（少なくとも）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 という名前の別のテストファイルを追加することを検討してください。