Repository navigation
Add #[nounroll] (or #[unroll(1)]) so ptxas keeps a loop rolled #1468
Description
What problem does this solve?
There's no way to keep a loop rolled through ptxas. ptxas fully unrolls small constant-trip-count loops on its own, even when the PTX keeps them rolled, and then schedules every iteration's loads at once. In speakrs's tensor-core convolution kernels (crates/speakrs-cuda-kernels/src/wideconv/tensor.rs, https://github.com/avencera/speakrs), that raised register pressure and hoisted later taps' weight loads, so we keep those loops rolled on purpose.
#[unroll(N)] only accepts factors in 2..=1024, so #[unroll(1)] isn't available. And per #1375, #[unroll(N)] doesn't emit loop metadata.
Proposed API / behavior
#[nounroll] // or #[unroll(1)]
while tap < TAPS { /* … */ tap += 1; }This would lower to llvm.loop.unroll.disable metadata on the loop. The NVPTX backend then emits .pragma "nounroll"; at the loop header, which ptxas honors. The libNVVM route should do the same.
Alternatives considered
We emit the pragma by hand at the top of the loop body:
#[inline(always)]
pub(crate) fn nounroll() {
unsafe { ptx_asm!(".pragma \"nounroll\";"); }
}This works, but it depends on where the asm lands relative to the loop header. dialect-ptx/src/scopes.rs already has to reason about the scope of a block-level .pragma "nounroll" that comes from inline asm.
Additional context
Related: #1375 (#[unroll(N)] emits no loop metadata, so the NVVM backend unrolls a second time).
Activity
nihalpasham commentedon Oct 10, 2026
Agreed, having an explicit #[nounroll] would be useful. I checked the metadata through LLVM and libNVVM, including LTO, and ptxas honors the resulting pragma.
A PR should preserve that loop policy through transformations and lowering, including debug builds.
For this one - we should cover nested loops and the main/remainder loops from partial unrolling; that’s also relevant to #1375.
There’s already a range-based for loop unrolling PR from @midagedev, #1346, which we’ll review. We should check how the two interact as well.
praveenperera commentedon Oct 10, 2026
Thanks. One correction to the issue text: I said
dialect-ptx/src/scopes.rsalready has to reason about a block-level.pragma "nounroll"that comes from inline asm. That isn't right.ScopeFlattenPlanis only used by the PTX raising indialect-ptx/src/raising/plan.rs, whichptx-scheduleuses forcargo oxide fuzz-schedule. It isn't on the kernel compile path, so it has no say in where a hand-writtenptx_asm!pragma ends up.hhdri commentedon Oct 11, 2026
A data point for the nested-loop case. In this register-MMA kernel, ptxas fully unrolls a 4-trip K loop even when its trip count is hidden behind inline asm. A hand-placed
.pragma "nounroll"keeps it rolled and removes most spills. This looks like the same failure that motivates the second part of #1235.Setup: cuda-oxide
bafffe45(currentmain), LLVM NVPTX route,ptxas13.4 (V13.4.92)-O3,sm_89(RTX 4060).The kernel is the 128×128 int8
mma.sync.m16n8k32Q4_K prefill tile from my comment on #1346 (#1346 (comment)), with 64 f32 accumulators per thread and a 128-register cap fromlaunch_bounds(256, 2). The shared-memory K loopwhile base < 8 { …; base += 2 }runs 4 trips. Each trip contains two 2-trip#[unroll]loops, one staging and one MMA. The MMA loop contains#[unroll]loops of 4 trips and of 16 trips, and the 16-trip loop has another 4-trip loop inside. Only the outer K loop's control changes between builds:mma.syncIMMABAR8ptx_asm!("mov.u32 %0, 8;")ptx_asm!(".pragma \"nounroll\";")first in the bodyIn the last row ptxas leaves the K loop rolled as well (3 barriers), presumably because it now contains a loop.
Timing. Each measurement covers the tile plus its activation quantization, run as a CUDA graph and compared with a native CUDA kernel in the same process. The GPU was shared with another job, so I kept only runs with no overlapping GPU process: 4 samples per shape per build. Compared with the opaque-count build:
In this kernel the measurable effect comes mostly from the visible-constant case doubling the spills. Removing the remaining spills gives only a small further gain.
Observations:
mma.sync). In SASS, the first two rows are fully unrolled (128IMMA, 9 barriers). That matches the issue description.mov.u32 %r, 8, so ptxas still unrolls 4×. We had used this pattern in our kernel, believing it kept the loop rolled. It changes the generated code enough to halve the spills compared with a visible constant, but it doesn't keep the loop rolled, so it can't replace#[nounroll].// => This Inner Loop Header: Depth=2), and ptxas honored it. Correct placement still depends on where LLVM puts the asm after loop rotation, so loop metadata would make it reliable. All rows produce bitwise-identical output: our checks compare against a native CUDA kernel and CPU references.forloops #1346: with that PR cherry-picked ontobafffe45, I wrote all ten#[unroll]loops in the kernel asfor x in 0..Nusizeand kept the pragma on the outer K loop. The build emits no warnings, and its PTX is identical to the counted-while+ pragma build apart from a shared-memory symbol hash, with the pragma still first in the K loop header. So feat(unroll): recognize rangeforloops #1346's canonicalization of inner loops doesn't disturb an outernounroll. This is only the hand-written pragma, though; an attribute-based#[nounroll]would need its own check against feat(unroll): recognize rangeforloops #1346.Relation to #1235. Its proposed counted register-MMA carrier is motivated partly by "libNVVM expanded all four K-steps before ptxas and introduced 168-216 bytes of spill traffic". That is the same failure as rows 1 and 2, happening in ptxas instead of libNVVM. A
#[nounroll]that reaches LLVM/libNVVM asllvm.loop.unroll.disableand ptxas as.pragma "nounroll"would likely cover the loop-expansion and spill part on both routes without the carrier. The u32 shared-address recurrence part of #1235 is separate and would still be needed.Prepared with help from an AI assistant (Claude Code).