feat(unroll): add #[nounroll] to keep loops rolled

[praveenperera]

Closes #1468.

What this adds

#[nounroll] on a loop in a #[kernel] or #[device] function keeps that loop rolled through LLVM, libNVVM and ptxas. ptxas otherwise fully unrolls small constant-trip-count loops even when the PTX keeps them rolled.

let mut tap = 0;
#[nounroll]
while tap < 4 {
    acc += w[tap] * x[i + tap];
    tap += 1;
}

The loop's latch branch carries !llvm.loop metadata with llvm.loop.unroll.disable. The NVPTX backend prints .pragma "nounroll"; at the loop header, and libNVVM reads the same metadata.

How it works

1. The loop visitor in cuda-macros strips #[nounroll] from for, while and loop and inserts a __nounroll_config() marker at the top of the body, as it does for #[unroll]. It rejects arguments and a loop carrying both #[unroll] and #[nounroll]. A literal #[unroll(1)] now points to #[nounroll].
2. mir-importer emits a mir.nounroll_hint op. The op has no SideEffects impl, so DCE keeps it, and the unroll pass ignores it: unrolling an enclosing loop copies the hint along with the inner loop.
3. mir_transforms::nounroll::resolve_nounroll_hints runs at the start of lower_mir_to_llvm, after WGMMA fusion, the last transform that reshapes loops. For each hint it finds the innermost loop, joins its back-edges through one latch with the existing merge_backedges, removes the hint, and returns the latch's mir.goto. Because it runs in lowering rather than MIR preparation, full-debug builds get it too.
4. mir-lower marks those gotos, and the branch conversion carries the mark to llvm.br.
5. The LLVM IR exporter appends , !llvm.loop !N to a marked branch and emits one distinct node per loop plus one shared !{!"llvm.loop.unroll.disable"} node, from the module's single metadata ID counter.

The single latch matters because LLVM reads loop metadata from the latch terminator and ignores it when several latches disagree. A branch that ends an inner loop and also jumps to the outer header would otherwise apply the request to both loops.

Interactions asked about in #1468

• Nested loops: a hint marks only the loop whose body holds it.
• Partial unrolling: #[unroll(4)] around a #[nounroll] loop leaves four inner copies in the main loop and one in the remainder loop, and all five are marked. The main and remainder loops themselves are not.
• Debug builds: the full-debug PTX has one .pragma "nounroll"; per annotated source loop.
• feat(unroll): recognize range for loops #1346 range for loops: #[nounroll] works on a range for loop, including one inside a fully unrolled loop.
• WGMMA counted K-loops: counted-loop fusion runs before the hints are resolved, so it now accepts a hint inside the loop it fuses. Without that, #[nounroll] on an F16 m64n64 K loop failed to compile, since F16 has no pointer fallback, and a BF16 one quietly lost the fused value-carrier form. The fused loop is one inline-PTX operation with no LLVM loop to mark, so fusion moves the request onto that operation, and its template starts the loop header with .pragma "nounroll";. This covers both the counted loop and the counted pipeline.
• [BUG]: #[unroll(N)] emits no loop metadata, so the NVVM backend unrolls the loops a second time #1375: on the llc route, LLVM still unrolls the remainder loop that #[unroll(N)] leaves behind, because that loop carries no metadata. Any #[nounroll] loop inside it is cloned and every clone stays rolled. The same latch metadata could close [BUG]: #[unroll(N)] emits no loop metadata, so the NVVM backend unrolls the loops a second time #1375 by marking the main and remainder loops; that changes existing #[unroll(N)] codegen, so it is left for a follow-up.

Testing

• mir-transforms/tests/nounroll.rs: latch placement for an inner and an outer loop of a nest, a loop with two continue paths, a #[nounroll] loop inside a full and a partial #[unroll] (3 and 5 marked copies, main and remainder unmarked), and a hint outside any loop. Each test checks the returned branch is the marked loop's only latch.
• mir-lower/tests/lowering_test/loops.rs: three sibling loops, two hinted, go through lower_mir_to_llvm and both the PTX and legacy NVVM exporters. Exactly two marked latches come out, each with its own distinct self-referencing node, one shared property node, and no metadata ID defined twice next to !nvvmir.version.
• wgmma_lowering.rs: the F16 counted K-loop still fuses with a hint in its latch, and the fused template's loop header starts with .pragma "nounroll"; only when the loop was hinted. Without the fusion change, the hinted loop failed with "invalid input program".
• wgmma.rs unit test: both counted templates put the pragma right after their loop label when asked, and omit it otherwise.
• collector.rs unit test: the marker check matches only cuda_device's __unroll_config and __nounroll_config, not a function like __nounroll_config_helper or one with the same name in another crate.
• A hint in unreachable code that branches into a loop marks nothing.
• cuda-macros unit tests: the marker goes first in while, for and loop bodies; arguments and #[unroll] + #[nounroll] are rejected.
• New nounroll example, with verify-code-shape.sh for the compile-only smoketest. It checks the exact count of marked latches per function in the pre-opt .ll, and at least that many pragmas in the PTX, with none in the control kernel. LLVM may clone a rolled loop, as in the [BUG]: #[unroll(N)] emits no loop metadata, so the NVVM backend unrolls the loops a second time #1375 case above.
• On an RTX 5090 (sm_120a, driver 595.99.02, CUDA 13.1, LLVM 21), the example passes in every route: llc PTX, libNVVM JIT, libNVVM cubin, and --device-debug. An sm_80 libNVVM cubin build (legacy LLVM 7 dialect) also compiles.
• SASS check. Without the annotation, ptxas unrolls the 4-tap control kernel to 8 loads. With it, each annotated loop stays a loop in SASS on the llc route and on both libNVVM cubins, sm_120a and sm_80:

Kernel	Rolled loops	Global loads
taps_default (control)	0	8
taps_rolled	1	2
rolled_k_unrolled_inner (K loop rolled, inner #[unroll])	1	4
rolled_inside_full_unroll	4	4
rolled_inside_partial_unroll (libNVVM: 5 inner, plus the main and remainder loops)	7	5
rolled_in_device_fn (loop in an #[inline(never)] #[device] fn)	1	1

• just check and scripts/sync-example-locks.sh --check on Linux with CUDA 13.1 and LLVM 21.
• cargo clippy --all-targets -D warnings and cargo fmt --check for the new example; scripts/smoketest.sh GPU, --compile-only and --full-debug lanes for nounroll, unroll_smoke and policy_config, and the GPU and --compile-only lanes for wgmma_mma_bf16; just book. (wgmma_mma_bf16 fails its --full-debug build with "invalid input program" on unmodified main at a25a5857 as well.)

The example's annotated loops all load from memory. With an arithmetic-only body and a constant trip count, LLVM or ptxas computes the result outright and deletes the loop, which leaves nothing to keep rolled.

Limits

Like #[unroll], the hint is mapped to the innermost loop around it. A loop whose body never repeats, such as loop { ...; break; }, produces no loop in MIR, so inside another loop the request applies to the enclosing loop. The book says so. Telling the two apart would need a second frontend marker that survives cloning, for both annotations.

Not in this PR

#[unroll(N)] on a loop that contains a range for loop warns and does not unroll on main today. The inner loop's Option match has an unreachable arm, which partial unrolling counts as a second exit. The example's partial-unroll kernel uses an inner while loop for that reason.

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