使用 x64 SIMD 进行蚕食改组

Nibble shuffling with x64 SIMD

我知道 byte shuffling 指令,但我想对半字节(4 位值)做同样的事情,具体来说,我想在 64 位字中打乱 16 个半字节.我的洗牌索引也存储为 16 个半字节。最有效的实施方式是什么?

使用必须以这种方式存储的控制向量进行任意洗牌?呃,很难合作。我想你必须将两者都解压以提供 SSSE3 pshufb 然后 re-pack 结果。

可能只是 punpcklbw 对 right-shifted 副本,然后进行 AND 掩码以仅保留每个字节中的低 4 位。然后pshufb.

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

但如果不这样做,re-packing 就是一个单独的问题。我想将相邻的字节对组合成低字节中的一个字,可能使用 pmaddubsw if throughput is more important than latency. Then you can packuswd (针对零或自身)或 pshufb (使用常量控制向量)。

如果您进行多次这样的洗牌,您可以将两个向量压缩为一个,以存储 movhps / movq。使用 AVX2,可以让所有其他指令在两个 128 位通道中处理两个独立的洗牌。

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

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

如果延迟是瓶颈而不是通常的吞吐量,请考虑将 pmaddubsw 替换为 right-shift 4、por 和 AND/pack。或者 pshufb 进行打包,同时忽略奇数字节中的垃圾。由于无论如何你都需要另一个常量,不妨将它设为 pshufb 常量而不是 and.

如果你有 AVX-512,移位和 bit-blend 与 vpternlogd 可以避免在洗牌之前需要屏蔽数据,并且 vpermb 而不是 vpshufb 会避免需要屏蔽控件,因此您可以完全避免使用 set1_epi8(0x0f) 常量。

clang 的 shuffle 优化器没有发现任何东西,只是像 GCC 那样编译它 as-written (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 register-copy 指令。)