使用 ymm 寄存器作为 "memory-like" 存储位置

Using ymm registers as a "memory-like" storage location

考虑 x86 中的以下循环:

; on entry, rdi has the number of iterations
.top:
; some magic happens here to calculate a result in rax
mov [array + rdi * 8], rax ; store result in output array
dec rdi
jnz .top

很简单:某些东西在 rax 中计算结果(未显示),然后我们将结果存储到一个数组中,与我们使用 rdi.

索引的顺序相反

我想转换上面的循环而不对内存进行任何写入(我们可以假设未显示的计算不写入内存)。

只要 rdi 中的循环计数有限,我就可以使用 ymm regs 提供的充足的 space(512 字节)来保存值,但是实际执行此操作似乎很尴尬,因为您不能 "index" 任意寄存器。

一种方法是始终将 ymm 寄存器的整个 "array" 随机排列一个元素,然后将该元素插入新释放的位置。

像这样:

vpermq  ymm3, ymm3, 10_01_00_11b ; left rotate ymm by qword
vpermq  ymm2, ymm2, 10_01_00_11b ; left rotate ymm by qword
vpermq  ymm1, ymm1, 10_01_00_11b ; left rotate ymm by qword
vpermq  ymm0, ymm0, 10_01_00_11b ; left rotate ymm by qword

vblenddd ymm3, ymm3, ymm2, 3     ; promote one qword of ymm2 to ymm3
vblenddd ymm2, ymm2, ymm1, 3     ; promote one qword of ymm1 to ymm2
vblenddd ymm1, ymm1, ymm0, 3     ; promote one qword of ymm0 to ymm1

pinsrq   xmm0, rax, 0  ; playing with mixed-VEX mode fire (see Peter's answer)

这显示只处理 16 个寄存器中的四个,所以显然要完成所有 16 个将需要大量代码(32 条指令)。

有没有更好的方法?

不可预测的分支是不可取的,但我们仍然可以考虑使用它们的解决方案。

您可以考虑的几个选项:

展开

如果展开循环(由于 ymm 寄存器中可用的存储空间限制为 64 个 qword,循环的迭代次数必然有限),您将有机会使用硬编码逻辑来插入您的来自 rax 直接在正确位置的结果,例如。 pinrsqmovq 有时与洗牌相结合,让您进入高车道。每次迭代可能只需要 1.25 条指令,比 32 条要好得多!

垂直车道

您当前的改组解决方案可以描述为通过寄存器的水平旋转,从 ymm N 的高 qword 进入 ymm N+1 的低 qword。也就是说,一个寄存器中的相邻元素在您的方案中逻辑上相邻。相反,您可以进行垂直旋转,其中给定 qword 通道中的元素在逻辑上与寄存器 ymm N-1ymm N+1 中同一通道中的元素相邻。这避免了任何水平洗牌的需要,并且大多数移位只需要一个寄存器-寄存器 mov。您只需要对第一个和最后一个寄存器进行特殊处理即可将您的元素包装到下一个通道中。

像这样:

; shift all lanes "up"
vmovdqa ymm15, ymm3
vmovdqa  ymm3, ymm2
vmovdqa  ymm2, ymm1
vmovdqa  ymm1, ymm0

; wrap from the top register back to ymm0, shifting to the left by 1
vpermq   ymm0, ymm15, 10_01_00_11b
; store new element
vpinsrq  ymm0, rax, 0

这与您将要获得的一般 "shift every element" 策略一样简单:每个使用的 ymm 寄存器一个 vmovdqa,加上附加指令来执行回绕和新元素插入。就向量操作而言,寄存器-寄存器移动比任何其他类型的操作都要快得多,因为它们可以被移动消除(0 延迟)并且每个周期可以执行 4 次。

这种方法确实需要一个临时寄存器(上例中的ymm15),我想不出一个简单的方法来消除它,所以你最多可以使用 15 个寄存器作为"queue".

的一部分

间接跳转

您可以根据迭代计数计算间接跳转到一个短序列(2-4 条指令),将元素放在正确的位置。基本上是 vpinsrq 并且在某些情况下需要一些额外的洗牌才能进入高车道。

这种类型的 table 可以是完全通用的,即允许以任何顺序写入任意索引,但是如果你知道你是按上面的顺序索引,你可以简化 table 使用该假设(即,您可以通过先写入低元素然后使用 vinserti128 或类似的东西在正确的时间将它们都移动到高车道来处理高车道。

这个 table 可能会在第一次时预测错误。之后,可能会也可能不会,这取决于间接分支预测器的模式和强度。

您不能 vpinsrq 进入 YMM 注册。只有 xmm 目标可用,因此它不可避免地会将完整 YMM 寄存器的上通道清零。它作为 128 位指令的 VEX 版本与 AVX1 一起引入。 AVX2 和 AVX512 没有将其升级到 YMM/ZMM 个目的地。我猜他们不想提供插入高通道的功能,提供仍然只查看 imm8 最低位的 YMM 版本会很奇怪。

您将需要一个暂存器,然后与 vpblendd 混合到 YMM 中。 或者(在 Skylake 或 AMD 上)使用旧版 SSE 版本保持高位字节不变! 在 Skylake 上,使用旧版 SSE 指令编写 XMM reg 对完整的错误依赖登记。您想要 这种错误的依赖。 (我还没有测试过这个;它可能会触发某种合并 uop)。但是你不希望它在 Haswell 上保存所有 YMM regs 的上半部分,进入“状态 C”。

显而易见的解决方案是给自己留一个临时注册以用于 vmovq+vpblendd(而不是 vpinsrq y,r,0)。那仍然是 2 微指令,但 vpblendd 不需要英特尔 CPU 上的端口 5,以防万一。 (movq 使用端口 5)。如果您真的对 space 感到吃力,可以使用 mm0..7 MMX 寄存器。


降低成本

通过嵌套循环,我们可以拆分工作。通过少量展开内循环,我们基本上可以去除那部分成本。

例如,如果我们有一个内部循环产生 4 个结果,我们可以在内循环中的 2 或 4 个寄存器上使用您的强力堆栈方法,在没有实际展开的情况下提供适度的开销(“魔术”负载出现只有一次)。 3 或 4 微指令,可选地没有循环携带的 dep 链。

; on entry, rdi has the number of iterations
.outer:
    mov       r15d, 3
.inner:
; some magic happens here to calculate a result in rax

%if  AVOID_SHUFFLES
    vmovdqa   xmm3, xmm2
    vmovdqa   xmm2, xmm1
    vmovdqa   xmm1, xmm0
    vmovq     xmm0, rax
%else
    vpunpcklqdq  xmm2, xmm1, xmm2        ; { high=xmm2[0], low=xmm1[0] }
    vmovdqa   xmm1, xmm0
    vmovq     xmm0, rax
%endif

    dec   r15d
    jnz   .inner

    ;; Big block only runs once per 4 iters of the inner loop, and is only ~12 insns.
    vmovdqa  ymm15, ymm14
    vmovdqa  ymm13, ymm12
    ...
    
    ;; shuffle the new 4 elements into the lowest reg we read here (ymm3 or ymm4)

%if  AVOID_SHUFFLES       ; inputs are in low element of xmm0..3
    vpunpcklqdq  xmm1, xmm1, xmm0     ; don't write xmm0..2: longer false dep chain next iter.  Or break it.
    vpunpcklqdq  xmm4, xmm3, xmm2
    vinserti128  ymm4, ymm1, xmm4, 1  ; older values go in the top half
    vpxor        xmm1, xmm1, xmm1     ; shorten false-dep chains

%else                     ; inputs are in xmm2[1,0], xmm1[0], and xmm0[0]
    vpunpcklqdq  xmm3, xmm0, xmm1     ; [ 2nd-newest,  newest ]
    vinserti128  ymm3, ymm2, xmm3, 1
    vpxor        xmm2, xmm2,xmm2   ; break loop-carried dep chain for the next iter
    vpxor        xmm1, xmm1,xmm1   ; and this, which feeds into the loop-carried chain
%endif

    sub   rdi, 4
    ja   .outer

好处:这只需要 AVX1(并且在 AMD 上更便宜,将 256 位向量保留在内部循环之外)。我们仍然得到 12 x 4 qwords 的存储而不是 16 x 4。无论如何这是一个任意数字。

展开有限

我们可以只展开内部循环,如下所示:

.top:
    vmovdqa     ymm15, ymm14
    ...
    vmovdqa     ymm3, ymm2           ; 12x movdqa
    vinserti128 ymm2, ymm0, xmm1, 1

    magic
    vmovq       xmm0, rax
    magic
    vpinsrq     xmm0, rax, 1
    magic
    vmovq       xmm1, rax
    magic
    vpinsrq     xmm1, rax, 1

    sub         rdi, 4
    ja          .top

当我们离开循环时,ymm15..2 and xmm1 和 0 充满了有价值的数据。如果它们在底部,它们将 运行 相同的次数,但 ymm2 将是 xmm0 和 1 的副本。 jmp 进入循环而不执行 vmovdqa 东西在第一个 iter 上是一个选项。

每 4x magic,端口 5(movq + pinsrq)需要 6 微指令,12 vmovdqa(无执行单元)和 1x vinserti128(又是端口 5)。所以这是每 4 magic 19 微指令,或 4.75 微指令。

您可以将 vmovdqa + vinsert 与第一个 magic 交错放置,或者在第一个 magic 之前/之后拆分它。您不能在 vinserti128 之前破坏 xmm0,但是如果您有备用整数 reg,则可以延迟 vmovq.

更多嵌套

另一个循环嵌套级别,或另一个展开,将大大减少vmovdqa 指令的数量。不过,将数据完全混入 YMM regs 的成本最低。 .

AVX512 可以给我们更便宜的 int->xmm。 (并且它将允许写入 YMM 的所有 4 个元素)。但我不认为它可以避免展开或嵌套循环以避免每次都接触所有寄存器。


PS:

我对洗牌累加器的第一个想法是将元素洗牌到左边。但后来我意识到这最终有 5 个状态元素,而不是 4 个,因为我们在两个 regs 中有 high 和 low,加上新写的 xmm0。 (并且可以使用 vpalignr。)

留在这里作为您可以使用 vshufpd 做什么的示例:在一个寄存器中从低位移动到高位,并将另一个寄存器的高位合并为新的低位。

    vshufpd   xmm2, xmm1,xmm2, 01b     ; xmm2[1]=xmm2[0], xmm2[0]=xmm1[1].  i.e. [ low(xmm2), high(xmm1) ]
    vshufpd   xmm1, xmm0,xmm1, 01b
    vmovq     xmm0, rax

AVX512:索引向量作为内存

对于将向量寄存器写入内存的一般情况,我们可以 vpbroadcastq zmm0{k1}, rax 并使用不同的 k1 掩码对其他 zmm 寄存器重复。带有合并掩码的广播(其中掩码设置了一位)为我们提供了一个索引存储到向量寄存器中,但我们需要一条指令用于每个可能的目标寄存器。

创建蒙版:

xor      edx, edx
bts      rdx, rcx          #  rdx = 1<<(rcx&63)
kmovq     k1, rdx
kshiftrq  k2, k1, 8
kshiftrq  k3, k1, 16
...

从 ZMM 寄存器读取

vpcompressq zmm0{k1}{z}, zmm1    ; zero-masking: zeros whole reg if no bits set
vpcompressq zmm0{k2},    zmm2    ; merge-masking
... repeat as many times as you have possible source regs

vmovq       rax, zmm0

(请参阅 vpcompressq 的文档:使用零屏蔽会将其写入的元素之上的所有元素归零)

要隐藏 vpcompressq 延迟,您可以将多个 dep 链放入多个 tmp 向量中,然后在最后 vpor xmm0, xmm0, xmm1。 (其中一个向量将全为零,另一个向量将包含所选元素。)

在 SKX 上它有 3c 延迟和 2c 吞吐量,according to this instatx64 report