diff --git a/cuda-oxide/crates/cuda-device/src/thread.rs b/cuda-oxide/crates/cuda-device/src/thread.rs index 0b3f33c435..2f67a5b820 100644 --- a/cuda-oxide/crates/cuda-device/src/thread.rs +++ b/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) { +/// 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; diff --git a/cuda-oxide/crates/cuda-macros/src/kernel/mod.rs b/cuda-oxide/crates/cuda-macros/src/kernel/mod.rs index 60f55f788a..d99c7bc164 100644 --- a/cuda-oxide/crates/cuda-macros/src/kernel/mod.rs +++ b/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::()` 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::()` 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(); } diff --git a/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs b/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs index b10f4cd455..60574d6bae 100644 --- a/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs +++ b/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) -> Option { + 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) -> Option { + 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::()` 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); diff --git a/cuda-oxide/crates/cuda-macros/src/lib.rs b/cuda-oxide/crates/cuda-macros/src/lib.rs index 0a3d738b55..feb6597148 100644 --- a/cuda-oxide/crates/cuda-macros/src/lib.rs +++ b/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 diff --git a/cuda-oxide/crates/cuda-macros/src/tests/launch.rs b/cuda-oxide/crates/cuda-macros/src/tests/launch.rs index 4cee9d382d..72b27eebf7 100644 --- a/cuda-oxide/crates/cuda-macros/src/tests/launch.rs +++ b/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::("0").is_err()); diff --git a/cuda-oxide/crates/dialect-mir/README.md b/cuda-oxide/crates/dialect-mir/README.md index c7c21b021b..48fde952d5 100644 --- a/cuda-oxide/crates/dialect-mir/README.md +++ b/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 diff --git a/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs b/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs index fc2d7f51ed..1dec37e083 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs @@ -489,6 +489,41 @@ impl Verify for MirUnrollHintOp { } } +/// A per-loop "keep this loop rolled" request, planted inside a loop body by +/// the `#[nounroll]` frontend. +/// +/// Like [`MirUnrollHintOp`], the op sits physically inside the loop it +/// annotates, so a pass can map it back to that loop through `LoopInfo`. The +/// loop-unroll pass ignores it: an unrolled outer loop copies it into every +/// body copy along with the inner loop it marks. Just before lowering, +/// `mir_transforms::nounroll` gives each marked loop a single latch, and the +/// lowering attaches `llvm.loop.unroll.disable` to that latch's branch. The +/// hint itself never reaches the LLVM dialect. It has no operands, results, or +/// attributes. +#[pliron_op( + name = "mir.nounroll_hint", + format, + interfaces = [ + pliron::builtin::op_interfaces::NOpdsInterface<0>, + pliron::builtin::op_interfaces::NResultsInterface<0>, + ] +)] +pub struct MirNoUnrollHintOp; + +impl MirNoUnrollHintOp { + /// Create a "keep this loop rolled" hint. + pub fn new(ctx: &mut Context) -> Self { + let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0); + MirNoUnrollHintOp { op } + } +} + +impl Verify for MirNoUnrollHintOp { + fn verify(&self, _ctx: &Context) -> pliron::result::Result<()> { + Ok(()) + } +} + /// Register control flow operations into the given context. pub fn register(ctx: &mut Context) { MirReturnOp::register(ctx); @@ -497,4 +532,5 @@ pub fn register(ctx: &mut Context) { MirAssertOp::register(ctx); MirUnreachableOp::register(ctx); MirUnrollHintOp::register(ctx); + MirNoUnrollHintOp::register(ctx); } diff --git a/cuda-oxide/crates/llvm-export/src/export/metadata.rs b/cuda-oxide/crates/llvm-export/src/export/metadata.rs index 1c46128ebd..d15a53b6ad 100644 --- a/cuda-oxide/crates/llvm-export/src/export/metadata.rs +++ b/cuda-oxide/crates/llvm-export/src/export/metadata.rs @@ -3,11 +3,14 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! NVVM annotations and version metadata. +//! NVVM annotations, version, and loop metadata. use std::collections::HashSet; use std::fmt::Write; +use pliron::context::Ptr; +use pliron::operation::Operation; + use super::state::{KernelBlockGeometry, ModuleExportState}; pub(super) fn needs_nvvm_annotations( @@ -196,6 +199,46 @@ pub(super) fn emit_nvvmir_version( .unwrap(); } +impl ModuleExportState<'_> { + /// The `, !llvm.loop !N` suffix for a latch branch marked by + /// [`set_loop_unroll_disable`](crate::ops::set_loop_unroll_disable), or an + /// empty string for any other branch. + /// + /// Each marked branch gets its own `distinct` node, because LLVM identifies + /// a loop by its node. Lowering marks only a loop's single latch, so one + /// branch is one loop. + pub(super) fn loop_metadata_suffix(&mut self, branch: Ptr) -> String { + if !crate::ops::loop_unroll_disable(self.ctx, branch) { + return String::new(); + } + if self.loop_unroll_disable_property.is_none() { + self.loop_unroll_disable_property = Some(self.alloc_metadata_id()); + } + let node = self.alloc_metadata_id(); + self.unroll_disabled_loops.push(node); + format!(", !llvm.loop !{node}") + } +} + +/// Emit the loop nodes referenced by [`ModuleExportState::loop_metadata_suffix`]. +/// +/// ```text +/// !7 = distinct !{!7, !6} +/// !6 = !{!"llvm.loop.unroll.disable"} +/// ``` +/// +/// A loop node's first operand must be the node itself. The syntax is valid in +/// both modern LLVM and the LLVM 7 dialect that legacy libNVVM reads. +pub(super) fn emit_loop_metadata(output: &mut String, state: &ModuleExportState) { + let Some(property) = state.loop_unroll_disable_property else { + return; + }; + for node in &state.unroll_disabled_loops { + writeln!(output, "!{node} = distinct !{{!{node}, !{property}}}").unwrap(); + } + writeln!(output, "!{property} = !{{!\"llvm.loop.unroll.disable\"}}").unwrap(); +} + #[cfg(test)] mod tests { use super::*; diff --git a/cuda-oxide/crates/llvm-export/src/export/module.rs b/cuda-oxide/crates/llvm-export/src/export/module.rs index 1a9e934851..6d80ebdb48 100644 --- a/cuda-oxide/crates/llvm-export/src/export/module.rs +++ b/cuda-oxide/crates/llvm-export/src/export/module.rs @@ -31,7 +31,9 @@ use super::{ ExportedModule, config::{DebugKind, ExportBackendConfig, NvvmIrDialect}, externs::{DeviceExternDecl, DeviceExternType}, - metadata::{emit_nvvm_annotations, emit_nvvmir_version, needs_nvvm_annotations}, + metadata::{ + emit_loop_metadata, emit_nvvm_annotations, emit_nvvmir_version, needs_nvvm_annotations, + }, state::{DebugSharedFunctionScope, GlobalSourceInfo, GlobalSymbolInfo, ModuleExportState}, }; @@ -729,7 +731,13 @@ pub(super) fn export_module_with_externs_impl( emit_nvvmir_version(&mut output, &mut state, config.nvvmir_version()); } - // 9. DWARF metadata (if requested and source locations exist) + // 9. Loop metadata for loops that must stay rolled (`#[nounroll]`) + if !state.unroll_disabled_loops.is_empty() { + writeln!(&mut output).unwrap(); + emit_loop_metadata(&mut output, &state); + } + + // 10. DWARF metadata (if requested and source locations exist) if state.has_debug_metadata() { writeln!(&mut output).unwrap(); state.emit_debug_metadata(&mut output); @@ -876,6 +884,12 @@ pub(super) fn export_module_to_string_with_config( emit_nvvmir_version(&mut output, &mut state, config.nvvmir_version()); } + // Emit loop metadata for loops that must stay rolled (`#[nounroll]`) + if !state.unroll_disabled_loops.is_empty() { + writeln!(&mut output).unwrap(); + emit_loop_metadata(&mut output, &state); + } + // Emit DWARF line-table metadata if debug export requested it and at least // one function had a real source location. if state.has_debug_metadata() { diff --git a/cuda-oxide/crates/llvm-export/src/export/ops.rs b/cuda-oxide/crates/llvm-export/src/export/ops.rs index 1c1216bc43..e9fe6c75f2 100644 --- a/cuda-oxide/crates/llvm-export/src/export/ops.rs +++ b/cuda-oxide/crates/llvm-export/src/export/ops.rs @@ -668,7 +668,7 @@ impl<'a> ModuleExportState<'a> { } fn emit_br( - &self, + &mut self, op: &ops::BrOp, block_labels: &FxHashMap, String>, output: &mut String, @@ -676,7 +676,8 @@ impl<'a> ModuleExportState<'a> { let op_ref = op.get_operation().deref(self.ctx); let dest = op_ref.successors().next().unwrap(); let label = block_labels.get(&dest).ok_or("Missing block label")?; - writeln!(output, " br label %{label}").unwrap(); + let loop_metadata = self.loop_metadata_suffix(op.get_operation()); + writeln!(output, " br label %{label}{loop_metadata}").unwrap(); Ok(()) } @@ -697,7 +698,12 @@ impl<'a> ModuleExportState<'a> { write!(output, " br i1 ").unwrap(); self.export_value(cond, value_names, output)?; - writeln!(output, ", label %{true_label}, label %{false_label}").unwrap(); + let loop_metadata = self.loop_metadata_suffix(op.get_operation()); + writeln!( + output, + ", label %{true_label}, label %{false_label}{loop_metadata}" + ) + .unwrap(); Ok(()) } diff --git a/cuda-oxide/crates/llvm-export/src/export/state.rs b/cuda-oxide/crates/llvm-export/src/export/state.rs index c9fe47022a..71d0b0cc8f 100644 --- a/cuda-oxide/crates/llvm-export/src/export/state.rs +++ b/cuda-oxide/crates/llvm-export/src/export/state.rs @@ -157,6 +157,12 @@ pub(super) struct ModuleExportState<'a> { /// used for NVVM annotations/version nodes; debug-info nodes will use the /// same counter so the exporter never has to guess which IDs are free. next_metadata_id: usize, + /// `distinct` `!llvm.loop` node of each loop whose latch branch asked to + /// stay rolled, in emission order. + pub(super) unroll_disabled_loops: Vec, + /// The `!{!"llvm.loop.unroll.disable"}` property node shared by every + /// entry of `unroll_disabled_loops`, allocated on first use. + pub(super) loop_unroll_disable_property: Option, /// Which debug metadata tier this export should emit. pub(super) debug_kind: DebugKind, /// Where function-local statics are retained (per the consuming LLVM). @@ -258,6 +264,8 @@ impl<'a> ModuleExportState<'a> { global_symbols: FxHashMap::default(), global_sources: FxHashMap::default(), next_metadata_id: 0, + unroll_disabled_loops: Vec::new(), + loop_unroll_disable_property: None, debug_kind, debug_function_local_static_placement, nvvm_ir_dialect, diff --git a/cuda-oxide/crates/llvm-export/src/lib.rs b/cuda-oxide/crates/llvm-export/src/lib.rs index ba5dc7d272..7c00805e10 100644 --- a/cuda-oxide/crates/llvm-export/src/lib.rs +++ b/cuda-oxide/crates/llvm-export/src/lib.rs @@ -1425,6 +1425,35 @@ pub mod ops { .map(|a| a.0) } + /// Op-attribute key on a loop's latch branch when that loop must stay + /// rolled. The exporter prints it as `!llvm.loop` metadata carrying + /// `llvm.loop.unroll.disable`. + const LOOP_UNROLL_DISABLE_KEY: &str = "cuda_oxide_loop_unroll_disable"; + + /// Mark a loop's latch branch so the loop stays rolled. + /// + /// The branch must be the loop's only back-edge. LLVM reads loop metadata + /// from the latch terminator and ignores it when several latches disagree. + /// A fused WGMMA K-loop operation carries the same mark, and its lowering + /// emits `.pragma "nounroll";` in the loop it generates. + pub fn set_loop_unroll_disable(ctx: &mut Context, op: Ptr) { + let key = + Identifier::try_new(LOOP_UNROLL_DISABLE_KEY.to_string()).expect("valid identifier"); + op.deref_mut(ctx) + .attributes + .set(key, pliron::builtin::attributes::UnitAttr); + } + + /// Whether [`set_loop_unroll_disable`] marked this branch. + pub fn loop_unroll_disable(ctx: &Context, op: Ptr) -> bool { + let key = + Identifier::try_new(LOOP_UNROLL_DISABLE_KEY.to_string()).expect("valid identifier"); + op.deref(ctx) + .attributes + .get::(&key) + .is_some() + } + /// Attach the source-facing name of a function to a MIR/LLVM function op. pub fn set_debug_function_name(ctx: &mut Context, op: Ptr, name: &str) { set_string_attr(ctx, op, DEBUG_FUNCTION_NAME_KEY, name.to_string()); diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs index d88f37d4b0..4150d93c15 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs @@ -68,8 +68,8 @@ use crate::translator::rvalue; use crate::translator::values::{ValueMap, maybe_ptr_coerce}; use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; use dialect_mir::ops::{ - MirAssertOp, MirCondBranchOp, MirConstantOp, MirEqOp, MirGotoOp, MirNotOp, MirReturnOp, - MirUnrollHintOp, + MirAssertOp, MirCondBranchOp, MirConstantOp, MirEqOp, MirGotoOp, MirNoUnrollHintOp, MirNotOp, + MirReturnOp, MirUnrollHintOp, }; use pliron::basic_block::BasicBlock; use pliron::builtin::op_interfaces::OperandSegmentInterface; @@ -1099,6 +1099,26 @@ fn translate_call( return Ok(helpers::emit_goto(ctx, target, hint, block_map, loc)); } + // Per-loop `#[nounroll]` marker. Like the unroll marker above, plant the + // hint where the call sits, inside the annotated loop's body, so a later + // pass can map it back to that loop. It survives MIR preparation and is + // turned into latch metadata just before lowering. + if is_cuda_device_const_marker(func, "__nounroll_config") { + let Some(target) = target_usize else { + return input_err!( + loc, + TranslationErr::invalid_op("a nounroll marker call has no return target") + ); + }; + let hint = MirNoUnrollHintOp::new(ctx).get_operation(); + hint.deref_mut(ctx).set_loc(loc.clone()); + match prev_op { + Some(prev) => hint.insert_after(ctx, prev), + None => hint.insert_at_front(block_ptr, ctx), + } + return Ok(helpers::emit_goto(ctx, target, hint, block_map, loc)); + } + // Handle DynamicSharedArray specially to extract the ALIGN const generic. // The crate anchor tightens the name-substring gate: a user fn merely // named like it falls through to ordinary call handling. Kept in diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs index 0f76c7a81a..6c1bd8fd6f 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs @@ -188,7 +188,10 @@ fn value_group_constraints(descriptor_count: usize) -> String { value_group_constraints_for_count(VALUE_ACCUMULATOR_COUNT, descriptor_count) } -fn counted_loop_template(input_kind: WgmmaInputKind) -> String { +/// The PTX for a fused counted K loop. With `nounroll`, the loop header starts +/// with `.pragma "nounroll";`, which ptxas honors as the header's first +/// statement, so a `#[nounroll]` K loop stays rolled after fusion. +fn counted_loop_template(input_kind: WgmmaInputKind, nounroll: bool) -> String { let accumulators = value_accumulator_operand_list(VALUE_ACCUMULATOR_COUNT); let descriptor_base = VALUE_ACCUMULATOR_COUNT * 2; let desc_a_base = descriptor_base; @@ -207,6 +210,9 @@ fn counted_loop_template(input_kind: WgmmaInputKind) -> String { template.push_str(" setp.eq.u64 %loop_more, %remaining, 0;\n"); template.push_str(" @%loop_more bra.uni L__wgmma_done_${:uid};\n"); template.push_str("L__wgmma_loop_${:uid}:\n"); + if nounroll { + template.push_str(" .pragma \"nounroll\";\n"); + } let input_types = input_kind.ptx_suffix(); template.push_str(&format!( " wgmma.mma_async.sync.aligned.m64n64k16.f32.{input_types} \ @@ -224,7 +230,9 @@ fn counted_loop_template(input_kind: WgmmaInputKind) -> String { template } -fn counted_pipeline_template(slot_count: usize, max_pending_groups: u8) -> String { +/// The PTX for a fused counted pipeline. `nounroll` works as in +/// [`counted_loop_template`]. +fn counted_pipeline_template(slot_count: usize, max_pending_groups: u8, nounroll: bool) -> String { let result_count = slot_count * VALUE_ACCUMULATOR_COUNT; let control_base = result_count * 2; let step_base = control_base + slot_count * 2; @@ -248,6 +256,9 @@ fn counted_pipeline_template(slot_count: usize, max_pending_groups: u8) -> Strin template.push_str(" setp.eq.u64 %loop_more, %remaining, 0;\n"); template.push_str(" @%loop_more bra.uni L__wgmma_pipeline_done_${:uid};\n"); template.push_str("L__wgmma_pipeline_loop_${:uid}:\n"); + if nounroll { + template.push_str(" .pragma \"nounroll\";\n"); + } for slot in 0..slot_count { let accumulators = pipeline_accumulator_operand_list(slot); @@ -555,7 +566,8 @@ fn convert_mma_loop_values_for_kind( ); } - let template = counted_loop_template(input_kind); + let nounroll = llvm_export::ops::loop_unroll_disable(ctx, op); + let template = counted_loop_template(input_kind, nounroll); let constraints = counted_loop_constraints(); let f32_ty = FP32Type::get(ctx); @@ -632,7 +644,8 @@ pub(crate) fn convert_mma_loop_pipeline_values( ); } - let template = counted_pipeline_template(slot_count, max_pending_groups); + let nounroll = llvm_export::ops::loop_unroll_disable(ctx, op); + let template = counted_pipeline_template(slot_count, max_pending_groups, nounroll); let constraints = counted_pipeline_constraints(result_count, control_count); let f32_ty = FP32Type::get(ctx); let struct_ty: TypeHandle = llvm_types::StructType::get_unnamed( @@ -894,7 +907,7 @@ mod tests { #[test] fn counted_loop_template_keeps_descriptor_recurrence_inside_one_scope() { - let template = counted_loop_template(WgmmaInputKind::Bf16); + let template = counted_loop_template(WgmmaInputKind::Bf16, false); assert!(template.contains("m64n64k16.f32.bf16.bf16")); assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); assert_eq!(template.matches("wgmma.mma_async").count(), 1); @@ -944,7 +957,7 @@ mod tests { #[test] fn f16_counted_loop_template_changes_only_the_input_element_types() { - let template = counted_loop_template(WgmmaInputKind::F16); + let template = counted_loop_template(WgmmaInputKind::F16, false); assert_eq!(template.matches("wgmma.mma_async").count(), 1); assert!(template.contains("m64n64k16.f32.f16.f16")); assert!(!template.contains(".bf16.bf16")); @@ -969,9 +982,25 @@ mod tests { assert!(!template.contains(".reg .f32")); } + #[test] + fn nounroll_pragma_opens_each_fused_loop_header() { + const PRAGMA: &str = ".pragma \"nounroll\";"; + for kind in [WgmmaInputKind::Bf16, WgmmaInputKind::F16] { + let rolled = counted_loop_template(kind, true); + assert_eq!(rolled.matches(PRAGMA).count(), 1); + assert!(rolled.contains("L__wgmma_loop_${:uid}:\n .pragma \"nounroll\";\n")); + assert!(!counted_loop_template(kind, false).contains(PRAGMA)); + } + + let rolled = counted_pipeline_template(2, 1, true); + assert_eq!(rolled.matches(PRAGMA).count(), 1); + assert!(rolled.contains("L__wgmma_pipeline_loop_${:uid}:\n .pragma \"nounroll\";\n")); + assert!(!counted_pipeline_template(2, 1, false).contains(PRAGMA)); + } + #[test] fn counted_pipeline_template_supports_two_and_three_slot_wait_depths() { - let two_slot = counted_pipeline_template(2, 1); + let two_slot = counted_pipeline_template(2, 1, false); assert_eq!(two_slot.matches("wgmma.fence.sync.aligned").count(), 1); assert_eq!(two_slot.matches("wgmma.mma_async").count(), 2); assert_eq!( @@ -994,7 +1023,7 @@ mod tests { assert!(two_slot.contains("add.u64 %desc_b1, %desc_b1, $135;")); assert!(two_slot.contains("mov.u64 %remaining, $136;")); - let three_slot = counted_pipeline_template(3, 2); + let three_slot = counted_pipeline_template(3, 2, false); assert_eq!(three_slot.matches("wgmma.fence.sync.aligned").count(), 1); assert_eq!(three_slot.matches("wgmma.mma_async").count(), 3); assert_eq!( diff --git a/cuda-oxide/crates/mir-lower/src/convert/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/mod.rs index fd0f44c956..94f4bcfc15 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/mod.rs @@ -43,6 +43,18 @@ pub(crate) fn preserve_location( lowered } +/// Carry a `#[nounroll]` latch mark from a MIR branch to the LLVM branch that +/// replaces it. +pub(crate) fn preserve_loop_unroll_disable( + ctx: &mut Context, + source: Ptr, + lowered: Ptr, +) { + if llvm_export::ops::loop_unroll_disable(ctx, source) { + llvm_export::ops::set_loop_unroll_disable(ctx, lowered); + } +} + pub(crate) mod enum_payload_storage; mod generated_intrinsics; pub mod interface_impls; diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs index bf31f820cf..12a0fcfb62 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs @@ -185,6 +185,7 @@ pub(crate) fn convert_cond_branch( let llvm_br = llvm::CondBrOp::new(ctx, cond, true_block, true_args, false_block, false_args); crate::convert::preserve_location(ctx, op, llvm_br.get_operation()); + crate::convert::preserve_loop_unroll_disable(ctx, op, llvm_br.get_operation()); rewriter.insert_operation(ctx, llvm_br.get_operation()); rewriter.erase_operation(ctx, op); @@ -305,6 +306,7 @@ pub(crate) fn convert_goto( let llvm_br = llvm::BrOp::new(ctx, dest, final_args); crate::convert::preserve_location(ctx, op, llvm_br.get_operation()); + crate::convert::preserve_loop_unroll_disable(ctx, op, llvm_br.get_operation()); rewriter.insert_operation(ctx, llvm_br.get_operation()); rewriter.erase_operation(ctx, op); diff --git a/cuda-oxide/crates/mir-lower/src/lib.rs b/cuda-oxide/crates/mir-lower/src/lib.rs index 82fa281c4b..973a7b7c49 100644 --- a/cuda-oxide/crates/mir-lower/src/lib.rs +++ b/cuda-oxide/crates/mir-lower/src/lib.rs @@ -400,7 +400,15 @@ pub fn lower_mir_to_llvm_with_options( // unsupported pointer shapes retain the deferred pointer-group fallback. // Run this while MIR control flow and unsigned constants are still intact. wgmma_deferred_accumulator::fuse_deferred_accumulators(ctx, module_op)?; - // WGMMA fusion is the final MIR-producing transform in this entry point. + // `#[nounroll]` loops: give each marked loop one latch and mark its branch + // for `llvm.loop.unroll.disable`. This runs after every transform that can + // clone or reshape loops, including the WGMMA K-loop fusion above, and in + // full-debug builds too, which skip MIR preparation. + for latch in mir_transforms::nounroll::resolve_nounroll_hints(module_op, ctx)? { + llvm_export::ops::set_loop_unroll_disable(ctx, latch); + } + // WGMMA fusion and nounroll latch placement are the final MIR-producing + // transforms in this entry point. // Reverify immediately before dialect conversion erases pointer kinds so a // future transform cannot bypass the provenance invariant. dialect_mir::verification::verify_pointer_kind_producers(ctx, module_op)?; diff --git a/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs b/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs index 80a9d68bbd..a14c2491e3 100644 --- a/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs +++ b/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs @@ -59,8 +59,8 @@ use dialect_mir::attributes::MirCastKindAttr; use dialect_mir::{ ops::{ MirAddOp, MirArrayElementAddrOp, MirCastOp, MirCondBranchOp, MirConstantOp, MirFuncOp, - MirGeOp, MirGotoOp, MirGtOp, MirLeOp, MirLoadOp, MirLtOp, MirNotOp, MirStorageDeadOp, - MirStorageLiveOp, MirStoreOp, MirSubOp, + MirGeOp, MirGotoOp, MirGtOp, MirLeOp, MirLoadOp, MirLtOp, MirNoUnrollHintOp, MirNotOp, + MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, MirSubOp, }, types::{MirArrayType, MirPtrType, address_space}, }; @@ -166,6 +166,8 @@ struct CountedLoopPlan { trip_count: u64, row_type: TypeHandle, element_type: TypeHandle, + /// The K loop carried `#[nounroll]`, so the fused loop must stay rolled + nounroll: bool, } struct PipelinedCountedLoopPlan { @@ -183,6 +185,8 @@ struct PipelinedCountedLoopPlan { trip_count: u64, slot_types: Vec<(TypeHandle, TypeHandle)>, max_pending_groups: u8, + /// The K loop carried `#[nounroll]`, so the fused loop must stay rolled + nounroll: bool, } fn collect_blocks(ctx: &Context, root: Ptr) -> Vec> { @@ -283,7 +287,10 @@ fn loop_values_escape(ctx: &Context, loop_blocks: &FxHashSet>) - } fn counted_loop_operation_is_supported(ctx: &Context, operation: Ptr) -> bool { - Operation::get_op::(operation, ctx).is_some() + // `#[nounroll]` on the K loop leaves a hint in it until lowering. Fusion + // replaces the loop and moves the request onto the fused operation + Operation::get_op::(operation, ctx).is_some() + || Operation::get_op::(operation, ctx).is_some() || Operation::get_op::(operation, ctx).is_some() || Operation::get_op::(operation, ctx).is_some() || Operation::get_op::(operation, ctx).is_some() @@ -300,6 +307,18 @@ fn counted_loop_operation_is_supported(ctx: &Context, operation: Ptr) || is_representation_only_ptr_retype(ctx, operation) } +/// Whether a matched K loop carried `#[nounroll]`. Its hint becomes +/// unreachable with the rest of the loop, so the fused operation carries the +/// request instead. +fn loop_has_nounroll_hint(ctx: &Context, blocks: &FxHashSet>) -> bool { + blocks.iter().any(|block| { + block + .deref(ctx) + .iter(ctx) + .any(|operation| Operation::get_op::(operation, ctx).is_some()) + }) +} + fn pipelined_counted_loop_operation_is_supported(ctx: &Context, operation: Ptr) -> bool { (counted_loop_operation_is_supported(ctx, operation) && wgmma_mma_kind(ctx, operation).is_none_or(|kind| kind == WgmmaMmaKind::Bf16M64N64)) @@ -730,6 +749,7 @@ fn match_pipelined_counted_loop( trip_count, slot_types, max_pending_groups, + nounroll: loop_has_nounroll_hint(ctx, &r#loop.blocks), })) } @@ -914,6 +934,7 @@ fn match_counted_loop( trip_count, row_type, element_type, + nounroll: loop_has_nounroll_hint(ctx, &r#loop.blocks), })) } @@ -1393,6 +1414,7 @@ fn apply_pipelined_counted_loop_plan(ctx: &mut Context, plan: PipelinedCountedLo trip_count, slot_types, max_pending_groups, + nounroll, } = plan; let slot_count = usize::from(max_pending_groups) + 1; @@ -1442,6 +1464,9 @@ fn apply_pipelined_counted_loop_plan(ctx: &mut Context, plan: PipelinedCountedLo max_pending_groups, ); pipeline.deref_mut(ctx).set_loc(loc.clone()); + if nounroll { + llvm_export::ops::set_loop_unroll_disable(ctx, pipeline); + } let accumulator_results = (0..slot_count * ACCUMULATOR_LEN) .map(|index| pipeline.deref(ctx).get_result(index)) .collect::>(); @@ -1492,6 +1517,7 @@ fn apply_counted_loop_plan(ctx: &mut Context, plan: CountedLoopPlan) { trip_count, row_type, element_type, + nounroll, } = plan; debug_assert_eq!( @@ -1538,6 +1564,9 @@ fn apply_counted_loop_plan(ctx: &mut Context, plan: CountedLoopPlan) { } }; group.deref_mut(ctx).set_loc(loc.clone()); + if nounroll { + llvm_export::ops::set_loop_unroll_disable(ctx, group); + } let accumulator_results = (0..ACCUMULATOR_LEN) .map(|index| group.deref(ctx).get_result(index)) .collect::>(); @@ -1575,7 +1604,8 @@ fn apply_pointer_fallback( } fn is_ignorable(ctx: &Context, op: Ptr) -> bool { - Operation::get_op::(op, ctx).is_some() + Operation::get_op::(op, ctx).is_some() + || Operation::get_op::(op, ctx).is_some() || Operation::get_op::(op, ctx).is_some() || Operation::get_op::(op, ctx).is_some() || Operation::get_op::(op, ctx).is_some() diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/loops.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/loops.rs new file mode 100644 index 0000000000..9375912182 --- /dev/null +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/loops.rs @@ -0,0 +1,186 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::ops as mir; +use llvm_export::export::{NvvmExportConfig, NvvmIrDialect}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::op_interfaces::OperandSegmentInterface; +use pliron::builtin::ops::ModuleOp; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::region::Region; +use pliron::value::Value; + +use super::common::{append_return, build_test_kernel, make_test_ctx}; +use super::wgmma_lowering::append_mir_unsigned_constant; + +/// Append an op of `$op` with one `$result` to `$block` and return the result. +/// A macro because pliron's op-info type cannot be named outside pliron. +macro_rules! append_op { + ($ctx:expr, $block:expr, $op:ty, $result:expr, $operands:expr) => {{ + let op = Operation::new( + $ctx, + <$op>::get_concrete_op_info(), + vec![$result], + $operands, + vec![], + 0, + ); + op.insert_at_back($block, $ctx); + op.deref($ctx).get_result(0) + }}; +} + +/// Append `while i < n { i += 1 }` after `from`, with a `mir.nounroll_hint` +/// first in the body when `nounroll` is set, and return the empty exit block. +fn append_counted_loop( + ctx: &mut Context, + region: Ptr, + from: Ptr, + n: Value, + nounroll: bool, +) -> Ptr { + let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let header = BasicBlock::new(ctx, None, vec![u32_ty.into()]); + header.insert_at_back(region, ctx); + let body = BasicBlock::new(ctx, None, vec![]); + body.insert_at_back(region, ctx); + let exit = BasicBlock::new(ctx, None, vec![]); + exit.insert_at_back(region, ctx); + + let start = append_mir_unsigned_constant(ctx, from, u32_ty, 0); + goto(ctx, from, header, vec![start]); + + let i = header.deref(ctx).get_argument(0); + let lt = append_op!(ctx, header, mir::MirLtOp, i1_ty.into(), vec![i, n]); + let done = append_op!(ctx, header, mir::MirNotOp, i1_ty.into(), vec![lt]); + let (operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![done], vec![], vec![]]); + let branch = Operation::new( + ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + operands, + vec![exit, body], + 0, + ); + Operation::get_op::(branch, ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(ctx, segment_sizes); + branch.insert_at_back(header, ctx); + + if nounroll { + mir::MirNoUnrollHintOp::new(ctx) + .get_operation() + .insert_at_back(body, ctx); + } + let one = append_mir_unsigned_constant(ctx, body, u32_ty, 1); + let next = append_op!(ctx, body, mir::MirAddOp, u32_ty.into(), vec![i, one]); + goto(ctx, body, header, vec![next]); + exit +} + +fn goto(ctx: &mut Context, from: Ptr, to: Ptr, operands: Vec) { + Operation::new( + ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + operands, + vec![to], + 0, + ) + .insert_at_back(from, ctx); +} + +/// Check the loop metadata in exported LLVM IR: `marked` latch branches, each +/// with its own `distinct` self-referencing node, one shared +/// `llvm.loop.unroll.disable` node, and no metadata ID defined twice. +fn assert_loop_metadata(ir: &str, marked: usize) { + let mut definitions: HashMap<&str, &str> = HashMap::new(); + for line in ir.lines() { + if let Some((id, value)) = line + .strip_prefix('!') + .and_then(|rest| rest.split_once(" = ")) + && id.chars().all(|c| c.is_ascii_digit()) + { + assert!( + definitions.insert(id, value).is_none(), + "!{id} defined twice:\n{ir}" + ); + } + } + + let loop_ids: Vec<&str> = ir + .lines() + .filter(|line| line.trim_start().starts_with("br ")) + .filter_map(|line| line.split_once("!llvm.loop !")) + .map(|(_, rest)| rest.split(|c: char| !c.is_ascii_digit()).next().unwrap()) + .collect(); + assert_eq!(loop_ids.len(), marked, "marked latch branches:\n{ir}"); + assert_eq!( + loop_ids.iter().collect::>().len(), + marked, + "every marked loop needs its own node:\n{ir}" + ); + + let properties: Vec<&str> = definitions + .iter() + .filter(|(_, value)| **value == "!{!\"llvm.loop.unroll.disable\"}") + .map(|(id, _)| *id) + .collect(); + assert_eq!(properties.len(), 1, "one shared property node:\n{ir}"); + for id in loop_ids { + assert_eq!( + definitions.get(id).copied(), + Some(format!("distinct !{{!{id}, !{}}}", properties[0]).as_str()), + "loop node !{id}:\n{ir}" + ); + } +} + +/// Lowering keeps the hinted loops' latch marks, and the exporter emits valid +/// loop metadata in both the PTX and the legacy NVVM dialects, alongside the +/// other metadata that shares the module's ID counter. +#[test] +fn nounroll_hints_become_loop_metadata_on_their_latches() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![u32_ty.into()]); + let n = entry.deref(&ctx).get_argument(0); + let region = entry + .deref(&ctx) + .get_parent_region() + .expect("entry in a region"); + + // hinted, unhinted, hinted + let mut tail = entry; + for nounroll in [true, false, true] { + tail = append_counted_loop(&mut ctx, region, tail, n, nounroll); + } + append_return(&mut ctx, tail); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ptx_ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + assert_loop_metadata(&ptx_ir, 2); + + let nvvm_ir = llvm_export::export::export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .map_err(|error| anyhow::anyhow!(error))?; + assert!(nvvm_ir.contains("!nvvmir.version"), "{nvvm_ir}"); + assert_loop_metadata(&nvvm_ir, 2); + Ok(()) +} diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/main.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/main.rs index 09ab58e43a..10f4f948cc 100644 --- a/cuda-oxide/crates/mir-lower/tests/lowering_test/main.rs +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/main.rs @@ -14,6 +14,7 @@ mod calls_and_values; mod common; mod cp_async; mod inline_ptx; +mod loops; mod math_conversions; mod matrix_memory; mod mma; diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs index 0650940e1a..714267bb04 100644 --- a/cuda-oxide/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs @@ -2379,6 +2379,18 @@ fn test_pointer_form_wgmma_counted_k_loop_stays_register_resident() -> Result<() #[test] fn test_f16_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow::Error> { + assert_f16_counted_k_loop_fuses(false) +} + +/// `#[nounroll]` on the K loop plants a hint in its latch. Counted-loop fusion +/// must still accept the loop, since F16 has no pointer fallback, and the fused +/// loop must stay rolled. +#[test] +fn test_f16_wgmma_counted_k_loop_with_nounroll_hint_still_fuses() -> Result<(), anyhow::Error> { + assert_f16_counted_k_loop_fuses(true) +} + +fn assert_f16_counted_k_loop_fuses(plant_nounroll: bool) -> Result<(), anyhow::Error> { use dialect_mir::types::{MirArrayType, MirPtrType}; use pliron::basic_block::BasicBlock; use pliron::builtin::op_interfaces::OperandSegmentInterface; @@ -2477,6 +2489,11 @@ fn test_f16_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow: branch.insert_at_back(header, &ctx); // latch: one WGMMA per K iteration and affine descriptor recurrences. + if plant_nounroll { + mir::MirNoUnrollHintOp::new(&mut ctx) + .get_operation() + .insert_at_back(latch, &ctx); + } append_pointer_wgmma_mma_f16(&mut ctx, latch, accumulator, desc_a, desc_b); let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); @@ -2568,6 +2585,12 @@ fn test_f16_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow: "the F16 counted-loop template must contain exactly one F16 MMA" ); assert!(!template.contains(".bf16.bf16")); + // `#[nounroll]` moves from the fused K loop onto the template's own loop + assert_eq!( + template.contains("L__wgmma_loop_${:uid}:\n .pragma \"nounroll\";\n"), + plant_nounroll, + "the fused loop header must start with the nounroll pragma exactly when the K loop carried the hint:\n{template}" + ); assert_eq!( template.matches("wgmma.commit_group.sync.aligned").count(), 1 diff --git a/cuda-oxide/crates/mir-transforms/README.md b/cuda-oxide/crates/mir-transforms/README.md index 1ab94371e0..b90693249b 100644 --- a/cuda-oxide/crates/mir-transforms/README.md +++ b/cuda-oxide/crates/mir-transforms/README.md @@ -12,6 +12,12 @@ The first pass is loop unrolling, requested by the `#[unroll]` / `#[unroll(N)]` annotation, which the importer records as a `mir.unroll_hint` operation inside the annotated loop. Further loop passes belong here too. +`nounroll.rs` handles the opposite request, `#[nounroll]`. The importer +records it as a `mir.nounroll_hint` operation, which the unroll pass leaves in +place. `mir-lower` calls `resolve_nounroll_hints` just before lowering: it +gives each annotated loop a single latch, removes the hint, and returns the +latch branch, which the LLVM exporter marks with `llvm.loop.unroll.disable`. + Two more sources ship with it. `canonicalize.rs` prepares a loop for the unroller: it merges multiple back-edges (several `continue`s) into one synthetic latch and forwards header-carried values out of each loop exit so diff --git a/cuda-oxide/crates/mir-transforms/src/lib.rs b/cuda-oxide/crates/mir-transforms/src/lib.rs index 187a556927..092080ca78 100644 --- a/cuda-oxide/crates/mir-transforms/src/lib.rs +++ b/cuda-oxide/crates/mir-transforms/src/lib.rs @@ -9,11 +9,13 @@ //! that promotes memory slots to plain SSA values) and before the IR is lowered //! to the LLVM dialect on its way to PTX. The first pass here is loop unrolling, //! switched on by the `#[unroll]` / `#[unroll(N)]` annotation (recorded as a -//! `mir.unroll_hint` operation inside the annotated loop). More loop passes can -//! live here too. +//! `mir.unroll_hint` operation inside the annotated loop). [`nounroll`] handles +//! the opposite request, `#[nounroll]`, by marking the loop's latch just before +//! lowering. More loop passes can live here too. pub mod analyses; mod canonicalize; pub mod forward_compiler_result_bundles; +pub mod nounroll; pub mod scalarize_borrowed_aggregate_reads; pub mod unroll; diff --git a/cuda-oxide/crates/mir-transforms/src/nounroll.rs b/cuda-oxide/crates/mir-transforms/src/nounroll.rs new file mode 100644 index 0000000000..870290e88b --- /dev/null +++ b/cuda-oxide/crates/mir-transforms/src/nounroll.rs @@ -0,0 +1,177 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `#[nounroll]`: keep an annotated loop rolled through LLVM, libNVVM, and +//! ptxas. +//! +//! The frontend records the request as a `mir.nounroll_hint` operation inside +//! the annotated loop's body. The hint has no side effects interface, so DCE +//! keeps it, and the unroll pass ignores it: when an enclosing `#[unroll]` loop +//! is unrolled, each body copy gets its own copy of the inner loop and of its +//! hint. +//! +//! [`resolve_nounroll_hints`] runs just before lowering, after every MIR +//! transform that can clone or reshape loops. For each hint it finds the +//! innermost loop that contains it, routes that loop's back-edges through one +//! latch block ending in `mir.goto header`, removes the hint, and returns that +//! goto. The lowering attaches `llvm.loop.unroll.disable` to the branch the goto +//! becomes. +//! +//! The loop needs its own latch because LLVM reads loop metadata from the +//! latch's terminator. With several latches, LLVM ignores the metadata unless +//! every latch carries the same node. A branch that ends an inner loop and also +//! jumps to the outer header would apply the request to both loops. + +use dialect_mir::ops::control_flow::MirNoUnrollHintOp; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::graph::ControlFlowGraph; +use pliron::graph::dominance::DomInfo; +use pliron::linked_list::ContainsLinkedList; +use pliron::operation::Operation; +use pliron::pass::AnalysisManager; +use pliron::region::Region; +use pliron::result::Result; +use rustc_hash::FxHashSet; + +use crate::analyses::loop_info::LoopInfo; +use crate::canonicalize::{CanonicalizeOutcome, merge_backedges}; +use crate::unroll::{collect_functions, hint_location}; + +/// Map every `mir.nounroll_hint` in `module` to its loop and return the latch +/// branch of each marked loop. +/// +/// Each returned operation is a `mir.goto` whose only successor is its loop's +/// header, and it is the only back-edge of that loop. Every hint is erased, +/// including hints that could not be applied. Hints in unreachable blocks are +/// dropped silently; the others that cannot be applied produce a warning on +/// stderr, in the same form as the unroll pass's warnings. +/// +/// Functions without hints are left untouched. In a function with hints, this +/// may insert one forwarding latch block per marked loop. +pub fn resolve_nounroll_hints( + module: Ptr, + ctx: &mut Context, +) -> Result>> { + let mut latches = Vec::new(); + for func_op in collect_functions(module, ctx) { + let region = func_op.deref(ctx).get_region(0); + + // A hint in dead code has no loop to keep rolled. WGMMA counted-loop + // fusion leaves a fused K loop's blocks unreachable this way, after + // moving the request onto the fused operation. Merging back-edges + // below only adds blocks, so reachability is computed once + let reachable = reachable_blocks(ctx, region); + for (hint, block) in collect_hints(ctx, region) { + if !reachable.contains(&block) { + Operation::erase(hint, ctx); + } + } + + // One loop per round. Merging back-edges changes the CFG, so loop + // facts are recomputed before the next hint is placed + loop { + let hints = collect_hints(ctx, region); + let Some(&(hint, block)) = hints.first() else { + break; + }; + + let info = loop_info(module, ctx, region)?; + let Some(loop_id) = info.innermost_loop(block) else { + warn_not_applied( + ctx, + hint, + "the annotation is not inside a recognizable loop", + ); + Operation::erase(hint, ctx); + continue; + }; + + match merge_backedges(ctx, &info, loop_id) { + CanonicalizeOutcome::Changed => continue, + CanonicalizeOutcome::Unsupported(reason) => { + warn_not_applied( + ctx, + hint, + &format!("could not normalize its back-edges: {reason}"), + ); + } + CanonicalizeOutcome::Unchanged => { + // `Unchanged` means the loop has exactly one back-edge, a + // `mir.goto` to the header + let latch = info.loops()[loop_id].latches[0]; + let Some(branch) = latch.deref(ctx).get_terminator(ctx) else { + return pliron::input_err_noloc!( + "a normalized loop latch has no terminator" + ); + }; + latches.push(branch); + } + } + + // a loop can hold several hints, for example after block merging + for (other, other_block) in hints { + if info.innermost_loop(other_block) == Some(loop_id) { + Operation::erase(other, ctx); + } + } + } + } + + Ok(latches) +} + +/// The blocks of `region` reachable from its entry block. +fn reachable_blocks(ctx: &Context, region: Ptr) -> FxHashSet> { + let mut reachable = FxHashSet::default(); + let mut worklist: Vec<_> = region.deref(ctx).iter(ctx).take(1).collect(); + while let Some(block) = worklist.pop() { + if reachable.insert(block) { + worklist.extend(region.successors(ctx, &block)); + } + } + reachable +} + +/// Compute fresh loop facts for `region`. +fn loop_info(module: Ptr, ctx: &mut Context, region: Ptr) -> Result { + let mut analyses = AnalysisManager::default(); + let mut dom_info = analyses.get_analysis_mut::(module, ctx)?; + let dom = dom_info.get_dom_tree(ctx, region); + Ok(LoopInfo::compute(ctx, region, dom)) +} + +/// The `mir.nounroll_hint` ops in `region`, each with the block it sits in. +fn collect_hints(ctx: &Context, region: Ptr) -> Vec<(Ptr, Ptr)> { + let mut out = Vec::new(); + for block in region.deref(ctx).iter(ctx) { + for op in block.deref(ctx).iter(ctx) { + if Operation::get_op::(op, ctx).is_some() { + out.push((op, block)); + } + } + } + out +} + +/// Report a `#[nounroll]` that could not be applied, with the annotation's +/// position on a `-->` line when the hint carries one. +fn warn_not_applied(ctx: &Context, hint: Ptr, reason: &str) { + eprintln!( + "{}", + not_applied_warning(reason, hint_location(ctx, hint).as_deref()) + ); +} + +/// The full text of one "not applied" warning. +fn not_applied_warning(reason: &str, location: Option<&str>) -> String { + let mut warning = + format!("warning: #[nounroll] requested but the loop may still be unrolled: {reason}"); + if let Some(location) = location { + warning.push_str("\n --> "); + warning.push_str(location); + } + warning +} diff --git a/cuda-oxide/crates/mir-transforms/src/unroll.rs b/cuda-oxide/crates/mir-transforms/src/unroll.rs index 1616439232..1361cd7f70 100644 --- a/cuda-oxide/crates/mir-transforms/src/unroll.rs +++ b/cuda-oxide/crates/mir-transforms/src/unroll.rs @@ -116,7 +116,7 @@ fn unroll_kind(factor: u32) -> String { /// A hint op's source position as `file:line:column`, the form editors and /// terminals turn into a link. Returns `None` when the op does not carry a /// single position in a file. -fn hint_location(ctx: &Context, op: Ptr) -> Option { +pub(crate) fn hint_location(ctx: &Context, op: Ptr) -> Option { match op.deref(ctx).loc() { Location::SrcPos { src: Source::File(path), @@ -1653,7 +1653,7 @@ fn append_cmp( } /// All function ops in the module (each `mir.func`). -fn collect_functions(module: Ptr, ctx: &Context) -> Vec> { +pub(crate) fn collect_functions(module: Ptr, ctx: &Context) -> Vec> { let mut out = Vec::new(); let module_region = module.deref(ctx).get_region(0); let blocks: Vec> = module_region.deref(ctx).iter(ctx).collect(); diff --git a/cuda-oxide/crates/mir-transforms/tests/nounroll.rs b/cuda-oxide/crates/mir-transforms/tests/nounroll.rs new file mode 100644 index 0000000000..18807638da --- /dev/null +++ b/cuda-oxide/crates/mir-transforms/tests/nounroll.rs @@ -0,0 +1,223 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tests for `#[nounroll]` latch placement. Each test plants a +//! `mir.nounroll_hint` where the frontend would, runs the unroll pass when the +//! case involves one, then runs `resolve_nounroll_hints` and checks which loop +//! latches came back. +//! +//! The property that matters for LLVM: every marked loop ends up with exactly +//! one latch, that latch is the returned branch, and no unmarked loop's latch +//! is returned. + +mod common; + +use common::{block, goto, mir_ctx, multi_latch_counted_loop, nested_counted_loop}; +use dialect_mir::ops::{MirGotoOp, MirNoUnrollHintOp, MirUnrollHintOp}; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::graph::dominance::DomInfo; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::pass::AnalysisManager; +use pliron::region::Region; +use rustc_hash::FxHashSet; + +use mir_transforms::analyses::loop_info::LoopInfo; +use mir_transforms::nounroll::resolve_nounroll_hints; +use mir_transforms::unroll::unroll_annotated_loops; + +fn loop_info(ctx: &Context, region: Ptr) -> LoopInfo { + let mut dom = DomInfo::default(); + let dt = dom.get_dom_tree(ctx, region); + LoopInfo::compute(ctx, region, dt) +} + +fn plant_nounroll(ctx: &mut Context, block: Ptr) { + MirNoUnrollHintOp::new(ctx) + .get_operation() + .insert_at_front(block, ctx); +} + +fn nounroll_hint_count(ctx: &Context, region: Ptr) -> usize { + region + .deref(ctx) + .iter(ctx) + .flat_map(|block| block.deref(ctx).iter(ctx)) + .filter(|&op| Operation::get_op::(op, ctx).is_some()) + .count() +} + +/// Run the resolver and check the shared contract: hints are gone, the module +/// verifies, and every returned branch is a `mir.goto` to a loop header that is +/// that loop's only back-edge. Returns the headers of the marked loops. +fn resolve( + ctx: &mut Context, + module: Ptr, + region: Ptr, +) -> FxHashSet> { + let latches = resolve_nounroll_hints(module, ctx).expect("resolve nounroll hints"); + pliron::operation::verify_operation(module, ctx).expect("valid IR after latch placement"); + assert_eq!(nounroll_hint_count(ctx, region), 0, "every hint is erased"); + + let info = loop_info(ctx, region); + let mut headers = FxHashSet::default(); + for &branch in &latches { + assert!( + Operation::get_op::(branch, ctx).is_some(), + "a marked latch ends in mir.goto" + ); + let latch = branch + .deref(ctx) + .get_parent_block() + .expect("branch in a block"); + let successors: Vec<_> = branch.deref(ctx).successors().collect(); + let [header] = successors[..] else { + panic!("a marked latch has one successor, got {}", successors.len()); + }; + let lp = info + .loops() + .iter() + .find(|lp| lp.header == header) + .expect("a marked latch jumps to a loop header"); + assert_eq!( + lp.latches, + [latch], + "the marked branch is the loop's only latch" + ); + assert!(headers.insert(header), "one marked latch per loop"); + } + headers +} + +/// A hint marks the loop whose body holds it, not the loop around it, and the +/// enclosing loop's latch is not returned. +#[test] +fn nounroll_marks_only_the_annotated_loop_of_a_nest() { + for annotate_inner in [true, false] { + let mut ctx = mir_ctx(); + let lp = nested_counted_loop(&mut ctx, 3, 2); + let (body, expected) = if annotate_inner { + (lp.inner_body, lp.inner_header) + } else { + (lp.outer_body, lp.outer_header) + }; + plant_nounroll(&mut ctx, body); + + let headers = resolve(&mut ctx, lp.module, lp.region); + + assert_eq!(headers.into_iter().collect::>(), [expected]); + assert_eq!( + loop_info(&ctx, lp.region).loops().len(), + 2, + "both loops remain" + ); + } +} + +/// Two `continue` paths are two back-edges. LLVM ignores loop metadata unless +/// every latch carries the same node, so the resolver routes both through one +/// new latch block and returns its branch. +#[test] +fn nounroll_merges_several_back_edges_into_one_latch() { + let mut ctx = mir_ctx(); + let lp = multi_latch_counted_loop(&mut ctx, 4, 1, 1); + assert_eq!(loop_info(&ctx, lp.region).loops()[0].latches.len(), 2); + plant_nounroll(&mut ctx, lp.normal_latch); + + let headers = resolve(&mut ctx, lp.module, lp.region); + + assert_eq!(headers.into_iter().collect::>(), [lp.header]); +} + +/// A full `#[unroll]` of an outer loop copies the inner `#[nounroll]` loop, +/// hint included, into each body copy. Every copy must stay rolled. +#[test] +fn nounroll_inside_a_fully_unrolled_loop_marks_every_copy() { + let mut ctx = mir_ctx(); + let lp = nested_counted_loop(&mut ctx, 3, 2); + MirUnrollHintOp::new(&mut ctx, 0) + .get_operation() + .insert_at_front(lp.outer_body, &ctx); + plant_nounroll(&mut ctx, lp.inner_body); + + unroll_annotated_loops(lp.module, &mut ctx, &mut AnalysisManager::default()) + .expect("unroll the outer loop"); + assert_eq!(nounroll_hint_count(&ctx, lp.region), 3, "one hint per copy"); + + let headers = resolve(&mut ctx, lp.module, lp.region); + + let info = loop_info(&ctx, lp.region); + assert_eq!(info.loops().len(), 3, "three inner copies, no outer loop"); + assert!( + info.loops() + .iter() + .all(|inner| headers.contains(&inner.header)) + ); +} + +/// A partial `#[unroll(4)]` of an outer loop leaves a main loop and a remainder +/// loop. The main loop holds four inner copies and the remainder holds one. All +/// five stay rolled, and the main and remainder loops themselves are unmarked. +#[test] +fn nounroll_inside_a_partially_unrolled_loop_marks_main_and_remainder_copies() { + let mut ctx = mir_ctx(); + let lp = nested_counted_loop(&mut ctx, 8, 2); + MirUnrollHintOp::new(&mut ctx, 4) + .get_operation() + .insert_at_front(lp.outer_body, &ctx); + plant_nounroll(&mut ctx, lp.inner_body); + + unroll_annotated_loops(lp.module, &mut ctx, &mut AnalysisManager::default()) + .expect("partially unroll the outer loop"); + + let headers = resolve(&mut ctx, lp.module, lp.region); + + let info = loop_info(&ctx, lp.region); + let outer: FxHashSet<_> = info + .top_level() + .iter() + .map(|&id| info.loops()[id].header) + .collect(); + assert_eq!(outer.len(), 2, "main loop + remainder loop"); + let inner: FxHashSet<_> = info + .loops() + .iter() + .map(|lp| lp.header) + .filter(|header| !outer.contains(header)) + .collect(); + assert_eq!(inner.len(), 5); + assert_eq!(headers, inner, "exactly the inner copies are marked"); +} + +/// A hint outside every loop has nothing to mark. It is dropped with a warning +/// instead of failing the compile. +#[test] +fn nounroll_outside_a_loop_marks_nothing() { + let mut ctx = mir_ctx(); + let lp = nested_counted_loop(&mut ctx, 3, 2); + plant_nounroll(&mut ctx, lp.preheader); + + let headers = resolve(&mut ctx, lp.module, lp.region); + + assert!(headers.is_empty()); +} + +/// Dead code that branches into a loop puts that loop's blocks in its reach +/// backwards, but a hint there has no loop to keep rolled. WGMMA counted-loop +/// fusion leaves a fused K loop's blocks unreachable like this. +#[test] +fn nounroll_in_unreachable_code_marks_nothing() { + let mut ctx = mir_ctx(); + let lp = nested_counted_loop(&mut ctx, 3, 2); + let dead = block(&mut ctx, lp.region, vec![]); + goto(&mut ctx, dead, lp.outer_latch, vec![]); + plant_nounroll(&mut ctx, dead); + + let headers = resolve(&mut ctx, lp.module, lp.region); + + assert!(headers.is_empty()); +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/Cargo.lock new file mode 100644 index 0000000000..004c2bde68 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/Cargo.lock @@ -0,0 +1,753 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nounroll" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/Cargo.toml new file mode 100644 index 0000000000..2dc9b12d85 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "nounroll" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" + +[workspace] + +# Keeps `#[nounroll]` loops rolled through LLVM, libNVVM, and ptxas. +# cargo oxide run nounroll + +[dependencies] +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/src/main.rs new file mode 100644 index 0000000000..b7f8759a36 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/src/main.rs @@ -0,0 +1,318 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `#[nounroll]` keeps a loop rolled through LLVM, libNVVM, and ptxas. +//! +//! ptxas fully unrolls small constant-trip-count loops on its own, even when the +//! PTX keeps them rolled. `#[nounroll]` lowers to `llvm.loop.unroll.disable` +//! metadata on the loop's latch, and the NVPTX backend turns that into +//! `.pragma "nounroll";` at the loop header, which ptxas honors. +//! +//! - `taps_rolled` keeps a constant 4-trip loop rolled. `taps_default` is the +//! same loop without the annotation. +//! - `rolled_k_unrolled_inner` keeps an outer K loop rolled while `#[unroll]` +//! expands the loop inside it. +//! - `rolled_inside_full_unroll` and `rolled_inside_partial_unroll` put a +//! `#[nounroll]` range loop inside a loop that `#[unroll]` / `#[unroll(4)]` +//! unrolls. Every copy of the inner loop, including the one in the partial +//! unroll's remainder loop, stays rolled. +//! - `rolled_in_device_fn` calls a `#[device]` function whose `loop` carries +//! the annotation. +//! +//! Every annotated loop loads from memory. With an arithmetic-only body and a +//! constant trip count, LLVM or ptxas computes the loop's result outright and +//! deletes the loop, which leaves nothing to keep rolled. +//! +//! `verify-code-shape.sh` counts the pragmas in each PTX entry. GPU-less CI runs +//! it through `scripts/smoketest.sh --compile-only`. The host checks every +//! result. +//! +//! Run: cargo oxide run nounroll + +use cuda_core::simt::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; +use cuda_device::{DisjointSlice, device, kernel, thread}; +use cuda_host::cuda_module; + +const TAPS: usize = 4; +const N: usize = 32; +/// Largest outer trip count the nested kernels may read inputs for. +const MAX_TRIP: u32 = 16; + +#[cuda_module] +mod kernels { + use super::*; + + /// `out[i] = sum(w[t] * x[i + t])` over `TAPS` taps, with the tap loop kept + /// rolled. + #[kernel] + pub fn taps_rolled(w: &[f32], x: &[f32], mut out: DisjointSlice) { + let tid = thread::index_1d(); + let i = tid.get(); + if let Some(slot) = out.get_mut(tid) { + let mut acc = 0.0f32; + let mut tap = 0usize; + #[nounroll] + while tap < TAPS { + acc += w[tap] * x[i + tap]; + tap += 1; + } + *slot = acc; + } + } + + /// The same loop as `taps_rolled`, left to the default unrolling. + #[kernel] + pub fn taps_default(w: &[f32], x: &[f32], mut out: DisjointSlice) { + let tid = thread::index_1d(); + let i = tid.get(); + if let Some(slot) = out.get_mut(tid) { + let mut acc = 0.0f32; + let mut tap = 0usize; + while tap < TAPS { + acc += w[tap] * x[i + tap]; + tap += 1; + } + *slot = acc; + } + } + + /// An outer K loop of 4 trips stays rolled while `#[unroll]` expands the + /// 4-lane loop inside it. `out[i] = sum(a[k * 4 + n] * (i + 1))` over every + /// `k < 8` step 2 and `n < 4`. + #[kernel] + pub fn rolled_k_unrolled_inner(a: &[u32], mut out: DisjointSlice) { + let tid = thread::index_1d(); + let scale = tid.get() as u32 + 1; + if let Some(slot) = out.get_mut(tid) { + let mut acc = [0u32; 4]; + let mut k = 0usize; + #[nounroll] + while k < 8 { + let mut n = 0usize; + #[unroll] + while n < 4 { + acc[n] = acc[n].wrapping_add(a[k * 4 + n].wrapping_mul(scale)); + n += 1; + } + k += 2; + } + *slot = acc[0] + .wrapping_add(acc[1]) + .wrapping_add(acc[2]) + .wrapping_add(acc[3]); + } + } + + /// Full `#[unroll]` copies the inner `#[nounroll]` loop into each of its + /// four body copies. `out[i] = i + sum(a[i_outer * 3 + j])` over + /// `i_outer < 4` and `j < 3`. + /// + /// The inner body loads from memory: LLVM would otherwise compute the sum + /// in closed form and delete the loop, leaving nothing to keep rolled. The + /// unchecked load and wrapping index keep panic exits out of the loop. + #[kernel] + pub fn rolled_inside_full_unroll(a: &[u32], mut out: DisjointSlice) { + let tid = thread::index_1d(); + let base = tid.get() as u32; + if let Some(slot) = out.get_mut(tid) { + let mut acc = base; + let mut i = 0u32; + #[unroll] + while i < 4 { + #[nounroll] + for j in 0..3u32 { + let index = i.wrapping_mul(3).wrapping_add(j) as usize; + // SAFETY: the host passes `3 * MAX_TRIP` elements, and + // `i < 4 <= MAX_TRIP` + acc = acc.wrapping_add(unsafe { *a.get_unchecked(index) }); + } + i += 1; + } + *slot = acc; + } + } + + /// Partial `#[unroll(4)]` leaves four inner copies in its main loop and one + /// in its remainder loop. `out[i] = i + sum(a[i_outer * 3 + j])` over + /// `i_outer < trip` and `j < 3`. + /// + /// Partial unrolling needs the loop test to be the only exit, so the body + /// avoids bounds checks and overflow checks the same way. The inner loop + /// is a `while`: a range `for` adds an unreachable `Option` arm, which + /// partial unrolling counts as a second exit. + #[kernel] + pub fn rolled_inside_partial_unroll(a: &[u32], mut out: DisjointSlice, trip: u32) { + let tid = thread::index_1d(); + let base = tid.get() as u32; + if let Some(slot) = out.get_mut(tid) { + let mut acc = base; + let mut i = 0u32; + #[unroll(4)] + while i < trip { + let mut j = 0u32; + #[nounroll] + while j < 3 { + let index = i.wrapping_mul(3).wrapping_add(j) as usize; + // SAFETY: the host passes `3 * MAX_TRIP` elements and + // `trip <= MAX_TRIP` + acc = acc.wrapping_add(unsafe { *a.get_unchecked(index) }); + j = j.wrapping_add(1); + } + i = i.wrapping_add(1); + } + *slot = acc; + } + } + + /// `x + a[0] + a[1] + a[2] + a[3]`, summed by a rolled `loop`, so the + /// annotation is covered on `loop` and in a `#[device]` function. + #[inline(never)] + #[device] + fn rolled_sum(a: &[u32], x: u32) -> u32 { + let mut acc = x; + let mut step = 0usize; + #[nounroll] + loop { + // SAFETY: the host passes at least four elements + acc = acc.wrapping_add(unsafe { *a.get_unchecked(step) }); + step = step.wrapping_add(1); + if step == 4 { + break; + } + } + acc + } + + /// `out[i] = i + a[0] + a[1] + a[2] + a[3]`, computed by `rolled_sum`. + #[kernel] + pub fn rolled_in_device_fn(a: &[u32], mut out: DisjointSlice) { + let tid = thread::index_1d(); + let base = tid.get() as u32; + if let Some(slot) = out.get_mut(tid) { + *slot = rolled_sum(a, base); + } + } +} + +fn taps_expected(w: &[f32], x: &[f32]) -> Vec { + (0..N) + .map(|i| (0..TAPS).fold(0.0, |acc, tap| acc + w[tap] * x[i + tap])) + .collect() +} + +fn nested_expected(base: u32, outer_trips: u32, a: &[u32]) -> u32 { + (0..outer_trips * 3).fold(base, |acc, index| acc.wrapping_add(a[index as usize])) +} + +fn main() { + println!("=== nounroll ===\n"); + + let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); + let module = kernels::load(&ctx).expect("Failed to load embedded CUDA module"); + let stream = ctx.default_stream(); + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (N as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut failures = 0usize; + + let w: Vec = (0..TAPS).map(|t| 0.5 + t as f32).collect(); + let x: Vec = (0..N + TAPS).map(|v| (v % 7) as f32 - 3.0).collect(); + let w_dev = DeviceBuffer::from_host(&stream, &w).unwrap(); + let x_dev = DeviceBuffer::from_host(&stream, &x).unwrap(); + let want_taps = taps_expected(&w, &x); + + let mut d_rolled = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: each thread reads `x[i..i + TAPS]`, which `x` covers, and writes + // its own output element. + unsafe { module.taps_rolled(stream.as_ref(), cfg, &w_dev, &x_dev, &mut d_rolled) } + .expect("launch taps_rolled"); + let mut d_default = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: as for `taps_rolled`. + unsafe { module.taps_default(stream.as_ref(), cfg, &w_dev, &x_dev, &mut d_default) } + .expect("launch taps_default"); + // small integers times halves: every partial sum is exact in f32 + for (name, got) in [ + ("taps_rolled", d_rolled.to_host_vec(&stream).unwrap()), + ("taps_default", d_default.to_host_vec(&stream).unwrap()), + ] { + if got != want_taps { + println!("FAIL {name}: {got:?} expected {want_taps:?}"); + failures += 1; + } + } + + let a: Vec = (0..32).map(|v| v * 3 + 1).collect(); + let a_dev = DeviceBuffer::from_host(&stream, &a).unwrap(); + let mut d_k = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: every thread reads `a[0..32]` and writes its own output element. + unsafe { module.rolled_k_unrolled_inner(stream.as_ref(), cfg, &a_dev, &mut d_k) } + .expect("launch rolled_k_unrolled_inner"); + let a_sum: u32 = (0..8) + .step_by(2) + .flat_map(|k| (0..4).map(move |n| k * 4 + n)) + .map(|index| a[index]) + .sum(); + let want_k: Vec = (0..N as u32).map(|i| a_sum * (i + 1)).collect(); + if d_k.to_host_vec(&stream).unwrap() != want_k { + println!("FAIL rolled_k_unrolled_inner"); + failures += 1; + } + + let nested: Vec = (0..3 * MAX_TRIP).map(|v| v * 7 + 2).collect(); + let nested_dev = DeviceBuffer::from_host(&stream, &nested).unwrap(); + let mut d_full = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: the kernel reads `nested[0..12]` and each thread writes its own + // output element. + unsafe { module.rolled_inside_full_unroll(stream.as_ref(), cfg, &nested_dev, &mut d_full) } + .expect("launch rolled_inside_full_unroll"); + let want_full: Vec = (0..N as u32) + .map(|i| nested_expected(i, 4, &nested)) + .collect(); + if d_full.to_host_vec(&stream).unwrap() != want_full { + println!("FAIL rolled_inside_full_unroll"); + failures += 1; + } + + // 10 = two trips of the main loop plus a remainder of two + let trip = 10u32; + let mut d_partial = DeviceBuffer::::zeroed(&stream, N).unwrap(); + assert!(trip <= MAX_TRIP); + // SAFETY: the kernel reads `nested[0..3 * trip]`, which `nested` covers, + // and each thread writes its own output element. + unsafe { + module.rolled_inside_partial_unroll(stream.as_ref(), cfg, &nested_dev, &mut d_partial, trip) + } + .expect("launch rolled_inside_partial_unroll"); + let want_partial: Vec = (0..N as u32) + .map(|i| nested_expected(i, trip, &nested)) + .collect(); + if d_partial.to_host_vec(&stream).unwrap() != want_partial { + println!("FAIL rolled_inside_partial_unroll"); + failures += 1; + } + + let mut d_device = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: the kernel reads `nested[0..4]` and each thread writes its own + // output element. + unsafe { module.rolled_in_device_fn(stream.as_ref(), cfg, &nested_dev, &mut d_device) } + .expect("launch rolled_in_device_fn"); + let want_device: Vec = (0..N as u32) + .map(|i| nested[..4].iter().fold(i, |acc, &v| acc.wrapping_add(v))) + .collect(); + if d_device.to_host_vec(&stream).unwrap() != want_device { + println!("FAIL rolled_in_device_fn"); + failures += 1; + } + + if failures == 0 { + println!("nounroll: PASS ({N} threads; every rolled loop computes the same results)"); + } else { + println!("nounroll: FAIL ({failures} kernels mismatched)"); + std::process::exit(1); + } +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/verify-code-shape.sh new file mode 100644 index 0000000000..c213bcf950 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nounroll/verify-code-shape.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +set -euo pipefail + +root="$(cd "$(dirname "$0")" && pwd)" +ptx="${1:-$root/nounroll.ptx}" +ll="${2:-$root/nounroll.ll}" + +for file in "$ptx" "$ll"; do + if [[ ! -f "$file" ]]; then + echo "artifact not found: $file" >&2 + exit 1 + fi +done + +python3 - "$ptx" "$ll" <<'PY' +import re +import sys + +ptx_path, ll_path = sys.argv[1], sys.argv[2] +ptx = open(ptx_path, encoding="utf-8").read() +ll = open(ll_path, encoding="utf-8").read() +# Strip comments so a mention in a comment cannot satisfy or hide a check. +ptx = re.sub(r"/\*.*?\*/|//[^\n]*", "", ptx, flags=re.S) +ll = re.sub(r";[^\n]*", "", ll) + +# Marked loops per function. The full unroll leaves four copies of its inner +# loop; the partial unroll leaves four in its main loop and one in its +# remainder loop. `taps_default` is the control. +expected = { + "taps_rolled": 1, + "taps_default": 0, + "rolled_k_unrolled_inner": 1, + "rolled_inside_full_unroll": 4, + "rolled_inside_partial_unroll": 5, + "rolled_sum": 1, +} + + +def body(text, header_pattern, what): + # PTX forward-declares a callee defined after its caller; a declaration + # reaches `;` before any `{`, so take the first header that opens a body + brace = -1 + for header in re.finditer(header_pattern, text): + candidate = text.find("{", header.end()) + semicolon = text.find(";", header.end()) + if candidate >= 0 and (semicolon < 0 or candidate < semicolon): + brace = candidate + break + if brace < 0: + sys.exit(f"error: missing {what}") + depth = 1 + end = brace + 1 + while depth and end < len(text): + depth += (text[end] == "{") - (text[end] == "}") + end += 1 + if depth: + sys.exit(f"error: unterminated {what}") + return text[brace + 1 : end - 1] + + +for name, want in expected.items(): + # a device function's symbol is mangled or path-qualified around its name + symbol = r"[\w$.]*" + re.escape(name) + r"[\w$.]*" + ll_body = body(ll, r"define\b[^\n]*@\"?" + symbol + r"\"?\(", f"LLVM function {name}") + ptx_body = body( + ptx, r"\.(?:entry|func)\s+(?:\([^)]*\)\s*)?" + symbol + r"\s*\(", f"PTX function {name}" + ) + # The exporter marks exactly one latch per annotated loop + marked = len(re.findall(r"^\s*br\b[^\n]*!llvm\.loop\s+!\d+", ll_body, flags=re.M)) + if marked != want: + sys.exit(f"error: {name} has {marked} latches with !llvm.loop in {ll_path}, expected {want}") + # LLVM may copy a rolled loop (for example when it unrolls the remainder + # loop around it), and each copy keeps its pragma, so PTX can hold more + pragmas = len(re.findall(r'\.pragma\s+"nounroll"\s*;', ptx_body)) + if pragmas < want or (want == 0 and pragmas): + sys.exit(f'error: {name} has {pragmas} `.pragma "nounroll";` in {ptx_path}, expected {want}') + print(f"{name}: {marked} marked latch(es), {pragmas} nounroll pragma(s) PASS") +print("nounroll code shape: PASS") +PY diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs index 97e5b7c45a..c5048b7388 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs @@ -720,17 +720,25 @@ fn is_ptx_asm_marker_path(fn_path: &str) -> bool { || has_arity_suffix(fn_path, "cuda_device::ptx::__ptx_asm_void_") } -/// Returns true for the hidden `__unroll_config::` marker function. +/// Returns true for the hidden `__unroll_config::` and +/// `__nounroll_config` marker functions. /// -/// `#[unroll]` / `#[unroll(N)]` on a loop makes the `#[kernel]` or `#[device]` -/// macro plant a call to this marker at the top of the loop body. The MIR -/// importer rewrites that call to a `mir.unroll_hint` op (consumed by the -/// loop-unroll pass) and never emits a real call. So the marker's empty body must -/// not be collected, or it would show up as a dead `.func` in the generated PTX. -/// Matches both the re-exported path (`cuda_device::__unroll_config`) and the -/// full path. +/// `#[unroll]` / `#[unroll(N)]` / `#[nounroll]` on a loop makes the `#[kernel]` +/// or `#[device]` macro plant a call to one of these markers at the top of the +/// loop body. The MIR importer rewrites that call to a `mir.unroll_hint` or +/// `mir.nounroll_hint` op and never emits a real call. So the marker's empty +/// body must not be collected, or it would show up as a dead `.func` in the +/// generated PTX. Matches both the re-exported path +/// (`cuda_device::__unroll_config`) and the full path, with or without generic +/// arguments, and nothing else: a user function whose name merely starts with +/// a marker's name must still be collected. fn is_unroll_marker_path(fn_path: &str) -> bool { - fn_path.contains("::__unroll_config") + let Some(path) = fn_path.strip_prefix("cuda_device::") else { + return false; + }; + let name = path.strip_prefix("thread::").unwrap_or(path); + let name = name.split_once("::<").map_or(name, |(name, _)| name); + matches!(name, "__unroll_config" | "__nounroll_config") } /// Returns true for zero-cost compile-time configuration markers planted by @@ -2980,8 +2988,8 @@ pub fn dump_device_mir_info<'tcx>(tcx: TyCtxt<'tcx>, functions: &[CollectedFunct mod tests { use super::{ device_runtime_checks_target, is_host_cpu_arch_def_path, is_kernel_entry_def_path, - is_spin_loop_expansion, is_std_float_inherent_method, truncate_path_for_box, - unsupported_codegen_protocol_root, + is_spin_loop_expansion, is_std_float_inherent_method, is_unroll_marker_path, + truncate_path_for_box, unsupported_codegen_protocol_root, }; use reserved_oxide_symbols::{ DEVICE_PREFIX, KERNEL_PREFIX, LEGACY_DEVICE_PREFIX, LEGACY_KERNEL_PREFIX, @@ -3052,6 +3060,27 @@ mod tests { assert!(!is_spin_loop_expansion("::hint::spin_loop")); } + #[test] + fn unroll_markers_match_only_cuda_device_marker_names() { + for path in [ + "cuda_device::__unroll_config::<4>", + "cuda_device::thread::__unroll_config::<0>", + "cuda_device::__nounroll_config", + "cuda_device::thread::__nounroll_config", + ] { + assert!(is_unroll_marker_path(path), "{path}"); + } + for path in [ + "cuda_device::thread::__nounroll_config_helper", + "cuda_device::__unroll_configure::<4>", + "my_kernels::__nounroll_config", + "my_kernels::cuda_device::__unroll_config::<4>", + "cuda_device::thread::inner::__nounroll_config", + ] { + assert!(!is_unroll_marker_path(path), "{path}"); + } + } + #[test] fn ptx_merge_marker_matches_only_the_final_path_component() { let marker = ptx_merge_required_marker("map"); diff --git a/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md index cc90fe89a5..7e40a0e1db 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md +++ b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md @@ -45,6 +45,7 @@ fn helper(x: f32) -> f32 { x * x } | `#[kernel]` | Mark a function as a GPU kernel entry point (`.entry` in PTX) | | `#[device]` | Mark a helper function or `extern "C"` block for device compilation | | `#[unroll]` / `#[unroll(N)]` | Request full unrolling, or unrolling by a factor `N >= 2` | +| `#[nounroll]` | Keep a loop rolled through LLVM, libNVVM, and ptxas | | `#[launch_bounds(max_threads, min_blocks)]` | Occupancy hints for register allocation | | `#[constant]` | Place a `ConstantMemory` static in constant memory, with a host `set_` | | `#[launch_contract(...)]` | Declare the launch shape a kernel requires, unlocking a safe (non-`unsafe`) launch | diff --git a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md index f568257006..ea7f0b4824 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md +++ b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md @@ -149,7 +149,7 @@ Anonymous promoted allocations remain unsupported. | For Loops (range, iterator, enumerate) | **Full** | Full iterator desugaring: range-based, `slice.iter()`, `enumerate()`, nested loops, `break`, `continue`. | | While Loops / If-Else | **Full** | Baseline control flow fully supported. | | Break and Continue | **Full** | `break` and `continue` in for/while loops, including early exit. | -| Loop Unroll Annotations | **Partial** | `#[unroll]` and `#[unroll(N)]` request unrolling of explicit counted `while` loops and range `for` loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, including `isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, are not yet recognized. Nested loops and multiple `continue` paths work; full unrolling preserves `break` paths and multiple exit targets, while partial unrolling requires a positive-step `<`/`<=` loop with an invariant limit and only the normal header exit. Requests are capped at 1,024 copies, 8,192 cloned blocks, and 65,536 cloned operations. | +| Loop Unroll Annotations | **Partial** | `#[unroll]` and `#[unroll(N)]` request unrolling of explicit counted `while` loops and range `for` loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, including `isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, are not yet recognized. Nested loops and multiple `continue` paths work; full unrolling preserves `break` paths and multiple exit targets, while partial unrolling requires a positive-step `<`/`<=` loop with an invariant limit and only the normal header exit. Requests are capped at 1,024 copies, 8,192 cloned blocks, and 65,536 cloned operations. `#[nounroll]` keeps any loop rolled through LLVM, libNVVM, and ptxas. | | Monomorphization-Dead Branches | **Partial** | Branches that become dead after generic specialization (e.g. the const-false arm of `if M::ENABLED`) are ignored by symbol collection, panic checks, and pointer address-space inference, so panic-only hooks in dead arms compile. Only switches rustc itself folds are pruned: a constant discriminant operand or a direct single-assignment constant. Multi-step constant copy chains keep both arms, matching rustc's host monomorphization; this is deliberate, not a general constant-propagation pass. | ## Compiler: Arithmetic and Casting diff --git a/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md b/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md index 3a2d6515af..2802d6d624 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md +++ b/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md @@ -329,6 +329,40 @@ Choose a factor that matches the structure of the algorithm—often a pipeline depth or another small natural group—and benchmark the result. Removing `#[unroll]` should always remain an easy comparison. +--- + +## Keeping a loop rolled with `#[nounroll]` + +The unroll pass only acts on `#[unroll]` requests, but LLVM, libNVVM, and ptxas +have their own unrollers. ptxas in particular fully unrolls small +constant-trip-count loops even when the PTX keeps them rolled. `#[nounroll]` +asks all three to leave a loop alone. It does not go through the unroll pass: + +1. The macro inserts a different marker, and `mir-importer` emits a + `mir.nounroll_hint` operation inside the loop. +2. The hint survives MIR preparation. The unroll pass ignores it, so unrolling + an enclosing loop copies the hint along with the inner loop. +3. Just before lowering, after the last transform that can reshape loops, each + hint is mapped to its innermost loop. The loop's back-edges are joined + through one latch block, and the hint is removed. +4. Lowering marks that latch's branch, and the LLVM IR exporter attaches + `!llvm.loop` metadata with `llvm.loop.unroll.disable` to it. The NVPTX + backend prints `.pragma "nounroll";` at the loop header; libNVVM reads the + same metadata. + +Step 3 also runs in full-debug builds, which skip MIR preparation. A hint that +is not inside any loop is dropped with a warning that names the attribute's +`file:line:column`. A hint in unreachable code is dropped silently. WGMMA +counted-loop fusion replaces a K loop with one operation that generates its own +PTX loop. It moves a `#[nounroll]` request onto that operation, whose loop +header then starts with `.pragma "nounroll";`, and leaves the original blocks +unreachable. + +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. If it sits inside another loop, the request applies to that enclosing +loop. + To inspect the pass's analysis, build verbosely: ```bash diff --git a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md index 19986fe39e..beefa81861 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md +++ b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md @@ -83,13 +83,13 @@ the GPU memory hierarchy: ### Operations -`dialect-mir` defines 62 operations across 12 categories, one per module under +`dialect-mir` defines 63 operations across 12 categories, one per module under `cuda-oxide/crates/dialect-mir/src/ops/`: | Category | Examples | Count | | :----------- | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----: | | Function | `mir.func` | 1 | -| Control flow | `mir.goto`, `mir.cond_br`, `mir.return`, `mir.assert`, `mir.unreachable`, `mir.unroll_hint` | 6 | +| Control flow | `mir.goto`, `mir.cond_br`, `mir.return`, `mir.assert`, `mir.unreachable`, `mir.unroll_hint`, `mir.nounroll_hint` | 7 | | Constants | `mir.constant`, `mir.float_constant`, `mir.undef` | 3 | | Memory | `mir.alloca`, `mir.load`, `mir.store`, `mir.assign`, `mir.ref`, `mir.ptr_offset`, `mir.memcpy`, `mir.memmove`, `mir.shared_alloc`, `mir.global_alloc`, `mir.extern_shared` | 11 | | Arithmetic | `mir.add`, `mir.sub`, `mir.mul`, `mir.div`, `mir.rem`, `mir.checked_add`, `mir.checked_sub`, `mir.checked_mul`, `mir.neg`, `mir.not`, `mir.shr`, `mir.shl`, `mir.bitand`, `mir.bitor`, `mir.bitxor` | 15 | diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md b/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md index 303334a50e..bcc6cc7952 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md @@ -389,6 +389,37 @@ Unrolling trades larger generated code for fewer branches and more optimization opportunities. Use it for small or performance-critical loops, and measure the result. +### Keeping a loop rolled + +ptxas fully unrolls small loops with a constant trip count on its own, even +when the PTX keeps them rolled, and then schedules every iteration's loads at +once. That can raise register pressure enough to spill. Put `#[nounroll]` on a +loop to keep it rolled: + +```rust +let mut tap = 0; +#[nounroll] +while tap < 4 { + acc += weight(tap) * input(i + tap); + tap += 1; +} +``` + +The compiler attaches `llvm.loop.unroll.disable` metadata to the loop, so LLVM +and libNVVM leave it rolled, and the NVPTX backend emits `.pragma "nounroll";` +at the loop header, which ptxas honors. This works on any loop the compiler +sees as a loop, including `for` loops over iterators, and in full-debug +builds. + +`#[nounroll]` applies only to the loop that carries it and takes no arguments. +A loop cannot carry both `#[unroll]` and `#[nounroll]`. A `#[nounroll]` loop +inside a loop that `#[unroll]` or `#[unroll(N)]` unrolls is copied into each +unrolled body, and every copy stays rolled. + +Put the annotation on a loop that repeats. A loop whose body always leaves it, +such as `loop { ...; break; }`, compiles to straight-line code. Inside another +loop, the request then applies to the enclosing loop. + :::{seealso} For how the compiler analyzes and rewrites annotated loops, including the stage-index peephole, see [Compiler Optimizations](../compiler/compiler-optimizations.md).