HotSpot JavaVM のベクトル化変換

最近の HotSpot JavaVM はスカラー演算の繰り返し処理をベクトル化し SIMD 命令に変換する最適化を行っています (SIMD とは何ぞやという話は後半参照)。実際に最適化が効くコードで試してみたところ 1.5～2.8 倍程度の速度向上が見られたので、大量の演算処理を行う (GPGPU に頼らない) Java ライブラリでうまく使うことが出来れば有効な最適化手段になるかもしれません。

HotSpot の SIMD 最適化は Superword-Level Parallelism (SLP) に基づいています (以降この最適化は SLP と呼びます)。元々この論文は時代を反映して SIMD 未対応の画像・音声処理をコンパイラやランタイムのレイヤーで SIMD を利用する命令に変換することを目的としていましたが、これは SIMD 命令のない Java バイトコードで SIMD を利用するのにも有効です。

HotSpot JavaVM のベクトル化変換

SLP の変換処理

以下のような単純なループ処理を想定します。

for(int i=0; i<1024; i++){
  a[i] = b[i] + c[i];
}

このループに対して loop unrolling (ループ展開) を行います。

for(int i=0; i<1024; i+=4){
  a[i] = b[i] + c[i];
  a[i+1] = b[i+1] + c[i+1];
  a[i+2] = b[i+2] + c[i+2];
  a[i+3] = b[i+3] + c[i+3];
}

これだけでもインタラクション数や条件判定回数、分岐予測のペナルティが減るので高速化が期待できます (このようなループ展開はコンパイラや仮想マシンでできるものも結構あります)。

さてここでループ内の処理は以下のようなベクトル演算と見なせます。

このベクトル演算を実際にベクトル演算命令に置き換えます。

for(int i=0; i<1000; i+=4){
  load b[i:i+3] to xmm0
  load c[i:i+3] to xmm1
  add xmm0, xmm1 and store result to xmm0
  load xmm0 to a[i:i+3]
}

この変換により単純な演算処理の繰り返しをベクトル演算化できます。最近の CPU はどれも SIMD 命令が使えますから、ソースやバイトコードの変更なしに SIMD を使ってくれるのは効果的ですね。

Java 8 での動作

実際に以下のソースを使用して SLP を有効/無効化し実行時間を比較します。計測は適当なので正確なエビデンスが必要なら各自行ってください。

import java.util.Arrays;

public class J8SIMD {
  private static final int SIZE = 1024 * 1024;
  private static final float[] a = new float[SIZE];
  private static final float[] b = new float[SIZE];
  static {
    Arrays.fill(a, 1);
    Arrays.fill(b, 2);
  }
  public static void vectorAdd(){
    for(int i=0; i<a.length; i++){
      a[i] += b[i];
    }
  }
  public static void main(String[] args){
    // warming up
    for(int i=0; i<100; i++) vectorAdd();
    // measure
    long t0 = System.currentTimeMillis();
    for(int i=0; i<10000; i++){
      vectorAdd();
    }
    long t1 = System.currentTimeMillis();
    System.out.printf("vectorAdd: %,d[msec]", t1 - t0);
  }
}

上記ソースを Core i7-6600U (MMX, SSE, SSE2, SSSE3, SSE4, AVX 対応) で実行してみると優位な差が出ます (ただし SIMD 命令を使用したかの確証がないため単に loop unrolling しただけの差の可能性もなきにしもあらず)。

ここでは float を使用しましたが int でも double でも優位な差となりますので試してみてください。ただしデータ幅のせいか double は int や float ほど差が大きくありません。

SLP は Java7u40 あたりからデフォルトで有効になっていて、無効化する場合は -XX:-UseSuperWord オプションを使用します。

java8simd>java -XX:+UseSuperWord J8SIMD
vectorAdd: 4,624[msec]

java8simd>java -XX:-UseSuperWord J8SIMD
vectorAdd: 6,928[msec]

java8simd>java -version
java version "1.8.0_141"
Java(TM) SE Runtime Environment (build 1.8.0_141-b15)
Java HotSpot(TM) 64-Bit Server VM (build 25.141-b15, mixed mode)

各演算の実行時間 (ミリ秒) は以下の通り。

%= は数分待っても終わらなかったのであきらめました (何でこんな遅い?)。加減乗算は 1.5 倍の速度向上が見られました。除算の速度向上が 3 倍近くと特に顕著ですが、SIMD 命令で実行すると専用の FPU 上で処理されるなどの動きがあるのでしょうか?

