OpenXLA Project

`-chlo-legalize-to-stablehlo`

CHLO 演算フローから StableHLO および Shape 演算への合法化

`-shape-legalize-to-stablehlo`

形状関連のオペレーションを StableHLO に合法化します。

形状関連のオペレーションを StableHLO オペレーションに合法化する試験運用パス。

オプションのパスを介して形状とデータ計算を統合することで、StableHLO エコシステムは、StableHLO オペレーションを使用して動的性をモデル化するコンパイルパイプラインを活用できるようになります。

`-stablehlo-canonicalize-dynamism`

動的 StableHLO オペレーションを静的オペレーションに正規化します。

これらのオペレーションのすべての動的要素が実際には定数である場合、DynamicReshapeOp などの動的 StableHLO オペレーションを、DynamicReshapeOp から ReshapeOp、DynamicBroadcastInDim から BroadcastInDim などの対応する静的オペレーションに置き換えます。

  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>

  ==>

  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>

`-stablehlo-check-shape-assertions`

_stablehlo.custom_call @shapeアサーションオペレーションを確認します。

shape_assertion カスタム呼び出しを検証します。

シェイプアサーションは、StableHLO の動的ディメンションの制約を検証します。たとえば、フレームワークで DimA < 2 の制約を適用する必要がある場合、次の IR が出力される可能性があります。

%dimA = <get_dimension_size or input arg> : tensor<i32>
%c2 = stablehlo.constant dense<2> : tensor<i32>
%is_lt = stablehlo.compare LT %dimA, %c2 : tensor<i1>
stablehlo.custom_call @shape_assertion(%is_lt) { error_message = "DimA must be less than 2" }

パスが成功すると、図形が正しければ stablehlo.custom_call が削除されます。

オプション

-enable-shape-assertions : Whether shape assertions may generate errors.

`-stablehlo-compatibility-expander`

StableHLO オペレーションの互換性拡張機能。

StableHLO オペレーションが更新されるか、最新バージョンで新しいオペレーションが導入されます。このオプトインパスは、新しい StableHLO オペレーションを古いバージョンでサポートされている同等のオペレーションに分解することで、古い StableHLO バージョンとの下位互換性を拡張します。

なぜこのパスはオプトイン式なのですか？

StableHLO op の拡張機能は、OpenXLA エコシステムの特定の一般的なパターンの処理を大幅に簡素化するために使用されることがあります。これには、フレームワークとコンパイラのサポートが充実している TanOp や、スライスを使用して表現できるものの、シャーディングを非常に困難にする gather/scatter バッチ処理のディメンションなどが含まれます。このカテゴリの新機能では、後続の最適化で使用される重要な情報が破棄される可能性があるため、自動ダウングレードは提供されません。このパスは、ターゲットバージョンに基づいてこれらのオペレーションを拡張するために使用できます。これにより、互換性を最大限に高めることができますが、コンパイルの最適性が低下する可能性があります。

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %1 = stablehlo.tan %arg0 : tensor<4xf64>
  func.return %1 : tensor<4xf64>
}

==>

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %0 = stablehlo.sine %arg0 : tensor<4xf64>
  %1 = stablehlo.cosine %arg0 : tensor<4xf64>
  %2 = stablehlo.divide %0, %1 : tensor<4xf64>
  return %2 : tensor<4xf64>
}

オプション

-target : The target version. Must be a version of the form #.#.#.

`-stablehlo-complex-math-expander`

StableHLO の複素数演算用のエクスパンダ。

StableHLO の複素数演算は、StableHLO の実数演算を使用した分解です。

このステートメントは、複素数や複素数演算をネイティブにサポートするハードウェアが存在しないという前提に基づいています。つまり、コンパイラが実装する可能性のある複雑な数学演算のフォールバックメカニズムは冗長です。このパスを有効にすると、すべての StableHLO 複素数演算が展開されます。

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  %1 = stablehlo.sqrt %arg0 : tensor<4xcomplex<f64>>
  func.return %1 : tensor<4xcomplex<f64>>
}

==>

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  TBD
  return %2 : tensor<4xcomplex<f64>>
}

`-stablehlo-convert-to-signless`

符号なし整数になるように IR を変換するパス。

