feat(unroll): recognize range for loops

[midagedev]

Summary

#[unroll] on for i in a..b does not unroll ("no recognized induction variable"), so arrays indexed by i stay in local memory. The loop exits through a match on the Option from the inlined Range::next. This adds the dialect-mir canonicalization mentioned in #559. No block is copied.

#559 asked for a real use. Our GDN kernel had a 192 B local depot from range for loops over lane arrays. We had to rewrite them as counted while loops.

Changes

• canonicalize.rs: thread_iterator_exit_test. Only for a hinted loop whose header does not exit, so loops that unroll today are unchanged.
• Other iterator loops (a..=b, .step_by) still warn and stay loops.
• Tests use the existing evaluate_integer; one unroll_smoke kernel; docs.

Testing

• New tests fail on main, and on variants that read the wrong payload or give a both-reach value to one arm. i8/u8 boundary sweeps.
• Repro: 32 B depot, 8 ld.local -> none.
• just check passes (5,506 tests)
• cargo oxide run <example> passes (unroll_smoke on RTX 3090, LLVM and NVVM paths, memcheck and synccheck 0 errors)
• New example added (if applicable) — one kernel added to unroll_smoke

Checklist

• All commits signed off (git commit -s)
• SPDX headers on new source files (no new files)

Written with help from an AI assistant. I ran every test above.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hhdri]

I tested this PR on the real-use question from #559 and found one gap: range loops with signed counters, including the default i32, are still rejected.

Setup: PR head 5a59773 cherry-picked cleanly onto bafffe45 (current main). RTX 4060 (sm_89), driver 615.78.08, CUDA 13.4 (ptxas V13.4.92), nightly-2026-08-28.

Real use: a Q4_K prefill MMA tile

The kernel is a 128×128 int8 mma.sync.m16n8k32 prefill tile for an LLM server, with launch_bounds(256, 2) and 10 #[unroll] loops (trip counts 2–16, nested up to three deep around the MMA). I built it twice, changing only the loop form: counted while versus for i in 0..N. Resource counts come from ptxas -O3 -v.

Build	#[unroll] warnings	PTX ld/st.local	mma.sync in PTX	stack / spill st / spill ld
counted while, main	0	0	32	72 / 104 / 96 B
range for, main	10	96	1	320 / 0 / 0 B
range for, this PR	3	0	16	0 / 0 / 0 B
range for with usize counters, this PR	0	0	32	same PTX as counted while

Timing covers the tile plus its activation-quantization kernel: 40 launches per CUDA graph after warmup, 8 samples per shape, with builds interleaved. Medians:

M×N×K	counted while	range for, main	range for, this PR
128×12288×4096	0.373 ms	1.360 ms (3.65×)	0.361 ms (0.97×)
128×1024×4096	0.098 ms	0.333 ms (3.39×)	0.111 ms (1.13×)*
128×4096×4096	0.173 ms	0.500 ms (2.89×)	0.165 ms (0.96×)
512×12288×4096	1.539 ms	5.475 ms (3.56×)	1.486 ms (0.97×)

*In those runs the paired native CUDA kernel was also about 10% slower. Against it, the ratio is 0.79 versus 0.76, so the difference is about +3.5%.

On main, idiomatic range loops make this kernel 2.9–3.65× slower. With the PR and usize counters, the generated PTX matches the counted-while version exactly, apart from a shared-memory symbol hash. The PR doesn't change the counted-while PTX either. All four builds pass our GPU checks: bitwise against a native CUDA kernel and against CPU references, including zero inputs and signed-scale extremes.

Gap: signed counters (i32 by default)

Three of the ten loops never use their counter as an index, only in casts and shift amounts, so Rust's integer fallback makes it i32. The PR still rejects those three:

warning: #[unroll] requested but the loop was not unrolled: its exit test is not in the loop header, and it is not a range `for` loop over `a..b`

Minimal reproducer. With acc[j] indexing (usize counters), the PR fixes this kernel: ld/st.local drops from 18 to 0, and the PTX matches the counted-while form. Casting the counter instead leaves it i32:

