使用 x64 SIMD 进行半字节洗牌

Question

我知道字节洗牌指令，但我想对半字节(4 位值)做同样的事情，具体来说我想在 64 位单词中洗牌 16 个半字节。我的洗牌索引也存储为 16 个半字节。最有效的实现方式是什么？

Answer 1

使用必须以这种方式存储的控制向量的任意洗牌？呃，很难合作。我想你必须解开两者才能pshufbSSSE3 输入，然后重新打包该结果。

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

有时奇数/偶数拆分比加宽每个元素更容易(因此位只是保留在其原始字节或单词内)。 在这种情况下，如果我们可以更改您的半字节索引编号，punpcklqdq可以将奇数或偶数放在高半部分，准备将它们放回并 OR。

但如果不这样做，重新包装是一个单独的问题。 我想将相邻的字节对组合成低字节中的一个单词，如果吞吐量比延迟更重要，也许会pmaddubsw。 然后，您可以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/时钟吞吐量意味着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寄存器复制指令。

Answer 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);
// Look up 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 "n", r1);
printf("0x%" PRIx64 "n", r2);
printf("0x%" PRIx64 "n", r3);
}

叮当产量 (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 位元素向量中随机排列半字节;这种方法也避免了加宽的需要。如果你的洗牌是恒定的，并且你停留在向量中，这种方法减少到可怜的四个指令：2xvpmultishiftqb、1xvpslld和 1xvpternlogd。计数 μops 表明 128 位和 256 位矢量的延迟为 5，吞吐量为每 2 个周期 1，在随机 μops 上存在瓶颈;512 位向量的吞吐量为 3，因为后两条指令的执行单元减少。