`-stablehlo-legalize-composite-to-call`

複合演算を分解の呼び出しに置き換えます。

複合オペレーションをその分解の呼び出しに置き換えます。例:

stablehlo.composite "my_namespace.my_op" %arg0, %arg1 {
  decomposition = @bar,
  version = 1,
  composite_attributes = {
    "my_attribute": "my_value"
  }
}

次のようになります。

func.call @bar(%arg0, %arg1)

「except」フラグを使用して、この変換から除外する複合のサブセットを指定できます。例:

stablehlo-opt --stablehlo-legalize-composite-to-call=except='foo.baz,foo.qux'

オプション

-except : Names of composites that should not be replaced with calls.

`-stablehlo-legalize-deprecated-ops`

非推奨のオペレーションを十分にサポートされているオペレーションに合法化します。

StableHLO v1.0 Opset Deprecations RFC（#2283）では、いくつかの冗長なオペレーションを削除することが提案されています。このパスは、これらの op の削除がさまざまなコンパイルパイプラインに与える影響を、長期サポートされる対応する op に合法化することで評価します。

オプション

-fail-on-unused : Fail on (mostly) unused ops that are deprecated without any fallback.

`-stablehlo-legalize-qdq-to-quantized-op`

融合（逆量子化、浮動小数点演算、量子化）パターンを StableHLO 量子化演算に

融合（逆量子化、浮動小数点演算、量子化）パターンを StableHLO 量子化オペレーションに統合します。注: このパスでは、既存の op は削除されません。たとえば、次のプログラム

func.func @add(%arg0: tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
  func.return %2 : tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
}

次のようになります。

func.func @add(%arg0: tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.abs %arg0 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  %3 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  return %2 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
}

`-stablehlo-legalize-quant-to-math`

StableHLO の量子化された演算から StableHLO のプリミティブな数学演算に変換します。

UniformQuantized 型を使用する StableHLO プログラムを、意味的に同等の整数演算に変換します。

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

次のようになります。

func.func @add(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
  %0 = stablehlo.convert %arg0 : (tensor<i8>) -> tensor<f32>
  %cst = stablehlo.constant dense<0.333333343> : tensor<f32>
  %1 = chlo.broadcast_multiply %0, %cst : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<f32>
  %2 = chlo.broadcast_add %1, %cst_0 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %3 = stablehlo.round_nearest_even %2 : tensor<f32>
  %4 = stablehlo.convert %3 : (tensor<f32>) -> tensor<i32>
  %5 = stablehlo.convert %arg1 : (tensor<i8>) -> tensor<f32>
  %cst_1 = stablehlo.constant dense<0.666666686> : tensor<f32>
  %6 = chlo.broadcast_multiply %5, %cst_1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_2 = stablehlo.constant dense<1.33333337> : tensor<f32>
  %7 = chlo.broadcast_add %6, %cst_2 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %8 = stablehlo.round_nearest_even %7 : tensor<f32>
  %9 = stablehlo.convert %8 : (tensor<f32>) -> tensor<i32>
  %c = stablehlo.constant dense<2> : tensor<i32>
  %10 = chlo.broadcast_add %4, %9 : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %11 = chlo.broadcast_subtract %10, %c : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %c_3 = stablehlo.constant dense<-128> : tensor<i32>
  %c_4 = stablehlo.constant dense<127> : tensor<i32>
  %12 = stablehlo.clamp %c_3, %11, %c_4 : tensor<i32>
  %13 = stablehlo.convert %12 : (tensor<i32>) -> tensor<i8>
  return %13 : tensor<i8>
}

`-stablehlo-legalize-quantized-op-to-qdq`

量子化された StableHLO オペレーションを（逆量子化、浮動小数点演算、量子化）パターンに分解します。

均一な量子化/逆量子化オペレーションを使用して、StableHLO 量子化プログラムを分解します。たとえば、次のプログラム

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

次のようになります。

func.func @add(%arg0: tensor<!quant.uniform<i8:f32, 1.000000e+00>>, %arg1: tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<!quant.uniform<i8:f32, 1.000000e+00>>) -> tensor<f32>
  %1 = stablehlo.uniform_dequantize %arg1 : (tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<f32>
  %2 = stablehlo.add %0, %1 : tensor<f32>
  %3 = stablehlo.uniform_quantize %2 : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
  return %3 : tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
}

