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

Efficient way of rotating a byte inside an AVX register

Summary/tl;dr: 除了进行 2 倍移位和混合结果之外,还有什么方法可以按位旋转 YMM 寄存器中的字节(使用 AVX)一起?

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

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

一定有比使用 18 次操作循环一个字节更快的方法来做到这一点!

此外,我需要assemble 新寄存器中的所有字节。为此,我使用 "set" 指令将 7 个掩码加载到寄存器中,这样我就可以从每个寄存器中提取正确的字节。 I AND 这些掩码与寄存器以从中提取正确的字节。之后我将单字节寄存器异或在一起以获得包含所有字节的新寄存器。 这总共需要 7+7+6 次操作,所以另外 20 次操作(每个寄存器)。

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

这给出了每个寄存器总共 38 次操作,这似乎不太适合在寄存器内以不同方式旋转 6 个字节。

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

[根据第一条评论和一些编辑,最终的解决方案略有不同。先介绍一下,下面留下原意]

这里的主要思想是使用乘以 2 的幂来完成移位,因为这些常量可以在向量中变化。 @harold 指出了下一个想法,即两个重复字节的相乘将自动将移出的位 "rotation" 返回到较低位。

  1. 将字节解包并复制为 16 位值[... d c b a] -> [... dd cc bb aa]
  2. 生成一个 16 位常量[128 64 32 16 8 4 2 1]
  3. 你想要的字节是每个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位重新打包成两个八位向量,一个高字节向量和一个低字节向量
  5. OR 这两个向量完成旋转。

我对可用的乘法选项和你的指令集限制有点模糊,但理想的是产生 16 位乘积的 8 位 x 8 位乘法。据我所知,它不存在,这就是为什么我建议先解包,但我已经看到其他用于执行此操作的简洁算法。

XOP instruction set does provide _mm_rot_epi8()(这不是 Microsoft 特有的;从 4.4 或更早版本开始,它在 GCC 中也可用,并且在最近的 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 的英特尔酷睿 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+orblendv。前者很有意义,因为它纯粹是一种优化,但后者可能(但是,在当前的 Intel CPUs 上!)实际上更快。

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

我认为如果我在这里列出变体是最有用的,这样任何需要这种功能的用户都可以对他们的实际工作负载进行一些基准测试,并测试是否存在任何可衡量的差异。

我也同意他使用 oddeven 而不是 highlow 的建议,但请注意,由于向量中的第一个元素已编号元素0,第一个元素是even,第二个是odd,依此类推

#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);
}

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