AVX2 complex kernels: fenced K step, fmaddsub 4x6 default, vector pack/store

[lkdvos]

AVX2 complex kernels: fenced K step, fmaddsub 4x6 default, vector pack/store

Summary

On AVX2 machines (Zen 2 / Rome, and Intel without AVX-512), complex contractions ran at only 50–60% of OpenBLAS. There were three causes:

• every complex microkernel spilled registers in its inner loop;
• the default tile was smaller than the register file allows;
• packing and storing complex data fell back to scalar loops.

This PR fixes all three and changes nothing on AVX-512 or NEON. On a Rome node against current main, the complex suite runs in about 0.6× the time and the GEMM kernels now reach 0.9–1.0× OpenBLAS.

Background: why AVX2 is different

AVX2 has 16 vector registers of 256 bits, against AVX-512's 32 registers of 512 bits. A complex microkernel holds a tile of C accumulators in registers. Each K step also needs a few registers for the A sliver and for broadcasts of B. With only 16 registers there is almost no slack, so two things that do not matter on AVX-512 decide performance here:

• how many live values LLVM's scheduler chooses to keep;
• which tile shape fits in the register file.

All of the engine's tuning had been done on AVX-512 machines, so the AVX2 path had never been looked at closely.

What changes, and why

1. A scheduling fence in the AVX2 complex K step (memory_fence, steps.jl)

Problem. For a planar or fmaddsub K step, LLVM's haswell/znver2 scheduler hoists every B broadcast of the step to the top of the loop body to get the loads started early. With 16 registers those broadcasts cannot all stay live next to the accumulators, so they are spilled to the stack and reloaded right away. Main's AVX2 default (planar 4x5) has 4 stack accesses per K step. The register-optimal fmaddsub 4x6 has 24 (table below).

Fix. memory_fence() is an empty inline asm sideeffect with a memory clobber. It emits no instruction, but loads cannot move across it. Placed after each column group of the K step, it keeps each broadcast next to the FMAs that use it, so a single broadcast register is live at a time. The FMAs themselves can still be reordered freely.

For fmaddsub on AVX2 the step also follows OpenBLAS's Zen zgemm ordering: the swap(a) * b_imag pass, a fence, then the a * b_real pass on a freshly loaded a. That way a and swap(a) are never live together.

Scope. Only 256-bit kernels (W * sizeof(real(T)) == 32) take the fenced order. The planar and fmaddsub generators each build the same terms once and emit them either in the fenced order or in main's order, so every other shape generates exactly main's code. The mixed-domain kernels (ComplexRealKernel, RealComplexKernel) run the real SIMDKernel step, which does not spill, so they are untouched.

2. The complex pack and store fast paths open to AVX2 (complex_fastpath_isa_eligible)

Problem. The vectorised complex pack (de-interleaving or copying slivers) and the vectorised C-tile stores were gated to AVX-512. On AVX2 every complex sliver was therefore packed element by element, and every tile was stored element by element. The mixed-domain kernels reuse these stores, so they were affected too. At small K, or with a thin M or N, this overhead dominates.

Fix. The gate now admits any target with at least 256-bit vectors. The AVX2 code is the AVX-512 path at half width; no new code is needed. NEON and unknown targets keep the scalar loops.

3. AVX2 complex default: fmaddsub 4x6 / 8x6 when C's rows allow its vector store (select_shape)

Problem. Main's AVX2 complex default is planar 4x5 (ComplexF64) / 8x5 (ComplexF32). That is 10 complex accumulators, below what the register file can carry, so too few FMAs are in flight to hide their latency.

Choice. The new default is fmaddsub 4x6/W4 (ComplexF64) and 8x6/W8 (ComplexF32). Its single interleaved accumulator plane holds 12 accumulators + 2 A + 1 broadcast = 15 of 16 registers. With the fence it is spill-free whether B is packed or read in place.

Alternative considered. The natural alternative, planar with NR = 6, needs all 16 registers. It is spill-free with packed B, but when B is read in place (M ≤ 256) it needs one more live address, and it then spills 13 values per K step. On Rome it collapsed to about 0.6× of planar 4x5 on exactly those shapes, so it is not offered.

Why the condition. fmaddsub stores C as lane pairs. When C's unit-stride run along M does not cover whole MR-row slivers (for example C permuted so that M is strided), every tile falls back to a scalar store. That scalar store costs about 1.3× planar's per element, and at short K the store dominates the tile's time. The TCCG ccsd_t_2/3/4 and ao2mo cases ran 1.1–1.3× slower with fmaddsub for this reason. So fmaddsub is chosen only when run == m_length || run % MR == 0; otherwise the planner keeps main's planar 4x5/8x5, which is now fenced. To evaluate this, plan_contract now computes C's run along M for complex T as well.

