在 AVX 寄存器中旋转字节的有效方法

Question

Summary/tl;dr:除了进行 2x 移位并将结果混合在一起之外，还有什么方法可以按位旋转 YMM 寄存器中的字节（使用 AVX）？

对于 YMM 寄存器中的每 8 个字节，我需要在其中左循环 7 个字节。每个字节需要比前者多向左旋转一位。因此，第 1 个字节应旋转 0 位，第 7 个字节应旋转 6 位。

目前，我已经通过[我在此处使用 1 位循环作为示例] 将寄存器 1 位向左移动，7 位分别向右移动来实现这一点。然后我使用混合操作（内在操作 _mm256_blend_epi16）从第一个和第二个临时结果中选择正确的位，以获得我的最终旋转字节。
每个字节总共需要 2 次移位操作和 1 次混合操作，并且需要旋转 6 个字节，因此每个字节需要 18 次操作（移位和混合的性能几乎相同）。

必须有比使用 18 次操作来旋转单个字节更快的方法！

此外，我需要在新寄存器中组装所有字节。我通过将带有“set”指令的 7 个掩码加载到寄存器中来做到这一点，这样我就可以从每个寄存器中提取正确的字节。我将这些掩码与寄存器相结合，以从中提取正确的字节。之后我将单字节寄存器 XOR 在一起以获得包含所有字节的新寄存器。 这总共需要 7+7+6 次操作，因此需要另外 20 次操作（每个寄存器）。

我可以使用提取内在函数 (_mm256_extract_epi8) 来获取单个字节，然后使用 _mm256_set_epi8 来组装新寄存器，但我还不知道这是否会更快。 （英特尔内在函数指南中没有列出这些函数的性能，所以我可能在这里误解了一些东西。）

这为每个寄存器提供了总共 38 次操作，这对于在寄存器内以不同方式旋转 6 个字节似乎不太理想。

我希望更精通 AVX/SIMD 的人可以在这里指导我——无论我是否以错误的方式处理这件事——因为我觉得我现在可能正在这样做。

Answer 1

[基于第一条评论和一些编辑，得到的解决方案有点不同。我先介绍一下，然后把原来的想法留在下面]

这里的主要思想是使用乘以 2 的幂来完成移位，因为这些常数可以在向量中变化。 @harold 指出了下一个想法，即两个重复字节的乘法将自动将移出的位“旋转”回低位。

1. 将字节解包并复制为 16 位值[... d c b a] -> [... dd cc bb aa]
2. 生成一个16位常量[128 64 32 16 8 4 2 1]
3. 相乘
4. 你想要的字节是每个 16 位值的前八位，所以右移并重新打包

假设 __m128i 源（你只有 8 个字节，对吗？）：

__m128i duped = _mm_unpacklo_epi8(src, src);
__m128i res = _mm_mullo_epi16(duped, power_of_two_vector);
__m128i repacked = _mm_packus_epi16(_mm_srli_epi16(res, 8), __mm_setzero_si128());

---

[保存这个最初的想法以供比较]

这个怎么样：使用 2 的幂来完成移位，使用 16 位产品。然后 OR 产品的上半部分和下半部分来完成旋转。

1. 将字节解压缩为 16 位字。
2. 生成 16 位 [128 64 32 16 8 4 2 1]
3. 乘以 16 位字
4. 将 16 位重新打包成两个 8 位向量，一个高字节向量和一个低字节向量
5. OR 这两个向量来完成旋转。

我对可用的乘法选项和您的指令集限制有点模糊，但理想的情况是产生 16 位乘积的 8 位乘 8 位乘法。据我所知它不存在，这就是为什么我建议先解包，但我已经看到了其他巧妙的算法来做到这一点。

Answer 2

XOP instruction set 确实提供了_mm_rot_epi8()（它不是 Microsoft 特定的；它在 GCC 4.4 或更早版本中也可用，并且在最近的 clang 中也应该可用）。它可用于以 128 位为单位执行所需的任务。不幸的是，我没有支持 XOP 的 CPU，所以我无法对其进行测试。

在 AVX2 上，将 256 位寄存器分成两半，一个包含偶数字节，另一个奇数字节右移 8 位，允许 16 位向量乘法来解决问题。给定常量（使用 GCC 64 位组件数组格式）