#[kernel]
pub unsafe fn range_i32(input: *const u32, output: *mut u32) {
    let mut acc = [0u32; 16];
    let mut iteration = 0usize;
    while iteration < 4 {
        #[unroll]
        for j in 0..16 { // `j: i32`: only used through `as`
            acc[j as usize] = acc[j as usize].wrapping_add(unsafe { *input.add(iteration * 16 + j as usize) });
        }
        iteration += 1;
    }
    #[unroll]
    for j in 0..16 {
        unsafe { *output.add(cuda_device::thread::threadIdx_x() as usize * 16 + j as usize) = acc[j as usize]; }
    }
}

With this PR it still emits 2 warnings, 18 ld/st.local and 3 branches, the same as main. An explicit 0..16u32 or 0..16usize range unrolls without warnings.

Cause. I added a debug print at each rejection in thread_iterator_exit_test. The check that fires is the requirement that each arm ends in goto join. In core on this nightly, signed Step::forward_unchecked is start.checked_add_unsigned(n as $unsigned).unwrap_unchecked(), while unsigned uses unchecked_add. That leaves an extra branch in the Some arm (trimmed dump):

header(i: si32):  c = mir.not(mir.lt(i, 4));  mir.cond_br c [^none, ^some]
^none:            o = construct_enum None;   mir.goto ^join(o, i)
^some:            t = mir.checked_add(i, 1);  ...;  mir.cond_br(not overflow) [^some2, ^ovf]
^some2:           o = construct_enum Some(i); mir.goto ^join(o, t.0)
^ovf:             (outside the loop) mir.goto ...

The overflow edge is dead given the header's i < b. Following the single-entry ^some → ^some2 chain, or folding that checked_add, would probably cover signed ranges. The signed i8 fixtures in tests/common/mod.rs use a plain mir.add in the Some arm, and unroll_smoke only uses usize/u32 counters. An inferred-i32 case in unroll_smoke would exercise the rustc shape. If signed ranges stay out of scope, the docs and warning should say so, because a counter never used as an index defaults to i32.

A side note that also affects main: the warning has no source span. Finding the three loops in this kernel meant converting them one at a time.

Prepared with help from an AI assistant (Claude Code).

[midagedev]

Thank you for testing this on a real kernel, and for finding the cause. I can reproduce the shape from your dump: signed Step::forward_unchecked leaves a checked_add branch in the Some arm. I am extending the match to follow that single-entry chain, with an inferred-i32 case in unroll_smoke and a signed fixture in the tests. I will push it to this PR.

[midagedev]

Pushed 20e30b6. Thank you for the clear report and the reproducer.

The cause was as you described: core's signed step goes through checked_add, with the tuple in a stack slot and a branch on the overflow flag. The pass now folds that step to a plain add (sub for .rev()) when the header's < test shows it cannot overflow and the overflow edge leaves the loop.

Your range_i32 kernel went from 2 warnings, 18 ld/st.local and 3 branches to 0, 0 and 0 (sm_80, LLVM 21.1.8), and its PTX equals the usize form. unroll_smoke has a new i32 kernel, and the existing kernels' PTX is unchanged. The tests add a signed fixture in the backend's real shape and two that must stay rejected (an overflow edge that re-enters the loop, and a step of 2).

The missing source span in the warning is a separate change; I left it out of this PR.

[hhdri]

Thanks, confirmed. 20e30b60 fixes the signed case.

Setup: both PR commits cherry-picked cleanly onto bafffe45 (current main). RTX 4060 (sm_89), CUDA 13.4 (ptxas V13.4.92), nightly-2026-08-28.

Kernel	main	PR at 20e30b60
range (usize counter)	2 warnings, 18 ld/st.local, 3 bra	0 warnings, 0 / 0
range_i32 (inferred i32)	2 warnings, 18 ld/st.local, 3 bra	0 warnings, 0 / 0
counted while	0 / 0	0 / 0

Both range kernels now have the same instruction counts as the counted-while version (64 ld.global, 16 st.global). Our production kernels (counted while loops, 64 #[unroll] sites) produce identical PTX with and without the PR, apart from shared-memory symbol hashes. cargo test -p mir-transforms passes on the cherry-pick.

Prepared with help from an AI assistant (Claude Code).