# About Super Vectorization スーパーベクトル化（Super-Vectorization）は、高度なベクトル化戦略を用いたコンパイラ最適化の手法である。この手法では、スーパーベクトルと呼ばれる抽象化を使用して、ベクトルタイプの単一でコンパクトな表現を提供している。スーパーベクトルは、ハードウェアが高レベルのPrimitiveなセットを効率的に実装できるような「良い」ベクトルサイズの倍数として定義されている（ただし、ハードウェアによっては必ずしも定数倍であるとは限らない）。例えば、vector<16xf32>とvector<2x8xf32>は、vector<8xf32>のハードウェアベクトルに対して有効なスーパーベクトルといえる。 これにより、ループスケジューリング、タイリング、ループ入れ替えなどの後続のパスによるベクトル化への影響を減らすことが期待される。 スーパーベクトル化の主な特徴としては以下のようなものがある: - ベクトル化はスーパーベクトルの単位で行われる - スーパーベクトルはハードウェアベクトルのタイルとして抽象化され、ハードウェアレジスタの整数倍ではない場合もある - ベクトル型は依存性のないベクトル化を阻害せず、静的サイズのスライスであることを示す - スーパーベクトル化は後続のパスに対して不透明性を持ち、ベクトル化の制約を保持する 特に最後の後続のパスに対してベクトル化の制約を伝搬させることは、後続のパスでベクトル化を妨げる変換を防ぐ意味で重要である。 ### vs. VPLAN/SLP in LLVM LLVMにはVPLANとSLPというVectorizationのPlan機能がある。 MLIRのスーパーベクトルとVPLANは典型的なPolyhedral compilation flow（依存性解析、マルチレベル（スケジューリング+タイリング）、高速メモリへのフットプリントの変換、通信合成、マッピング、レジスタ最適化など）の場合ではベクトル化はアンロール前の遅い段階での適応となり、このケースにおいてはVPLANとは目的が同じであると言える。ただし、MLIRのAffine dialectはLLVMよりも高レベルなIRであるので、LLVMよりもシンプルな実装と柔軟性が実現されている。 SLP VectorizationとMLIRのスーパーベクトルは目的が違っており、特に扱う情報が違う。MLIRのSuperVectorizationは最も変化の速い次元と再起的にネストするパターンをキャプチャすることによって、メモリ参照の連続性を調べるということをやっている。一方で、SLPは単一のアンロールされたループ内部のフラットなパターンマッチングをしていて、その後Load/Storeを1次元ベクトルに結合する、という動きを取っている。この二つは目的が違うものであるため、相互補完的な位置付けである。 ## Size of Super Vectorization 特定のターゲットにおいて、最小のn次元ベクトルサイズの概念が指定されると、スーパーベクトル化はこれらの倍数を対象とする。以下では、「k .」が「の倍数」として使用している（例：vector<16 x k . 128>は、vector<16 x 128>、vector<16 x 256>、vector<16 x 1024>）。 - CPU： `vector<k . HW_vector_size>` - `vector<k' . core_count x k . HW_vector_size>` - `vector<socket_count x k' . core_count x k . HW_vector_size>` - GPU: `vector<k . warp_size>` - `vector<k . warp_size x float2>` - `vector<k . warp_size x float4>` - `vector<k . warp_size x 4 x 4x 4>` （テンソルコアサイズ用） ベクトルに対するループと操作はこれらのスーパーベクトル形状で生成され、後続のLoweringでは実際のHWベクトルサイズに具現化される。 ## Transfer vector 高レベルのIRにおけるベクトル化のためのロードは以下のようにして`vector.transfer`を利用して行われる: ``` affine.for %i = ? to ? step ? { %v_a = vector.transfer_read A[%i] : memref<?xf32>, vector<128xf32> } ``` ## Where places are super-vectorizable in pass pipeline スーパーベクトルを適応できる場所は大きく6つある: - 1. MLIRへの変換直後 - 2. 依存性解析の後 - Polyhederalスケジューリングの前 - 3. Polyhederal Stypeのスケジューリング直後 - 4. Polyhederal Stypeのスケジューリング + Tiling直後 - 5. スケジューリング + Tiling + 再スケジューリング直後 - 6. VPLANに対するMLIRベースの代替手段 ### Pattern 1 & 2 - Pattern 1: MLIRへの変換直後 - Pattern 2: 依存性解析の後 - Polyhederalスケジューリングの前 1のユーザー言語・ライブラリからMLIRへ変換した直後では、言語やライブラリレベルのアノテーション情報が利用可能である。ループタイプのアノテーション（並列、リダクション、スキャン、依存距離ベクトル、ベクトル化可能なループなど）やメモリアクセスのアノテーション（エイリアスを持たない書き込みが保証されたもの、構築による順列の間接アクセスなど）などが利用できるため、これらを積極的に活用できるタイミングである。 2の依存性解析後でかつ、Polyhederalスケジューリングの前の適応では、MLIRでは伝統的な依存性解析が利用可能であり、ベクトル化とコストモデルを駆動するために使用される。 この1と2においては、メリットとリスクが存在する: - メリット： - ベクトル化は型システムに組み込まれ、ループスケジューリング、タイリング、ループ入れ替えなどの影響から保護される。 - 後続のパスがベクトルタイプで動作する場合、ベクトルの形状、関連するループのイテレータプロパティ、アライメント、最も変化の速い次元の連続性は、スーパーベクトルタイプをLoweringするまで保持される。これにより、フェーズ順序のマイナスの影響が大幅に抑制されると期待される。 - リスク： - a. スーパーベクトル化の後続のすべてのパスは、基本的なベクトルタイプで動作する必要がある（これは常に真実とは限らないが、ベクトル化が適用される場所に依存する） - b. ベクトル化制約を早すぎる段階で課すことが、ループの統合、タイリング、その他の変換に全体的に悪影響を及ぼす可能性がある。なぜなら、基本的なベクトルタイプで操作する際に依存距離が粗くなるためである。そのため、パターンの収益性分析には、特定のパターンにおける最大量の融合もキャプチャする要素を含める必要がある。 ### Pattern 3 - Pattern 3: Polyhederal Stypeのスケジューリング直後 Polyhederal Stypeのスケジューリング直後のアルゴリズムは、局所性と並列性を向上させ、構成可能（最大融合、スマート融合など）であることが知られている。これらはベクトル化に必要な連続性プロパティにも逆効果を及ぼす可能性があるが、vector.transferのコピー、リシェイプ、パッド、転置の抽象化により、これらのプロパティを取り戻すことが期待される。 ### Pattern 4 & 5 - Pattern 4: Polyhederal Stypeのスケジューリング + Tiling直後 - Pattern 5: スケジューリング + Tiling + 再スケジューリング直後 4と5のポイントはおそらく最も有望な場所を示している。タイリングの適用により、再スケジューリングは局所性よりも並列性と分散（たとえば、最小融合）に関心を持つことができるようになる。 これらの段階では、おそらく多くの言語/ユーザー/ライブラリレベルのアノテーションを失ってしまったかもしれない一方で、スケジューリングとタイリングを通じて並列性と局所性を得た可能性がある。ただし、タイリングがスーパーベクトル化で使用されている完全タイルのみの抽象化と互換性があることを確認する必要がある。そうでないと、その結果を受け入れる必要が生じます。 ### Pattern 6 - Pattern 6: VPLANに対するMLIRベースの代替手段 ## Maching and Rewrite Algorithm スーパーベクトル化アルゴリズムは以下の手順で処理される: ### 1. パターンマッチング スーパーベクトル化パターンを定義し、それらをAffineForOpのツリーにマッチングする。スーパーベクトル化パターンは、再帰的なデータ構造であり、ネストされた非完全ネストループをマッチングしてキャプチャする。 これらのループは、以下のような条件がある: - a. 適合するループのアノテーション（例：parallel、reduction、vectorizableなど）が付属していることと、 - b. 指定されたMinor次元に沿ったすべての連続したロード/ストア操作を持つこと ### 2. パターンの利益分析 定義されたパターンごとの効果を分析する。 ### 3. ループとOperatorの書き換え 次に、各パターンに対して以下の順序で反復的にループとそのネストされた操作を書き換える。 - a. トポロジカルオーダーでループとそのネストされたすべての操作を書き換える。書き換えは、ループを粗くし、OperatorとOperandをベクトル形式に変換することによって実装される。このトポロジカルオーダーは、あるOperatorのすべてのOperandがその操作自体よりも前に単一のトラバースでベクトル化されていることを保証する（ただし、ループネストの外で定義されたOperandは除く）。アルゴリズムは以下の操作をベクトル形式に変換できる： - Affineロードおよびストア操作はOpaqueなベクトル転送読み取りおよび書き込み操作に変換される - スカラー定数操作/オペランドはベクトル定数操作（Splat）に変換される - Uniform Operand（ベクトル次元にマッピングされていないループの帰納変数、または現時点ではループネストの外で定義されたOperand）はベクトルにブロードキャストされる - `iter_args`を持つAffine For操作は、その `iter_args` Operandと結果をベクトル化する - ループネスト内の残りの操作は、スカラータイプをベクトルタイプに拡張してベクトル化される - b. 現在のパターンのルートであるAffineForOpの下にあるすべてが適切にベクトル化されていれば、そのループをIRにコミットしてスカラーループを削除する。 - そうでない場合は、ベクトル化されたループを破棄して元のスカラーループを保持する - c. 次のパターンでベクトル化が適用される - パターン干渉回避はまだ実装されていないことと、既にベクトルロードに対してさらなるベクトル化をサポートしないため、パターンがまだベクトル化可能であることを再検証する必要がある ## Example `@vector_add_2d` という関数を例にスーパーベクトル化を考える: ```mlir func.func @vector_add_2d(%M : index, %N : index) -> f32 { %A = alloc (%M, %N) : memref<?x?xf32, 0> %B = alloc (%M, %N) : memref<?x?xf32, 0> %C = alloc (%M, %N) : memref<?x?xf32, 0> %f1 = arith.constant 1.0 : f32 %f2 = arith.constant 2.0 : f32 affine.for %i0 = 0 to %M { affine.for %i1 = 0 to %N { // non-scoped %f1 affine.store %f1, %A[%i0, %i1] : memref<?x?xf32, 0> } } affine.for %i2 = 0 to %M { affine.for %i3 = 0 to %N { // non-scoped %f2 affine.store %f2, %B[%i2, %i3] : memref<?x?xf32, 0> } } affine.for %i4 = 0 to %M { affine.for %i5 = 0 to %N { %a5 = affine.load %A[%i4, %i5] : memref<?x?xf32, 0> %b5 = affine.load %B[%i4, %i5] : memref<?x?xf32, 0> %s5 = arith.addf %a5, %b5 : f32 // non-scoped %f1 %s6 = arith.addf %s5, %f1 : f32 // non-scoped %f2 %s7 = arith.addf %s5, %f2 : f32 // diamond dependency. %s8 = arith.addf %s7, %s6 : f32 affine.store %s8, %C[%i4, %i5] : memref<?x?xf32, 0> } } %c7 = arith.constant 7 : index %c42 = arith.constant 42 : index %res = load %C[%c7, %c42] : memref<?x?xf32, 0> return %res : f32 } ``` `-affine-super-vectorize` passに対してInnermost-loopのベクトル化を実施するパラメータを指定する: ``` -affine-super-vectorize="virtual-vector-size=256 test-fastest-varying=0" ``` すると、innermost-loopで扱うデータのみ`vector.transfer_write` で`vector<256xf32>`とベクトル化されているのがわかる: ```mlir func.func @vector_add_2d(%arg0 : index, %arg1 : index) -> f32 { %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %1 = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %2 = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %cst = arith.constant 1.0 : f32 %cst_0 = arith.constant 2.0 : f32 affine.for %i0 = 0 to %arg0 { affine.for %i1 = 0 to %arg1 step 256 { %cst_1 = arith.constant dense<vector<256xf32>, 1.0> : vector<256xf32> vector.transfer_write %cst_1, %0[%i0, %i1] : vector<256xf32>, memref<?x?xf32> } } affine.for %i2 = 0 to %arg0 { affine.for %i3 = 0 to %arg1 step 256 { %cst_2 = arith.constant dense<vector<256xf32>, 2.0> : vector<256xf32> vector.transfer_write %cst_2, %1[%i2, %i3] : vector<256xf32>, memref<?x?xf32> } } affine.for %i4 = 0 to %arg0 { affine.for %i5 = 0 to %arg1 step 256 { %3 = vector.transfer_read %0[%i4, %i5] : memref<?x?xf32>, vector<256xf32> %4 = vector.transfer_read %1[%i4, %i5] : memref<?x?xf32>, vector<256xf32> %5 = arith.addf %3, %4 : vector<256xf32> %cst_3 = arith.constant dense<vector<256xf32>, 1.0> : vector<256xf32> %6 = arith.addf %5, %cst_3 : vector<256xf32> %cst_4 = arith.constant dense<vector<256xf32>, 2.0> : vector<256xf32> %7 = arith.addf %5, %cst_4 : vector<256xf32> %8 = arith.addf %7, %6 : vector<256xf32> vector.transfer_write %8, %2[%i4, %i5] : vector<256xf32>, memref<?x?xf32> } } %c7 = arith.constant 7 : index %c42 = arith.constant 42 : index %9 = load %2[%c7, %c42] : memref<?x?xf32> return %9 : f32 } ``` 次にoutermostとinntermostの両方をベクトル化する: ``` -affine-super-vectorize="virtual-vector-size=32,256 test-fastest-varying=1,0" ``` すると、outer/inner-mostで扱うデータがそれぞれ指定されたサイズのベクトルとなっているのがわかる: ```mlir func.func @vector_add_2d(%arg0 : index, %arg1 : index) -> f32 { %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %1 = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %2 = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %cst = arith.constant 1.0 : f32 %cst_0 = arith.constant 2.0 : f32 affine.for %i0 = 0 to %arg0 step 32 { affine.for %i1 = 0 to %arg1 step 256 { %cst_1 = arith.constant dense<vector<32x256xf32>, 1.0> : vector<32x256xf32> vector.transfer_write %cst_1, %0[%i0, %i1] : vector<32x256xf32>, memref<?x?xf32> } } affine.for %i2 = 0 to %arg0 step 32 { affine.for %i3 = 0 to %arg1 step 256 { %cst_2 = arith.constant dense<vector<32x256xf32>, 2.0> : vector<32x256xf32> vector.transfer_write %cst_2, %1[%i2, %i3] : vector<32x256xf32>, memref<?x?xf32> } } affine.for %i4 = 0 to %arg0 step 32 { affine.for %i5 = 0 to %arg1 step 256 { %3 = vector.transfer_read %0[%i4, %i5] : memref<?x?xf32> vector<32x256xf32> %4 = vector.transfer_read %1[%i4, %i5] : memref<?x?xf32>, vector<32x256xf32> %5 = arith.addf %3, %4 : vector<32x256xf32> %cst_3 = arith.constant dense<vector<32x256xf32>, 1.0> : vector<32x256xf32> %6 = arith.addf %5, %cst_3 : vector<32x256xf32> %cst_4 = arith.constant dense<vector<32x256xf32>, 2.0> : vector<32x256xf32> %7 = arith.addf %5, %cst_4 : vector<32x256xf32> %8 = arith.addf %7, %6 : vector<32x256xf32> vector.transfer_write %8, %2[%i4, %i5] : vector<32x256xf32>, memref<?x?xf32> } } %c7 = arith.constant 7 : index %c42 = arith.constant 42 : index %9 = load %2[%c7, %c42] : memref<?x?xf32> return %9 : f32 } ``` ## Reduction Reductionのスーパーベクトル化は以下の条件が満たされると実行できる: - サポートされているリダクションの種類であること - ベクトル化が1次元であること - ループのステップサイズが1であること 非ベクトル次元の場合と比較して、このようなループのベクトル化中に次の2つの追加的な処理が行われる: 1. ループから返されるベクトルは、vector.reduceを使用してスカラーにReductionされる 2. ループの最後に生成されたベクトルにマスクを適用して、Accumulatorにゴミの値が書き込まれるのを防ぐ Reductionのベクトル化はデフォルトでは無効になっており、ループからReductionへのマップをユーティリティ関数に渡すか、ベクトル化パスに`vectorize-reductions=true`を指定することで有効にすることができる。 Example: ```mlir func @vecred(%in: memref<512xf32>) -> f32 { %cst = arith.constant 0.000000e+00 : f32 %sum = affine.for %i = 0 to 500 iter_args(%part_sum = %cst) -> (f32) { %ld = affine.load %in[%i] : memref<512xf32> %cos = math.cos %ld : f32 %add = arith.addf %part_sum, %cos : f32 affine.yield %add : f32 } return %sum : f32 } ``` `-affine-super-vectorize`に`vectorize-reductions=true`のオプションを与える: ``` -affine-super-vectorize="virtual-vector-size=128 test-fastest-varying=0 \ vectorize-reductions=true" ``` 以下のIRが生成される: ```mlir #map = affine_map<(d0) -> (-d0 + 500)> func.func @vecred(%arg0: memref<512xf32>) -> f32 { %cst = arith.constant 0.000000e+00 : f32 %cst_0 = arith.constant dense<0.000000e+00> : vector<128xf32> %0 = affine.for %arg1 = 0 to 500 step 128 iter_args(%arg2 = %cst_0) -> (vector<128xf32>) { // %2 is the number of iterations left in the original loop. %2 = affine.apply #map(%arg1) %3 = vector.create_mask %2 : vector<128xi1> %cst_1 = arith.constant 0.000000e+00 : f32 %4 = vector.transfer_read %arg0[%arg1], %cst_1 : memref<512xf32>, vector<128xf32> %5 = math.cos %4 : vector<128xf32> %6 = arith.addf %arg2, %5 : vector<128xf32> // We filter out the effect of last 12 elements using the mask. %7 = select %3, %6, %arg2 : vector<128xi1>, vector<128xf32> affine.yield %7 : vector<128xf32> } %1 = vector.reduction <add>, %0 : vector<128xf32> into f32 return %1 : f32 } ``` この例では128のサイズでベクトル化されているのがわかる。`affine_map`と`vector.create_mask`でマスクを作っているのはAccumlatorにゴミが入らないようにするためである（selectOpでマスクをかけている）。ループを抜けたあとの最後のReductionによって、指定されたOp（この例ではAdd）が実行され、返り値はスカラー値になっている。 --- ## Reference - [Dialect/Affine/Transforms/SuperVectorize.cpp#L41-L570](https://github.com/llvm/llvm-project/blob/release/16.x/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp#L41-L570) - [2013 LLVM Developers’ Meeting: “Vectorization in LLVM”](https://www.youtube.com/watch?v=TVV5v5R43nA)