前の記事であるUTF-8のコードポイントはどうやってもっと高速に数えるかで AVX2 での高速化をやりましたが、今ではさらにベクタの長い AVX-512 というものがあるので、これでもやってみようかと思います。なお、想定するプロセッサは Skylake-X です。

まず、zmm レジスタ向けの水平加算関数を先に書いておきます。

inline int32_t avx512_horizontal_sum_epi8(__m512i x)
{
    __m512i sumhi = _mm512_unpackhi_epi8(x, _mm512_setzero_si512());
    __m512i sumlo = _mm512_unpacklo_epi8(x, _mm512_setzero_si512());
    __m512i sum16x32 = _mm512_add_epi16(sumhi, sumlo);
    __m256i sum16x16 = _mm256_add_epi16(_mm512_castsi512_si256(sum16x32), _mm512_extracti64x4_epi64(sum16x32, 1));
    __m128i sum16x8 = _mm_add_epi16(_mm256_castsi256_si128(sum16x16), _mm256_extracti128_si256(sum16x16, 1));
    __m128i sum16x4 = _mm_add_epi16(sum16x8, _mm_srli_si128(sum16x8, 8));
    uint64_t tmp = _mm_cvtsi128_si64(sum16x4);
    tmp += (tmp >> 32);
    tmp += (tmp >> 16);
    return tmp & 0xffff;
}

前置き

AVX-512 が AVX2 までと異なる点として、ベクタ比較命令の結果やブレンド命令の制御オペランドなどの、実質的に bool のベクタであるような値を保持するためにマスクレジスタというものを使うようになった、という点が挙げられます。

前の記事では最終的に VPCMPGTB 命令の結果を VPSUBB に直接渡すという手法に着地しましたが、これは VPCMPGTB の結果が ymm レジスタに入るからできたことです。 AVX-512 では結果がマスクレジスタに入るせいでこの手法は直接使えないため、別の方法を考える必要があります。

POPCNT に帰着する

AVX-512 の比較命令の結果として得られるマスクレジスタは前述のとおり bool のベクタですが、最終的にやりたいことは true の数を数えることです。これはつまり立っているビットを数えることに他ならないので、汎用レジスタにコピーして POPCNT する、というのが自然な発想でしょう。というわけでこんなコードが出来上がります。

size_t avx512_cmpgt_popcnt(const char *p, size_t sz)
{
    size_t result = 0;
    for (size_t i = 0; i < sz; i += 64) {
        __m512i s = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i));
        __mmask64 m = _mm512_cmpgt_epi8_mask(s, _mm512_set1_epi8(-0x41));
        result += _mm_popcnt_u64(m);
    }
    return result;
}

ループが一重で済むので見た目がシンプルなのが利点です。

VPMOVM2B を使う

マスクレジスタの値を AVX2 風のベクタの値に変換する VPMOVM2B という命令があります。これを使うと1命令増えるものの AVX2 の時と全く同じ手法が使えます。

size_t avx512_cmpgt_movm2b_sub(const char *p, size_t sz)
{
    size_t result = 0;
    for (size_t i = 0; i < sz;) {
        __m512i sum = _mm512_setzero_si512();
        size_t j = 0;
        size_t limit = std::min<size_t>(255 * 64, sz - i);
        for (; j < limit; j += 64) {
            __m512i s = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i + j));
            __mmask64 m = _mm512_cmpgt_epi8_mask(s, _mm512_set1_epi8(-0x41));
            sum = _mm512_sub_epi8(sum, _mm512_movm_epi8(m));
        }
        i += j;
        result += avx512_horizontal_sum_epi8(sum);
    }
    return result;
}

VPMOVDQU8 を使う

前掲の VPMOVM2B 命令ですが、みなさんお世話になっている Agner Fog さんが出している最適化情報リソースの命令表 (4. Instruction tables) を見ると、いかにも単純な命令なのになんとレイテンシが3もあります。一方、VPMOVDQU8 というマスク付き mov でも同じことができて、こちらは普通にレイテンシが1です(正確には命令表には 8/16 の命令についてはは掲載されていないのですが、 32/64 の命令は掲載されていてそれと異なる理由はなさそうなのでこう書いています)。というわけで VPMOVDQU8 を使うと VPMOVM2B よりレイテンシが改善されそうに思われます。

