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

cuda-oxide/crates/cuda-device/src/thread.rs

@@ -1411,6 +1411,42 @@ const fn validate_unroll_factor(factor: u32) {
     );
 }
 
+/// Compile-time "keep this loop rolled" marker (internal, do not call directly).
+///
+/// The `#[kernel]` and `#[device]` macros insert this marker at the start of a
+/// loop body annotated with `#[nounroll]`. The MIR importer turns it into a
+/// `mir.nounroll_hint` operation. Just before lowering, the hint marks the
+/// loop's latch with `llvm.loop.unroll.disable` metadata. LLVM and libNVVM then
+/// leave the loop rolled, and the NVPTX backend emits `.pragma "nounroll";` at
+/// the loop header so ptxas keeps it rolled too. It generates no runtime code.
+///
+/// # Usage
+///
+/// Put the annotation directly on the loop:
+///
+/// ```rust,ignore
+/// #[kernel]
+/// pub fn my_kernel(mut output: DisjointSlice<f32>) {
+///     let mut acc = 0.0;
+///     let mut tap = 0;
+///     #[nounroll]
+///     while tap < 4 {
+///         acc += tap as f32;
+///         tap += 1;
+///     }
+///     // ...
+/// }
+/// ```
+///
+/// The marker applies to the loop that carries the annotation, not to loops
+/// nested in its body. A `#[nounroll]` loop inside a loop that `#[unroll]`
+/// unrolls is copied into each unrolled body, and every copy stays rolled.
+#[inline(never)]
+pub fn __nounroll_config() {
+    // This function is detected at compile time and removed.
+    // No runtime code is generated.
+}
+
 #[cfg(test)]
 mod tests {
     use super::block_linear_thread_id;

cuda-oxide/crates/cuda-macros/src/kernel/mod.rs

@@ -204,13 +204,13 @@ pub(crate) fn kernel_entry(attr: TokenStream, item: TokenStream) -> TokenStream
         .into();
     }
 
-    // Consume any `#[unroll]` / `#[unroll(N)]` attributes written directly on
-    // loops inside the kernel body. We strip the attribute from the loop
-    // expression (so rustc never sees an expression attribute, keeping us off
-    // nightly `stmt_expr_attributes`) and inject an
-    // `__unroll_config::<FACTOR>()` marker into the loop body. If the visitor
-    // hits a malformed attribute it records the error so we can surface it as a
-    // compile error below.
+    // Consume any `#[unroll]` / `#[unroll(N)]` / `#[nounroll]` attributes
+    // written directly on loops inside the kernel body. We strip the attribute
+    // from the loop expression (so rustc never sees an expression attribute,
+    // keeping us off nightly `stmt_expr_attributes`) and inject an
+    // `__unroll_config::<FACTOR>()` or `__nounroll_config()` marker into the
+    // loop body. If the visitor hits a malformed attribute it records the
+    // error so we can surface it as a compile error below.
     if let Err(err) = rewrite_loop_unroll_attrs(&mut input) {
         return err.to_compile_error().into();
     }

cuda-oxide/crates/cuda-macros/src/launch_attrs.rs

