[llvm] [AMDGPU] Price an add reduction over a vector of i1 (PR #217327)

[via llvm-commits]

michaelselehov wrote:

---

## What this changes

`GCNTTIImpl::getArithmeticReductionCost()` now charges 4 instructions per element for an add reduction over a vector of i1. Nothing outside AMDGPU changes. The patch is 12 lines of code and two tests.

## Why the current cost is wrong

The SLP vectorizer emits an unsigned add reduction over `<n x i1>` as `ctpop(bitcast <n x i1> to iN)`. See `SLPVectorizer.cpp`, in `emitReduction()`: the ctpop form is used whenever the reduction kind is `Add`, the vector element type is i1, and the destination type is wider than i1.

Every element of that vector is a separate mask. The bitcast has to pack them: the target selects a 0 or a 1 per element, shifts it into place and merges it. The generic model prices the reduction as a shuffle tree instead, which is far cheaper than the packing. So the vectorizer replaces a cheap scalar chain of adds with an expensive packed form and calls it a win.

## The measurement

I compile two kernels that differ only in the reduction, and I subtract a baseline that does the same loads and compares. The operand is divergent. I exclude `s_delay_alu` and `s_wait_alu`, which GFX11 and GFX12 need between dependent operations and which no cost model counts. Script and raw counts are reproducible with any `llc`.

| target | instructions for 8 elements | per element |
|---|---|---|
| gfx900 | 33 | 4.1 |
| gfx942 | 33 | 4.1 |
| gfx1030 | 33 | 4.1 |
| gfx1100, `real-true16` off | 33 | 4.1 |
| gfx1100, `real-true16` on | 41 | 5.1 |
| gfx1201, `real-true16` off | 33 | 4.1 |
| gfx1201, `real-true16` on | 41 | 5.1 |

At 16 elements the numbers per element are 4.2 and 5.4. The packing is the same sequence everywhere: 8 compares, 8 selects, 7 shifts, 7 ors and one bit count. `real-true16` adds one `v_mov_b16` per element, which I do not charge, so the estimate stays below the measurement on every subtarget and the cost needs no subtarget check.

## The effect

The regression that led me here is rocSPARSE `csrgemm` on gfx1201, Navi48. Measured against a ROCm build without the regression, 29 matrices from the rocSPARSE QA set, 5 repetitions per arm, arms alternate per repetition:

| quantity | median | worst |
|---|---|---|
| the regression | 6.7% more time | 32.1% on `torso1` |
| with this patch, against the build without the regression | 0.5% less time | 6.9% on `torso1` |

28 of the 29 matrices land within 3% of the build that has no regression. For `torso1` I also built the same source with `-vectorize-slp=false`: it runs in 11.08 ms, the patch runs in 11.09 ms, so the patch already takes everything the vectorizer can give there and the rest has another cause.

## Answers to the questions I expect

**Why not `getCastInstrCost()`, since the bitcast is what costs?** I tried exactly that, alone: the bitcast of `<8 x i1>` then prices at 32 instead of 8, the add reduction keeps its generic cost, and the generated code does not change at all, 516 packing sites in the module both ways. SLP takes the `getArithmeticReductionCost()` branch for this tree, and `emitReduction()` then builds the bitcast without asking anyone for a price. A cast cost is invisible to this decision. It is still worth fixing for the and and or reductions, which really do lower to a packed compare, but that is a separate change.

**Why only `Add`?** Because only the add reduction is emitted as `ctpop(bitcast ...)`. The and and or reductions over i1 are also too cheap today, by about four times, but they lower to a packed compare rather than a bit count, so the same number does not describe them.

**This is a middle end bug, so fix the middle end.** I agree that SLP prices one form and emits another, and I filed ISSUE-NUMBER for it. The cost side cannot move first. Here is what every target says today about an add reduction over `<8 x i1>`, next to what the hardware needs:

| target | cost SLP asks today | cost of the form SLP emits | real instructions, parity form | real instructions, packed form |
|---|---|---|---|---|
| x86-64-v3 | 2 | 2 | 10 | 2 |
| aarch64 | 2 | 6 | 1 | 1 |
| thumbv8.1m.main+mve | 394 | 6 | 7 | 40 |
| riscv64+v | 2 | 43 | 0 | 0, one `vcpop.m` |
| gfx1201 | 15 | 9 | 15 | 41 |

Route the query to the form that SLP emits and MVE drops from 394 to 6 while the code needs 40 instructions, so it starts emitting the slow sequence. RISCV climbs from 2 to 43 while one `vcpop.m` does the job, so it stops emitting the fast one. AMDGPU still sees 9 against 41 and needs a target number anyway. Every target needs its own cost before the dispatch can change.

**You are overpricing a real parity reduction.** True. An add reduction that keeps the i1 result is a parity computation, about 2 instructions per element, and I charge 4. The hook sees the type, not the form, so it cannot tell the two apart. It does not change a decision on this target. A chain of `add i1` returning i1 is refused by SLP at cost 1 without this patch, at 4, 8, 16 and 32 elements, so the patch only widens a margin that already pointed the same way. I also checked the other clients of this hook with a patched and an unpatched `opt`: the loop vectorizer, VPlan, VectorCombine and LowerMatrixIntrinsics all produce identical output. The loop vectorizer does build `reduce.add(<n x i1>)`, but the reduction sits in the middle block and the loop cost does not include it, so the reported cost for VF 2 is the same with and without the patch. Only `-prefer-inloop-reductions` exposes the number, and then the loop still vectorizes and only the interleave count changes. That flag is off by default and only ARM opts into in-loop reductions.

**Make the backend cheap instead of teaching the cost model.** That is the real fix, and it would remove the need for this number: the target could count the bits with a chain of adds over the masks and never pack them. It is a much bigger change, and it does not help the other targets in the table above.

## Testing

`check-llvm` is clean. The new cost model test uses one check prefix for gfx900, gfx1030, gfx1201 and gfx1201 with `-real-true16` off, which is how it states that the cost does not depend on the subtarget. It also covers code size, a single element vector, 64 elements and a scalable vector. The SLP test shows the scalar chain that the vectorizer now leaves alone, and a forced arm that shows the packed form it would otherwise build. Both tests are generated with the update scripts.


https://github.com/llvm/llvm-project/pull/217327