size_t avx512_cmpgt_movdqu8_add(const char *p, size_t sz)
{
    size_t result = 0;
    for (size_t i = 0; i < sz;) {
        __m512i sum = _mm512_setzero_si512();
        size_t j = 0;
        size_t limit = std::min<size_t>(255 * 64, sz - i);
        for (; j < limit; j += 64) {
            __m512i s = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i + j));
            __mmask64 m = _mm512_cmpgt_epi8_mask(s, _mm512_set1_epi8(-0x41));
            sum = _mm512_add_epi8(sum, _mm512_maskz_mov_epi8(m, _mm512_set1_epi8(1)));
        }
        i += j;
        result += avx512_horizontal_sum_epi8(sum);
    }
    return result;
}

マスク付き加算を使う

上の2つではマスクレジスタをベクタレジスタに変換してから加算してましたが、加算に対してもマスクが使えるのだから素直にマスク付き加算をしてしまえば話が単純です。

size_t avx512_cmpgt_maskadd(const char *p, size_t sz)
{
    size_t result = 0;
    for (size_t i = 0; i < sz;) {
        __m512i sum = _mm512_setzero_si512();
        size_t j = 0;
        size_t limit = std::min<size_t>(255 * 64, sz - i);
        for (; j < limit; j += 64) {
            __m512i s = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i + j));
            __mmask64 m = _mm512_cmpgt_epi8_mask(s, _mm512_set1_epi8(-0x41));
            sum = _mm512_mask_add_epi8(sum, m, sum, _mm512_set1_epi8(1));
        }
        i += j;
        result += avx512_horizontal_sum_epi8(sum);
    }
    return result;
}

ただし、この手法を取る場合、加算のところに注意が必要かもしれません。マスクなし命令とマスク付き命令とでレイテンシが異なる可能性があるからです。実際、サイボウズ・ラボの光成滋生さんは AVX-512 詳解の発表においてマスクを使うと「少し低速」と言っています

とはいえ追加のレイテンシが複数クロックかかるとも思えないので、レイテンシが1から2になるものとしてそれを回避するように手動で2倍にアンローリングをかけてみましょう(とりあえずこの部分にはコンパイラによるアンローリングは期待しないものとします)

size_t avx512_cmpgt_maskadd128(const char *p, size_t sz)
{
    size_t result = 0;
    for (size_t i = 0; i < sz;) {
        __m512i sum0 = _mm512_setzero_si512();
        __m512i sum1 = _mm512_setzero_si512();
        size_t j = 0;
        size_t limit = std::min<size_t>(255 * 128, sz - i);
        for (; j < limit; j += 128) {
            __m512i s0 = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i + j));
            __m512i s1 = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i + j + 64));
            __mmask64 m0 = _mm512_cmpgt_epi8_mask(s0, _mm512_set1_epi8(-0x41));
            __mmask64 m1 = _mm512_cmpgt_epi8_mask(s1, _mm512_set1_epi8(-0x41));
            sum0 = _mm512_mask_add_epi8(sum0, m0, sum0, _mm512_set1_epi8(1));
            sum1 = _mm512_mask_add_epi8(sum1, m1, sum1, _mm512_set1_epi8(1));
        }
        i += j;
        result += avx512_horizontal_sum_epi8(sum0);
        result += avx512_horizontal_sum_epi8(sum1);
    }
    return result;
}

VPSHUFB に回帰する

今までマスクレジスタをどうやってベクタレジスタに変換するかを考えていましたが、大元の記事に書いてある手法である VPSHUFB を使った方法なら結果は最初からベクタレジスタで得られるので、単純に足すだけです。とりあえずこれも実装しておきましょう。

余談ですが、intrinsic としては \_mm_setr_epi8 や \_mm128_setr_epi8 はあっても何故か \_mm512_setr_epi8 はなくて \_mm512_set_epi8 しかなく、面食らいました。Intel の intrinsic guide にもないんですが、なんでですかね…?

size_t avx512_pshufb_add(const char *p, size_t sz)
{
    size_t result = 0;
    for (size_t i = 0; i < sz;) {
        __m512i sum = _mm512_setzero_si512();
        size_t j = 0;
        size_t limit = std::min<size_t>(255 * 64, sz - i);
        for (; j < limit; j += 64) {
            const __m512i table = _mm512_set_epi8(
                1, 1, 1, 1,             // 0xF .. 0xC
                0, 0, 0, 0,             // 0xB .. 0x8
                1, 1, 1, 1, 1, 1, 1, 1, // 0x7 ..
                1, 1, 1, 1,             // 0xF .. 0xC
                0, 0, 0, 0,             // 0xB .. 0x8
                1, 1, 1, 1, 1, 1, 1, 1, // 0x7 ..
                1, 1, 1, 1,             // 0xF .. 0xC
                0, 0, 0, 0,             // 0xB .. 0x8
                1, 1, 1, 1, 1, 1, 1, 1, // 0x7 ..
                1, 1, 1, 1,             // 0xF .. 0xC
                0, 0, 0, 0,             // 0xB .. 0x8
                1, 1, 1, 1, 1, 1, 1, 1  // 0x7 ..
            );
            __m512i s = _mm512_load_si512(reinterpret_cast<const __m512i *>(p + i + j));
            s = _mm512_and_si512(_mm512_srli_epi16(s, 4), _mm512_set1_epi8(0x0F));
            s = _mm512_shuffle_epi8(table, s);
            sum = _mm512_add_epi8(sum, s);
        }
        i += j;
        result += avx512_horizontal_sum_epi8(sum);
    }
    return result;
}

速いのはどれか

で、結局どれが速いのかですが…実機が無いので分かりません。が、それで終わるのもどうかと思うので、ひとまずコンパイラが吐いたバイナリを眺めてみます。コンパイルオプションは /O2 /clang:-march=skylake-avx512 /clang:-mtune=skylake-avx512 です。

まず、avx512_cmpgt_popcnt をループアンローリングしてくれていません。 #pragma unroll を付けると、オプティマイザがアンロールを実行できなかったと警告を出してきます。これを他のルーチンと同様な形の二重ループにするとアンローリングしてくれるのですが、二重ループじゃないとアンローリングできないとかじゃないよなぁ…

マスクレジスタを元にベクタ加算する手法(avx512_cmpgt_movm2b_sub、avx512_cmpgt_movdqu8_add、avx512_cmpgt_maskadd、avx512_cmpgt_maskadd128 の4つ)を見てみると、全てマスクを VMOVDQU8 でベクタに変換して加算するバイナリ(つまり avx512_cmpgt_movdqu8_add で期待したのと等価な処理)に変換されていました。VPMOVM2B が VMOVDQU8 に変換されるところまではある程度予測の範囲内でしたが、マスク付き加算のコードを置き換えてくるとは思いませんでした。下手な考え休むに似たりということのようです。

まとめ

  • AVX-512 でのちょっとした最適化手法について書きましたが、あまりうまく行った気はしません。
  • Clang すごいっすね(10日ぶり2回目)

Appendix

ソースコード(Windows + Visual Studio 2017 + LLVM integration 向けなので Unix 向けには多少修正する必要があります)

実機は持ってませんが Intel SDE で正しい結果が得られることは確認してあります。

Trackback

