AVX2 的矢量移位的 AVX 替代方案?

AVX alternative of AVX2's vector shift?

在 AVX2 中,我们有 _mm256_srlv_epi32(a, b)_mm256_sllv_epi32(a, b) 用于将 'a' 中的一组 8 个值移动 'b' 中的 8 个值。是否有使用 AVX 的有效替代方案,以便我可以留在 AVX 而不必吐出标量代码?

AVX1 没有 256b 整数运算,只有 FP。所以我假设您真的在寻找 __m128i _mm_srlv_epi32() 的替代品。使用 extractf128 / insertf128,您可以轻松地为 256b 向量执行此操作,但最好只使用更多 128b loads/stores,尤其是。如果您的 AVX2 版本可以在支持 AVX2 的 CPU 上 运行。 (现有的 AVX1-only CPU 恰好都有 128b load/store 数据路径,所以 256b loads/stores 几乎没有优势。)

从矢量到标量的往返行程非常昂贵(在标量存储后重新加载时存储转发停顿,或者很多 movd / pextrd / pinsrd),因此,即使是非常笨重的代码也可能比整数代码更好,这取决于吞吐量或延迟在您使用它的代码中是否更重要。

我的最佳想法基本上是矢量 regs 中的标量:4 个班次(每个班次对应一个班次)和 3 个立即混合以组合结果。

update:思路2:左移32位乘以2count。请参阅此答案的结尾。