static const __m256i epi16_highbyte = { 0xFF00FF00FF00FF00ULL,
                                        0xFF00FF00FF00FF00ULL,
                                        0xFF00FF00FF00FF00ULL,
                                        0xFF00FF00FF00FF00ULL };
static const __m256i epi16_lowbyte  = { 0x00FF00FF00FF00FFULL,
                                        0x00FF00FF00FF00FFULL,
                                        0x00FF00FF00FF00FFULL,
                                        0x00FF00FF00FF00FFULL };
static const __m256i epi16_oddmuls  = { 0x4040101004040101ULL,
                                        0x4040101004040101ULL,
                                        0x4040101004040101ULL,
                                        0x4040101004040101ULL };
static const __m256i epi16_evenmuls = { 0x8080202008080202ULL,
                                        0x8080202008080202ULL,
                                        0x8080202008080202ULL,
                                        0x8080202008080202ULL };

旋转操作可以写成

__m256i byteshift(__m256i value)
{
    return _mm256_or_si256(_mm256_srli_epi16(_mm256_mullo_epi16(_mm256_and_si256(value, epi16_lowbyte), epi16_oddmuls), 8),
                           _mm256_and_si256(_mm256_mullo_epi16(_mm256_and_si256(_mm256_srai_epi16(value, 8), epi16_lowbyte), epi16_evenmuls), epi16_highbyte));
}

这已经过验证，可以在使用 GCC-4.8.4 的 Intel Core i5-4200U 上产生正确的结果。例如，输入向量（作为单个 256 位十六进制数）

88 87 86 85 84 83 82 81 38 37 36 35 34 33 32 31 28 27 26 25 24 23 22 21 FF FE FD FC FB FA F9 F8

旋转成

44 E1 D0 58 24 0E 05 81 1C CD C6 53 A1 CC 64 31 14 C9 C4 52 21 8C 44 21 FF BF BF CF DF EB F3 F8

其中最左边的八位字节向左旋转 7 位，接下来的 6 位，依此类推；第七个八位字节不变，第八个八位字节循环 7 位，依此类推，对于所有 32 个八位字节。

我不确定上述函数定义是否能编译成最佳机器码——这取决于编译器——但我当然对它的性能感到满意。

由于您可能不喜欢上述简洁的函数格式，因此这里是程序化的扩展形式：

static __m256i byteshift(__m256i value)
{
    __m256i low, high;
    high = _mm256_srai_epi16(value, 8);
    low = _mm256_and_si256(value, epi16_lowbyte);
    high = _mm256_and_si256(high, epi16_lowbyte);
    low = _mm256_mullo_epi16(low, epi16_lowmuls);
    high = _mm256_mullo_epi16(high, epi16_highmuls);
    low = _mm256_srli_epi16(low, 8);
    high = _mm256_and_si256(high, epi16_highbyte);
    return _mm256_or_si256(low, high);
}

---

在评论中，Peter Cordes 建议将srai+and 替换为srli，并可能将最后的and+or 替换为blendv。前者很有意义，因为它纯粹是一种优化，但后者可能不会（在当前的 Intel CPU 上！）实际上更快。

我尝试了一些微基准测试，但无法获得可靠的结果。我通常在 x86-64 上使用 TSC，并使用存储在数组中的输入和输出进行数十万次测试的中位数。

我认为，如果我仅在此处列出变体是最有用的，因此任何需要此类功能的用户都可以对其实际工作负载进行一些基准测试，并测试是否有任何可测量的差异。

我也同意他的建议，用odd和even代替high和low，但是请注意，由于向量中的第一个元素编号为元素0，所以第一个元素是偶数，第二个奇数，以此类推。

#include <immintrin.h>

static const __m256i epi16_oddmask  = { 0xFF00FF00FF00FF00ULL,
                                        0xFF00FF00FF00FF00ULL,
                                        0xFF00FF00FF00FF00ULL,
                                        0xFF00FF00FF00FF00ULL };
static const __m256i epi16_evenmask = { 0x00FF00FF00FF00FFULL,
                                        0x00FF00FF00FF00FFULL,
                                        0x00FF00FF00FF00FFULL,
                                        0x00FF00FF00FF00FFULL };