Menu clean-up. The fmaddsub NR = 5 AVX2 tiles (4x5/W4, 8x5/W8) existed only because NR = 6 used to spill; nothing selects them now, so they are dropped from the menu. The AVX2 shape override covers OneMKernel and FMAddSubKernel alike, which also gives an explicitly requested OneMKernel ComplexF32 an AVX2 shape (8x6/W8) instead of an AVX-512 one.

How it was checked

Codegen and plan checks are deterministic and were run on the development machine; performance was measured on Slurm nodes only.

• Nothing changes off AVX2.

  • code_llvm/code_native of add_tile were compared under -C icelake-server, znver4, znver2 and haswell. The set is every planar, fmaddsub, 1m and mixed-domain menu shape, with B packed and with B in place. All 57 non-256-bit variants are identical to main.
  • select_shape and full plans match main on icelake-server, cascadelake, znver4 and znver5 profiles. The plans cover the 490 suite contract cases, same-eltype and both mixed orientations: kernel, blocking, line packing and execution path.
  • On a znver2 profile, 419 same-eltype plans change, and only in their kernel and the blocking derived from it.

• No spills. Hot-loop stack accesses per K step, -C znver2 (haswell identical), B packed / B in place:

  kernel	main	this PR
  planar 4x5/W4, 8x5/W8	4 / 4	0 / 0
  fmaddsub 4x6/W4, 8x6/W8 (new default)	24 / 24	0 / 0
  planar 4x6/W4, 8x6/W8 (not offered)	14 / 12	0 / 13
  mixed-domain 256-bit shapes	0	0

• Tests. Pkg.test() passes. The suite also passes with QS_FAKE_ISA=avx2 and julia_args = ["-C", "znver2"], apart from the host-detection test, which cannot pass under a forced ISA. New tests pin:

  • the AVX2 shape selection, with and without a vector-store run;
  • the fast-path gate;
  • zero spills and the FMA count of the fenced AVX2 tiles.

• Compile latency and allocations.

  • The first ComplexF64 plan + execute! takes 1.79 s on main and 1.88 s here natively, and 1.62 vs 1.78 s under a forced AVX2 profile. About 40 ms of that is the shared planar generator, which now carries the fenced order. On AVX2 the first call also compiles fmaddsub instead of planar.
  • Specialisation counts are unchanged natively.
  • A reused plan's execute! allocates 0 B everywhere.

Performance

Rome (EPYC 7742, Zen 2), one core, benchmark/submit_ab.sh alternating main (ad94b73) and this branch for two rounds on the same node. Best of the two rounds, with no correctness mismatches against StridedBLAS.

time, this PR / main (geomean)	ComplexF64	ComplexF32
bench_kernels.jl, default kernel (10 shapes)	0.59	0.54
bench_to_suite.jl contract (≈155 cases)	0.63	0.60
bench_to_suite.jl network: mps, ctmrg, trg (20)	0.57	0.52

Every kernel shape is faster. Large GEMMs go from 0.51–0.57× to 0.88–0.99× of OpenBLAS (ComplexF64), and from 0.52–0.56× to 0.92–1.04× (ComplexF32). No suite case is consistently slower.

These ratios are larger than when the same changes were measured on 4a77066, before the cleanup in #22. There they gave 0.87 / 0.79 on kernels and 0.83 / 0.78 on the suite. The difference is that main itself got slower on AVX2 complex in #22, while this branch did not:

	main 4a77066 → ad94b73	this PR, before → after #22
default ComplexF64 kernel	36–42 → 25–27 GF/s	unchanged
dim63_2_2_2_gemm_ready	11.9 → 18.6 s	10.0 → 9.9 s
mps_2site_D256	261 → 399 ms	217 → 214 ms

These come from different nodes, so only the size of the gap is meaningful. The pre-cleanup run also showed a few consistent losses (dim32/96_1_3_1_contract_scrambled, ComplexF64 intensli_7_dim24); they are gone in this run. Whether #22 also slowed real eltypes on AVX2 is not covered here.

🤖 Generated with Claude Code

[codecov]

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

ℹ️ You can also turn on project coverage checks and project coverage reporting on Pull Request comment

Thanks for integrating Codecov - We've got you covered ☂️