Cranelift x64: a shuffle whose mask reads one operand still depends on the other

Summary

On x64, a shuffle gets a single pshufb only when both operands are the same
value (shuffle a a). If the operands differ but the mask only selects lanes from
one of them (every index < 16, or every index >= 16), lowering falls through to:

• vpermi2b mask, a, b when AVX512VL + AVX512VBMI are available, or
• pshufb a + pshufb b (zeroing mask) + por otherwise.

Both read the operand the mask never selects. That costs instructions, and it
also adds a false data dependency on that operand.

Wasm producers emit this pattern. In core::arch::wasm32::i8x16_shuffle::<0, 4, 1, 5, …>(x, x), LLVM marks the second operand undefined, and the wasm backend fills
it with whatever local is handy. In the case below that local is the loop's
accumulator, so the shuffle waits on the previous iteration's add and the loop
runs serially. V8 lowers the same module to a single pshufb.

Reproducer

CLIF (test compile precise-output, target x86_64 ssse3 has_avx512vl has_avx512vbmi):

function %shuffle_lhs_only(i8x16, i8x16) -> i8x16 {
block0(v0: i8x16, v1: i8x16):
    v2 = shuffle v0, v1, [0 4 1 5 2 6 3 7 8 12 9 13 10 14 11 15]
    return v2
}

Observed (main at dcfbea1):

  movdqa %xmm0, %xmm5
  movdqu (%rip), %xmm0
  movdqa %xmm5, %xmm6
  vpermi2b %xmm1, %xmm6, %xmm0

With target x86_64 sse41 it is two pshufbs and a por. Expected in both
cases: pshufb (%rip), %xmm0. The mirror case, every index >= 16, behaves the
same way with v1 and v0 swapped.

Wasm loop showing the false dependency (shuffle-false-dep.wat, attached):

(loop $l
  (local.set $v
    (i8x16.shuffle 0 4 1 5 2 6 3 7 8 12 9 13 10 14 11 15
      (v128.load …) (local.get $acc)))          ;; $acc is never selected
  (local.set $acc (i32x4.add (local.get $acc) (i32x4.dot_i16x8_s …)))
  (br_if $l …))

Impact

shuffle-false-dep.wat, i7-1165G7 (has AVX512VBMI):

n = 1_000_000_000, wasmtime run --invoke run shuffle-false-dep.wat 1000000000,
pinned to one core; two runs each agree to the hundredth of a second. The loop
on 49.0.3:

vmovdqu (%rax, %rcx), %xmm1
subl    $1, %edx
vmovdqu 0x38(%rip), %xmm2
vpermi2b %xmm0, %xmm1, %xmm2      ; %xmm0 is $acc, never selected
vpmovsxbw %xmm2, %xmm1
vpalignr $8, %xmm2, %xmm2, %xmm2
vpmovsxbw %xmm2, %xmm2
vpmaddwd %xmm2, %xmm1, %xmm1
vpaddd  %xmm1, %xmm0, %xmm0       ; $acc

With vpshufb 0x27(%rip), %xmm1, %xmm2 in place of the vmovdqu 0x38(%rip) /
vpermi2b pair, the shuffle no longer waits on the previous iteration's vpaddd.

In a real int8 matrix-vector kernel (tract's wasm_i32_4x1, the inner loop of a
transformer decoder's weight matmuls), patching only this operand in the module
(passing the loaded vector twice, so that shuffle a a matches) takes a
512×512 call from 56.7 µs to 27.2 µs on wasmtime 49.0.3: 2.1×, level with V8
(26 µs).

Where it comes from

cranelift/codegen/src/isa/x64/lower.isle, shuffle rules: rule 3 is
(shuffle _ a a …). Rules 2 and 1 (vpermi2b) and the por fallback take
(shuffle _ a b …) with no check on which operands the mask reads. Several
special cases already match "lhs only" or "rhs only" masks (pshufd_lhs_imm /
pshufd_rhs_imm, pshuflw_*, pshufhw_*), but no general byte shuffle does.
Neither the mid-end nor the wasm translator canonicalizes shuffle a b to
shuffle a a.

Related

• v8x16.shuffle optimizations needed #2531: shuffle optimizations needed. Closed once the common patterns got
  single-instruction lowerings; single-operand masks with distinct operands were
  not among them.
• x64: lower i8x16.shuffle to VPERMI2B when possible #2930: added the vpermi2b lowering for shuffles of two different registers.
  It is the rule that now also catches these single-operand masks.
• x64: Lower shuffle and swizzle in ISLE #4772: port of the x64 shuffle lowering to ISLE (the rules above).
• x64: Add a smattering of lowerings for shuffle specializations #5930, x64: Add shuffle specialization for palignr #5999: more shuffle specializations (pshufd/pshuf{l,h}w lhs/rhs,
  palignr); they show the lhs/rhs-only pattern this needs.
• x64: Add non-SSSE3 lowerings of pshufb #6606: non-SSSE3 pshufb lowering (the X86Pshufb libcall).