StableHLO でのダイナミズム

動的性の現在の状態は、動的性 RFC でより正式に説明されています。このページでは、RFC の概要を説明し、動的プログラムとやり取りするための重要な API とツールについて説明します。

ダイナミズムの用語とサポートの概要

まず、このドキュメントに記載されている用語と、StableHLO でのサポートの概要について説明します。

動的な次元

動的ディメンションとは、ディメンションのサイズが不明なディメンションを指します。StableHLO では、動的ディメンションを ?（つまり tensor<16x?xf32>）を使用して表します。

バウンドダイナミズム

境界付き動的とは、値に既知の上限がある動的ディメンションを指します。通常、これは実行中にテンソルをパディングするのに役立ちます。StableHLO では、境界付きの動的性をテンソルエンコードとして #stablehlo.bounds を使用して表します。つまり、1 つの動的ディメンションが 16 で制限され、もう 1 つが制限されていないランク 2 テンソルは tensor<?x?xf32, #stablehlo.bounds<16, ?>> として表すことができます。

StableHLO は有界動的性を表現できますが、フレームワークのサポートは TensorFlow に由来するものが限られており、PyTorch/XLA で一部サポートされています。

無制限の動的性

無制限の動的ディメンションとは、その名のとおり、サイズに既知の境界がない動的ディメンションを指します。この種の動的性は StableHLO で非常に一般的であり、JAX、PyTorch/XLA、TF のサポートにより、動的なバッチサイズやシーケンス長を持つモデルのエクスポートによく使用されます。

StableHLO では、この形式の動的性に対して境界エンコードを省略します（tensor<?x?xf32>）。

シェイプのポリモーフィズム

形状ポリモーフィズムは、JAX から継承した用語です。

シェイプポリモーフィズムには、次の 2 つの重要な意味があります。

プログラムのすべての動的性は、入力引数にまで遡ります。
すべての動的性は、テンソルの形状のみに関係します。つまり、データに依存しません。

これらの 2 つのルールにより、プログラムの静的形状がわかれば、動的プログラムを取得して、コンパイル用の静的プログラムに完全に絞り込むことができます（「動的プログラムを絞り込むためのコンパイラパス」を参照）。

一般に、シェイプポリモーフィズムは無制限の動的性を使用します。既知の引数のシェイプが完全に静的なプログラムにつながる場合は、値をバインドする方法を推測する必要はありません。

データ依存の動的性

データ依存の動的性とは、テンソル内のデータに関連する動的ディメンションサイズを指します。標準的な例は、テンソル値で 0 であるすべての要素のインデックスを返す nonzeros 関数です。形状はデータを評価しないとわかりませんが、多くの場合、有界動的性を使用してコンパイルし、潜在的な出力テンソルサイズに余分なメモリを費やすことができます。

多くのデータ依存の動的オペレーションは、有界動的性を使用してモデル化できます。この場合、テンソルサイズの上限が指定され、ハードウェアは通常、テンソルパディングを介してこれを実装します。現在、PyTorch/XLA と TensorFlow ではデータ依存の動的性がサポートされていますが、JAX ではデータ依存の動的性につながるオペレーションはトレースされません。

動的ディメンションを含むプログラムをエクスポートする

動的バッチサイズまたはシーケンス長でプログラムをエクスポートする方法については、StableHLO チュートリアルをご覧ください。

動的プログラムを改良するためのコンパイラパス

ダイナミズムパスパイプラインを削除

シェイプの調整に役立つパスがいくつかあります。これらはすべて、パスパイプライン createStablehloRemoveDynamismPipeline にバンドルされています。

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

ダイナミズムを調整するための個々のパス

形状の調整に役立つパスは次のとおりです。

stablehlo-refine-arguments: 入力引数を具体的なテンソル型に置き換えます。
stablehlo-refine-shapes を使用して、新しい入力引数の形状情報をプログラム全体に伝播します。
stablehlo-canonicalize-dynamism: 動的オペレーションを静的バリアントに置き換えます。
stablehlo-check-shape-assertions: シェイプアサーションのカスタム呼び出しをチェックして削除します。

最新情報と例については、リンク先のドキュメントをご覧ください。

例: ダイナミズムはどのように役立ちますか？また、どのように使用できますか？

ダイナミズムには多くの用途がありますが、ここでは主に、シェイプポリモーフィズムの一般的なユースケースである、柔軟なエクスポートモデル表現の作成に焦点を当てます。これは通常、動的なバッチサイズやシーケンス長を表すために使用されます。

静的 add_one モデル

次の簡単な add_one モデルを使用して、これを示します。

def add_one(x):
  return x + 1

tensor<4xf32> を使用してトレースすると、次の StableHLO プログラムが得られます。

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

このモデルは、tensor<4xf32> 形状の入力引数に対してのみ機能します。バッチサイズやシーケンス長を変更した場合は、ソースコードを再トレースして StableHLO に再変換する必要があります。また、ソースコードにアクセスできる保証もありません。

動的 add_one モデル

ここで、シェイプのポリモーフィックな動的性が重要になります。代わりに、JAX と PyTorch/XLA は、次のように動的に有効な IR を持つ add_one モデルを出力し、動的入力形状に合わせて定数をブロードキャストできます。

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

このモデル表現ははるかに柔軟で、バッチサイズやシーケンス長などの値の遅延指定が可能です。このモデルは、動的形状をサポートするプラットフォーム（AI Edge など）にデプロイできます。また、このドキュメントで説明する動的パスを使用して調整することもできます。

動的モデルの改良

たとえば、次のパスの順序付けでこのプログラムを完全に絞り込むことができます。

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

プログラムは次のように段階的に変換されます。

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
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) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}

StableHLO でのダイナミズム コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。