feat(unroll): recognize range for loops

[midagedev]

What this adds

#[unroll] on for i in a..b and for i in (a..b).rev() now unrolls. After inlining, such a loop exits through a match on the Option returned by Range::next, so the unroll pass found no header exit test and no induction variable; arrays indexed by i stayed in local memory. A new dialect-mir canonicalization threads that match into a header exit test, and the signed-counter checked_add/checked_sub step core emits is folded to a plain add/sub where the header test proves it cannot overflow. Casts between identical dialect types (isize/i64, usize/u64) are elided in the importer so isize loops look like i64 loops.

nested_lane_range (16x16 over a 16-lane accumulator), PTX:
before: 17 ld.local, 17 st.local, .local .b8 __local_depot[64]
after:   0 ld.local,  0 st.local, no depot

let mut acc = [0u32; 16];
#[unroll]
for i in 0..16 {
    #[unroll]
    for j in 0..16 { acc[j] = acc[j].wrapping_add(mix(input[i * 16 + j])); }
}

What we fixed

• Range for loops over any integer counter, including isize, with full #[unroll] and partial #[unroll(N)], and (a..b).rev() with full unroll.
• The signed step: core's checked_add overflow branch is dead under the header's strict lo < hi test and is folded away; the fold refuses any other step or an overflow edge that re-enters the loop.
• No-op casts no longer hide the counter from the pass.
• a..=b, .step_by(k) and custom iterators still warn and stay loops; the warning text now distinguishes an unrecognized step from a loop that is not an a..b range.

Implementation

• canonicalize.rs: thread_iterator_exit_test runs only for a hinted loop whose header cond_br has both successors inside the loop. It requires single-entry arms ending in goto to a two-predecessor join that reads the discriminant of one argument, each arm building its value with mir.construct_enum, and evaluates the match blocks with a small evaluator that accepts only storage markers, constants, the discriminant read, IntToInt casts and eq/ne. Replacements are planned before any mutation; a match value used where both targets reach is refused. Dead blocks are removed by the next simplify_cfg round in unroll.rs.
• canonicalize.rs: fold_checked_step matches core's checked_add/checked_sub by constant 1 on the counter the header orders strictly, with the overflow edge leaving the loop.
• analyses/induction.rs: match_cmp, unwrap_not, const_i128 and CmpPred::negate made pub(crate) for reuse.
• mir-importer translator/rvalue/expr.rs: a cast whose operand type equals the result type returns the operand and emits nothing, beside the existing reborrow shortcut. Pointer-kind, signedness and width changes still emit mir.cast.
• unroll_smoke: new kernels range_for_loops, signed_range_for_loops, signed_partial_range_for, isize_range_for, nested_lane_range; verify-code-shape.sh fails the build when nested_lane_range or isize_range_for uses local memory. nested_lane_range keeps a runtime inner while so LLVM cannot unroll it alone; it is the discriminating gate. isize_range_for is a control, since LLVM already unrolls it on main.
• Docs: #[unroll] macro docs, thread.rs, supported-features and the three book pages list the same supported forms.

Verification

Checked at 4d9b8295:

• ✅ cargo test -p mir-transforms: 91 passed. New tests fail on main; i8/u8 sweeps over every boundary pair, interpreter over every end value for factors 2, 3 and 4, negative cases for a re-entering overflow edge, a step of 2 and a value read after both exits.
• ✅ just test: 3986 passed; just book builds; fmt, clippy, rustdoc and all 18 check-*.sh guards pass; sync-example-locks.sh --check OK after merging current main.
• ✅ cargo oxide run unroll_smoke on RTX 5090 (sm_120), LLVM and libNVVM routes: PASS; memcheck and synccheck 0 errors.
• ✅ nested_lane_range with main's backend: 17 ld.local, 17 st.local, 64-byte depot; here 0/0/none. The 17 pre-existing kernels are instruction-identical to main.
• ✅ Scratch kernels: u8/i8/u16/i64/u32::MAX boundaries, empty ranges, negative isize bounds, rev(), partial remainders for several n, break with value, continue, runtime bounds, mixed casts; all correct, 0 local memory, memcheck 0 errors.

Caveats

