分岐命令なしに 256 ビット整数の下位 k ビットの population count を求める
複数ワードで構成される一つの整数の下位 k ビットに対する population count を分岐命令なしに実現する方法について Java で考えてみました。
はじめに
ランククエリを実現するランク辞書実装の中には、複数個 (n 個) のワードで構成される大きなビット数の整数に対して下位 k ビットだけの population count (a.k.a. hamming weight) の計算を必要とする実装が存在します。具体的には Rank-1L や Poppy がそれに相当し、例えば Rank-1L であれば、64 * 4 = 256 ビットの整数に対する下位 k ビットの population count の計算が必要となります。
単に 1 ワードだけを対象とした下位 k ビットの population count であれば、以下のようにビットシフト、除算、ビット積、POPCNT 命令をそれぞれ 1 下位ずつ実行するだけで容易に実現できます。
public interface SingleLongPopulationCount {
static int popcnt(long value, int k) {
return Long.bitCount(value & ((1L << k) - 1));
}
}
しかしながら 2 ワード以上を対象とする場合は事情が異なります。ワードをまたいで一括でビット演算を実現するのは (AVX ほげほげを利用すれば実現できるかもしれませんが) Java では難しく、愚直に実装するならば以下のようにループを用いたコードのようになります1。しかしこの実装の場合は、ループカウンタの部分で 1 回以上の分岐命令 の実行が必要になります。
public interface LongsPopulationCount {
static int popcnt(long[] values, int k) {
int l = k / 64;
int result = Long.bitCount(values[l] & ((1L << k) - 1));
for (int i = 0; i < l; i++) { // ループカウンタのチェックで分岐命令が実行される
result += Long.bitCount(values[i]);
}
return result;
}
}
分岐命令は Making Your Code Faster by Taming Branches の記事にもあるように、最近の CPU に備わっている分岐予測や投機的実行によってその速度パフォーマンスはある程度よいものが期待できるものの、分岐予測に失敗したときの時間的なコストはそれなりに高くつきます。そのため、速度パフォーマンスをよりよくする場合においては、可能な限り分岐命令の実行回数を削減する、もしくは完全に分岐命令を排除することが望まれます。
もちろん、ランク辞書の速度パフォーマンスにおいては 一般的に知られているレイテンシの差 にあるように、分岐予測ミスよりも辞書を参照する際の L1/L2 キャッシュミスの方が支配的となります。ゆえに分岐命令の有無が速度パフォーマンスに 深刻な 影響を与えることはないのですが、分岐命令を削減・排除することで速度パフォーマンスの向上が見込めることに違いはない… ということで、このエントリでは分岐命令を排した 256 ビット整数下位 k ビットの population count を計算する Java 実装を紹介し、またその速度パフォーマンスをベンチマーク計測して評価してみます。
分岐命令を排除する方法
前述したように Java は複数ワードを一括してビット演算するのは難しいため、ここではワードそれぞれに対して異なるビットマスクを用意し、そのワードの個数だけビット演算をする方法を考えてみます。
つまり 256 ビット (64 ビット * 4 ワード) の整数に対して、例えば下位 130 ビットの population count を求めようと思えば、1 ワード目と 2 ワード目はすべてのビットが 1 のビットマスクを用意し、3 ワード目は 0x0000_0000_0000_0003 のビットマスク、4 ワード目はすべてのビットが 0 のビットマスクを用意して、それぞれのワードとのビット積を求めた後に 4 回ほど Long.bitCount() で POPCNT 命令を実行する、という処理になります。
これを分岐命令なしに実現する具体的な実装はいくつか考えられそうですが、今回取り上げるのはこちらです。
public interface Popcnt256 {
static int branchless(long[] values, int offset, int bitIndex) {
long b = 0b1110_1101_1011_0111__0000_0111_0011_0001L >>> (bitIndex >> 6);
long rightShiftBits = 255L - bitIndex;
long mask0 = (0x7fff_ffff_ffff_ffffL | (b << 35)) >> rightShiftBits;
long mask1 = ((0x7fff_ffff_ffff_ffffL | (b << 39)) + (b & 1)) >> rightShiftBits;
long mask2 = ((0x7fff_ffff_ffff_ffffL | (b << 43)) + ((b >>> 4) & 1)) >> rightShiftBits;
long mask3 = ((0x7fff_ffff_ffff_ffffL | (b << 47)) + ((b >>> 8) & 1)) >> rightShiftBits;
return Long.bitCount(values[offset] & mask0)
+ Long.bitCount(values[offset + 1] & mask1)
+ Long.bitCount(values[offset + 2] & mask2)
+ Long.bitCount(values[offset + 3] & mask3);
}
}
この実装では、最上位ビットが立っている場合の算術右シフトの挙動 (最上位ビットに 1 が fill され続ける) を利用して、すべてのビットが 1 となるケースと下位の一部のビットだけが 1 となるケースを作り分けたり、すべてのビットが 1 であるビットマスクに +1 してすべてが 0 のビットマスクを作り出すといった泥臭いことをしています。
さてこちらの実装、確かに分岐命令を排除できてはいるものの、多数のビット演算・加算や必ず 4 回実行される POPCNT 命令があることから一見しただけではループで愚直に実装した場合よりも速度パフォーマンスが悪くなりそうな印象があります。そこで次は、ベンチマークを用いて実際の速度パフォーマンスを計測してみることにします。
ベンチマーク
今回ベンチマークで計測するタスクは以下になります (都合2 により、乱数のサンプリング処理も計測対象に含まれています)。
- 64 ビット整数の乱数 4 つで構成される 256 ビットの乱数と、0 以上 256 未満の整数の乱数 k を毎回サンプリングして生成する
- その 256 ビットの乱数に対して下位 k ビットの population count を計算をする
ベンチマーク計測に用いたコードは こちら です。今回はループによる実装 (loop) と前述した分岐命令を排除した実装 (branchless) を計測します。加えて、乱数のサンプリング処理に要する時間を測るために、乱数生成のみを計測対象とするベンチマーク (noop) を用意しています。こちらを用いて、乱数サンプリングの影響を除外したスループットを推定してみることにします。
ベンチマーク計測環境はこちらです。
- Intel(R) Core(TM) i5-6287U CPU @ 3.10GHz
- macOS 10.14.6（18G3020）
- OpenJDK Runtime Environment 18.9 (build 11.0.2+9)
そして、ベンチマークの計測結果は以下になります。
| 実装 | スループット [ops/ms] |
スループット (乱数生成を除外) [ops/ms] |
|---|---|---|
| loop | 48,145.78 | 88,412.69 |
| branchless | 55,907.20 (vs. loop +16.12%) |
118,664.43 (vs. loop +34.22%) |
| (noop) | 105,712.05 | - |
上記より明らかなように、分岐命令の排除は速度パフォーマンスの向上に寄与することがわかります。具体的には、ループによる実装に対して少なくとも 16%、おそらくは 30% 超ぐらいは速度性能が向上すると見込まれます。
まとめ
複数ワードからなる整数の下位 k ビットのみの population count を、ループを使わずに分岐命令なしに実現する方法について考えてみました。結果として、256 ビットの整数を対象とした場合にループを用いる実装よりも 少なくとも 16% 以上 (おそらくは 30% 以上は期待できるほど) の速度性能を向上させることができました。
なお今回は 256 ビットの整数のみを対象にベンチマークを計測しましたが、128 ビットもしくは 512 ビットの整数を対象とする場合はまた状況が異なる可能性があることに注意が必要です。また、Intel 製以外の CPU の場合にも異なる結果となる可能性が十分にあります。
-
実際に C による Poppy のリファレンス実装では、ループによる実装が採用されています。 ↩
-
JMH には、
@Setup(Level.Invocation)アノテーションを利用することでベンチマーク対象メソッドを呼び出す度にセットアップ処理を実行することができるのですが、このLevel.Invocationはベンチマーク対象のメソッドが 1ms を超えるケースでの利用にのみ適しており、今回のように μs オーダーで実行される処理で利用するには適していません。 ↩