如果移位计数不是编译时常量,您将需要解压缩移位计数向量,以便将每个移位计数作为向量的 64b。 (非可变移位指令可以在寄存器中获取它们的计数,但它们会查看整个低位 64b。并且不像标量移位那样屏蔽(模字大小),它们会饱和。

将 xmm 寄存器的 4 个元素中的每一个元素都隔离在一个否则为零的目标中是很棘手的。您不能只是将它们向下移动到底部,因为那样会从第二个元素中留下非零字节。

由于这是针对没有 AVX2 的 AVX,我假设您有一个单独的 AVX2 CPU 版本。所以对于Intel来说,这个版本会用在SnB/IvB上。这意味着您有两个 128b 洗牌单元,而不是在 Haswell 和更高版本上只有一个。

## 4 shift-counts in the elements of   xmm0 = [ D C B A ].  element 1 isolated in xmm1, etc.
vpsrlq      xmm2, xmm0, 32           ; xmm2 = [ 0 D 0 B ]
vpunpckhqdq xmm4, xmm2, xmm0         ; xmm4 = [ D C 0 D ]
vpshufd     xmm3, xmm4, 0b01010110   ; xmm3 = [ 0 0 0 C ]
vblendps    xmm1, xmm2, xmm0, 0b0001 ; xmm1 = [ 0 D 0 A ]
; or
vpblendw     xmm1, xmm2, xmm0, 0b00000011 ; xmm1 = [ 0 D 0 A ]

vblendps 运行s 在 p0/5 上 SnB/IvB。 p1/p5 上 SnB/IvB 上的等效 vpblendw 运行s。在 Haswell/SKL 上,它是 p015 与 p5,因此 blendps 更好(与 PAND 相同的端口选择)。对于 SnB,可以使用两者的组合来混合移位结果。对于内在函数,在整数数据上使用 FP 指令需要大量转换,这使得源代码难看且难以阅读。除非您打算使用性能计数器和微基准对其进行调整以使其最适合周围的代码,否则只需将 pblendw 用于 SnB/IvB。否则只需投射并使用 blendps.

如果您有可用的 [ 0 -1 0 -1 ] 掩码,则可以使用矢量 AND 在更多端口上 运行,并缩短 xmm3 的依赖链。这不足以证明加载或生成掩码的合理性,因此更喜欢使用 shifts/shuffles/blends.

完成所有操作的先前版本
vpcmpeqw   xmm5, xmm5,xmm5            ; all-ones
vpsrlq     xmm5, xmm5, 32             ; [ 0 -1  0 -1 ]: generate the mask on the fly if desired

vpand       xmm1, xmm5, xmm0           ; [ 0 C 0 A ]
vpsrlq      xmm2, xmm0, 32             ; [ 0 D 0 B ]
vpunpckhqdq xmm3, xmm1,xmm1            ; [ 0 C 0 C ]  ; saves 1B vs. the equivalent pshufd: no imm8 byte
vpunpckhqdq xmm4, xmm2,xmm2            ; [ 0 D 0 D ]

旁注:奇怪的是,在 Skylake 上,VPSRLVD ymm,ymm,ymmPSRLD xmm,xmm,xmm(2 微指令)便宜(1 微指令)。不过,立即数 PSRLD 只有 1 uop。 (来自 Agner Fog's insn tables)。

@BeeOnRope 的测试证实 Agner 的延迟数是从数据输入到数据输出,移位计数不在关键路径上。从移位计数输入到数据输出的延迟是 2c(xmm) 或 4c(ymm),通常 1c 是车道内广播,而 3c 是车道交叉广播。


uop 计数:

使用编译时常量移位计数的标量代码,整个过程可能如下所示:

movaps    [rsp - 16], xmm0
shr       [rsp - 16], 3         ; 3 uops with a memory-destination.  5 uops for variable count with a memory destination
shr       [rsp - 12], 1
shr       [rsp -  8], 4
shr       [rsp -  4], 1
movaps    xmm0, [rsp - 16]      ; store-forwarding stall here from the 4x 32b stores to the 128b load

或者对于可变计数:

## data in xmm0,  shift counts in xmm1, results in xmm2
vmovd      eax, xmm0      ; 1 uop
vmovd      ecx, xmm1      ; 1 uop
shr        eax, cl        ; 3 uops because of CISC stupidity
vmovd      xmm2, eax      ; 1 uop

vpextrd    eax, xmm0, 1   ; 2 uops
vpextrd    ecx, xmm1, 1   ; 2 uops
shr        eax, cl        ; 3 uops because of CISC stupidity
vpinsrd    xmm2, eax, 1   ; 2 uops

... repeat twice more, for indices 2 and 3    

因此可变计数移位的全寄存器方式为6uops + 9uops * 3,共33 uops。


内存目标版本是 14 个融合域微指令,因为我计算了一个将移位计数作为编译时常量的版本。将计数加载或 pextring 计数到 ecx 中会更多,因为每个可变计数移位比立即计数移位多 2 微指令。


因此,尽管 SSE/AVX 版本非常糟糕,但并没有那么糟糕。全变向量版本还是

  • 4 微指令解包计数
  • 四个 vpsrld xmm,xmm insns 的 8 微指令
  • 3 微指令用于 vpblendwvblendps 合并这些结果。
  • 总计 = 15 个完全可变 AVX1 的融合域微指令

所以全可变矢量版本只和全常量存储/标量洗牌/重载版本一样糟糕,并且其中有一个存储转发停顿。

请注意,仅计算融合域微指令并不总是唯一相关的事情。延迟可能很重要,未融合域中的执行端口压力可能很重要。


比较:

  • Skylake:vpsrlvd ymm, ymm, ymm 是 1 uop,1c 延迟,每 0.5c 吞吐量一个。
  • Haswell/BDW: vpsrlvd ymm, ymm, ymm 是 3 微指令,2c 延迟,每 2c 吞吐量一个。

请记住,这是针对 256b 向量的。我所做的所有计数都是针对 128b 向量的。

在 Haswell(而不是 SnB/IvB)上,我的 SSE 版本可能会在 shuffle 端口吞吐量上出现瓶颈。延迟也会更糟,因为资源冲突限制了它可以利用的 insn 级并行性的数量。


使用SSE4.1左移pmulld乘以2的幂

在 SnB/IvB 上,SSE4.1 pmulld 是 1 uop,5c 延迟,每 1c 吞吐量一个。
在 Haswell 上,它是 2 微指令,10c 延迟,每 2c 吞吐量一个。 (Skylake 的吞吐量翻倍,因为它的 uops 在 p1 和 p0 上可以 运行)

诀窍是将移位计数变成 2c。一种方法是使用可变移位。如果你可以重复使用 2c 的取幂向量来移动多个其他向量,这很好,否则就是先有鸡还是先有蛋的问题。

如果移位计数的范围很小(即 0..7),您可以使用 SSSE3 pshufb 作为 LUT 将计数向量映射到 2^c 的向量。每个元素低字节的0必须变成1(20),但其他字节的0必须保持为零。

##           1<<8 or higher is 0, in an 8bit element
## xmm5 = _mm_set_epi8(0, 0, ..., 1<<7, ..., 1<<2, 1<<1, 1<<0);
## xmm4 = _mm_set1_epi32(0x000000ff);        
## data in xmm0, shift counts in xmm1
movdqa    xmm2, xmm5           ; avoid this with AVX
pshufb    xmm2, xmm5           ; 2^count
pand      xmm2, xmm4           ; zero all but the low byte in each element
pmulld    xmm0, xmm2           ; data * 2^count

Intel SnB/IvB:3 uops(不包括 AVX 不需要的 movdqa)。从轮班计数到结果的延迟:7c。从移位数据到结果的延迟:5c。吞吐量:每 1c 一个(因为所有三个 uops 都可以 运行 在不同的端口上)。

使用 Haswell 及更高版本:延迟增加 5c。 Penryn/Nehalem pmulld 也比 SnB 需要更多的 uops,但没有像 Haswell 那样糟糕的延迟。


LUT 在上部 64b 中全为零,但说服编译器只存储相关部分并使用 movq 加载它并非易事。我不会在这里讨论。

为了处理更大的移位计数,我们可以使用多个 LUT 从 [ D-8 C-8 B-8 A-8 ] 进行查找,以获取每个 32b 元素的第二个字节的值,等等。请注意 C-8 有符号如果 C<8BLENDVB 根据设置的符号位合并。但是,它很昂贵,因此一系列合并可能并不比仅使用早期的 shift/blend-immediate 方法更好。


除了屏蔽 pshufb 结果之外,您还可以添加 set1_epi32(1) 的向量。那么 LUT 中具有非零字节的索引范围将为 1..8,移位计数向量中的填充 0 字节将查找 LUT 的低位元素(应为 0)。这样做将使即时常量生成更加可行:

## xmm5 = _mm_set_epi8(0, 0, ..., 1<<7, ..., 1<<2, 1<<1, 1<<0, 0);
## data in xmm0, shift counts in xmm1
pcmpeqw   xmm4,xmm4            ; all-ones

psubd     xmm1, xmm4           ; shift_counts -= -1
movdqa    xmm2, xmm5
pshufb    xmm2, xmm1           ; 2^count
pmulld    xmm0, xmm2           ; data * 2^count

这没有任何优势,除非你真的关心在一个更少的 insn 中动态生成一个常量。 (set1_epi32(0xff) 使用 pcmpeqw / psrld 24 可以快速生成,但编译器通常只能在一个 insn 中即时生成。)


更新:

OP 在聊天中澄清问题实际上 简单得多 :被移动的数据是一个编译时常量(特别是 0xF)。另外,只需要结果的低8位。

这使得只需将 PSHUFB 作为 LUT 即可实现,无需乘法。请参阅此答案的前一部分,该部分使用 pshufb 来执行 2<<count.

如果您想要 32b 的结果,您可以生成 [ 0 0 D+8 D | 0 0 C+8 C | ... ] 用作控制掩码。在 LUT 的每一半中使用正确的数据,将产生正确的两个字节。

只是在混合中提出另一个想法,如果移位很小(在这种情况下 <= 4),那么一系列 compare/mask/add 操作并不会太低效并且仅使用 SSE2 指令:

__m128i mm_sllv_4_epi32(__m128i v, __m128i vcount)
{
    const __m128i vone = _mm_set1_epi32(1);
    __m128i vtest, vmask;

    vtest = _mm_set1_epi32(0);
    vmask = _mm_cmpgt_epi32(vcount, vtest);
    v = _mm_add_epi32(v, _mm_and_si128(v, vmask));

    vtest = _mm_add_epi32(vtest, vone);
    vmask = _mm_cmpgt_epi32(vcount, vtest);
    v = _mm_add_epi32(v, _mm_and_si128(v, vmask));

    vtest = _mm_add_epi32(vtest, vone);
    vmask = _mm_cmpgt_epi32(vcount, vtest);
    v = _mm_add_epi32(v, _mm_and_si128(v, vmask));

    vtest = _mm_add_epi32(vtest, vone);
    vmask = _mm_cmpgt_epi32(vcount, vtest);
    v = _mm_add_epi32(v, _mm_and_si128(v, vmask));

    return v;
}

显然,您仍然需要将此应用于 AVX 向量的每一半。