`-stablehlo-legalize-to-vhlo`

StableHLO を VHLO に合法化。

StableHLO を VHLO の最新バージョンの ops に合法化します。これらのオペレーションは、VhloToVersionPass を使用して上位互換性を確保するために、古いバージョンの VHLO にダウングレードできます。

stablehlo.exponential %[[ARG0]] <{result_accuracy = DEFAULT}> : tensor<f32>
# ====>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.DEFAULT_v1}> : !vhlo.tensor_v1<!vhlo.f32_v1>

VHLO を使用して前方互換性と後方互換性を維持する方法について詳しくは、vhlo.md > VHLO 言語をご覧ください。

オプション

-allow-other-dialects : Allow serialization to use other (potentially unstable) dialects, inserts unrealized casts between dialects.

`-stablehlo-refine-arguments`

メイン関数の引数シェイプを絞り込みます。

入力型シグネチャを使用して、メイン関数の引数を変更します。custom_call @stablehlo.shape_refinement_operand_wrapper で引数をラップして、形状の絞り込みが実行される前に IR が有効になるようにします。

func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  ...
}

==>

func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
    : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
}

refinedTypesOption を使用して、絞り込み型のリストを指定できます。これは、MLIR で --types='tensor<...>,tensor<...>' を使用して指定するか、パス作成メソッドに渡すことができます。絞り込み型リストでは、絞り込まれる main メソッドのすべての引数の型を指定する必要があります。

オプション

-types : The new types to be used for the main function's arguments, specified as an MLIR TypeRange 'tensor<1x2xf32>, ...'

`-stablehlo-refine-shapes`

StableHLO プログラム全体でシェイプを調整します。

ops 内の形状を調整する StableHLO プログラムについて説明します。

このパスの主なユースケースは、動的にシェイプされたプログラムを静的なシェイプに特化させることです。動的にシェイプされた StableHLO プログラムが正しい構造を持っている場合、その引数型を動的シェイプから静的シェイプに更新してこのパスを実行すると、静的シェイプがプログラム全体に伝播されます。

このパスは、結果の使用をオペランドで直接置き換えることで custom_call @shape_refinement_operand_wrapper を削除し、静的形状をプログラム全体に伝播します。

  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
      : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  %1 = stablehlo.add %0, %0 : tensor<?xf32>

  ==>

  %1 = stablehlo.add %arg0, %arg0 : tensor<16xf32>

形状の調整に有効なモジュールには、次のプロパティが必要です。

すべての動的形状は入力形状にのみ依存します（入力配列の内容には形状の依存関係はありません）。入力形状（stablehlo.get_dimension_size で指定されたものなど）またはグローバル定数（シンボリック整数の解決値、つまり tensor : A = 5 など）にのみ推移的に依存するオペレーションを dimension オペレーションと呼びます。すべてのディメンション値は、プロシージャ間の定数フォールディングによって定数に解決できます。
中間関数は、引数リストの先頭に複数のトークン引数（!stablehlo.token 型）を取り、その後に、シンボリック整数の解決値（tensor : A = 5）などの定数整数スカラーであるグローバル定数引数を取ることができます。
一部の中間関数は、グローバル定数（symint 値の floordiv など）の計算を返すことがあります。これらの関数は、絞り込み後に定数値を返すことのみで示されます。これらの関数はインライン化されます。
単一の関数へのすべての呼び出しは同じ引数形状に解決され、再帰関数呼び出しや相互再帰関数呼び出しは行われません。

`-stablehlo-wrap-in-composite`

非複合 StableHLO オペレーションを複合オペレーションでラップします。

StableHLO オペレーションを stablehlo.composite オペレーションでラップします。

たとえば、次のような簡単な StableHLO プログラムを考えてみましょう。

