使用 x64 SIMD 进行半字节洗牌答案

【问题标题】：Nibble shuffling with x64 SIMD使用 x64 SIMD 进行半字节洗牌
【发布时间】：2022-10-23 03:44:29
【问题描述】：

我知道byte shuffling 指令，但我想对半字节（4 位值）做同样的事情，具体来说，我想在一个 64 位字中随机播放 16 个半字节。我的洗牌索引也存储为 16 个半字节。最有效的实现是什么？

【问题讨论】：

标签： x86-64 simd sse

【解决方案1】：

具有必须以这种方式存储的控制向量的任意洗牌？嗯，很难相处。我想你必须解压两者来提供 SSSE3 pshufb，然后重新打包结果。

可能只是 punpcklbw 对右移副本，然后使用 AND 掩码以仅保留每个字节中的低 4 位。然后pshufb。

有时，奇数/偶数拆分比扩大每个元素更容易（因此位仅保留在其原始字节或字内）。在这种情况下，如果我们可以更改您的半字节索引编号，punpcklqdq 可以将奇数或偶数半字节放在高半部分，准备好将它们放回原位和 OR。

但如果不这样做，重新包装是一个单独的问题。我猜想将相邻的字节对组合成低字节中的一个单词，如果吞吐量比延迟更重要，可能会使用pmaddubsw。然后您可以packuswd（针对零或自身）或pshufb（使用恒定控制向量）。

如果你在做多个这样的洗牌，你可以把两个向量压缩成一个，用movhps/movq存储。使用 AVX2，可能有可能让所有其他指令在两个 128 位通道中的两个独立 shuffle 上工作。

// UNTESTED, requires only SSSE3
#include <stdint.h>
#include <immintrin.h>

uint64_t shuffle_nibbles(uint64_t data, uint64_t control)
{
  __m128i vd = _mm_cvtsi64_si128(data);    // movq
  __m128i vd_hi = _mm_srli_epi32(vd, 4);   // x86 doesn't have a SIMD byte shift
  vd = _mm_unpacklo_epi8(vd, vd_hi);       // every nibble at the bottom of a byte, with high garbage
  vd = _mm_and_si128(vd, _mm_set1_epi8(0x0f));  // clear high garbage for later merging

  __m128i vc = _mm_cvtsi64_si128(control);
  __m128i vc_hi = _mm_srli_epi32(vc, 4);
  vc = _mm_unpacklo_epi8(vc, vc_hi);

  vc = _mm_and_si128(vc, _mm_set1_epi8(0x0f));  // make sure high bit is clear, else pshufb zeros that element.
       //  AVX-512VBMI  vpermb doesn't have that problem, if you have it available
  vd = _mm_shuffle_epi8(vd, vc);

       // left-hand input is the unsigned one, right hand is treated as signed bytes.
  vd = _mm_maddubs_epi16(vd, _mm_set1_epi16(0x1001));  // hi nibbles << 4 (*= 0x10), lo nibbles *= 1.

  // vd has nibbles merged into bytes, but interleaved with zero bytes
  vd = _mm_packus_epi16(vd, vd);  // duplicate vd into low & high halves.
  //  Pack against _mm_setzero_si128() if you're not just going to movq into memory or a GPR and you want the high half of the vector to be zero.
  return _mm_cvtsi128_si64(vd);
}

在 shuffle 之前（而不是之后）使用 0x0f 屏蔽数据允许在具有两个 shuffle 单元的 CPU 上进行更多 ILP。至少如果它们已经在向量寄存器中有 uint64_t 值，或者如果数据和控制值来自内存，那么两者都可以在同一个周期中加载。如果来自 GPR，vmovq xmm, reg 的 1/clock 吞吐量意味着 dep 链之间存在资源冲突，因此它们不能在同一个周期中启动。但是由于我们的数据可能在控制之前就已经准备好了，因此尽早屏蔽会使其远离控制->输出延迟的关键路径。

