AVX2 矢量化 256 位查找表（32 个无符号字符）

Posted 2023-02-16

【中文标题】AVX2 矢量化 256 位查找表（32 个无符号字符）

【英文标题】：AVX2 vectorized 256 bit lookup table (32 unsigned chars)

【发布时间】：2017-05-04 19:19:36

【问题描述】：

我是 AVX 内在函数（以及一般的 AVX）的新手，我正在尝试加快一些使用由 32 个无符号字符组成的 256 位查找表的代码。目前代码（带有虚拟数据）是这样编写的：

unsigned char lookup_table[32] =  0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 ;
unsigned char result[8];
unsigned char indices[8] =  0, 4, 8, 12, 16, 20, 24, 28;
for(int i = 0; i < 8; i++)

    result[i] = lookup_table[indices[i]];

效果很好并导致将以下内容放入“结果”中：

0, 4, 8, 12, 16, 20, 24, 28

为了加快速度，我已将上述代码替换为以下 AVX 指令：

unsigned char lookup_table[32] =  0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 ;
unsigned char result[8];
unsigned char indices[8] =  0, 4, 8, 12, 16, 20, 24, 28;
__m256i avxTable = _mm256_loadu_si256((__m256i*)&table);
__m256i avxIndices = _mm256_loadu_si256((__m256i*)&indices);

__m256i avxResult= _mm256_shuffle_epi8(avxTable , avxIndices);

这会导致以下输出：

0, 4, 8, 12, 0, 4, 8, 12

我收集到的是 _mm256_shuffle_epi8 内部将索引与 0X0F （根据 https://software.intel.com/en-us/node/524017 的伪代码），有效地使任何大于 16 的索引再次“环绕”，因此重复 (0, 4 , 8, 12)。

我是否使用了错误的 AVX 调用？我认为这应该有效吗？

【问题讨论】：

这行不通。您可以尝试使用收集指令，但它们至少加载 32 位块，因此效率值得怀疑，尤其是在 Haswell 上，收集指令无论如何都很慢。

总体思路是合理的，但您需要注意，shuffle 实际上是 2 x 128 位操作，而不是正确的 256 位 shuffle（像许多其他 AVX 指令一样）。该解决方案将比上述代码复杂得多，但仍应比标量代码更高效。

您实际上最好进行 2 x 128 位 SSE 洗牌，并使用位 4 从两个 16 路查找中选择最终输出。

【参考方案1】：

这是一个使用 SSE 而不是 AVX 的解决方案。请注意，它并行执行 16 次查找（对于 128 位 SIMD 和 8 位元素，您不能比这更少）：

#include <stdio.h>
#include <smmintrin.h> // SSE 4.1

int main()

    unsigned char lookup_table[32] =  0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
                                       16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31 ;

    unsigned char result[16];
    unsigned char indices[16] =  0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30 ;

    __m128i vIndices, vSelect, vTable0, vTable1, vResult0, vResult1, vResult;

    vIndices = _mm_loadu_si128((__m128i *)&indices);
    vSelect = _mm_cmpgt_epi8(vIndices,  _mm_set1_epi8(15));
    vTable0 = _mm_loadu_si128((__m128i *)&lookup_table[0]);
    vTable1 = _mm_loadu_si128((__m128i *)&lookup_table[16]);
    vResult0 = _mm_shuffle_epi8(vTable0, vIndices);
    vResult1 = _mm_shuffle_epi8(vTable1, vIndices);
    vResult = _mm_blendv_epi8(vResult0, vResult1, vSelect);
    _mm_storeu_si128((__m128i *)result, vResult);

    printf("%vd\n", vResult);
    return 0;

编译测试：

$ gcc -Wall test_lut.c -msse4 && ./a.out 
0 4 8 12 16 20 24 28 2 6 10 14 18 22 26 30

【讨论】：

谢谢，保罗，这非常有帮助（而且效果很好！）。根据您之前的 cmets，我已经意识到我需要形成一个掩码来指示哪些索引大于 16，但不知道 blendv 指令在那里，这对我来说是缺失的部分。再次感谢您的帮助！

@Paul R 如果我有 2048 个元素的查找表，你有什么建议？

@StNickolay：如果你有 AVX2，那么使用聚集的负载，否则你可能不走运。

@Paul R 你的意思是，用 SSE 实现复杂的逻辑并只使用普通的 C 代码没有意义吗？

@StNickolay：是的，对于较大的 LUT 操作，随着表大小的增长，您从 SIMD 中获得的好处将越来越少。对于 2048 年，您可能已经过了临界点。