func.func @main(%arg0 : tensor<2xf32>, %arg1 : tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

このパスを適用して stablehlo.add オペレーションをラップすると、次のプログラムが生成されます。

func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.composite "stablehlo.add" %arg0, %arg1 {decomposition = @stablehlo.add.impl} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  return %0 : tensor<2xf32>
}
func.func private @stablehlo.add.impl(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

注:

生成された stablehlo.composite オペレーションの name 属性は、常にラップされた元のオペレーションの名前と同じになります（たとえば、stablehlo.add オペレーションをラップすると、複合オペレーションの名前も "stablehlo.add" になります）。
元のオペレーション（stablehlo.composite オペレーションの decomposition 属性で参照される）をカプセル化するプライベート関数は、<op_name>.impl[.N] パターンを使用して命名されます。ここで、<op_name> は元のオペレーションの名前、N はモジュール内の名前の競合を防ぐために生成される一意の整数識別子です。

このパスは次の 2 つの方法で使用できます。

モード 1: コマンドラインの使用

このモードは、生成された stablehlo.composite オペレーションの属性に対する制御が最小限であるため、デバッグまたはテストを目的としています。op-names（オペレーション名のカンマ区切りリスト）オプションを使用して指定されたオペレーションのすべてのインスタンスをラップします。新しく作成された stablehlo.composite オペレーションの属性は、元のオペレーションの属性と同じになります。

使用例:

stablehlo-opt input.mlir --stablehlo-wrap-in-composite=op-names='stablehlo.add,stablehlo.mul' -o output.mlir

モード 2: カスタマイズされた属性処理によるプログラムによるモジュール全体のラッピング

このモードでは、プログラムによるラッピングがモジュール全体に拡張され、どのオペレーションがラッピングされるか、その属性をきめ細かく制御できます。これは、CompositeAttributeProviderMap を引数として受け取る createStablehloWrapInCompositePass API を使用することで実現できます。

CompositeAttributeProviderMap は、どのオペレーションをラップの対象とするか、またその属性をどのように処理するかを指定するマップです。セマンティクスは次のとおりです。

キー（mlir::TypeID）: MLIR オペレーションの TypeID。オペレーションの TypeID がマップ内のキーと一致する場合、そのオペレーションはラッピングの候補になります。
値（ラムダ関数）: 型 std::function<std::optional<NamedAttrList>(Operation*)> のラムダ関数。この関数は、候補オペレーションごとに適用されます。
- 入力: mlir::Operation*。これは、TypeID キーに対応するオペレーションタイプのインスタンスです。
- 戻り値: std::optional<NamedAttrList>。
  - ラムダが NamedAttrList（std::optional でラップ）を返すと、オペレーションは stablehlo::composite オペレーションでラップされ、返された属性は複合の属性の設定に使用されます。
  - ラムダが std::nullopt を返す場合、オペレーションはラップされません。これにより、カスタム条件に基づいて選択的にラップできます。

例（C++）:


stablehlo::CompositeAttributeProviderMap compositeAttributeProviderMap;

compositeAttributeProviderMap[mlir::TypeID::get<mlir::stablehlo::AddOp>()] =
  [](mlir::Operation* op) -> std::optional<mlir::NamedAttrList> {
  // Custom logic to determine if and how to wrap the operation.
  // Example: Only wrap if it's on a specific type.
  if (mlir::isa<mlir::Float32Type>(op->getOperand(0).getType())) {
    return mlir::NamedAttrList(op->getAttrs());
  }
  return std::nullopt; // Do not wrap.
};

pm.addPass(createStablehloWrapInCompositePass(compositeAttributeProviderMap, compositeVersion));
if (mlir::failed(pm.run(module))) {
  return;
}

オプション

-op-names : The names of the ops to wrap.
-version  : The version number of the composite op.

`-vhlo-legalize-to-stablehlo`

VHLO を StableHLO に合法化。

`-vhlo-to-version`

互換性のために VHLO のバージョンを変換します。

IR のアップグレードとダウングレードのために VHLO のバージョン間で変換し、上位互換性と下位互換性を維持します。

"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>
# ==( -target=1.0.0 )==>
"vhlo.exponential_v1"(%[[ARG0]])
# ==( -target=1.9.0 )==>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>

VHLO を使用して前方互換性と後方互換性を維持する方法について詳しくは、vhlo.md > VHLO 言語をご覧ください。

オプション

-target : The target version. Must be a version of the form #.#.# .