使用 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

你说得对，Gather 比 Haswell 上的 PINSRD 循环慢。在 Broadwell 上，它可能几乎是收支平衡的。 （另请参阅x86 标签 wiki 以获得性能链接，尤其是 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 pshufbs。在 AVX512F 之前没有比这更好的了：__m512i _mm512_permutex2var_epi32 (__m512i a, __m512i idx, __m512i b)。有字节（vpermi2b 在 AVX512VBMI）、字（vpermi2w 在 AVX512BW）、dword（这个，vpermi2d 在 AVX512F）和 qword（vpermi2q 在 AVX512F）元素大小版本。这是一个完整的跨通道混洗，索引到两个连接的源寄存器。 （就像 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...) 行，而且没有 cmets。 :( 如果您使用健全的变量名称，而不是 avx0，那么阅读起来会更容易...对小于 4B 的元素使用 x86 收集指令当然需要烦人的屏蔽。但是，您可以使用移位而不是 pack或。

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

---

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

---

可能的加速：手动设计随机播放

如果您愿意为g_tables 的特定值手动设计一个随机播放，那么这种方式可能会加快速度。从src 加载一个向量，使用编译时常量pshufb 或pshufd 对其进行混洗，然后一次性存储任何连续的块。 （可能使用pextrd 或pextrq，或者更好的是从矢量底部开始movq。甚至是全矢量movdqu）。

实际上，使用shufps 可以加载多个src 向量并在它们之间进行洗牌。它在整数数据上运行良好，除了在 Nehalem 上（也可能在 Core2 上）没有减速。 punpcklwd / dq / qdq（以及对应的punpckhwd等）可以交错向量的元素，并为数据移动提供与shufps不同的选择。

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

如果g_tables 可以采用太多可能的值，则可以 JIT 编译自定义 shuffle 函数。不过，这可能真的很难做好。