llvm / llvm/llvm-project

[NVPTX] Branch folding from #191889 makes ptxas stop unrolling loops (26% regression on a GEMM inner loop)

Open
#219,778 0 comments 1 reaction 0 assignees View on GitHub
backend:NVPTX
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

We maintain a Rust-to-PTX compiler (cuda-oxide) and hit this when we moved from LLVM 22 to 23: one of our GEMM kernels quietly dropped 26%. Same IR in, same math out, bit-identical results. The only thing that changed was how llc lays out the loop branches.

Since #191889 (NVPTX reverseBranchCondition), branch folding and block placement collapse the classic two-jump loop idiom into a single negated conditional. Perfectly legal, one instruction smaller:

```text
llc-22: llc-23:
setp.eq.b32 %p11, %r3, %r26; setp.eq.b32 %p10, %r3, %r26;
@%p11 bra EXIT; @!%p10 bra LOOP;
bra.uni LOOP; // exit falls through
```

The problem is what ptxas does with it. Its SASS-level loop unroller apparently recognizes hot loops by that old branch layout, and it just stops unrolling the new form. Same module, CUDA 13.3 ptxas, sm_120:

```text
llc-22 llc-23
registers 38 26
SASS FFMA 18 4
SASS LDG.E 35 7 <- loads in flight collapse
runtime 0.297 ms 0.376 ms (1024^3 SGEMM inner loop,
7218 GFLOPS 5708 GFLOPS identical numerics)
```

We hand-edited the PTX to isolate exactly what ptxas keys on. It needs BOTH of these, and either one in the folded form kills the unrolling:

1. the loop guard's conditional branch must have the loop preheader as its TAKEN target (llc-23 falls through into the preheader instead)
2. the latch must be "conditional exit + `bra.uni` back to the loop"

Negation itself is innocent by the way: `@!%p bra PREHEADER; bra.uni SKIP;` still unrolls fine. It's purely about which edge is the taken target.

I looked around a bit to see if this was a known issue and found the following:

- On the parent PR (#185755), @Artem-B asked whether this helps at the SASS level and suggested "forwarding to ptxas folks". This is the concrete downside measurement for that conversation. We're filing the ptxas side internally at NVIDIA in parallel, since ideally ptxas should recognize both layouts.
- Someone else likely hit the same class already, unanswered: https://forums.developer.nvidia.com/t/truncated-loop-exit-compare-in-ptx-prevents-ptxas-from-unrolling-loops/375849

Our workaround: we now pass `-disable-branch-fold -disable-block-placement` to llc. That restores the unrollable layout and full performance (we swept ~700 kernels; every register change was either the unroller coming back or a return to the LLVM-22 numbers). But it's a blunt instrument. Something scoped would be much nicer, e.g. a TTI hook the way relative lookup tables are handled, or simply not folding the loop guard/latch pair on NVPTX until ptxas catches up.

Contributor guide

Open the contributing guide

Research direction

Start at the NVPTX branch-folding and block-placement changes from #191889, using llc-22 and llc-23 with the reported GEMM PTX as the reproduction. Compare the generated branch layouts and test the scoped workaround against the -disable-branch-fold and -disable-block-placement baseline; done means NVPTX preserves a ptxas-unrollable loop layout without disabling those passes globally.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
compilers, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.