only 1 comment untill now

  1. 公開いただいているコードを流させていただきました。ご指摘のように、avx512_cmpgt_movm2b_sub、avx512_cmpgt_movdqu8_add、avx512_cmpgt_maskadd、avx512_cmpgt_maskadd128の4つはほぼほぼ同じ速度のようで、かつこれが最速となるようです。vmovdqu8が最速だと見抜いてる(?)clangはたしかにすごい…

    サイズの大きいものは、メモリ帯域に引っかかってしまっているように思います。

    Win10 x64 / 電源オプション 高パフォーマンス
    i9 7980XE @ 4.0GHz (AVX Offsetなし)
    DDR4-3600, 4ch
    VC++2019 + LLVM integration (LLVM+clang 8.0)
    /O2 /clang:-march=skylake-avx512 /clang:-mtune=skylake-avx512

    size 16384 loop 1
    avx_count_utf8_codepoint : 2.370us 6.9GB/s 18.52clk/32B
    avx_count_utf8_codepoint_loopend: 0.395us 41.5GB/s 3.09clk/32B
    opt_innermost_content : 0.395us 41.5GB/s 3.09clk/32B
    opt_innermost_content_loopend : 0.395us 41.5GB/s 3.09clk/32B
    avx512_cmpgt_popcnt : 0.790us 20.7GB/s 6.17clk/32B
    avx512_cmpgt_popcnt_nestedloop : 0.000us infGB/s 0.00clk/32B
    avx512_cmpgt_movm2b_sub : 0.000us infGB/s 0.00clk/32B
    avx512_cmpgt_movdqu8_add : 0.395us 41.5GB/s 3.09clk/32B
    avx512_cmpgt_maskadd : 0.395us 41.5GB/s 3.09clk/32B
    avx512_cmpgt_maskadd128 : 0.395us 41.5GB/s 3.09clk/32B
    avx512_pshufb_add : 0.395us 41.5GB/s 3.09clk/32B
    size 16384 loop 10000000
    avx_count_utf8_codepoint : 0.406us 40.3GB/s 3.17clk/32B
    avx_count_utf8_codepoint_loopend: 0.199us 82.5GB/s 1.55clk/32B
    opt_innermost_content : 0.313us 52.3GB/s 2.45clk/32B
    opt_innermost_content_loopend : 0.140us 117.3GB/s 1.09clk/32B
    avx512_cmpgt_popcnt : 0.169us 97.0GB/s 1.32clk/32B
    avx512_cmpgt_popcnt_nestedloop : 0.104us 157.4GB/s 0.81clk/32B
    avx512_cmpgt_movm2b_sub : 0.111us 147.7GB/s 0.87clk/32B
    avx512_cmpgt_movdqu8_add : 0.112us 146.9GB/s 0.87clk/32B
    avx512_cmpgt_maskadd : 0.111us 148.1GB/s 0.86clk/32B
    avx512_cmpgt_maskadd128 : 0.110us 149.6GB/s 0.86clk/32B
    avx512_pshufb_add : 0.138us 118.7GB/s 1.08clk/32B
    size 229376 loop 1000000
    avx_count_utf8_codepoint : 5.980us 38.4GB/s 3.34clk/32B
    avx_count_utf8_codepoint_loopend: 3.395us 67.6GB/s 1.89clk/32B
    opt_innermost_content : 4.616us 49.7GB/s 2.58clk/32B
    opt_innermost_content_loopend : 2.253us 101.8GB/s 1.26clk/32B
    avx512_cmpgt_popcnt : 2.730us 84.0GB/s 1.52clk/32B
    avx512_cmpgt_popcnt_nestedloop : 1.732us 132.4GB/s 0.97clk/32B
    avx512_cmpgt_movm2b_sub : 1.562us 146.9GB/s 0.87clk/32B
    avx512_cmpgt_movdqu8_add : 1.518us 151.1GB/s 0.85clk/32B
    avx512_cmpgt_maskadd : 1.528us 150.2GB/s 0.85clk/32B
    avx512_cmpgt_maskadd128 : 1.496us 153.3GB/s 0.83clk/32B
    avx512_pshufb_add : 1.899us 120.8GB/s 1.06clk/32B
    size 6291456 loop 40000
    avx_count_utf8_codepoint : 195.752us 32.1GB/s 3.98clk/32B
    avx_count_utf8_codepoint_loopend: 194.159us 32.4GB/s 3.95clk/32B
    opt_innermost_content : 194.098us 32.4GB/s 3.95clk/32B
    opt_innermost_content_loopend : 193.907us 32.4GB/s 3.95clk/32B
    avx512_cmpgt_popcnt : 193.948us 32.4GB/s 3.95clk/32B
    avx512_cmpgt_popcnt_nestedloop : 186.218us 33.8GB/s 3.79clk/32B
    avx512_cmpgt_movm2b_sub : 177.776us 35.4GB/s 3.62clk/32B
    avx512_cmpgt_movdqu8_add : 178.904us 35.2GB/s 3.64clk/32B
    avx512_cmpgt_maskadd : 178.943us 35.2GB/s 3.64clk/32B
    avx512_cmpgt_maskadd128 : 179.061us 35.1GB/s 3.64clk/32B
    avx512_pshufb_add : 179.052us 35.1GB/s 3.64clk/32B
    size 134217728 loop 1000
    avx_count_utf8_codepoint : 6781.612us 19.8GB/s 6.47clk/32B
    avx_count_utf8_codepoint_loopend: 6597.984us 20.3GB/s 6.29clk/32B
    opt_innermost_content : 6481.035us 20.7GB/s 6.18clk/32B
    opt_innermost_content_loopend : 6331.704us 21.2GB/s 6.04clk/32B
    avx512_cmpgt_popcnt : 6300.549us 21.3GB/s 6.01clk/32B
    avx512_cmpgt_popcnt_nestedloop : 6262.211us 21.4GB/s 5.97clk/32B
    avx512_cmpgt_movm2b_sub : 6229.497us 21.5GB/s 5.94clk/32B
    avx512_cmpgt_movdqu8_add : 6204.488us 21.6GB/s 5.92clk/32B
    avx512_cmpgt_maskadd : 6204.954us 21.6GB/s 5.92clk/32B
    avx512_cmpgt_maskadd128 : 6249.845us 21.5GB/s 5.96clk/32B
    avx512_pshufb_add : 6316.001us 21.3GB/s 6.02clk/32B

Add your comment now