C++
SIMD
AVX2

AVX2を使ったテーブル引き処理について

More than 1 year has passed since last update.

AVX2 intrinsics 関数を使ったテーブル引き処理について解説します。

1バイトインデックスによる1バイト要素のテーブル引き処理

__m256i型のサイズは256ビットで1バイト型(char, unsigned char)の情報を32個格納する事が出来ます。

__m256i型の変数の32個のそれぞれの 0~255 の値でテーブル引きをした値を、__m256i型の値で返す処理の実装を紹介します。

単純な実装

以下は intrinsincs 関数を使わずに書いたものです。
VC++では使えますが、GCCではコンパイルエラーが起きました。
__m256i型の共用体のメンバーアクセスは移植性が低いみたいですね。

#ifdef _MSC_VER
#define YESINLINE __forceinline
#else
#define YESINLINE __attribute__((always_inline)) inline
#endif

static YESINLINE __m256i
ymm_u8lookup_naive(const uint8_t table[256], __m256i idx)
{
    __m256i ret;
    for (int i = 0; i<32; ++i) {
        ret.m256i_u8[i] = table[idx.m256i_u8[i]];
    }
    return ret;
}

_mm256_i32gather_epi32 関数を使った実装

次に _mm256_i32gather_epi32 関数を使った実装を紹介します。

static YESINLINE __m256i
ymm_u8lookup_avx2gather(const uint8_t* lut, __m256i vindex) {

    __m256i lo = _mm256_unpacklo_epi8(vindex, _mm256_setzero_si256());
    __m256i hi = _mm256_unpackhi_epi8(vindex, _mm256_setzero_si256());
    __m256i idx0 = _mm256_unpacklo_epi16(lo, _mm256_setzero_si256());
    __m256i idx1 = _mm256_unpackhi_epi16(lo, _mm256_setzero_si256());
    __m256i idx2 = _mm256_unpacklo_epi16(hi, _mm256_setzero_si256());
    __m256i idx3 = _mm256_unpackhi_epi16(hi, _mm256_setzero_si256());

    const int* base = (const int*)(lut - 3);
    __m256i nidx0 = _mm256_i32gather_epi32(base, idx0, 1);
    __m256i nidx1 = _mm256_i32gather_epi32(base, idx1, 1);
    __m256i nidx2 = _mm256_i32gather_epi32(base, idx2, 1);
    __m256i nidx3 = _mm256_i32gather_epi32(base, idx3, 1);

    nidx0 = _mm256_srli_epi32(nidx0, 24);
    nidx1 = _mm256_srli_epi32(nidx1, 24);
    nidx2 = _mm256_srli_epi32(nidx2, 24);
    nidx3 = _mm256_srli_epi32(nidx3, 24);

    nidx0 = _mm256_packus_epi32(nidx0, nidx1);
    nidx2 = _mm256_packus_epi32(nidx2, nidx3);
    nidx0 = _mm256_packus_epi16(nidx0, nidx2);

    __m256i ret = nidx0;
    return ret;
}

動作確認を行ったところ、Haswellではintrinsicsを使わない単純な実装に比べると少し速く、Skylakeでは倍ぐらい速くなっていました。とはいえかなり遅めな処理である事には変わりがありません。

_mm256_shuffle_epi8 関数等を使った実装

次に紹介する方法は 16要素単位でシャッフルする処理を16回繰り返す事によって1バイトのテーブル引きを行います。

template <unsigned N>
YESINLINE __m256i
ymm_u8lookup_avx2shuffle(
//  const __m128i* lut,
    const __m256i* lut,
    __m256i vindex,
    __m256i m256i_u8_all_16,
    __m256i m256i_u8_all_112
) {
    static_assert(N != 0, "N must not be 0.");
    static_assert(N <= 16, "N must be less than or equal to 16.");

    // a heck a lot of instructions needed...
    //LOOKUP(0)
//  __m256i t = _mm256_broadcastsi128_si256(lut[0]);
    __m256i t = _mm256_loadu_si256(lut + 0);
    __m256i tmp = _mm256_adds_epu8(vindex, m256i_u8_all_112);
    __m256i s = _mm256_sub_epi8(vindex, m256i_u8_all_16);
    __m256i ret = _mm256_shuffle_epi8(t, tmp);
    if (N == 1) return ret;

#define LOOKUP(idx) \
    t = _mm256_loadu_si256(lut + idx);\
    tmp = _mm256_adds_epu8(s, m256i_u8_all_112);\
    s = _mm256_sub_epi8(s, m256i_u8_all_16);\
    tmp = _mm256_shuffle_epi8(t, tmp);\
    ret = _mm256_or_si256(ret, tmp); \
    if (idx + 1 == N) return ret;

    LOOKUP(1)
    LOOKUP(2)
    LOOKUP(3)
    LOOKUP(4)
    LOOKUP(5)
    LOOKUP(6)
    LOOKUP(7)
    LOOKUP(8)
    LOOKUP(9)
    LOOKUP(10)
    LOOKUP(11)
    LOOKUP(12)
    LOOKUP(13)
    LOOKUP(14)
    LOOKUP(15)
#undef LOOKUP
}

