使用 AVX2 更快地查找表
Faster lookup tables using AVX2
我正在尝试加快执行一系列查找表的算法。我想使用 SSE2 或 AVX2。我试过使用 _mm256_i32gather_epi32 命令,但速度慢了 31%。有人对任何改进或不同的方法有任何建议吗?
时间:
C 代码 = 234
聚集 = 340
static const int32_t g_tables[2][64]; // values between 0 and 63
template <int8_t which, class T>
static void lookup_data(int16_t * dst, T * src)
{
const int32_t * lut = g_tables[which];
// Leave this code for Broadwell or Skylake since it's 31% slower than C code
// (gather is 12 for Haswell, 7 for Broadwell and 5 for Skylake)
#if 0
if (sizeof(T) == sizeof(int16_t)) {
__m256i avx0, avx1, avx2, avx3, avx4, avx5, avx6, avx7;
__m128i sse0, sse1, sse2, sse3, sse4, sse5, sse6, sse7;
__m256i mask = _mm256_set1_epi32(0xffff);
avx0 = _mm256_loadu_si256((__m256i *)(lut));
avx1 = _mm256_loadu_si256((__m256i *)(lut + 8));
avx2 = _mm256_loadu_si256((__m256i *)(lut + 16));
avx3 = _mm256_loadu_si256((__m256i *)(lut + 24));
avx4 = _mm256_loadu_si256((__m256i *)(lut + 32));
avx5 = _mm256_loadu_si256((__m256i *)(lut + 40));
avx6 = _mm256_loadu_si256((__m256i *)(lut + 48));
avx7 = _mm256_loadu_si256((__m256i *)(lut + 56));
avx0 = _mm256_i32gather_epi32((int32_t *)(src), avx0, 2);
avx1 = _mm256_i32gather_epi32((int32_t *)(src), avx1, 2);
avx2 = _mm256_i32gather_epi32((int32_t *)(src), avx2, 2);
avx3 = _mm256_i32gather_epi32((int32_t *)(src), avx3, 2);
avx4 = _mm256_i32gather_epi32((int32_t *)(src), avx4, 2);
avx5 = _mm256_i32gather_epi32((int32_t *)(src), avx5, 2);
avx6 = _mm256_i32gather_epi32((int32_t *)(src), avx6, 2);
avx7 = _mm256_i32gather_epi32((int32_t *)(src), avx7, 2);
avx0 = _mm256_and_si256(avx0, mask);
avx1 = _mm256_and_si256(avx1, mask);
avx2 = _mm256_and_si256(avx2, mask);
avx3 = _mm256_and_si256(avx3, mask);
avx4 = _mm256_and_si256(avx4, mask);
avx5 = _mm256_and_si256(avx5, mask);
avx6 = _mm256_and_si256(avx6, mask);
avx7 = _mm256_and_si256(avx7, mask);
sse0 = _mm_packus_epi32(_mm256_castsi256_si128(avx0), _mm256_extracti128_si256(avx0, 1));
sse1 = _mm_packus_epi32(_mm256_castsi256_si128(avx1), _mm256_extracti128_si256(avx1, 1));
sse2 = _mm_packus_epi32(_mm256_castsi256_si128(avx2), _mm256_extracti128_si256(avx2, 1));
sse3 = _mm_packus_epi32(_mm256_castsi256_si128(avx3), _mm256_extracti128_si256(avx3, 1));
sse4 = _mm_packus_epi32(_mm256_castsi256_si128(avx4), _mm256_extracti128_si256(avx4, 1));
sse5 = _mm_packus_epi32(_mm256_castsi256_si128(avx5), _mm256_extracti128_si256(avx5, 1));
sse6 = _mm_packus_epi32(_mm256_castsi256_si128(avx6), _mm256_extracti128_si256(avx6, 1));
sse7 = _mm_packus_epi32(_mm256_castsi256_si128(avx7), _mm256_extracti128_si256(avx7, 1));
_mm_storeu_si128((__m128i *)(dst), sse0);
_mm_storeu_si128((__m128i *)(dst + 8), sse1);
_mm_storeu_si128((__m128i *)(dst + 16), sse2);
_mm_storeu_si128((__m128i *)(dst + 24), sse3);
_mm_storeu_si128((__m128i *)(dst + 32), sse4);
_mm_storeu_si128((__m128i *)(dst + 40), sse5);
_mm_storeu_si128((__m128i *)(dst + 48), sse6);
_mm_storeu_si128((__m128i *)(dst + 56), sse7);
}
else
#endif
{
for (int32_t i = 0; i < 64; i += 4)
{
*dst++ = src[*lut++];
*dst++ = src[*lut++];
*dst++ = src[*lut++];
*dst++ = src[*lut++];
}
}
}
你说得对,收集比 Haswell 上的 PINSRD
循环慢。它可能在 Broadwell 上接近收支平衡。 (另见 x86 tag wiki for perf links, especially Agner Fog's insn tables, microarch pdf, and optimization guide)
如果您的索引很小,或者您可以将它们切片,pshufb
可以用作具有 4 位索引的并行 LUT。它为您提供 16 个 8 位 table 条目,但您可以使用 punpcklbw 之类的东西将字节结果的两个向量合并为一个 16 位结果向量。 (将 tables 用于 LUT 条目的高半部分和低半部分,具有相同的 4 位索引)。
这种技术用于伽罗华域乘法,当您想要将 GF16 值的大缓冲区的每个元素乘以相同的值时。 (例如,对于 Reed-Solomon 纠错码。)正如我所说,利用这一点需要利用您的用例的特殊属性。
AVX2 可以在 256b 向量的每个通道中并行执行两个 128b pshufb
s。没有比 AVX512F 更好的了:__m512i _mm512_permutex2var_epi32 (__m512i a, __m512i idx, __m512i b)
。有byte(vpermi2b
in AVX512VBMI)、word(vpermi2w
in AVX512BW)、dword(这一个,vpermi2d
in AVX512F)、qword(vpermi2q
in AVX512F)元素尺寸版本。这是一个完整的跨通道洗牌,索引到两个串联的源寄存器中。 (就像 AMD XOP 的 vpperm
)。
一个内部指令 (vpermt2d
/ vpermi2d
) 背后的两个不同指令让您可以选择用结果覆盖 table 或覆盖索引向量。编译器将根据重复使用的输入进行选择。
您的具体情况:
*dst++ = src[*lut++];
查找-table 实际上是 src
,而不是您调用的变量 lut
。 lut
实际上是遍历一个数组,该数组用作 src
.
的随机播放控制掩码
您应该将 g_tables
设为 uint8_t
的数组以获得最佳性能。条目仅为 0..63,因此适合。零扩展加载到完整寄存器与正常加载一样便宜,因此它只是减少了缓存占用空间。要将其用于 AVX2 收集,请使用 vpmovzxbd
。内在函数很难用作负载,因为没有采用 int64_t *
的形式,只有 __m256i _mm256_cvtepu8_epi32 (__m128i a)
采用 __m128i
。这是 IMO 内在函数的主要设计缺陷之一。
我没有任何加快循环的好主意。标量代码可能是去这里的方式。我猜 SIMD 代码将 64 int16_t 个值洗牌到一个新的目的地。我花了一段时间才弄明白,因为我没有立即找到 if (sizeof...)
行,而且没有评论。 :( 如果你使用合理的变量名,而不是 avx0
,它会更容易阅读......对小于 4B 的元素使用 x86 收集指令当然需要烦人的屏蔽。但是,你可以代替 pack
使用 shift 和 OR。
您可以为 sizeof(T) == sizeof(int8_t)
或 sizeof(T) == sizeof(int16_t)
制作一个 AVX512 版本,因为所有 src 都将适合一个或两个 zmm
寄存器。
如果 g_tables
被用作 LUT,AVX512 可以轻松完成,vpermi2b
。但是,如果没有 AVX512,你会遇到困难,因为 64 字节 table 对于 pshufb
来说太大了。为每个输入通道使用 pshufb
的四个通道 (16B) 可以工作:屏蔽掉 0..15 之外的索引,然后屏蔽 16..31 之外的索引,等等,使用 pcmpgtb
或其他东西。然后你必须将所有四个车道放在一起。所以这很糟糕。
可能的加速:手动设计随机播放
如果您愿意为 g_tables
的特定值手动设计随机播放,这样可能会加速。从 src
加载一个向量,用编译时常量 pshufb
或 pshufd
对其进行混洗,然后一次存储任何连续的块。 (也许使用 pextrd
或 pextrq
,或者更好的是从向量底部开始 movq
。甚至是完整向量 movdqu
)。
实际上,shufps
可以加载多个 src
向量并在它们之间进行混洗。它在整数数据上运行良好,除了在 Nehalem 上(也许在 Core2 上也一样)没有减速。 punpcklwd
/ dq
/ qdq
(以及相应的 punpckhwd
等)可以交错向量的元素,并为数据移动提供与 shufps 不同的选择。
如果构建几个完整的 16B 向量不需要太多指令,那么你的状态很好。
如果 g_tables
可以采用太多可能的值,则可以 JIT 编译自定义随机播放函数。不过,这可能真的很难做好。
我正在尝试加快执行一系列查找表的算法。我想使用 SSE2 或 AVX2。我试过使用 _mm256_i32gather_epi32 命令,但速度慢了 31%。有人对任何改进或不同的方法有任何建议吗?
时间: C 代码 = 234 聚集 = 340
static const int32_t g_tables[2][64]; // values between 0 and 63
template <int8_t which, class T>
static void lookup_data(int16_t * dst, T * src)
{
const int32_t * lut = g_tables[which];
// Leave this code for Broadwell or Skylake since it's 31% slower than C code
// (gather is 12 for Haswell, 7 for Broadwell and 5 for Skylake)
#if 0
if (sizeof(T) == sizeof(int16_t)) {
__m256i avx0, avx1, avx2, avx3, avx4, avx5, avx6, avx7;
__m128i sse0, sse1, sse2, sse3, sse4, sse5, sse6, sse7;
__m256i mask = _mm256_set1_epi32(0xffff);
avx0 = _mm256_loadu_si256((__m256i *)(lut));
avx1 = _mm256_loadu_si256((__m256i *)(lut + 8));
avx2 = _mm256_loadu_si256((__m256i *)(lut + 16));
avx3 = _mm256_loadu_si256((__m256i *)(lut + 24));
avx4 = _mm256_loadu_si256((__m256i *)(lut + 32));
avx5 = _mm256_loadu_si256((__m256i *)(lut + 40));
avx6 = _mm256_loadu_si256((__m256i *)(lut + 48));
avx7 = _mm256_loadu_si256((__m256i *)(lut + 56));
avx0 = _mm256_i32gather_epi32((int32_t *)(src), avx0, 2);
avx1 = _mm256_i32gather_epi32((int32_t *)(src), avx1, 2);
avx2 = _mm256_i32gather_epi32((int32_t *)(src), avx2, 2);
avx3 = _mm256_i32gather_epi32((int32_t *)(src), avx3, 2);
avx4 = _mm256_i32gather_epi32((int32_t *)(src), avx4, 2);
avx5 = _mm256_i32gather_epi32((int32_t *)(src), avx5, 2);
avx6 = _mm256_i32gather_epi32((int32_t *)(src), avx6, 2);
avx7 = _mm256_i32gather_epi32((int32_t *)(src), avx7, 2);
avx0 = _mm256_and_si256(avx0, mask);
avx1 = _mm256_and_si256(avx1, mask);
avx2 = _mm256_and_si256(avx2, mask);
avx3 = _mm256_and_si256(avx3, mask);
avx4 = _mm256_and_si256(avx4, mask);
avx5 = _mm256_and_si256(avx5, mask);
avx6 = _mm256_and_si256(avx6, mask);
avx7 = _mm256_and_si256(avx7, mask);
sse0 = _mm_packus_epi32(_mm256_castsi256_si128(avx0), _mm256_extracti128_si256(avx0, 1));
sse1 = _mm_packus_epi32(_mm256_castsi256_si128(avx1), _mm256_extracti128_si256(avx1, 1));
sse2 = _mm_packus_epi32(_mm256_castsi256_si128(avx2), _mm256_extracti128_si256(avx2, 1));
sse3 = _mm_packus_epi32(_mm256_castsi256_si128(avx3), _mm256_extracti128_si256(avx3, 1));
sse4 = _mm_packus_epi32(_mm256_castsi256_si128(avx4), _mm256_extracti128_si256(avx4, 1));
sse5 = _mm_packus_epi32(_mm256_castsi256_si128(avx5), _mm256_extracti128_si256(avx5, 1));
sse6 = _mm_packus_epi32(_mm256_castsi256_si128(avx6), _mm256_extracti128_si256(avx6, 1));
sse7 = _mm_packus_epi32(_mm256_castsi256_si128(avx7), _mm256_extracti128_si256(avx7, 1));
_mm_storeu_si128((__m128i *)(dst), sse0);
_mm_storeu_si128((__m128i *)(dst + 8), sse1);
_mm_storeu_si128((__m128i *)(dst + 16), sse2);
_mm_storeu_si128((__m128i *)(dst + 24), sse3);
_mm_storeu_si128((__m128i *)(dst + 32), sse4);
_mm_storeu_si128((__m128i *)(dst + 40), sse5);
_mm_storeu_si128((__m128i *)(dst + 48), sse6);
_mm_storeu_si128((__m128i *)(dst + 56), sse7);
}
else
#endif
{
for (int32_t i = 0; i < 64; i += 4)
{
*dst++ = src[*lut++];
*dst++ = src[*lut++];
*dst++ = src[*lut++];
*dst++ = src[*lut++];
}
}
}
你说得对,收集比 Haswell 上的 PINSRD
循环慢。它可能在 Broadwell 上接近收支平衡。 (另见 x86 tag wiki for perf links, especially Agner Fog's insn tables, microarch pdf, and optimization guide)
如果您的索引很小,或者您可以将它们切片,pshufb
可以用作具有 4 位索引的并行 LUT。它为您提供 16 个 8 位 table 条目,但您可以使用 punpcklbw 之类的东西将字节结果的两个向量合并为一个 16 位结果向量。 (将 tables 用于 LUT 条目的高半部分和低半部分,具有相同的 4 位索引)。
这种技术用于伽罗华域乘法,当您想要将 GF16 值的大缓冲区的每个元素乘以相同的值时。 (例如,对于 Reed-Solomon 纠错码。)正如我所说,利用这一点需要利用您的用例的特殊属性。
AVX2 可以在 256b 向量的每个通道中并行执行两个 128b pshufb
s。没有比 AVX512F 更好的了:__m512i _mm512_permutex2var_epi32 (__m512i a, __m512i idx, __m512i b)
。有byte(vpermi2b
in AVX512VBMI)、word(vpermi2w
in AVX512BW)、dword(这一个,vpermi2d
in AVX512F)、qword(vpermi2q
in AVX512F)元素尺寸版本。这是一个完整的跨通道洗牌,索引到两个串联的源寄存器中。 (就像 AMD XOP 的 vpperm
)。
一个内部指令 (vpermt2d
/ vpermi2d
) 背后的两个不同指令让您可以选择用结果覆盖 table 或覆盖索引向量。编译器将根据重复使用的输入进行选择。
您的具体情况:
*dst++ = src[*lut++];
查找-table 实际上是 src
,而不是您调用的变量 lut
。 lut
实际上是遍历一个数组,该数组用作 src
.
您应该将 g_tables
设为 uint8_t
的数组以获得最佳性能。条目仅为 0..63,因此适合。零扩展加载到完整寄存器与正常加载一样便宜,因此它只是减少了缓存占用空间。要将其用于 AVX2 收集,请使用 vpmovzxbd
。内在函数很难用作负载,因为没有采用 int64_t *
的形式,只有 __m256i _mm256_cvtepu8_epi32 (__m128i a)
采用 __m128i
。这是 IMO 内在函数的主要设计缺陷之一。
我没有任何加快循环的好主意。标量代码可能是去这里的方式。我猜 SIMD 代码将 64 int16_t 个值洗牌到一个新的目的地。我花了一段时间才弄明白,因为我没有立即找到 if (sizeof...)
行,而且没有评论。 :( 如果你使用合理的变量名,而不是 avx0
,它会更容易阅读......对小于 4B 的元素使用 x86 收集指令当然需要烦人的屏蔽。但是,你可以代替 pack
使用 shift 和 OR。
您可以为 sizeof(T) == sizeof(int8_t)
或 sizeof(T) == sizeof(int16_t)
制作一个 AVX512 版本,因为所有 src 都将适合一个或两个 zmm
寄存器。
如果 g_tables
被用作 LUT,AVX512 可以轻松完成,vpermi2b
。但是,如果没有 AVX512,你会遇到困难,因为 64 字节 table 对于 pshufb
来说太大了。为每个输入通道使用 pshufb
的四个通道 (16B) 可以工作:屏蔽掉 0..15 之外的索引,然后屏蔽 16..31 之外的索引,等等,使用 pcmpgtb
或其他东西。然后你必须将所有四个车道放在一起。所以这很糟糕。
可能的加速:手动设计随机播放
如果您愿意为 g_tables
的特定值手动设计随机播放,这样可能会加速。从 src
加载一个向量,用编译时常量 pshufb
或 pshufd
对其进行混洗,然后一次存储任何连续的块。 (也许使用 pextrd
或 pextrq
,或者更好的是从向量底部开始 movq
。甚至是完整向量 movdqu
)。
实际上,shufps
可以加载多个 src
向量并在它们之间进行混洗。它在整数数据上运行良好,除了在 Nehalem 上(也许在 Core2 上也一样)没有减速。 punpcklwd
/ dq
/ qdq
(以及相应的 punpckhwd
等)可以交错向量的元素,并为数据移动提供与 shufps 不同的选择。
如果构建几个完整的 16B 向量不需要太多指令,那么你的状态很好。
如果 g_tables
可以采用太多可能的值,则可以 JIT 编译自定义随机播放函数。不过,这可能真的很难做好。