使用AVX2快速查找表

Question

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

Answer 1

你说得对，集合比哈斯韦尔上的PINSRD循环慢。很可能在布罗德威尔几乎收支相抵。(有关perf链接，请参见x86标记wiki，特别是Agner Fog's insn表、microarch和优化指南)

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

这种技术用于Galois乘法，当您想要将GF16值的大缓冲区中的每个元素乘以相同的值时。(例如里德-所罗门纠错码)就像我说的，利用这一点需要利用用例的特殊属性。

AVX2可以在256 b矢量的每条车道上并行执行两个128 B pshufbs。在AVX512F：__m512i _mm512_permutex2var_epi32 (__m512i a, __m512i idx, __m512i b)之前没有什么比这更好的了。有字节(vpermi2b in AVX512VBMI)、word (AVX512BW中的vpermi2w)、dword (这里是AVX512F中的vpermi2d )和qword (vpermi2q in AVX512F)的元素大小版本。这是一个完全的交叉车道洗牌，索引为两个连接的源寄存器.(比如AMD的vpperm)。

一个内部(vpermt2d / vpermi2d)后面的两个不同的指令为您提供了一个选择:用结果覆盖表或覆盖索引向量。编译器将根据输入的可重用性来选择。

你的具体案例：

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

查找表实际上是src，而不是称为lut的变量。实际上，lut正在遍历一个数组，该数组用作src的洗牌控制掩码。

为了获得最佳性能，您应该使g_tables成为一个uint8_t数组。条目只有0..63，所以它们很合适。零扩展到全寄存器的负载和普通负载一样便宜，因此它只会减少缓存占用。若要与AVX2 use一起使用它，请使用vpmovzxbd。内部很难作为负载使用，因为没有任何形式需要int64_t *，只有使用__m128i的__m256i _mm256_cvtepu8_epi32 (__m128i a)。这是一个主要的设计缺陷与本质，海事组织。

我没有什么好主意来加速你的循环。标量代码可能是这里的方法。我猜，SIMD代码将64 int16_t值调入一个新的目的地。我花了一段时间才弄明白这一点，因为我没有立刻找到if (sizeof...)行，也没有评论。)(如果你使用的是理智的变量名，而不是avx0.对小于4B的元素使用x86收集指令当然需要恼人的掩蔽。但是，您可以使用shift和OR代替pack。

您可以为sizeof(T) == sizeof(int8_t)或sizeof(T) == sizeof(int16_t)制作一个zmm版本，因为所有的src都可以放入一个或两个zmm寄存器中。

如果g_tables被用作LUT，AVX512可以轻松地使用vpermi2b。但是，如果没有AVX512，您将遇到困难，因为64字节表对于pshufb来说太大了。对每个输入车道使用pshufb的四个车道(16B)可以工作:屏蔽0.15以外的索引，然后用pcmpgtb或其他什么方式屏蔽16.31以外的索引。那你就得把四条车道合起来。所以这太糟糕了。

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

如果您愿意为g_tables的特定值手工设计洗牌，则可能会出现这种加速。从src加载向量，用编译时常量pshufb或pshufd对其进行洗牌，然后一次性存储任何连续块。(可能是pextrd或pextrq，或者更好的矢量底部的movq )。甚至是全矢量movdqu)。

实际上，使用src可以加载多个shufps向量并在它们之间进行切换。它可以很好地工作在整数数据上，除了Nehalem (也可能在Core2上)之外，没有任何减速。punpcklwd / dq / qdq (以及相应的punpckhwd等)可以交错向量元素，并给出与shufps不同的数据移动选择。

如果不需要太多的指令来构造几个完整的16B向量，那么你就处于良好的状态。

如果g_tables可以接受太多可能的值，那么就有可能编译一个定制的洗牌函数。不过，这可能真的很难做好。