Skip to content

[llvm] Stop failed SLP ordered reductions from starving the block - #1297

Open
Phil-amd wants to merge 1 commit into
mainfrom
fix/issue-1273-slp-ordered-reduction-budget
Open

Phil-amd wants to merge 1 commit into
mainfrom
fix/issue-1273-slp-ordered-reduction-budget

Conversation

@Phil-amd

Copy link
Copy Markdown
Member

Fixes #1273.

What. A new LLVM extension, slp-ordered-reduction-budget-refund, stops SLP's failed ordered-reduction attempts from using up the per-block scheduling budget. Since llvm/llvm-project@94e366ef2060, which is in our pin, SLP tries strict fadd chains as ordered reductions. In a large unrolled kernel every rejected attempt still charges slp-schedule-budget (100000 per block, never given back). Once that budget is spent, SLP can no longer schedule the profitable vectorization that comes after the attempts.

On MI355X, Primus-Turbo's rmsnorm_mxfp4_fused_kernel (32768×2880 bf16):

Time VGPR+AGPR Spilled VGPRs
Before 297 us 512 346
After 152 us 220 0

The outputs are bit-identical.

The switch -fly-slp-ordered-reduction-refund is on by default. To turn it off for one kernel:

compile_hints={"llvm_options": {"fly-slp-ordered-reduction-refund": False}}

Why not disable-binop-extract-shuffle, the workaround proposed in #1273. It turns off a general VectorCombine fold, has to be set per kernel, and leaves the cause in place. Any large unrolled kernel can still lose the vectorization that follows a failed ordered reduction. The two combine well, though:

On this kernel Time VGPRs
This change alone 152 us 220
The workaround alone 140 us 238
Both 140 us 242 (same as 0.2.4)

With this change alone, a few cross-load packs remain, and that fold rewrites them into shuffles. That is why the workaround is still faster here.

Why the switch is on by default. CONTRIBUTING.md asks for an extension switch that is off by default. Keeping this one off would need FlyDSL to:

  • set a global cl::opt on every compile;
  • check that the LLVM in use registers the switch, which FLYDSL_LLVM_NO_EXT=1 builds and external LLVMs do not;
  • accept that code calling gpu-module-to-binary directly gets no fix.

It can still be turned off without rebuilding LLVM, which is what the rule protects.

Key points

  • Narrow trigger. It acts only when an ordered attempt in HorizontalReduction::tryToReduce vectorizes nothing. On gfx942, none of the 520 kernels compiled by tests/kernels change, and compile time is unchanged. With the switch off, the ISA is byte-identical to the 0.3.4.1 wheel's.
  • Only the root is marked. A failed root is not retried. This follows llvm/llvm-project@1b10cd91, which also treats only the root of a failed ordered reduction as a dead end; its leaves stay valid roots. Marks are cleared at each new block.
  • Not covered: the tryToReduceOrdered path.
  • Upstream: not yet submitted, although CONTRIBUTING.md asks for that first. Upstream main still regresses on this kernel (116 spilled VGPRs).

@Phil-amd Phil-amd self-assigned this Oct 15, 2026
llvm/llvm-project@94e366ef2060 lets SLP try strict fadd chains as ordered
reductions. In a large unrolled kernel each attempt schedules trees spanning
thousands of instructions and charges them to the per-block scheduling budget
(slp-schedule-budget, 100000), which a discarded tree never gives back. In
Primus-Turbo's rmsnorm_mxfp4_fused_kernel one root's 31 rejected bundles use up
the whole budget, so the profitable vectorization that follows cannot be
scheduled: on gfx950 the kernel needs 512 VGPR+AGPR, spills 346 VGPRs and runs
2.2x slower than with the LLVM in 0.2.4 (#1273). Upstream main still regresses.

The extension slp-ordered-reduction-budget-refund handles an ordered attempt
in HorizontalReduction::tryToReduce that vectorizes nothing: it gives the
attempt's budget back and marks the root as analyzed so it is not retried.
llvm/llvm-project@1b10cd91 likewise treats only the root of a failed ordered
reduction as a dead end. The tryToReduceOrdered path is not covered. Not yet
submitted upstream.

Its switch, -fly-slp-ordered-reduction-refund, is on by default, so everything
that uses the bundled LLVM gets the fix, including code that runs
gpu-module-to-binary itself. A kernel that regresses turns it off with
compile_hints={"llvm_options": {"fly-slp-ordered-reduction-refund": False}}.
This departs from the "(default off)" wording in CONTRIBUTING.md. Keeping the
switch off would need FlyDSL to set it on every compile, which means a global
cl::opt set and restored per compile, a registration check so that
FLYDSL_LLVM_NO_EXT=1 and external LLVM builds still compile, and no fix on
paths outside FlyDSL's pipeline.

On MI355X the kernel drops from 297 us to 152 us and from 512 VGPR+AGPR with
346 spilled to 220 VGPRs with none; its outputs are bit-identical. With the
switch off its ISA is byte-identical to the 0.3.4.1 wheel's. On gfx942, none of
the 520 kernels compiled by tests/kernels change, and compile time is
unchanged.
@Phil-amd
Phil-amd force-pushed the fix/issue-1273-slp-ordered-reduction-budget branch from eec9b5d to f4eb1d3 Compare October 16, 2026 10:04

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant