或るプログラマの一生

前の記事である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 で正しい結果が得られることは確認してあります。