@@ -391,7 +391,11 @@ impl Parse for UnrollArgs {
             1 => {
                 let factor = values.pop().unwrap();
                 match factor.literal_value {
-                    Some(0 | 1) => Err(syn::Error::new_spanned(
+                    Some(1) => Err(syn::Error::new_spanned(
+                        &factor.expr,
+                        "an unroll factor of 1 does not unroll; use #[nounroll] to keep the loop rolled",
+                    )),
+                    Some(0) => Err(syn::Error::new_spanned(
                         &factor.expr,
                         "partial unroll factor must be at least 2; use #[unroll] for full unrolling",
                     )),
@@ -410,8 +414,9 @@ impl Parse for UnrollArgs {
     }
 }
 
-/// `VisitMut` pass that consumes `#[unroll]` / `#[unroll(N)]` attributes written
-/// directly on loops inside a `#[kernel]` (or `#[device]`) function body.
+/// `VisitMut` pass that consumes `#[unroll]` / `#[unroll(N)]` and `#[nounroll]`
+/// attributes written directly on loops inside a `#[kernel]` (or `#[device]`)
+/// function body.
 ///
 /// This mirrors how CubeCL's `#[cube]` macro handles per-loop `#[unroll]`: the
 /// enclosing function macro owns the whole body AST, so it can strip the inner
@@ -427,12 +432,14 @@ impl Parse for UnrollArgs {
 ///    statement of that loop's body block.
 ///
 /// `FACTOR` follows [`UnrollArgs`]: bare `#[unroll]` => `0` (full unroll),
-/// `#[unroll(N)]` => `N`. The importer reads the marker from the block it lands
-/// in, so the request applies to that loop only.
+/// `#[unroll(N)]` => `N`. `#[nounroll]` inserts
+/// `cuda_device::thread::__nounroll_config();` instead. The importer reads the
+/// marker from the block it lands in, so the request applies to that loop only.
+/// A loop may carry one of the two requests, not both.
 ///
 /// The visitor recurses through nested blocks/loops/ifs via the default
 /// `visit_mut` traversal, so an annotated loop anywhere in the function is
-/// handled. Loops without an `#[unroll]` attribute are left untouched, so a
+/// handled. Loops without either attribute are left untouched, so a
 /// kernel with no per-loop annotations expands byte-identically to before.
 #[derive(Default)]
 pub(crate) struct LoopUnrollAttrVisitor {
@@ -484,6 +491,49 @@ impl LoopUnrollAttrVisitor {
         Some((factor, attr.span()))
     }
 
+    /// If `attrs` contains an outer `#[nounroll]`, remove it and return its
+    /// span. Records an error and returns `None` if it carries arguments.
+    fn take_nounroll(&mut self, attrs: &mut Vec<syn::Attribute>) -> Option<proc_macro2::Span> {
+        let idx = attrs
+            .iter()
+            .position(|attr| attr.path().is_ident("nounroll"))?;
+        let attr = attrs.remove(idx);
+        if !matches!(attr.meta, syn::Meta::Path(_)) {
+            if self.error.is_none() {
+                self.error = Some(syn::Error::new_spanned(
+                    &attr,
+                    "nounroll takes no arguments: #[nounroll]",
+                ));
+            }
+            return None;
+        }
+        Some(attr.span())
+    }
+
+    /// Take this loop's `#[unroll]` / `#[nounroll]` request and build the
+    /// marker statement for it. Both on one loop contradict each other, so
+    /// that records an error instead of picking one.
+    fn take_loop_marker(&mut self, attrs: &mut Vec<syn::Attribute>) -> Option<Stmt> {
+        let unroll = self.take_unroll_factor(attrs);
+        let nounroll = self.take_nounroll(attrs);
+        match (unroll, nounroll) {
+            (Some(_), Some(span)) => {
+                if self.error.is_none() {
+                    self.error = Some(syn::Error::new(
+                        span,
+                        "a loop cannot carry both #[unroll] and #[nounroll]",
+                    ));
+                }
+                None
+            }
+            (Some((factor, span)), None) => Some(Self::marker_stmt(&factor, span)),
+            (None, Some(span)) => Some(syn::parse_quote_spanned! {span=>
+                cuda_device::thread::__nounroll_config();
+            }),
+            (None, None) => None,
+        }
+    }
+
     /// Build the `__unroll_config::<FACTOR>()` marker statement. The call
     /// carries `span`, the `#[unroll]` attribute's, so diagnostics about the
     /// loop point at the attribute. The braces around the factor keep the
@@ -501,32 +551,23 @@ impl LoopUnrollAttrVisitor {
 
 impl VisitMut for LoopUnrollAttrVisitor {
     fn visit_expr_mut(&mut self, expr: &mut Expr) {
-        // Pull the unroll factor (if any) off this loop expr and inject the
+        // Pull the unroll request (if any) off this loop expr and inject the
         // marker into its body. We act only on loop expressions; everything
         // else falls through to the default recursion below.
         match expr {
             Expr::ForLoop(for_loop) => {
-                if let Some((factor, span)) = self.take_unroll_factor(&mut for_loop.attrs) {
-                    for_loop
-                        .body
-                        .stmts
-                        .insert(0, Self::marker_stmt(&factor, span));
+                if let Some(marker) = self.take_loop_marker(&mut for_loop.attrs) {
+                    for_loop.body.stmts.insert(0, marker);
                 }
             }
             Expr::While(while_loop) => {
-                if let Some((factor, span)) = self.take_unroll_factor(&mut while_loop.attrs) {
-                    while_loop
-                        .body
-                        .stmts
-                        .insert(0, Self::marker_stmt(&factor, span));
+                if let Some(marker) = self.take_loop_marker(&mut while_loop.attrs) {
+                    while_loop.body.stmts.insert(0, marker);
                 }
             }
             Expr::Loop(loop_expr) => {
-                if let Some((factor, span)) = self.take_unroll_factor(&mut loop_expr.attrs) {
-                    loop_expr
-                        .body
-                        .stmts
-                        .insert(0, Self::marker_stmt(&factor, span));
+                if let Some(marker) = self.take_loop_marker(&mut loop_expr.attrs) {
+                    loop_expr.body.stmts.insert(0, marker);
                 }
             }
             _ => {}
@@ -538,8 +579,9 @@ impl VisitMut for LoopUnrollAttrVisitor {
     }
 }
 
-/// Consume per-loop unroll annotations in a kernel or device function and
-/// replace them with the marker calls understood by the MIR importer.
+/// Consume per-loop `#[unroll]` / `#[nounroll]` annotations in a kernel or
+/// device function and replace them with the marker calls understood by the
+/// MIR importer.
 pub(crate) fn rewrite_loop_unroll_attrs(input: &mut ItemFn) -> syn::Result<()> {
     let mut visitor = LoopUnrollAttrVisitor::default();
     visitor.visit_block_mut(&mut input.block);

cuda-oxide/crates/cuda-macros/src/lib.rs

@@ -472,6 +472,27 @@ pub fn cuda_module(attr: TokenStream, item: TokenStream) -> TokenStream {
 /// at most 1,024 body copies, 8,192 cloned basic blocks, and 65,536 cloned
 /// operations. Factors above 1,024 are rejected; unsupported loop shapes warn
 /// and are not unrolled.
+///
+/// # Keeping a loop rolled
+///
+/// Put `#[nounroll]` on a loop to keep it rolled. ptxas otherwise fully unrolls
+/// small constant-trip-count loops on its own, which can raise register
+/// pressure. The annotation becomes `llvm.loop.unroll.disable` metadata on the
+/// loop, so LLVM and libNVVM leave it rolled and ptxas sees
+/// `.pragma "nounroll";` at its header.
+///
+/// ```ignore
+/// #[kernel]
+/// pub fn example(out: *mut f32) {
+///     let mut acc = 0.0;
+///     let mut tap = 0;
+///     #[nounroll]
+///     while tap < 4 { acc += load(tap); tap += 1; }
+/// }
+/// ```
+///
+/// `#[nounroll]` applies to the annotated loop only and takes no arguments. A
+/// loop cannot carry both `#[unroll]` and `#[nounroll]`.
 #[proc_macro_attribute]
 pub fn kernel(attr: TokenStream, item: TokenStream) -> TokenStream {
     kernel::kernel_entry(attr, item)
@@ -849,7 +870,7 @@ pub fn cooperative_launch(attr: TokenStream, item: TokenStream) -> TokenStream {
 /// - Return values (unlike kernels which must return `()`)
 /// - Be called from kernels and other device functions
 /// - Use generics (each monomorphization becomes a separate device function)
-/// - Use per-loop `#[unroll]` and `#[unroll(N)]` annotations
+/// - Use per-loop `#[unroll]`, `#[unroll(N)]`, and `#[nounroll]` annotations
 ///
 /// # Loop unrolling
 ///
@@ -866,6 +887,9 @@ pub fn cooperative_launch(attr: TokenStream, item: TokenStream) -> TokenStream {
 /// blocks, and 65,536 cloned operations. Factors above 1,024 are rejected;
 /// unsupported loop shapes warn and are not unrolled.
 ///
+/// `#[nounroll]` keeps a loop rolled through LLVM, libNVVM, and ptxas, as in
+/// kernels.
+///
 /// # Example: Device Function Definition
 ///
 /// ```ignore

cuda-oxide/crates/cuda-macros/src/tests/launch.rs

@@ -367,6 +367,73 @@ fn malformed_unroll_attr_records_error() {
     );
 }
 
+#[test]
+fn nounroll_injects_marker_first_in_each_loop_kind() {
+    let func: ItemFn = parse_quote! {
+        fn k(n: u32) {
+            let mut i = 0u32;
+            #[nounroll]
+            while i < 4 { i += 1; }
+            #[nounroll]
+            for j in 0..n { let _ = j; }
+            #[nounroll]
+            loop { break; }
+        }
+    };
+    let out = run_loop_unroll_visitor(func);
+    let marker = "cuda_device::thread::__nounroll_config();";
+    assert_eq!(
+        out.matches(marker).count(),
+        3,
+        "one marker per loop:\n{out}"
+    );
+    for body_start in ["whilei<4{", "forjin0..n{", "loop{"] {
+        assert!(
+            out.contains(&format!("{body_start}{marker}")),
+            "marker must be the first statement after `{body_start}`:\n{out}"
+        );
+    }
+    assert!(
+        !out.contains("nounroll]"),
+        "attribute must be stripped:\n{out}"
+    );
+    assert!(!out.contains("__unroll_config"), "no unroll marker:\n{out}");
+}
+
+#[test]
+fn nounroll_rejects_arguments_and_unroll_on_the_same_loop() {
+    for func in [
+        parse_quote! {
+            fn k() {
+                #[nounroll(2)]
+                for i in 0..4 { let _ = i; }
+            }
+        },
+        parse_quote! {
+            fn k() {
+                #[unroll]
+                #[nounroll]
+                for i in 0..4 { let _ = i; }
+            }
+        },
+        parse_quote! {
+            fn k() {
+                #[nounroll]
+                #[unroll(4)]
+                for i in 0..4 { let _ = i; }
+            }
+        },
+    ] {
+        let mut func: ItemFn = func;
+        let err = rewrite_loop_unroll_attrs(&mut func).expect_err("must be rejected");
+        let message = err.to_string();
+        assert!(
+            message.contains("no arguments") || message.contains("both #[unroll] and #[nounroll]"),
+            "unexpected error: {message}"
+        );
+    }
+}
+
 #[test]
 fn partial_unroll_factor_must_be_at_least_two() {
     assert!(syn::parse_str::<UnrollArgs>("0").is_err());

cuda-oxide/crates/dialect-mir/README.md

@@ -220,12 +220,12 @@ Pointers carry an NVPTX address space independently of their source kind:
 
 ## Operations
 
-62 operations across 12 modules, one row per file in `src/ops/`:
+63 operations across 12 modules, one row per file in `src/ops/`:
 
 | Module         | Ops | Description                                                                                              |
 |----------------|-----|----------------------------------------------------------------------------------------------------------|
 | `function`     | 1   | `MirFuncOp` -- function definition                                                                       |
-| `control_flow` | 6   | return, goto, cond_branch, assert, unreachable, unroll_hint                                              |
+| `control_flow` | 7   | return, goto, cond_branch, assert, unreachable, unroll_hint, nounroll_hint                               |
 | `memory`       | 11  | alloca, load, store, ref, assign, ptr_offset, memcpy, memmove, shared_alloc, global_alloc, extern_shared |
 | `constants`    | 3   | integer, float, and undef constants                                                                      |
 | `arithmetic`   | 15  | add/sub/mul/div/rem, checked variants, bitwise, shifts                                                   |
@@ -300,7 +300,7 @@ src/
 ├── ops/
 │   ├── mod.rs                   # Op module registry + re-exports
 │   ├── function.rs              # MirFuncOp
-│   ├── control_flow.rs          # Terminators, branches, unroll hints
+│   ├── control_flow.rs          # Terminators, branches, loop hints
 │   ├── memory.rs                # Load, store, alloc, memcpy, shared memory
 │   ├── constants.rs             # Integer and float literals
 │   ├── arithmetic.rs            # Math, bitwise, shifts, checked ops