• #[unroll(N)] on (a..b).rev() or on a loop with a runtime start still warns, as it does for counted while loops.
• a..=b and .step_by(k) are not recognized.
• The fold matches core's current step sequence; a future change there turns into a warning, which nested_lane_range in the compile-only CI lane would surface.
• #[unroll(N)] still emits no loop metadata ([BUG]: #[unroll(N)] emits no loop metadata, so the NVVM backend unrolls the loops a second time #1375), so NVVM may unroll the remainder again.
• #[unroll] rejects a loop bound computed from a const generic (N / 8) #1494 (const-generic bounds) and #[unroll] over the clone budget leaves an array-indexing loop rolled and moves the array to .local, with no way to raise the budget #1478 (clone budget) are not addressed.

[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).

[nihalpasham]

Thanks for the PR, it looks good: no soundness problems but three things before it lands:

• isize counters still bail. for i in 0isize..8 warns and stays a loop; every other integer type unrolls. The cause is small: core adds 1 to an isize by casting to i64 and back, both are the same type in our IR, so the step block carries two casts that change nothing and match_checked_step does not recognize it. Fix it in the importer, not the pass: when a cast's source and target translate to the same type, return the operand and emit nothing. The transmute path in the same function already does this. The isize block then looks like the i64 one, and every other pass stops seeing no-op casts too.
• The example cannot show the pass working. The new kernels are small enough that LLVM unrolls them on its own; built with main's backend they already have 0 ld/st.local. Add a kernel LLVM leaves alone, say a 16x16 nested range loop over a 16-lane accumulator (main: 17 ld.local, 17 st.local and a local depot; this PR: 0/0), plus isize loops (a..b, .rev(), #[unroll(4)]), and a verify-code-shape.sh for unroll_smoke that asserts no .local in those kernels so CI catches a regression. Several examples have one to copy.
• Docs: appendix/supported-features.md still says counted while loops only; the macro and thread.rs comments and the three book pages should list the supported forms the same way, isize included once it works; and the warning for an unrecognized step must not call a real range loop "not a range loop".

Also merge main in (conflict free): the branch pins nightly-2026-08-28 and CI runs on the head sha, so today it would test the wrong toolchain.

In short, to land this: we'll need the importer no-op cast fix with isize loops in unroll_smoke, the heavy kernel with its verify-code-shape.sh, the docs fixes, main merged. I will vet CI on that head.

[nihalpasham]

Reviewed 20e30b60: changes requested. The details and the list to land this are in my comment above: #1346 (comment)

[midagedev]

Thanks for the review. All four are done on feat/unroll-range-for.

• Importer: a cast between identical types now returns the operand, like the reborrow path. isize loops over a..b, (a..b).rev() and #[unroll(4)] now unroll with no warnings.
• Example: added nested_lane_range (16x16 over a 16-lane accumulator) and isize_range_for. Without #[unroll], nested_lane_range has 17 ld.local, 17 st.local and a 64-byte local depot. With it, 0, 0 and no depot. On the current LLVM a plain 16x16 loop gets unrolled anyway, so the inner loop has a runtime bound to keep it from doing that. verify-code-shape.sh fails CI on any .local in these two kernels.
• Docs: supported-features.md, the macro and thread.rs comments, and the three book pages list the same forms. The warning on a real range loop no longer says "not a range loop".
• Merged upstream/main, so CI runs on nightly-2026-10-11. cargo test -p mir-transforms passes (91 tests).

Ran unroll_smoke on an RTX 3090 (sm_86). Results match (unroll_smoke: PASS, 32 threads).

[nihalpasham]

/ok to test 4d9b829

[nihalpasham]

/ok to test 660873c

[nihalpasham]

Checked at 660873c0 on the 5090:

• All four items from the last review are in. The importer drops a cast only between identical dialect types (isize/i64, usize/u64, char/u32); signedness, width and pointer-kind casts still emit mir.cast.
• nested_lane_range with main's backend: 17 ld.local, 17 st.local, a 64-byte depot. With this branch: none. The 17 pre-existing unroll_smoke kernels are instruction-identical to main.
• isize loops over a..b, (a..b).rev() and #[unroll(4)] unroll with no warning. The two remaining warnings come from the deliberate 0..=4 and step_by(3) negatives.
• I pushed one fix on top: the new verify-code-shape.sh used a POSIX character class Python does not support, so its depot check never matched. It does now, and the gate still fails on main's PTX.
• The example passes on the LLVM and libNVVM routes with memcheck and synccheck at 0 errors; cargo test -p mir-transforms 91 passed, just test 3,986 passed. CI is green on this head.

#1494 (const-generic bounds) and #1478 (clone budget) stay open; neither path is touched here.