使用AVX2更快的查找表

时间:2016-03-04 07:05:31

标签: algorithm performance optimization sse simd

我正在尝试加速执行一系列查找表的算法。我想使用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++];
        }
    }
}

1 个答案:

答案 0 :(得分:10)

你认为聚集比Haswell的PINSRD循环要慢。在布罗德威尔,它几乎可以收支平衡。 (另请参阅标记wiki以获取perf链接,尤其是Agner Fog's insn tables, microarch pdf, and optimization guide

如果您的索引很小,或者您可以将它们切片,pshufb可以用作具有4位索引的并行LUT 。它为您提供了16个8位表条目,但您可以使用punpcklbw之类的东西将两个字节结果向量组合成一个16位结果向量。 (LUT条目的高半部分和低半部分的单独表格,具有相同的4位索引)。

当你想要将GF16值的大缓冲区的每个元素乘以相同的值时,这种技术被用于伽罗瓦域乘法。 (例如,对于Reed-Solomon纠错码。)就像我说的,利用这一点需要利用你的用例的特殊属性。

AVX2可以在256b矢量的每个通道中并行执行两个128b pshufb s。在AVX512F之前没有什么比这更好的了:__m512i _mm512_permutex2var_epi32 (__m512i a, __m512i idx, __m512i b)。在AVX512BBMI中有字节(vpermi2b),AVX512BW中的字(vpermi2w),dword(这一个,AVX512F中的vpermi2d)和qword(AVX512F中的vpermi2q)元素大小版本。这是一个完整的跨通道shuffle,索引到两个连接的源寄存器。 (就像AMD XOP的vpperm)。

一个内在函数(vpermt2d / vpermi2d)背后的两个不同指令可让您选择用结果覆盖表,或覆盖索引向量。编译器将根据重用的输入进行选择。

您的具体案例:

*dst++ = src[*lut++];

查找表实际上是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而不是sizeof(T) == sizeof(int8_t)使用班次和OR。

您可以为sizeof(T) == sizeof(int16_t)zmm制作AVX512版本,因为所有src都适合一个或两个g_tables寄存器。

如果vpermi2b被用作LUT,AVX512可以使用pshufb轻松完成。但是,你很难用AVX512,因为64字节的表对于pshufb来说太大了。对每个输入通道使用pcmpgtb的四个通道(16B)可以工作:屏蔽掉0..15之外的指数,然后屏蔽16..31之外的指数等,使用g_tables或其他东西。然后你必须将所有四个通道组合在一起。所以这很糟糕。

可能的加速:手动设计洗牌

如果你愿意为src的具体值手动设计一个洗牌,那么就有可能加速。从pshufb加载一个向量,用编译时常量pshufdpextrd对其进行混洗,然后一次性存储任何连续的块。 (也许使用pextrqmovq,或者更好的movdqu来自矢量底部。甚至是完整的矢量src)。

实际上,shufps可以加载多个punpcklwd向量并在它们之间进行混洗。它在整数数据上工作正常,除了Nehalem(也可能在Core2上)之外没有减速。 dq / qdq / punpckhwd(和相应的g_tables等)可以交错向量元素,并为shufps提供不同的数据移动选择。

如果构建一些完整的16B向量不需要太多指令,那么你的状态就会很好。

如果{{1}}可以采用太多可能的值,则可以JIT编译自定义shuffle函数。不过,这可能很难做得很好。