突然に64bitのレジスタ上でselectの操作をしたくなるときありますね (ありますね)。これはビット演算テクニックAdvent Calendar 19日目の記事です。
selectとは
select_q(k) は、配列x上を先頭から見たとき、k番目に出現する値qのインデクスを返す関数として定義されます。たとえば、配列 [0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1] について、(左を先頭にして) select_1(3) = 6 となります。
このselectの演算ですが、64bitのレジスタを64要素の {0, 1} の配列と見て、この上で実行したくなるときがあります。これはsuccinct data structureの文脈で出てきたりします。(自分は別の目的でした)
そこにpdepがあらわれた
x86_64にintelが導入したBMI2命令セットに、pdepという命令があります。
こんな感じに、マスクレジスタを使ってsrcの下位に並んだビットをscatterします。これを使って64bitのレジスタv上でselect_1(k)を計算します。
1をkだけ左にシフトし、これをselectしたいベクタをマスクにしてpdepし、結果のレジスタをtzcnt (下位から0を数える) します。また、vのうち1の個数がk個以下の場合は、pdepで1が追い出されるので結果は64となります。シフト、pdep、tzcntともに3クロックかかるのでいうほど速くありませんが、分岐がないのは良いです。めでたしめでたし。
おわりに
pdepってチェスとか将棋にしか使いどころがないのではと思っていたのですが、まさか自分が使うことになるとはと驚き、こんな使い方があったかと発見に嬉しくなりました。この驚きと喜びを誰かに伝えたかった...