transformations: (x86) fold vector loads into FMA memory operands

[manishpaulish]

Depends on #6359, which is the first commit here. Opening as a draft until that lands, then I will rebase so this is a clean single-commit diff.

x86 FMA instructions can take a multiply operand straight from memory, and on AVX-512 can apply an embedded broadcast to it. The RSM forms of vfmadd231pd and vfmadd231ps are defined and have emission tests, but nothing generates them, so every FMA currently pays for a separate load or broadcast plus the vector register to hold the result.

This adds x86-fold-memory-operands, which folds a single-use load or broadcast-load into the FMA that consumes it.

The fold is rejected when:

• the loaded value has more than one use, since folding would duplicate the memory access rather than remove it
• a memory write sits between the load and the FMA, since the fold sinks the read past it
• the target is not AVX-512 and the load is a broadcast, since the embedded broadcast modifier needs EVEX

Both multiply operands are considered, as multiplication commutes.

On tests/filecheck/projects/libxsmm at arch=avx512 this removes all eight broadcasts:

	before	after
instructions	28	20
broadcast instructions	8	0

and eight fewer live vector values reach the register allocator. The generated assembly still assembles, links against the existing main.c and matches the naive reference, checked on hardware with avx512f, avx512vl and fma.

Run it before x86-allocate-registers so the freed registers are available to the allocator.

The last commit is unrelated housekeeping I hit while debugging this: raising DiagnosticException from inside except KeyError chained the dict lookup onto the traceback, so a vector too wide for the target printed KeyError: 512 above the actual explanation. Happy to split it out if you would rather.

[codecov]

Codecov Report

❌ Patch coverage is 93.47826% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.94%. Comparing base (b642df9) to head (1c36d9a).

Files with missing lines	Patch %	Lines
xdsl/transforms/x86_fold_memory_operands.py	91.17%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6360      +/-   ##
==========================================
+ Coverage   86.93%   86.94%   +0.01%     
==========================================
  Files         439      440       +1     
  Lines       65902    65988      +86     
  Branches     7490     7502      +12     
==========================================
+ Hits        57289    57372      +83     
+ Misses       7036     7034       -2     
- Partials     1577     1582       +5     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[superlopuh]

This is cool but also touches on some tricky things, where technically this is not always an optimisation, and is an instruction selection and scheduling concern. In that sense it does feel like it might be worth having the rewrite pattern somewhere, but I'm not 100% sure that it would be pulling its weight as for the time being I'd expect this kind of choice to be made before the IR is in x86. What motivated you to make this change? If there's a reason to add it I'd be open to it, then it would be good to discuss what the objective of the pass should be.

Also yes, please open a new PR for the KeyError fix, with a dedicated test.

[manishpaulish]

Fair points, and I think you are right on the architecture.

On what motivated it: I was measuring the libxsmm matmul pipeline and noticed the RSM forms of vfmadd231pd/ps are defined, registered and asm-tested, but no pass ever generates them, so they are effectively dead code in the dialect. What I was actually chasing was register pressure. In that kernel at arch=avx512 the eight broadcasts each hold a vector register for exactly one use, and folding them frees eight registers going into allocation. That mattered to me because while measuring for #6358 I found the allocator gives up with Error allocating op rather than spilling once it runs out, so registers freed before allocation are worth more than the instruction count suggests.

On "not always an optimisation", you are right that the pass has no basis for the claim. Two things worth separating:

The regression I would have expected, a load hoisted out of a loop getting sunk back in and reloading every iteration, does not actually happen, because _safe_to_sink requires the load and the FMA to be in the same block. But that is incidental. I wrote that restriction to keep the memory-write check simple, not because I had reasoned about loop hoisting, and it would stop being true the moment anyone relaxed it.

The case I have no answer for is scheduling. Whether the fused form wins depends on whether the kernel is load-port bound or register bound, and the pass has no cost model to reason about that at all. It just always folds.

So I would not argue for this landing as an unconditional pass. The two coherent homes seem to be either in lowering, where convert-vector-to-x86 selects the folded form directly, or as a register-pressure-driven decision late in the backend, roughly where LLVM does spill folding. Both of those want the target information that #6358 would add, so I would rather park this behind that work than push it in as a peephole with no cost model.

Happy to close this if you would rather not have it sitting open. Otherwise I will leave it as a draft and come back once the arch plumbing exists and there is something to make the decision with.

KeyError fix is split out with a dedicated test, opening that separately.