PMD_MT 高速化版+6
久しぶりに更新して、AVX2版の高速化と、AVX512版を追加、さらにいくつかの最適化を積み上げて高速化。環境とオプションによっては、少しは速くなったかも?
※注意点: すみませんがSSE版は全くさわっていません…。
高速化以外の変更点
・想定動作環境の変更。 (Win 8.1 と Win10)
・どのSIMDを使用しているかを設定画面のウィンドウタイトルに表示するように。
ダウンロード>>
ダウンロード (ミラー) >>
ソースはこちら
高速化のためにやったこと
大雑把に言うと、6つの最適化をやってみた。ひとつずつは大したことなくても、これらをを積み上げたので、それなりに高速化できているかも。
(1) gauss縦横の統合
(2) AVX512BW対応
(3) vpmulld削減
(4) AVX512VBMI/VNNI対応
(5) vpgatherの置き換え (AVX512)
(6) vpgatherをやめる (Ryzen)
(1)(3)は修正PMD:オンの場合の変更、(5)(6)はuseExp:オンの場合の変更。
まあ、細かい高速化の中身は後で書くとして、まずはどのくらい高速化できたかから。
テスト環境と高速化結果
AVX512はとにかく種類が多い…。なんかほかにもあった気がするけど省略。
今回は1920x1080の動画の5120フレーム分の処理時間から、1フレームの平均処理時間を測った。
PMD_MTは設定によって、処理内容が変わってくるけど、設定は、
・強さ:100
・閾値:100
・回数:2
・useExp:オン
・修正PMD:オン
でテストしている。
修正PMD:オンの場合、PMD_MTの処理は大きく分けて、ノイズ除去の重みを決めるためにガウシアンフィルタをかける箇所と、実際にノイズ除去する部分に分かれるので、これらを分けて時間を測った。
まずは、Skylake-X(i9 7980XE)での結果から。従来1フレームに1.92msかかっていたのが、高速化(1)(2)(3)(5)で1.25msまで短縮できていて、かなり高速化できている。高速化(1)と(5)の効果が大きめ。
次に、IceLake(i5 1035G7)の結果。従来1フレームに6.64msかかっていたのが、高速化(1)(2)(3)(4)(5)で、4.15msまで短縮できている。こちらも高速化(1)と(5)の効果が大きめ。
RyzenはAVX512に対応していないので、AVX512関係の最適化は適用できず、あまり大きな高速化にはなっていない。高速化(1)と(6)が強く効いてるけど、高速化(6)はテーブルから表引きするときに、AVX2のgather命令を使わずに、普通にひとつづつ表引きするというもので、正直まあ、これで速くなるというのはRyzenのgatherがとても遅いからである…。
(1) gauss縦横の統合
PMD_MTでは、ノイズ除去の重みを決める際に、5x5のガウシアンフィルタを適用した値とオリジナルの値との差分を計算する。
5ｘ5のガウシアンフィルタの重みは下の図のような感じで、中心にある自分のピクセルの周辺5x5のピクセルに下のような重みをかけて足し、その後重みの総和(256)で割る。
これをまじめにやると、1ピクセル計算するのに5x5=25回分掛けたり足したりしないといけなくて、時間がかかってしまう。
ところがよく見るとこのガウシアンフィルタの重みは縦横の重みが共通(1,4,6,4,1)で対称なため、一度、横方向のみにフィルタをかけてから次に縦方向のみにフィルタをかけることで(丸め誤差を無視すれば)同じ計算結果になるし、計算回数は5+5=10回分で済むので、計算回数を削減できる。
そこでPMD_MTでは、従来からこのように縦横を分離して計算してきた。
この方法の問題点は、縦横2回をそれぞれ計算しているので、メモリアクセスがそのたびに発生してしまうこと。SIMD化するとメモリ帯域の制限が厳しくなりがちなので、メモリアクセスはなるべく減らしたい。
そこで今回は、横方向フィルタ→縦方向フィルタの処理関数を統合し、一気に行うように実装を変更した。これでメモリアクセスは縦横の処理それぞれ発生していたのが1回分に減ったはず。
やってみると、予想通りそれなりに効果がある感じ。メモリアクセスを減らすのはやはり大事なようだ。
なお、コードはだいぶ複雑化したけど、ここではそういうのは全く気にしない（笑）
(2) AVX512BW対応
せっかく手元にSkylake-Xのマシンがあるので、これを有効活用してみようということで、処理をそのままAVX2からAVX512に拡張した。
が、あとの結果をみてもらえればわかるけど、それだけだと高速化はわずかなのが悲しい。
いろいろ理由は考えられるけど、AVX512はその名の通り512bit幅(=64byte)で、キャッシュラインと同じサイズ。Aviutlのフレームは64byteでアライメントが取れているとは限らないので、メモリアクセスのたびにキャッシュラインをまたいでアクセスする羽目になる、というのが原因の一つかもしれない(けどいかんともしがたい)。
AVX2→AVX512で大変だったのがやっぱりYC48の取り扱い。YC48は、輝度と色差x2が16bit幅でY,Cb,Crの順に並び、これがピクセル数分続く形になっている。下の図で白が輝度、赤と青が色差、数字がピクセル番号を表すとして、上段のような感じ。
ただ、この形はSIMDで計算するには相性が悪いので、下の図の下段のように並び変えて計算し、最後には再度元に戻してあげる必要がある。
そのためにはひたすらshuffle命令だのvperm命令だのを使ってあげる必要があって、これを考えるのが結構大変だった。
結局、YC48→Y,Cb,Cr(上図の上段→下段)はこんな感じ。
で、その逆のY,Cb,Cr→YC48はこんな感じ。
結局便利なvpermt2w命令(_mm512_mask_permutexvar_epi16)を乱発して誤魔化してるけど、なんかもっとシンプルな方法ありそう…。
(3) vpmulld削減
vpmulldというのはSIMDレジスタで32bit同士の掛け算(a×b)をする命令でとても便利なのだけど、なぜだか知らないけどIntel系ではとても遅いらしい。Ryzenはそんなでもなさそうなのになぜなのだ…。
で、今回はよく考えると16bit同士の掛け算をするvpmaddwd(こっちはそこまで遅くない)で十分計算できるのに気づいたので、データの並べ替えが必要になる分だけ命令数は増えるけど置き換えてみた。するとほんのちょっとだけど速くなった。
実際の変更箇所はこんな感じ。
vpmaddwdが使えるようにするデータの並べ替えを考えるのが面倒だからと言って、あまり遅い命令を使うものではない…という反省。
(4) AVX512VNNI/AVX512VBMI対応
せっかくIceLakeが手元にあるので、新しくIceLakeで追加されたAVX512命令を使ってみようと思って対応してみた。
けど、まあほんの一部書き換えただけなので、そんなに速くならなかった。
ちなみにVNNIはVector Neural Network Instructions、VBMIはVector Bit Manipulationの略らしい。
(5) vpgatherの置き換え (AVX512)
以前も書いたけど、pmd_mtでノイズ除去するとき、データ構造と重み関数を以下のように定義して、
あるピクセルは以下のように計算する。(cb, crも同様)
というわけで、1ピクセル計算するのにfunc_weightは4回x3 (y, cb,cr)=12回呼ばれることになり、結構重い。さらにuseExp:オンの場合、その名の通りfunc_weightの中にexpが使用されている。expを12回もやっていると、まあ、だいぶ重い計算になる。
もともとpmd_mtはこの重い重み計算を回避するため、最初にとりうるすべてのx(-4500～4500)についてあらかじめ重み計算を行ってテーブルを作っておき(なのでテーブルサイズはおよそ9000)、あとはテーブル参照するだけというコードになっていた。
こうすることで、重み計算を毎回行わなくて済み、だいぶ高速化できる。さらに、これをSIMDでやる場合には、表引き(というか間接参照)を一気にやってくれるvpgatherという命令があるので、これを使う。(ただし、HaswellやRyzenでは遅い)
これを、
とやるとvpgatherdd命令1発でやってくれる。このvpgatherdd命令は、Skylake以降ではそれなりに高速だが、やはり要素数分のメモリアクセスが必要なのは変わらないので、どちらかというと遅い部類の命令で、避けられるなら避けたい。
ここでひとつ気になってくるのがテーブルのサイズ。もともとはテーブルの入力としてありうる -4500 ～ 4500 のだいたい9000のテーブルサイズになっているけど、本当にこんなに大きなテーブルが必要なのかな? ということで、実際のuseExp:オンのときのテーブルの値を確認すると(下図)、0を頂点にして山型の形になっていて、山の幅がPMD_MTのパラメータthresholdの値によって変わる感じの値になっている。
特徴としては、
・最大値は16384、16bit整数で事足りる
・0をはさんで左右対称→絶対値をとれば負側の値は要らない (テーブルサイズを半分にできる)
・あることろから急激に値が下がって1を下回る→テーブルは整数で持っているので、あるところからは0が続くだけ (ということはそれ以降はテーブルは要らない)
となっていて、thresholdの値に左右されるものの、例えばデフォルトのthreshold=100であれば、必要なテーブルサイズは100程度ともともとの9000に比べると小さいことがわかる。
AVX512BWにはvpermt2w/vpermi2wという、zmmレジスタふたつに入った合計64の16bit値から、16bit単位で任意の値を32並列で引っ張ってこれるという強力なshuffle命令が追加されている。ということは、これを使えばテーブルサイズが64までなら一発で表引きできるし、テーブルサイズが128までなら2回、テーブルサイズが256でも4回この命令を使えば表引きできてしまう。
たとえば、コードにするとこんな感じ。
さすがにテーブルサイズが大きすぎると、このvpermi2w命令を何回も呼ばなければならず、テーブルサイズが大きくても1回で済むvpgatherのほうが速い、ということになってしまうが、さっき紹介したように今回はある程度テーブルサイズが小さく済ませることができる。例えばthreshold=100ならテーブルサイズが100程度なので、vpermi2wを2回実行すれば済んでしまうので、遅い部類の命令であるvpgatherddを1回使うより速くなりそう…ということで試してみた。
結果はさっきの表のとおりで、IceLake/Skylake-Xともにかなり高速化できている。
なお、気を付けないといけないのは、useExp:オフのときや、thresholdの値が大きいときはテーブルサイズもかなり大きくなり、1000以上になることもあること。こういう場合は、従来通りvpgatherを使うようにしてある。
表引きのあたりは、そもそも表にするべきかどうかなども含めて、いろいろ工夫の余地があって面白い。
(6) vpgatherをやめる (@Ryzen)
RyzenはAVX2までの対応でAVX512は実装されていないし、AVX2にはvpermi2wのような強力なshuffle命令は存在しないので、さっきの高速化は使えない。
なお悪いことに、Ryzenのgather命令は遅いらしい。
というわけでRyzenでは表引きの際gatherを使うのではなく、愚直にひとつひとつ値をとってくるようにすると、そこそこ速くなった。えええ゛…。
とはいえ、gather命令は結構便利なので、Ryzenでも今後はgatherが速くなるといいな…。
というわけでいろいろな最適化を試すことで、全体としてもそこそこ高速化することができた。
メモリアクセスを減らせば高速化できる、というのはまあいつも通り。
一方でこれまでAVX2を単純にAVX512に置き換えるだけではあまり速くならないことが多くて悲しかったが、今回はAVX512で追加された強力なshuffle命令を活用してまあまあ速くすることができた。
AVX512は単に256bit→512bitとbit幅が倍になっただけじゃなくて、便利な新命令が増えていたり、なにより大変便利なマスクレジスタという仕組みも追加されていて、非常に面白い拡張になっている。なんか無駄にいろいろな拡張の名前がある(AVX512F/AVX512DQ/AVX512BW/AVX512VNNI/AVX512VBMI/AVX512VBMI2/........と無数にある…)のは、もうちょっと整理してよね！という感じでいただけないけど、今後はIceLakeやその後継のTigerLake、あと噂ではRocketLakeなどにも搭載されて使えるCPUが広がってくるはずなので、もっと活用して速くなるようになっていくといいなと思う。
ただまあ、そもそもIntel 10nmはいつデスクトップにやってくるんですかねえ…。(RocketLakeは14nm?)
※注意点: すみませんがSSE版は全くさわっていません…。
高速化以外の変更点
・想定動作環境の変更。 (Win 8.1 と Win10)
・どのSIMDを使用しているかを設定画面のウィンドウタイトルに表示するように。
ダウンロード>>
ダウンロード (ミラー) >>
ソースはこちら
高速化のためにやったこと
大雑把に言うと、6つの最適化をやってみた。ひとつずつは大したことなくても、これらをを積み上げたので、それなりに高速化できているかも。
(1) gauss縦横の統合
(2) AVX512BW対応
(3) vpmulld削減
(4) AVX512VBMI/VNNI対応
(5) vpgatherの置き換え (AVX512)
(6) vpgatherをやめる (Ryzen)
(1)(3)は修正PMD:オンの場合の変更、(5)(6)はuseExp:オンの場合の変更。
まあ、細かい高速化の中身は後で書くとして、まずはどのくらい高速化できたかから。
テスト環境と高速化結果
| CPU | i9 7980XE | i5 1035G7 | Ryzen7 3700X |
|---|---|---|---|
| CPU世代 | Skylake-X | IceLake | Zen2 |
| コア | 18C/36T | 4C/8T | 8C/16T |
| クロック | 4.0GHz固定 | 定格 | 定格 |
| SIMD | AVX2 AVX512F AVX512DQ AVX512BW | AVX2 AVX512F AVX512DQ AVX512BW AVX512VBMI AVX512VNNI | AVX2 |
| メモリ | DDR4-3600, 4ch | LPDDR4-3733, 2ch | DDR4-3600, 2ch |
| メモリ容量 | 16GB | 8GB | 16GB |
| OS | Win10 x64 | ||
AVX512はとにかく種類が多い…。なんかほかにもあった気がするけど省略。
今回は1920x1080の動画の5120フレーム分の処理時間から、1フレームの平均処理時間を測った。
PMD_MTは設定によって、処理内容が変わってくるけど、設定は、
・強さ:100
・閾値:100
・回数:2
・useExp:オン
・修正PMD:オン
でテストしている。
修正PMD:オンの場合、PMD_MTの処理は大きく分けて、ノイズ除去の重みを決めるためにガウシアンフィルタをかける箇所と、実際にノイズ除去する部分に分かれるので、これらを分けて時間を測った。
まずは、Skylake-X(i9 7980XE)での結果から。従来1フレームに1.92msかかっていたのが、高速化(1)(2)(3)(5)で1.25msまで短縮できていて、かなり高速化できている。高速化(1)と(5)の効果が大きめ。
次に、IceLake(i5 1035G7)の結果。従来1フレームに6.64msかかっていたのが、高速化(1)(2)(3)(4)(5)で、4.15msまで短縮できている。こちらも高速化(1)と(5)の効果が大きめ。
RyzenはAVX512に対応していないので、AVX512関係の最適化は適用できず、あまり大きな高速化にはなっていない。高速化(1)と(6)が強く効いてるけど、高速化(6)はテーブルから表引きするときに、AVX2のgather命令を使わずに、普通にひとつづつ表引きするというもので、正直まあ、これで速くなるというのはRyzenのgatherがとても遅いからである…。
(1) gauss縦横の統合
PMD_MTでは、ノイズ除去の重みを決める際に、5x5のガウシアンフィルタを適用した値とオリジナルの値との差分を計算する。
5ｘ5のガウシアンフィルタの重みは下の図のような感じで、中心にある自分のピクセルの周辺5x5のピクセルに下のような重みをかけて足し、その後重みの総和(256)で割る。
これをまじめにやると、1ピクセル計算するのに5x5=25回分掛けたり足したりしないといけなくて、時間がかかってしまう。
ところがよく見るとこのガウシアンフィルタの重みは縦横の重みが共通(1,4,6,4,1)で対称なため、一度、横方向のみにフィルタをかけてから次に縦方向のみにフィルタをかけることで(丸め誤差を無視すれば)同じ計算結果になるし、計算回数は5+5=10回分で済むので、計算回数を削減できる。
そこでPMD_MTでは、従来からこのように縦横を分離して計算してきた。
この方法の問題点は、縦横2回をそれぞれ計算しているので、メモリアクセスがそのたびに発生してしまうこと。SIMD化するとメモリ帯域の制限が厳しくなりがちなので、メモリアクセスはなるべく減らしたい。
そこで今回は、横方向フィルタ→縦方向フィルタの処理関数を統合し、一気に行うように実装を変更した。これでメモリアクセスは縦横の処理それぞれ発生していたのが1回分に減ったはず。
やってみると、予想通りそれなりに効果がある感じ。メモリアクセスを減らすのはやはり大事なようだ。
なお、コードはだいぶ複雑化したけど、ここではそういうのは全く気にしない（笑）
(2) AVX512BW対応
せっかく手元にSkylake-Xのマシンがあるので、これを有効活用してみようということで、処理をそのままAVX2からAVX512に拡張した。
が、あとの結果をみてもらえればわかるけど、それだけだと高速化はわずかなのが悲しい。
いろいろ理由は考えられるけど、AVX512はその名の通り512bit幅(=64byte)で、キャッシュラインと同じサイズ。Aviutlのフレームは64byteでアライメントが取れているとは限らないので、メモリアクセスのたびにキャッシュラインをまたいでアクセスする羽目になる、というのが原因の一つかもしれない(けどいかんともしがたい)。
AVX2→AVX512で大変だったのがやっぱりYC48の取り扱い。YC48は、輝度と色差x2が16bit幅でY,Cb,Crの順に並び、これがピクセル数分続く形になっている。下の図で白が輝度、赤と青が色差、数字がピクセル番号を表すとして、上段のような感じ。
ただ、この形はSIMDで計算するには相性が悪いので、下の図の下段のように並び変えて計算し、最後には再度元に戻してあげる必要がある。
そのためにはひたすらshuffle命令だのvperm命令だのを使ってあげる必要があって、これを考えるのが結構大変だった。
結局、YC48→Y,Cb,Cr(上図の上段→下段)はこんな感じ。
char *src; //YC48データへのポインタ
alignas(64) static const uint16_t PACK_YC48_SHUFFLE_AVX512[32] = {
0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45,
48, 51, 54, 57, 60, 63, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29
};
__m512i z0 = _mm512_load_si512((__m512i *)PACK_YC48_SHUFFLE_AVX512);
__m512i z5 = _mm512_loadu_si512((__m512i *)(src + 0));
__m512i z4 = _mm512_loadu_si512((__m512i *)(src + 64));
__m512i z3 = _mm512_loadu_si512((__m512i *)(src + 128));
__m512i z1, z2;
__m512i z6 = _mm512_ternarylogic_epi64(
_mm512_setzero_si512(), _mm512_setzero_si512(),
_mm512_setzero_si512(), 0xff); //-1
__mmask32 k7 = 0xffc00000;
__mmask32 k6 = 0xffe00000;
z1 = z0;
z1 = _mm512_permutex2var_epi16(z5/*a*/, z1/*idx*/, z4/*b*/);
z1 = _mm512_mask_permutexvar_epi16(z1/*src*/, k7, z0/*idx*/, z3);
z0 = _mm512_sub_epi16(z0, z6);
z2 = z0;
z2 = _mm512_permutex2var_epi16(z5/*a*/, z2/*idx*/, z4/*b*/);
z2 = _mm512_mask_permutexvar_epi16(z2/*src*/, k6, z0/*idx*/, z3);
z0 = _mm512_sub_epi16(z0, z6);
z6 = z0;
z6 = _mm512_permutex2var_epi16(z5/*a*/, z6/*idx*/, z4/*b*/);
z6 = _mm512_mask_permutexvar_epi16(z6/*src*/, k6, z0/*idx*/, z3);
y = z1; //出力 - Y のみ
cb = z2; //出力 - Cbのみ
cr = z6; //出力 - Crのみ
で、その逆のY,Cb,Cr→YC48はこんな感じ。
__m512i zY; //入力 - Y成分のみ
__m512i zU; //入力 - Cb成分のみ
__m512i zV; //入力 - Cr成分のみ
__m512i z0, z1, z2;
__m512i zOffset10 = _mm512_set1_epi16(10);
__m512i zOffset21 = _mm512_set1_epi16(21);
alignas(64) static const uint16_t shuffle_yc48[] = {
0, 32, 64, 1, 33, 65, 2, 34, 66, 3, 35, 67, 4, 36, 68, 5,
37, 69, 6, 38, 70, 7, 39, 71, 8, 40, 72, 9, 41, 73, 10, 42
};
__mmask32 mask1 = 0xDB6DB6DBu;
__mmask32 mask2 = 0xB6DB6DB6u;
__mmask32 mask1not = ~mask1;
__mmask32 mask2not = ~mask2;
__m512i zShuffle = _mm512_load_si512((__m512i *)shuffle_yc48);
z0 = _mm512_permutex2var_epi16(zY/*a*/, zShuffle/*idx*/, zU/*b*/);
z0 = _mm512_mask_mov_epi16(z0, mask1not,
_mm512_permutexvar_epi16(zShuffle/*idx*/, zV));
__m512i zShuffleM21 = _mm512_sub_epi16(zShuffle, zOffset21);
__m512i zShuffleP10 = _mm512_add_epi16(zShuffle, zOffset10);
z1 = _mm512_permutex2var_epi16(zY/*a*/, zShuffleM21/*idx*/, zU/*b*/);
z1 = _mm512_mask_mov_epi16(z1, mask2not,
_mm512_permutexvar_epi16(zShuffleP10/*idx*/, zV));
__m512i zShuffleP21 = _mm512_add_epi16(zShuffle, zOffset21);
__m512i zShuffleM42 = _mm512_sub_epi16(zShuffleM21, zOffset21);
z2 = _mm512_permutex2var_epi16(zU/*a*/, zShuffleP21/*idx*/, zV/*b*/);
z2 = _mm512_mask_mov_epi16(z2, mask1not,
_mm512_permutexvar_epi16(zShuffleM42/*idx*/, zY));
z0; //出力 - YC48(1)
z1; //出力 - YC48(2)
z2; //出力 - YC48(3)
結局便利なvpermt2w命令(_mm512_mask_permutexvar_epi16)を乱発して誤魔化してるけど、なんかもっとシンプルな方法ありそう…。
(3) vpmulld削減
vpmulldというのはSIMDレジスタで32bit同士の掛け算(a×b)をする命令でとても便利なのだけど、なぜだか知らないけどIntel系ではとても遅いらしい。Ryzenはそんなでもなさそうなのになぜなのだ…。
で、今回はよく考えると16bit同士の掛け算をするvpmaddwd(こっちはそこまで遅くない)で十分計算できるのに気づいたので、データの並べ替えが必要になる分だけ命令数は増えるけど置き換えてみた。するとほんのちょっとだけど速くなった。
実際の変更箇所はこんな感じ。
vpmaddwdが使えるようにするデータの並べ替えを考えるのが面倒だからと言って、あまり遅い命令を使うものではない…という反省。
(4) AVX512VNNI/AVX512VBMI対応
せっかくIceLakeが手元にあるので、新しくIceLakeで追加されたAVX512命令を使ってみようと思って対応してみた。
けど、まあほんの一部書き換えただけなので、そんなに速くならなかった。
ちなみにVNNIはVector Neural Network Instructions、VBMIはVector Bit Manipulationの略らしい。
(5) vpgatherの置き換え (AVX512)
以前も書いたけど、pmd_mtでノイズ除去するとき、データ構造と重み関数を以下のように定義して、
typedef struct {
short y; // 画素(輝度 )データ ( 0 ～ 4096 )
short cb; // 画素(色差(青))データ ( -2048 ～ 2048 )
short cr; // 画素(色差(赤))データ ( -2048 ～ 2048 )
} PIXEL_YC;
PIXEL_YC *dst; //出力フレーム
PIXEL_YC *src; //入力フレーム
PIXEL_YC *gref; //srcにガウスぼかしを適用したフレーム
int func_weight(int x) {
const double range = 4.0;
const double strength2 = strength/100.0;
const double inv_threshold2 = 1.0 / pow(2.0,threshold*0.1);
return int((1<<16)/range * strength2
* exp(-x*x * inv_threshold2));
}
あるピクセルは以下のように計算する。(cb, crも同様)
dst->y = src->y + (short)((
((src-max_w)->y - src->y)*func_weight((gref-max_w)->y - gref->y) + //上
((src-1 )->y - src->y)*func_weight((gref-1)->y - gref->y) + //左
((src+1 )->y - src->y)*func_weight((gref+1)->y - gref->y) + //右
((src+max_w)->y - src->y)*func_weight((gref+max_w)->y - gref->y) //下
) >>16 );
というわけで、1ピクセル計算するのにfunc_weightは4回x3 (y, cb,cr)=12回呼ばれることになり、結構重い。さらにuseExp:オンの場合、その名の通りfunc_weightの中にexpが使用されている。expを12回もやっていると、まあ、だいぶ重い計算になる。
もともとpmd_mtはこの重い重み計算を回避するため、最初にとりうるすべてのx(-4500～4500)についてあらかじめ重み計算を行ってテーブルを作っておき(なのでテーブルサイズはおよそ9000)、あとはテーブル参照するだけというコードになっていた。
int *pmdp; //重みの計算結果を格納したテーブル
dst->y = src->y + (short)((
((src-max_w)->y - src->y)*pmdp[(gref-max_w)->y - gref->y] + //上
((src-1 )->y - src->y)*pmdp[(gref-1)->y - gref->y] + //左
((src+1 )->y - src->y)*pmdp[(gref+1)->y - gref->y] + //右
((src+max_w)->y - src->y)*pmdp[(gref+max_w)->y - gref->y] //下
) >>16 );
こうすることで、重み計算を毎回行わなくて済み、だいぶ高速化できる。さらに、これをSIMDでやる場合には、表引き(というか間接参照)を一気にやってくれるvpgatherという命令があるので、これを使う。(ただし、HaswellやRyzenでは遅い)
int *pDst, *pTable, *pIndex;
for (int i = 0; i < 8; i++) {
pDst[i] = pTable[pIndex[i]];
}
これを、
int *pTable;
__m256i yDst, yIndex;
yDst = _mm256_i32gather_epi32(pTable, yIndex, 4);
とやるとvpgatherdd命令1発でやってくれる。このvpgatherdd命令は、Skylake以降ではそれなりに高速だが、やはり要素数分のメモリアクセスが必要なのは変わらないので、どちらかというと遅い部類の命令で、避けられるなら避けたい。
ここでひとつ気になってくるのがテーブルのサイズ。もともとはテーブルの入力としてありうる -4500 ～ 4500 のだいたい9000のテーブルサイズになっているけど、本当にこんなに大きなテーブルが必要なのかな? ということで、実際のuseExp:オンのときのテーブルの値を確認すると(下図)、0を頂点にして山型の形になっていて、山の幅がPMD_MTのパラメータthresholdの値によって変わる感じの値になっている。
特徴としては、
・最大値は16384、16bit整数で事足りる
・0をはさんで左右対称→絶対値をとれば負側の値は要らない (テーブルサイズを半分にできる)
・あることろから急激に値が下がって1を下回る→テーブルは整数で持っているので、あるところからは0が続くだけ (ということはそれ以降はテーブルは要らない)
となっていて、thresholdの値に左右されるものの、例えばデフォルトのthreshold=100であれば、必要なテーブルサイズは100程度ともともとの9000に比べると小さいことがわかる。
AVX512BWにはvpermt2w/vpermi2wという、zmmレジスタふたつに入った合計64の16bit値から、16bit単位で任意の値を32並列で引っ張ってこれるという強力なshuffle命令が追加されている。ということは、これを使えばテーブルサイズが64までなら一発で表引きできるし、テーブルサイズが128までなら2回、テーブルサイズが256でも4回この命令を使えば表引きできてしまう。
たとえば、コードにするとこんな感じ。
template< int pmdc > // pmdc * 64 = テーブルのサイズ
static __forceinline void lut_pmdp16(
const int16_t *pmdp16, //テーブル(16bit)
__m512i& z0 //表引きしたいインデックス
//表引きの結果はz0に上書きして返す
) {
if (pmdc <= 1) { //テーブルサイズが64以下の時
__m512i zpmdp16_0 = _mm512_load_si512((__m512i *)(pmdp16 + 0));
__m512i zpmdp16_1 = _mm512_load_si512((__m512i *)(pmdp16 + 32));
z0 = _mm512_permutex2var_epi16(zpmdp16_0, z0, zpmdp16_1);
} else if (pmdc <= 2) { //テーブルサイズが128以下の時
__m512i zpmdp16_0 = _mm512_load_si512((__m512i *)(pmdp16 + 0));
__m512i zpmdp16_1 = _mm512_load_si512((__m512i *)(pmdp16 + 32));
__mmask32 m0 =_mm512_cmplt_epi16_mask(z0, _mm512_set1_epi16(64));
z0 = _mm512_mask2_permutex2var_epi16(
zpmdp16_0, z0, m0, zpmdp16_1);
__m512i zpmdp16_2 = _mm512_load_si512((__m512i *)(pmdp16 + 64));
__m512i zpmdp16_3 = _mm512_load_si512((__m512i *)(pmdp16 + 96));
z0 = _mm512_mask2_permutex2var_epi16(
zpmdp16_2, z0, _knot_mask32(m0), zpmdp16_3);
} else if .... //
さすがにテーブルサイズが大きすぎると、このvpermi2w命令を何回も呼ばなければならず、テーブルサイズが大きくても1回で済むvpgatherのほうが速い、ということになってしまうが、さっき紹介したように今回はある程度テーブルサイズが小さく済ませることができる。例えばthreshold=100ならテーブルサイズが100程度なので、vpermi2wを2回実行すれば済んでしまうので、遅い部類の命令であるvpgatherddを1回使うより速くなりそう…ということで試してみた。
結果はさっきの表のとおりで、IceLake/Skylake-Xともにかなり高速化できている。
なお、気を付けないといけないのは、useExp:オフのときや、thresholdの値が大きいときはテーブルサイズもかなり大きくなり、1000以上になることもあること。こういう場合は、従来通りvpgatherを使うようにしてある。
表引きのあたりは、そもそも表にするべきかどうかなども含めて、いろいろ工夫の余地があって面白い。
(6) vpgatherをやめる (@Ryzen)
RyzenはAVX2までの対応でAVX512は実装されていないし、AVX2にはvpermi2wのような強力なshuffle命令は存在しないので、さっきの高速化は使えない。
なお悪いことに、Ryzenのgather命令は遅いらしい。
というわけでRyzenでは表引きの際gatherを使うのではなく、愚直にひとつひとつ値をとってくるようにすると、そこそこ速くなった。えええ゛…。
とはいえ、gather命令は結構便利なので、Ryzenでも今後はgatherが速くなるといいな…。
というわけでいろいろな最適化を試すことで、全体としてもそこそこ高速化することができた。
メモリアクセスを減らせば高速化できる、というのはまあいつも通り。
一方でこれまでAVX2を単純にAVX512に置き換えるだけではあまり速くならないことが多くて悲しかったが、今回はAVX512で追加された強力なshuffle命令を活用してまあまあ速くすることができた。
AVX512は単に256bit→512bitとbit幅が倍になっただけじゃなくて、便利な新命令が増えていたり、なにより大変便利なマスクレジスタという仕組みも追加されていて、非常に面白い拡張になっている。なんか無駄にいろいろな拡張の名前がある(AVX512F/AVX512DQ/AVX512BW/AVX512VNNI/AVX512VBMI/AVX512VBMI2/........と無数にある…)のは、もうちょっと整理してよね！という感じでいただけないけど、今後はIceLakeやその後継のTigerLake、あと噂ではRocketLakeなどにも搭載されて使えるCPUが広がってくるはずなので、もっと活用して速くなるようになっていくといいなと思う。
ただまあ、そもそもIntel 10nmはいつデスクトップにやってくるんですかねえ…。(RocketLakeは14nm?)
スポンサーサイト