以下は int 版に変更してビット操作を行った結果です。ビットシフト演算は SLP で最適化されないようです。

SLP が有効に機能したかは顕著な数値で出るので -XX:-UseSuperWord オプションの有無で判断できそうです。

SLP の発動条件

デバッグコンパイルされた OpenJDK で -XX:+PrintCompilation -XX:+TraceLoopOpts オプションを使用するとループがどう認識されているかが、また -XX:+PrintAssembly で HotSpot が生成したアセンブリが表示されるそうです。

ただ連休つぶしてそこまでやってられないので Vectorization in HotSpot JVM に載っていた Java 8/9ea の条件をかいつまんで載せます。将来のバージョンでより発動条件や制限が緩和されてゆくと思いますので時事で追ってください。

• Trip-Counted Loop であること。つまりカウントループかつ、ループ中に limit が変化せず step がコンパイル時に決定していること。

for(int i=0; i<100; i++){ } // OK
for(int i=5; i<100; i++){ } // OK
for(int i=t0; i<t1; i++){ } // OK
for(int i=0; i<100; i+=4){ } // OK
for(int i=0; i<100; i+=d){ } // NG
for(int i=0; i<100; i*=2){ } // NG
for(int i=0; i<t1; i++){     // NG
  ...
  if(...) t1 ++;
}

• ループカウンタは int, short, char のみ。

for(byte i=0; i<100; i++){ }  // NG
for(char i=0; i<100; i++){ }  // OK
for(short i=0; i<100; i++){ } // OK
for(int i=0; i<100; i++){ }   // OK
for(long i=0; i<100; i++){ }  // NG

• 関数コールが入っていると loop unrolling が行われない。

for(int i=0; i<100; i++){  // NG
  ...
  f();
}

• 対象のデータが配列であること。

for(int i=0; i<100; i++){  // OK
  a[i] += b[i];
}
for(int i=0; i<100; i++){  // NG
  a.set(i, a.get(i) + b.get(i));
}

• 単純な四則演算であること。

a[i] ++;  // OK
a[i] += 2; // OK
a[i] *= b[i]; // OK
a[i] = b[i] + c[i]; // OK

• インデックスが連続していること。

for(int i=0; i<100; i++){  // NG
  a[i] += b[2 * i];
}

• 集約 (aggregation) は Java 9 から一部出来るかも知れない

int r = 0;
for(int i=0; i<100; i++){
  r += a[i] * b[i];
}

動作の特徴と Java 9 の拡張

当然ですが Java (および Scala や Kotlin 等の JavaVM 言語) でのコレクション API や Stream 操作は SIMD 命令に変換可能な操作ではありません。それらを使って大量の数値演算を行っている部分がボトルネックであれば、配列 & ループカウントへの置き換えを検討する価値があるかもしれません。

Unsafe はメソッド呼び出しとなるため、速度優先のリスクテイクで Unsafe を使用しているコードでこの最適化は期待できません。

ベクトル化は 64 バイトのアラインメント制約があり、配列前方の端数はベクトル化されていない pre-loop で行われます。また配列終端の 4 の倍数になっていない要素もベクトル化されていない post-loop が行われます。ただし post-loop は Java 9 でベクトル化されるそうです。

Java 8 では 128bit レジスタ (xmm) を使用するため 2 または 4 要素ごとにベクトル化されますが Java 9 で AVX2 なら 32 まで増やせるとか何とか。

SIMD

SIMD (Single Instruction Multiple Data) はたくさん並べた数値データに一度に同じ命令を適用する仕組み。

流体シミュレーションのような科学技術系の演算処理、画像処理、音声処理、暗号化、機械学習の処理は「配列 $a$ に配列 $b$ の値を掛ける」といった線形代数的な演算が頻繁かつ大量に出てきます。

大量の演算を行うために無駄な回路を省いて低コスト、省電力、省スペース化した演算コアを大量に集積したものが GPU やスパコンのベクトル演算ユニットです。これらは1つの命令をレジスタに乗っている大量のデータに一度に適用できます。

最近の CPU は数個～数十個程度のベクトル演算を行う CPU 命令が用意されています。画像処理で大量の演算を扱う GPU は演算専用のプロセッサを数千個積んでいてそれらに一度に演算命令を出します。どちらも SIMD ですが HotSpot の SLP で最適化されるのは CPU 命令の方です。

参考資料

• Exploiting Superword Level Parallelism with Multimedia Instruction Sets
• Virtualization in HotSpot JVM
• JAVA and SIMD