static const __m256i epi16_evenmuls = { 0x4040101004040101ULL,
                                        0x4040101004040101ULL,
                                        0x4040101004040101ULL,
                                        0x4040101004040101ULL };
static const __m256i epi16_oddmuls  = { 0x8080202008080202ULL,
                                        0x8080202008080202ULL,
                                        0x8080202008080202ULL,
                                        0x8080202008080202ULL };

/* Original version suggested by Nominal Animal. */
__m256i original(__m256i value)
{
    return _mm256_or_si256(_mm256_srli_epi16(_mm256_mullo_epi16(_mm256_and_si256(value, epi16_evenmask), epi16_evenmuls), 8),
                           _mm256_and_si256(_mm256_mullo_epi16(_mm256_and_si256(_mm256_srai_epi16(value, 8), epi16_evenmask), epi16_oddmuls), epi16_oddmask));
}

/* Optimized as suggested by Peter Cordes, without blendv */
__m256i no_blendv(__m256i value)
{
    return _mm256_or_si256(_mm256_srli_epi16(_mm256_mullo_epi16(_mm256_and_si256(value, epi16_evenmask), epi16_evenmuls), 8),
                           _mm256_and_si256(_mm256_mullo_epi16(_mm256_srli_epi16(value, 8), epi16_oddmuls), epi16_oddmask));
}

/* Optimized as suggested by Peter Cordes, with blendv.
 * This is the recommended version. */
__m256i optimized(__m256i value)
{
    return _mm256_blendv_epi8(_mm256_srli_epi16(_mm256_mullo_epi16(_mm256_and_si256(value, epi16_evenmask), epi16_evenmuls), 8),
                              _mm256_mullo_epi16(_mm256_srli_epi16(value, 8), epi16_oddmuls), epi16_oddmask);
}

以下是以显示各个操作的方式编写的相同函数。虽然它根本不影响理智的编译器，但我已经标记了函数参数和每个临时值const，以便很明显如何将每个插入到后续表达式中，以将函数简化为上述简洁形式。

__m256i original_verbose(const __m256i value)
{
    const __m256i odd1  = _mm256_srai_epi16(value, 8);
    const __m256i even1 = _mm256_and_si256(value, epi16_evenmask);
    const __m256i odd2  = _mm256_and_si256(odd1, epi16_evenmask);
    const __m256i even2 = _mm256_mullo_epi16(even1, epi16_evenmuls);
    const __m256i odd3  = _mm256_mullo_epi16(odd3, epi16_oddmuls);
    const __m256i even3 = _mm256_srli_epi16(even3, 8);
    const __m256i odd4  = _mm256_and_si256(odd3, epi16_oddmask);
    return _mm256_or_si256(even3, odd4);
}

__m256i no_blendv_verbose(const __m256i value)
{
    const __m256i even1 = _mm256_and_si256(value, epi16_evenmask);
    const __m256i odd1  = _mm256_srli_epi16(value, 8);
    const __m256i even2 = _mm256_mullo_epi16(even1, epi16_evenmuls);
    const __m256i odd2  = _mm256_mullo_epi16(odd1, epi16_oddmuls);
    const __m256i even3 = _mm256_srli_epi16(even2, 8);
    const __m256i odd3  = _mm256_and_si256(odd2, epi16_oddmask);
    return _mm256_or_si256(even3, odd3);
}

__m256i optimized_verbose(const __m256i value)
{
    const __m256i even1 = _mm256_and_si256(value, epi16_evenmask);
    const __m256i odd1  = _mm256_srli_epi16(value, 8);
    const __m256i even2 = _mm256_mullo_epi16(even1, epi16_evenmuls);
    const __m256i odd2  = _mm256_mullo_epi16(odd1, epi16_oddmuls);
    const __m256i even3 = _mm256_srli_epi16(even2, 8);
    return _mm256_blendv_epi8(even3, odd2, epi16_oddmask);
}

我个人确实最初以上述详细形式编写我的测试函数，因为形成简洁版本是一组微不足道的复制粘贴。但是，我确实测试了这两个版本以验证不会引入任何错误，并保持详细版本可访问（作为评论左右），因为简洁版本基本上是只写的。编辑详细版本，然后将其简化为简洁形式，比尝试编辑简洁版本要容易得多。