AVX2技术可以加速查找表操作

Question

AVX2技术可以加速查找表操作

7

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

- ChipK

2

参见：在什么情况下，AVX2 gather指令比单独加载数据更快？。 - Paul R

我更新了我的回答。我认为你最好的希望是为特定的洗牌（g_tables内容）专门编写代码。通过一些shufps在向量之间移动数据并同时进行洗牌，以及pshufb，你可能能够设置一些向量存储。 - Peter Cordes

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Peter Cordes · Accepted Answer

你说的没错，对于Haswell架构来说，gather比PINSRD循环慢。在Broadwell架构上可能几乎持平。(另请参见x86标签维基页面中的性能链接，特别是Agner Fog的指令表、微架构pdf和优化指南)

如果您的索引很小，或者可以将它们分片，pshufb可以作为带有4位索引的并行LUT使用。它为您提供了16个8位表条目，但您可以使用诸如punpcklbw之类的内容将两个字节结果向量合并为一个16位结果向量。(对于LUT条目的高半部分和低半部分分别使用不同的表，具有相同的4位索引)。

当您想要将大缓冲区中的每个GF16值的元素乘以相同的值时，这种技术被用于Galois Field multiplies。(例如，用于Reed-Solomon错误纠正码)。就像我说过的，利用它需要利用您的用例的特殊属性。

AVX2可以在256b向量的每个通道中并行执行两个128b的pshufb操作。直到AVX512F出现之前，没有更好的选择: __m512i _mm512_permutex2var_epi32 (__m512i a, __m512i idx, __m512i b)。其中包括按字节（AVX512VBMI中的vpermi2b）、按字（AVX512BW中的vpermi2w）、按双字（即vpermi2d 在AVX512F中）和按四字（即AVX512F中的vpermi2q）元素大小版本。这是一个完整的跨通道混洗，索引到两个连接的源寄存器中（类似于AMD XOP中的vpperm）。

一个内在函数(vpermt2d / vpermi2d)背后的两个不同指令，让你可以选择用结果覆盖表格，还是用索引向量覆盖。编译器会根据哪些输入被重复使用来进行选择。

您的具体情况：

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

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

为了获得最佳性能，您应该将g_tables设置为uint8_t数组。条目仅为0..63，因此它们适合这样做。零扩展负载到完整寄存器与正常负载一样便宜，因此它只减少了缓存占用。要使用AVX2 gather，使用vpmovzxbd。内部函数很难使用作为负载，因为没有接受int64_t *的形式，只有__m256i _mm256_cvtepu8_epi32 (__m128i a)，它接受一个__m128i。这是内部函数的主要设计缺陷之一，以我个人的看法。

我没有任何加速您的循环的好思路。在这里，标量代码可能是最好的选择。SIMD代码将64个int16_t值洗牌到一个新的目标中，我猜。我花了一段时间才弄明白，因为我没有立即找到if (sizeof...)行，并且没有注释。:( 如果你使用合理的变量名而不是avx0，阅读起来会更容易。对于小于4B的元素，使用x86 gather指令肯定需要烦人的掩码。然而，您可以使用移位和OR而不是pack。

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

如果被用作LUT，AVX512可以轻松地使用< vpermi2b >完成操作。但是，如果没有AVX512，会很困难，因为64字节的表对于< pshufb >来说太大了。每个输入通道使用四条车道（16B）的< pshufb >可能有效：使用< pcmpgtb >或其他方式掩码超出0..15之外的索引，然后掩码超出16..31之外的索引，以此类推。接着必须将所有四个车道进行逻辑或运算。所以这非常糟糕。

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

如果您愿意为特定的g_tables值手动设计洗牌，那么有可能会有加速效果。从src加载一个向量，使用编译时常量pshufb或pshufd进行洗牌，然后一次性存储任何连续的块。(可以使用pextrd或pextrq，甚至更好的是从向量底部使用movq，或者甚至是全向量movdqu）。

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

如果构建几个完整的16B向量不需要太多指令，那么你就很棒了。

如果可以采用过多的可能值，可能可以JIT编译自定义洗牌函数。尽管这很难做到很好。