Edited at

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