如果延迟是瓶颈而不是通常的吞吐量，请考虑将 pmaddubsw 替换为右移 4、por 和 AND/pack。或 pshufb 打包，同时忽略奇数字节中的垃圾。由于无论如何您都需要另一个常量，因此不妨将其设为pshufb 常量而不是and。

如果您有 AVX-512，与 vpternlogd 的移位和位混合可以避免在洗牌之前需要屏蔽数据，而 vpermb 而不是 vpshufb 将避免需要屏蔽控件，所以你会避免set1_epi8(0x0f) 完全不变。

clang 的 shuffle 优化器没有发现任何东西，只是像 GCC (https://godbolt.org/z/xz7TTbM1d) 那样编译它，即使使用 -march=sapphirerapids。没有发现它可以使用vpermb 而不是vpand / vpshufb。

shuffle_nibbles(unsigned long, unsigned long):
        vmovq   xmm0, rdi
        vpsrld  xmm1, xmm0, 4
        vpunpcklbw      xmm0, xmm0, xmm1        # xmm0 = xmm0[0],xmm1[0],xmm0[1],xmm1[1],xmm0[2],xmm1[2],xmm0[3],xmm1[3],xmm0[4],xmm1[4],xmm0[5],xmm1[5],xmm0[6],xmm1[6],xmm0[7],xmm1[7]
        vmovq   xmm1, rsi
        vpsrld  xmm2, xmm1, 4
        vpunpcklbw      xmm1, xmm1, xmm2        # xmm1 = xmm1[0],xmm2[0],xmm1[1],xmm2[1],xmm1[2],xmm2[2],xmm1[3],xmm2[3],xmm1[4],xmm2[4],xmm1[5],xmm2[5],xmm1[6],xmm2[6],xmm1[7],xmm2[7]
        vmovdqa xmm2, xmmword ptr [rip + .LCPI0_0] # xmm2 = [15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15]
        vpand   xmm0, xmm0, xmm2
        vpand   xmm1, xmm1, xmm2
        vpshufb xmm0, xmm0, xmm1
        vpmaddubsw      xmm0, xmm0, xmmword ptr [rip + .LCPI0_1]
        vpackuswb       xmm0, xmm0, xmm0
        vmovq   rax, xmm0
        ret

（没有 AVX，它需要 2 个额外的 movdqa 寄存器复制指令。）

【讨论】：

诅咒。我想我会想出一个聪明的方法来在洗牌后重新组合半字节 - OR'ing 移位和最终洗牌到 [63:0] ......这要好得多。随着 SIMD 指令集的不断扩展，我越来越难以记住它们。
在这种情况下，_mm_srli_epi32 而不是 _mm_srli_epi64 是否有特殊原因？
@BrettHale：我认为32位元素大小可能至少与任何CPU上的其他任何东西一样快，如果有任何差异，除了乘法。对于某些操作（如 paddq 和 pcmpgtq），64 位元素大小具有更长的操作码和/或在某些 CPU 上速度较慢，所以我从不使用 epi64，而其他大小也同样好（除了随机播放）其中较大的元素较少是快点）。为了实现一个全寄存器，IIRC GCC 选择了pcmpgtd，我认为这就是我开始选择 32 作为“默认值”的原因，因为任何大小都有效。
@BrettHale：因此，当您必须模拟其他尺寸时，32 对我来说似乎是一个很好的默认值。（在这种情况下 ..._epi8 通过移位和屏蔽）。最接近我们想要的可用大小是..._epi16，这使它成为另一个直观的选择。我避免这种情况的部分原因是我想提醒初学者，这种仿真技术并不依赖于移位宽度仅比我们想要仿真的移位宽度“大一号”，16 位对此没有什么特别的.
我已经包含了一个test，以及你的代码序言/尾声。一些测试向量包括：https://godbolt.org/z/qMca4sPbh

【解决方案2】：

我今天遇到了这个问题。在 AVX-512 中，您可以使用 vpmultishiftqb (1)，这是 Ice Lake 及之后（显然在 Zen 4 中，根据维基百科）提供的一个有趣的指令，以更快地洗牌。它的强大之处在于它以非对齐方式排列字节的能力：它在每个 64 位元素中获取八个 8 位块并选择未对齐来自相应元素的 8 位块。下面是一个实现。

#include <immintrin.h>
#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>

// Convention: (a & (0xf << (4 * i))) >> (4 * i) is the ith nibble of a
// (i.e., lowest-significant is 0)
uint64_t shuffle_nibbles(uint64_t data, uint64_t indices) {
#if defined(__AVX512VBMI__) && defined(__AVX512VL__)
    // If your data is already in vectors, then this method also works in parallel
    const __m128i lo_nibble_msk = _mm_set1_epi8(0x0f);
    __m128i v_data = _mm_cvtsi64_si128(data);
    __m128i v_indices = _mm_cvtsi64_si128(indices);

    __m128i indices_lo = _mm_and_si128(lo_nibble_msk, v_indices);
    __m128i indices_hi = _mm_andnot_si128(lo_nibble_msk, v_indices);
    indices_lo = _mm_slli_epi32(indices_lo, 2);
    indices_hi = _mm_srli_epi32(indices_hi, 2);

    // Lookup unaligned bytes
    __m128i shuffled_hi = _mm_multishift_epi64_epi8(indices_hi, v_data);
    __m128i shuffled_lo = _mm_multishift_epi64_epi8(indices_lo, v_data);

    shuffled_hi = _mm_slli_epi32(shuffled_hi, 4);
    // msk ? lo : hi
    __m128i shuffled = _mm_ternarylogic_epi32(lo_nibble_msk, shuffled_lo, shuffled_hi, 202);

    return _mm_cvtsi128_si64(shuffled);
#else
    // Fallback scalar implementation (preferably Peter Cordes's SSE solution--this is as an example)
    uint64_t result = 0;
    for (int i = 0; i < 16; ++i) {
        indices = (indices >> 60) + (indices << 4);

        int idx = indices & 0xf;
        result <<= 4;
        result |= (data >> (4 * idx)) & 0xf;
    }

    return result;
#endif
}

int main() {
        // 0xaa025411fe034102
        uint64_t r1 = shuffle_nibbles(0xfedcba9876543210, 0xaa025411fe034102);
        // 0x55fdabee01fcbefd
        uint64_t r2 = shuffle_nibbles(0x0123456789abcdef, 0xaa025411fe034102);
        // 0xaaaa00002222aaaa
        uint64_t r3 = shuffle_nibbles(0xaa025411fe034102, 0xeeee11110000ffff);

        printf("0x%" PRIx64 "
", r1);
        printf("0x%" PRIx64 "
", r2);
        printf("0x%" PRIx64 "
", r3);
}

Clang 产生 (2):

.LCPI0_0:
        .zero   16,60
shuffle_nibbles(unsigned long, unsigned long):
        vmovq   xmm0, rdi
        vmovq   xmm1, rsi
        vpslld  xmm2, xmm1, 2
        vpsrld  xmm1, xmm1, 2
        vmovdqa xmm3, xmmword ptr [rip + .LCPI0_0] # xmm3 = [60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60]
        vpand   xmm1, xmm1, xmm3
        vpmultishiftqb  xmm1, xmm1, xmm0
        vpand   xmm2, xmm2, xmm3
        vpmultishiftqb  xmm0, xmm2, xmm0
        vpslld  xmm1, xmm1, 4
        vpternlogd      xmm1, xmm0, dword ptr [rip + .LCPI0_1]{1to4}, 216
        vmovq   rax, xmm1

就我而言，我在 64 位元素向量中洗牌；这种方法也避免了加宽的需要。如果您的 shuffle(s) 是恒定的并且您停留在向量中，则此方法将减少到只有四个指令：2x vpmultishiftqb、1x vpslld 和 1x vpternlogd。对于 128 位和 256 位向量，计算 µops 表明延迟为 5，吞吐量为每 2 个周期 1 个，在 shuffle µops 上成为瓶颈；由于后两条指令的执行单元减少，512 位向量的吞吐量为 3。

【讨论】：