[NVPTX] Branch folding from #191889 makes ptxas stop unrolling loops (26% regression on a GEMM inner loop)
- 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
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