基于标量整数条件的 AVX 向量寄存器的条件移动(cmov)?

Conditional move (cmov) for AVX vector registers based on scalar integer condition?

对于 64 位寄存器,有 CMOVcc A, B 指令,如果条件 cc 满足,则只将 B 写入 A

; Do rax <- rdx iff rcx == 0
test rcx, rcx
cmove rax, rdx

但是,我找不到任何与 AVX 等效的东西。我仍然想根据 RFLAGS 的值移动,只是操作数更大:

; Do ymm1 <- ymm2 iff rcx == 0
test rcx, rcx
cmove ymm1, ymm2  (invalid)

是否有 cmov 的 AVX 等效项?如果没有,如何以无分支的方式实现这个操作?

虽然 cmov 没有向量化版本,但可以使用位掩码和 blending.

实现等效功能

假设我们有两个 256 位向量 value1value2,它们驻留在相应的向量寄存器 ymm1ymm2 中:

align 32
value1: dq 1.0, 2.0, 3.0, 4.0
value2: dq 5.0, 6.0, 7.0, 8.0
; Operands for our conditional move
vmovdqa ymm1, [rel value1]
vmovdqa ymm2, [rel value2]

我们要比较两个寄存器 rcxrdx:

; Values to compare
mov rcx, 1
mov rdx, 2

如果它们相等,我们希望将 ymm2 复制到 ymm1(因此 select value2),否则我们希望保留 ymm1因此 value1.

使用 cmov 的等效(无效)表示法:

cmp rcx, rdx
cmove ymm1, ymm2  (invalid)

首先,我们将 rcxrdx 加载到向量寄存器中并 broadcast 它们,因此它们被复制到相应寄存器的所有 64 位块中(. 描述了一个连接):

vmovq xmm0, rcx          ; xmm0 <- 0 . rcx
vpbroadcastq ymm1, xmm0  ; ymm1 <- rcx . rcx . rcx . rcx
vmovq xmm0, rdx          ; xmm0 <- 0 . rdx
vpbroadcastq ymm2, xmm0  ; ymm2 <- rdx . rdx . rdx . rdx

然后,我们使用 vpcmpeqq:

生成掩码
; If rcx == rdx:  ymm0 <- ffffffffffffffff.ffffffffffffffff.ffffffffffffffff.ffffffffffffffff
; If rcx != rdx:  ymm0 <- 0000000000000000.0000000000000000.0000000000000000.0000000000000000
vpcmpeqq ymm0, ymm1, ymm2

最后,我们blendymm2进入ymm1,使用掩码ymm0:

; If rcx == rdx: ymm1 <- ymm2
; If rcx != rdx: ymm1 <- ymm1
vpblendvb ymm1, ymm1, ymm2, ymm0

感谢 @fuz,他在评论中概述了这种方法!

鉴于此分支代码(如果条件预测良好,这将是有效的):

    cmp rcx, rdx
    jne  .nocopy
     vmovdqa  ymm1, ymm2       ;; copy if RCX==RDX
.nocopy:

我们可以通过基于比较条件创建一个 0 / -1 向量并在其上混合来实现无分支。一些优化与其他答案:

  • Broadcast after XMM compare,所以你不需要广播两个输入。保存一条指令,并只比较 XMM(在 Zen1 上保存一个 uop)。
  • 如果成本较低,请将整数输入减少为一个整数。所以你只需要将一个东西从整数复制到 XMM regs。标量 xor 可以在任何执行端口上 运行,而 vmovd/q xmm, reg 只能在英特尔的单个执行端口上 运行:端口 5,与 vpbroadcastq ymm, xmm 等向量洗牌所需的端口相同.

除了总共节省 1 条指令外,它还使其中一些指令更便宜(对同一执行端口的竞争更少,例如,标量异或根本不是 SIMD)并且脱离了关键路径(异或归零)。并且在循环中,你可以在循环外准备一个归零向量。

;; inputs: RCX, RDX.  YMM1, YMM2
;; output: YMM0

   xor      rcx, rdx        ; 0 or non-0.
   vmovq    xmm0, rcx
         vpxor xmm3, xmm3, xmm3   ; can be done any time, e.g. outside a loop
   vcmpeqq  xmm0, xmm0, xmm3      ; 0 if RCX!=RDX,  -1 if RCX==RDX

   vpbroadcastq ymm0, xmm0
   vpblendvb    ymm0, ymm1, ymm2, ymm0   ; ymm0 = (rcx==rdx) ? ymm2 : ymm1

销毁旧的 RCX 意味着您可能需要 mov,但这仍然值得。

rcx >= rdx(无符号)这样的条件可以用cmp rdx, rcx/sbb rax,rax来实现一个0/-1整数(你可以广播而不 需要 vpcmpeqq).

有符号大于条件更痛苦;您可能最终想要 vpcmpgtq 的 2x vmovq,而不是 cmp/setg/vmovd / vpbroadcastb。特别是如果您没有方便的注册到 setg 以避免可能的错误依赖。 setg al / 读取 EAX 对于部分寄存器停顿不是问题:足够新的 CPU 可以拥有 AVX2 。 (只有 Intel 曾经这样做过,而 Haswell 没有。)所以无论如何,你 可以 只是 setcc 进入你的 cmp 输入之一的低字节.

注意vblendvpsvblendvpd只关心每个dword或qword元素的高字节。如果你有两个正确的符号扩展整数, 并且减去它们不会溢出c - d 将直接用作你的混合控件,只需广播它。整数 SIMD 指令之间的 FP 混合 vpaddd 在具有 AVX2 的英特尔 CPU 上(在 AMD 上可能类似)在输入和输出上有额外的 1 个周期的旁路延迟,但您保存的指令也会有延迟。

对于无符号的 32 位数字,您可能已经将它们零扩展为整数 regs 中的 64 位。在这种情况下,sub rcx, rdx 可以将 RCX 的 MSB 设置为与 cmp ecx, edx 设置 CF 的方式相同。 (请记住 jb / cmovbFLAGS conditionCF == 1

;; unsigned 32-bit compare, with inputs already zero-extended
   sub   rcx, rdx               ; sets MSB = (ecx < edx)
   vmovq xmm0, rcx
   vpbroadcastq   ymm0, xmm0

   vblendvpd      ymm0, ymm1, ymm2, ymm0   ; ymm0 = ecx<edx ? ymm2 : ymm1

但是如果您的输入已经是 64 位的,并且您不知道它们的范围是有限的,您需要一个 65 位的结果来完全捕获一个64位减法结果。

这就是为什么 jl 的条件是 SF != OF,而不仅仅是 a-b < 0,因为 a-b 是用 t运行 数学运算完成的。 jb 的条件是 CF == 1 (而不是 MSB)。