LoginSignup
12
11

More than 5 years have passed since last update.

@beru(湯浅 勝久)

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

Last updated at Posted at 2016-10-05

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_16 と m256i_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

12
11
2

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
12
11