引数の m256i_u8_all_16m256i_u8_all_112 には下記のようにして設定した値を渡して実行します。

__m256i m256i_u8_all_16 = _mm256_set1_epi8(0x10);
__m256i m256i_u8_all_112 = _mm256_set1_epi8(112);

テーブル引きする値の範囲が 0~255 ではなくもっと小さい範囲の場合は、テンプレートパラメータの N を調整する事により、実行される命令数を減らして処理速度を上げる事が出来ます。

なお上記の実装では 256バイトのテーブルをそのまま使えません。下記の関数で別のテーブル(__m256iレジスタ16個分の512バイトのテーブル)に変換して使う必要があります。

__forceinline void setLUT(__m256i lut[16], const uint8_t table[256]) {
    __m128i val;
    for (size_t i = 0; i < 16; ++i) {
        val = _mm_loadu_si128((__m128i*)(table + i * 16));
        lut[i] = _mm256_broadcastsi128_si256(val);
    }
}

テーブル引き処理の中で _mm256_loadu_si256 関数ではなくて _mm256_broadcastsi128_si256 関数を使えば、いちいち変形した別のテーブルを用意する必要は有りませんが、そのintrinsic関数を使うとVC++では不要であるはずの vmovdqu 命令が余分に生成されてしまい残念な事に遅くなってしまいます。
https://connect.microsoft.com/VisualStudio/Feedback/Details/790352
https://connect.microsoft.com/VisualStudio/Feedback/Details/1287053
https://developercommunity.visualstudio.com/content/problem/1496/calling-avx2-intrinsics-function-mm256-broadcastsi.html

Haswellではこの方法が他の方法に比べると速いようですが、Skylakeでは _mm256_i32gather_epi32 関数を使った方法の方が少し速いようです。
ただしケースバイケースで周りのコードに依存するのか?(知識不足で原因が分かりません)Skylakeでもgatherを使うと遅い事があります。

後書き

AVX2 を使って工夫してもテーブル参照は遅い印象を受けます。この部分が足を引っ張って処理速度を上げられないコードは結構世の中にあると思います。

今回紹介した1バイトインデックスによる1バイト要素のテーブル引き処理だと、用途によっては精度が足りない事が多いと思うので、下記の処理も必要になるのではないかと思います。

  • 1バイトインデックスによる2バイト要素のテーブル引き処理
  • 1バイトインデックスによる4バイト要素のテーブル引き処理
  • 2バイトインデックスによる1バイト要素のテーブル引き処理
  • 2バイトインデックスによる2バイト要素のテーブル引き処理
  • 2バイトインデックスによる4バイト要素のテーブル引き処理
  • 4バイトインデックスによる1バイト要素のテーブル引き処理
  • 4バイトインデックスによる2バイト要素のテーブル引き処理
  • 4バイトインデックスによる4バイト要素のテーブル引き処理

どのテーブル引きにしろ _mm256_i32gather_epi32 intrinsic 関数を使って出来ますが、あまり速くないので悩みどころです。

参考にした(ような気がする)情報

http://stackoverflow.com/a/35790337/4699324
http://qiita.com/tanakmura/items/094757927592b6a54237
http://qiita.com/tanakmura/items/6c8b76f5228a2b772d85
http://www.agner.org/optimize/instruction_tables.pdf
https://msdn.microsoft.com/en-us/library/gg466498(v=vs.100).aspx
https://chessprogramming.wikispaces.com/XOP#Instructions-Packed%20Permute%20Bytes