diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/common.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/common.rs index a5df8d3a17..90d283bd77 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/common.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/common.rs @@ -30,11 +30,11 @@ pub(super) fn anyhow_to_pliron(e: anyhow::Error) -> pliron::result::Error { /// type alone can look under-aligned: `{ i8, [7 x i8] }` says "align 1" /// to LLVM, while Rust may require 8. /// -/// The alloca lands at the use site, same as -/// [`convert_extract_array_element`](super::array_extract::convert_extract_array_element); -/// the standard `opt -O2` run (SROA) -/// removes it again. Hoisting these into the function's entry block is a -/// known follow-up for the unoptimized (`CUDA_OXIDE_NO_OPT=1`) path. +/// The alloca and store initially land at the use site, same as +/// [`convert_extract_array_element`](super::array_extract::convert_extract_array_element). +/// NVVM normalization hoists eligible fixed slots into the true entry block, +/// leaving initialization here. This exposes static storage to NVVM even +/// when the optional `opt -O2` / SROA pass is disabled. pub(super) fn spill_enum_value( ctx: &mut Context, rewriter: &mut DialectConversionRewriter, diff --git a/cuda-oxide/crates/nvvm-transforms/README.md b/cuda-oxide/crates/nvvm-transforms/README.md index 75a77caf16..4202bc895d 100644 --- a/cuda-oxide/crates/nvvm-transforms/README.md +++ b/cuda-oxide/crates/nvvm-transforms/README.md @@ -11,5 +11,10 @@ selected libNVVM will take: - Blackwell and newer keep modern operations, apart from the NVVM-wide compatibility rewrites that apply everywhere. +After those rewrites, fixed scratch allocations move into each function's entry +block so NVVM can use static storage. Initialization stays in its original block. +Dynamic counts and loop slots whose addresses may escape remain in place. The +normalization updates the LLVM-dialect IR before verification and text export. + Ordinary PTX builds skip this stage: it exists for the NVVM path, where libNVVM rather than `llc` produces the final code. diff --git a/cuda-oxide/crates/nvvm-transforms/src/entry_allocas.rs b/cuda-oxide/crates/nvvm-transforms/src/entry_allocas.rs new file mode 100644 index 0000000000..532b481249 --- /dev/null +++ b/cuda-oxide/crates/nvvm-transforms/src/entry_allocas.rs @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Normalize fixed scratch slots after LLVM dialect legalization for NVVM. +//! +//! Non-entry allocas can introduce dynamic-stack restoration when helpers are +//! inlined. Moving eligible storage into entry lets NVVM use static slots. +//! Initialization and memory operations stay at their original control-flow sites. + +use std::collections::HashSet; + +use llvm_export::ops::{self, FuncOp}; +use pliron::{ + builtin::attributes::IntegerAttr, + context::{Context, Ptr}, + graph::traversals::region::sccs_in_topological_order, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, + value::Value, +}; + +/// Normalize every LLVM function in the module, leaving declarations alone. +/// The function-local analysis runs after all other NVVM compatibility rewrites +/// and updates the same IR that the pipeline subsequently verifies and exports. +pub(super) fn hoist_fixed_allocas(ctx: &mut Context, module: Ptr) { + let functions: Vec<_> = module + .deref(ctx) + .regions() + .flat_map(|region| region.deref(ctx).iter(ctx).collect::>()) + .flat_map(|block| block.deref(ctx).iter(ctx).collect::>()) + .filter_map(|op| Operation::get_op::(op, ctx)) + .collect(); + for function in functions { + hoist_function_allocas(ctx, &function); + } +} + +/// A fixed slot outside a cycle executes at most once; its address stays unique. +/// Inside a cycle, reuse is allowed only when no iteration can retain or observe +/// the address. Dynamic counts and potentially escaping loop slots stay in place. +fn hoist_function_allocas(ctx: &mut Context, function: &FuncOp) { + let Some(entry) = function.get_entry_block(ctx) else { + return; + }; + let region = function.get_operation().deref(ctx).get_region(0); + let mut candidates = Vec::new(); + for block in region.deref(ctx).iter(ctx) { + if block == entry { + continue; + } + for allocation in block.deref(ctx).iter(ctx) { + if Operation::get_op::(allocation, ctx).is_none() { + continue; + } + let Some(count) = allocation.deref(ctx).get_operand(0).defining_op() else { + continue; + }; + let Some(constant) = Operation::get_op::(count, ctx) else { + continue; + }; + let attribute = constant.get_value(ctx); + let Some(count) = + (&*attribute as &dyn pliron::attribute::Attribute).downcast_ref::() + else { + continue; + }; + candidates.push((block, allocation, count.clone())); + } + } + if candidates.is_empty() { + return; + } + let cyclic_blocks: HashSet<_> = sccs_in_topological_order(ctx, ®ion) + .into_iter() + .filter(|component| component.is_cyclic) + .flat_map(|component| component.nodes) + .collect(); + let eligible: Vec<_> = candidates + .into_iter() + .filter(|(block, op, _)| { + !cyclic_blocks.contains(block) + || pointer_is_local_scratch(ctx, op.deref(ctx).get_result(0)) + }) + .collect(); + let anchor = entry.deref(ctx).iter(ctx).next(); + for (_, allocation, count) in eligible { + // The count's old definition may be in a conditional block. Give the + // moved allocation its own entry constant rather than violating SSA + // dominance or moving a definition that other body operations still use. + let constant = ops::ConstantOp::new(ctx, Box::new(count)); + if let Some(anchor) = anchor { + constant.get_operation().insert_before(ctx, anchor); + } else { + constant.get_operation().insert_at_back(entry, ctx); + } + let count = constant.get_operation().deref(ctx).get_result(0); + Operation::replace_operand(allocation, ctx, 0, count); + // Move the original operation so alignment, source location, and debug + // provenance survive. Stores and lifetime operations are not moved. + allocation.unlink(ctx); + if let Some(anchor) = anchor { + allocation.insert_before(ctx, anchor); + } else { + allocation.insert_at_back(entry, ctx); + } + } +} + +/// Follow pointer views, accepting only loads and stores through the address. +/// Calls, block arguments, comparisons/conversions, returns, and storing the +/// pointer as data can reveal allocation identity across iterations and stop us. +fn pointer_is_local_scratch(ctx: &Context, pointer: Value) -> bool { + let mut pending = vec![pointer]; + let mut visited = HashSet::new(); + while let Some(pointer) = pending.pop() { + if !visited.insert(pointer) { + continue; + } + for usage in pointer.uses(ctx) { + let operation = usage.user_op(); + let index = usage.find_index(ctx); + let object = Operation::get_op_dyn(operation, ctx); + let op = object.as_ref(); + if (op.is::() && index == 0) || (op.is::() && index == 1) { + continue; + } + if index == 0 + && (op.is::() + || op.is::() + || op.is::()) + { + pending.push(operation.deref(ctx).get_result(0)); + continue; + } + return false; + } + } + true +} diff --git a/cuda-oxide/crates/nvvm-transforms/src/lib.rs b/cuda-oxide/crates/nvvm-transforms/src/lib.rs index ba0e8ce56b..cb0b3b9f33 100644 --- a/cuda-oxide/crates/nvvm-transforms/src/lib.rs +++ b/cuda-oxide/crates/nvvm-transforms/src/lib.rs @@ -12,6 +12,7 @@ #![warn(missing_docs)] +mod entry_allocas; mod helpers; mod legalize; @@ -30,6 +31,8 @@ use pliron::{ /// compatibility pass. Modern modules retain modern operations but still /// rewrite integer widths that libNVVM does not accept in bit-manipulation /// intrinsics. +/// Finally, eligible fixed scratch allocations move into each function's entry +/// block so NVVM can use static storage without dynamic-stack restoration. /// /// `capability` is the numeric compute capability of the build target /// (`52`, `86`, `90`, ...). The legacy pass uses it to reject rewrites whose @@ -45,5 +48,7 @@ pub fn legalize_for_nvvm( match dialect { NvvmIrDialect::LegacyLlvm7 => legalize::legalize_for_legacy_nvvm(ctx, module, capability), NvvmIrDialect::Modern => legalize::legalize_nvvm_bit_intrinsics(ctx, module), - } + }?; + entry_allocas::hoist_fixed_allocas(ctx, module); + Ok(()) } diff --git a/cuda-oxide/crates/nvvm-transforms/tests/entry_allocas.rs b/cuda-oxide/crates/nvvm-transforms/tests/entry_allocas.rs new file mode 100644 index 0000000000..dc10964a25 --- /dev/null +++ b/cuda-oxide/crates/nvvm-transforms/tests/entry_allocas.rs @@ -0,0 +1,554 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + attributes::ICmpPredicateAttr, + export::{NvvmExportConfig, NvvmIrDialect, export_module_to_string_with_config}, + op_interfaces::CastOpInterface, + ops::{ + AddrSpaceCastOp, AllocaOp, BitcastOp, BrOp, CallOp, CondBrOp, ConstantOp, FuncOp, GepIndex, + GetElementPtrOp, ICmpOp, LoadOp, PtrToIntOp, ReturnOp, SelectOp, StoreOp, + }, + types::{FuncType, PointerType, VoidType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::IntegerAttr, + op_interfaces::CallOpCallable, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + common_traits::Verify, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, + printable::Printable, + r#type::Typed, + utils::apint::APInt, + value::Value, +}; +use std::num::NonZero; + +/// Find the module-level block that owns LLVM function definitions. +/// Fixtures insert their declarations and definitions in this shared scope. +fn module_top_block(ctx: &Context, module: &ModuleOp) -> Ptr { + module + .get_operation() + .deref(ctx) + .get_region(0) + .deref(ctx) + .iter(ctx) + .next() + .unwrap() +} + +/// Uses that can reveal a scratch allocation's identity across iterations. +/// The same uses are safe to hoist when the allocation executes at most once. +#[derive(Clone, Copy, Debug)] +enum AddressUse { + /// Direct memory accesses through a chain of pointer views are local. + /// This is the positive control beside every rejection case. + Local, + /// Passing the address to a callee may retain it or compare it. + /// Reject calls even when the declaration has no visible implementation. + Call, + /// A backedge can retain the previous iteration's pointer as a block argument. + /// This is the LLVM-dialect representation of a pointer PHI. + Phi, + /// A pointer equality test can observe repeated allocation identity. + /// It must not become a comparison against one reused slot. + Compare, + /// Converting an address to an integer exposes its numeric identity. + /// The numeric result may outlive the pointer-typed use chain. + Convert, + /// Returning an address exposes it outside this function's use graph. + /// Conservative classification does not assume anything about the caller. + Return, + /// Storing the pointer as data can keep it alive across loop iterations. + /// A store through the pointer is accepted, but storing the pointer is not. + Store, + /// A select can forward an address to later operations or iterations. + /// Unsupported pointer-producing operations must fail closed. + Select, +} + +/// Build a scratch slot followed by each accepted pointer-view operation. +/// Its last view is used both for ordinary loads/stores and the chosen escape. +fn address_use_fixture( + ctx: &mut Context, + usage: AddressUse, + cyclic: bool, + multi_block: bool, +) -> (ModuleOp, Ptr, Ptr, Ptr) { + let module = ModuleOp::new(ctx, "address_use".try_into().unwrap()); + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); + let pointer_ty = PointerType::get(ctx, 0); + let global_ty = PointerType::get(ctx, 1); + let ty = FuncType::get( + ctx, + pointer_ty.into(), + vec![i1_ty.into(), pointer_ty.into()], + false, + ); + let func = FuncOp::new(ctx, "scratch_views".try_into().unwrap(), ty); + let entry = func.get_or_create_entry_block(ctx); + func.get_operation() + .insert_at_back(module_top_block(ctx, &module), ctx); + let repeat = entry.deref(ctx).get_argument(0); + let destination = entry.deref(ctx).get_argument(1); + let region = func.get_operation().deref(ctx).get_region(0); + let body = BasicBlock::new(ctx, None, vec![pointer_ty.into()]); + let exit = BasicBlock::new(ctx, None, vec![]); + body.insert_at_back(region, ctx); + exit.insert_at_back(region, ctx); + BrOp::new(ctx, body, vec![destination]) + .get_operation() + .insert_at_back(entry, ctx); + let count = one(ctx, body); + let slot = AllocaOp::new(ctx, i32_ty.into(), count, 0); + slot.get_operation().insert_at_back(body, ctx); + let pointer = slot.get_operation().deref(ctx).get_result(0); + let gep = GetElementPtrOp::new(ctx, pointer, vec![GepIndex::Constant(0)], i32_ty.into()); + gep.get_operation().insert_at_back(body, ctx); + let view = gep.get_operation().deref(ctx).get_result(0); + let cast = BitcastOp::new(ctx, view, pointer_ty.into()); + cast.get_operation().insert_at_back(body, ctx); + let view = cast.get_operation().deref(ctx).get_result(0); + let global = AddrSpaceCastOp::new(ctx, view, global_ty.into()); + global.get_operation().insert_at_back(body, ctx); + let view = global.get_operation().deref(ctx).get_result(0); + let generic = AddrSpaceCastOp::new(ctx, view, pointer_ty.into()); + generic.get_operation().insert_at_back(body, ctx); + let view = generic.get_operation().deref(ctx).get_result(0); + StoreOp::new(ctx, count, view) + .get_operation() + .insert_at_back(body, ctx); + LoadOp::new(ctx, view, i32_ty.into()) + .get_operation() + .insert_at_back(body, ctx); + match usage { + AddressUse::Call => { + let callee_ty = FuncType::get( + ctx, + VoidType::get(ctx).into(), + vec![pointer_ty.into()], + false, + ); + let name: pliron::identifier::Identifier = "observe".try_into().unwrap(); + FuncOp::new(ctx, name.clone(), callee_ty) + .get_operation() + .insert_at_back(module_top_block(ctx, &module), ctx); + CallOp::new(ctx, CallOpCallable::Direct(name), callee_ty, vec![view]) + .get_operation() + .insert_at_back(body, ctx); + } + AddressUse::Compare => { + ICmpOp::new(ctx, ICmpPredicateAttr::EQ, view, destination) + .get_operation() + .insert_at_back(body, ctx); + } + AddressUse::Convert => { + PtrToIntOp::new(ctx, view, i64_ty.into()) + .get_operation() + .insert_at_back(body, ctx); + } + AddressUse::Store => { + StoreOp::new(ctx, view, destination) + .get_operation() + .insert_at_back(body, ctx); + } + AddressUse::Select => { + SelectOp::new(ctx, repeat, view, destination) + .get_operation() + .insert_at_back(body, ctx); + } + AddressUse::Local | AddressUse::Phi | AddressUse::Return => {} + } + let forwarded = if matches!(usage, AddressUse::Phi) { + view + } else { + destination + }; + if cyclic { + let latch = if multi_block { + let latch = BasicBlock::new(ctx, None, vec![]); + latch.insert_at_back(region, ctx); + BrOp::new(ctx, latch, vec![]) + .get_operation() + .insert_at_back(body, ctx); + latch + } else { + body + }; + CondBrOp::new(ctx, repeat, body, vec![forwarded], exit, vec![]) + .get_operation() + .insert_at_back(latch, ctx); + } else { + BrOp::new(ctx, exit, vec![]) + .get_operation() + .insert_at_back(body, ctx); + } + let returned = if matches!(usage, AddressUse::Return) { + view + } else { + destination + }; + ReturnOp::new(ctx, Some(returned)) + .get_operation() + .insert_at_back(exit, ctx); + (module, slot.get_operation(), entry, body) +} + +#[test] +fn loop_address_uses_fail_closed_through_pointer_views_in_both_dialects() { + for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { + for multi_block in [false, true] { + for usage in [ + AddressUse::Local, + AddressUse::Call, + AddressUse::Phi, + AddressUse::Compare, + AddressUse::Convert, + AddressUse::Return, + AddressUse::Store, + AddressUse::Select, + ] { + let mut ctx = Context::new(); + let (module, slot, entry, body) = + address_use_fixture(&mut ctx, usage, true, multi_block); + normalize(&mut ctx, &module, dialect); + let expected = if matches!(usage, AddressUse::Local) { + entry + } else { + body + }; + assert_eq!( + slot.deref(&ctx).get_parent_block(), + Some(expected), + "{dialect:?}: {usage:?}, multi_block={multi_block}" + ); + } + } + } +} + +#[test] +fn once_only_slots_can_hoist_even_when_their_addresses_are_observable() { + for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { + for usage in [ + AddressUse::Call, + AddressUse::Compare, + AddressUse::Convert, + AddressUse::Return, + AddressUse::Store, + AddressUse::Select, + ] { + let mut ctx = Context::new(); + let (module, slot, entry, _) = address_use_fixture(&mut ctx, usage, false, false); + normalize(&mut ctx, &module, dialect); + assert_eq!( + slot.deref(&ctx).get_parent_block(), + Some(entry), + "{dialect:?}: {usage:?}" + ); + } + } +} + +/// Exercise the public normalization entry point and verify the mutated IR. +/// Running twice also checks that already-normalized slots and counts stay put. +fn normalize(ctx: &mut Context, module: &ModuleOp, dialect: NvvmIrDialect) { + module.get_operation().deref(ctx).verify(ctx).unwrap(); + nvvm_transforms::legalize_for_nvvm(ctx, module.get_operation(), dialect, 90).unwrap(); + module.get_operation().deref(ctx).verify(ctx).unwrap(); + let once = module.get_operation().deref(ctx).disp(ctx).to_string(); + nvvm_transforms::legalize_for_nvvm(ctx, module.get_operation(), dialect, 90).unwrap(); + module.get_operation().deref(ctx).verify(ctx).unwrap(); + assert_eq!( + once, + module.get_operation().deref(ctx).disp(ctx).to_string(), + "normalization must be idempotent" + ); +} + +/// Build a fixed count at the allocation site, outside the entry block. +/// Normalization must provide a dominating count without moving this definition. +fn one(ctx: &mut Context, block: Ptr) -> Value { + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let op = ConstantOp::new(ctx, Box::new(attr)); + op.get_operation().insert_at_back(block, ctx); + op.get_operation().deref(ctx).get_result(0) +} + +#[test] +fn fixed_slot_precedes_body_phis_but_initialization_stays_in_body() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "entry_alloca".try_into().unwrap()); + let top = module_top_block(&ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let ty = FuncType::get(&ctx, i32_ty.into(), vec![i32_ty.into()], false); + let func = FuncOp::new(&mut ctx, "scratch".try_into().unwrap(), ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let input = entry.deref(&ctx).get_argument(0); + let region = func.get_operation().deref(&ctx).get_region(0); + let body = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + body.insert_at_back(region, &ctx); + BrOp::new(&mut ctx, body, vec![input]) + .get_operation() + .insert_at_back(entry, &ctx); + let value = body.deref(&ctx).get_argument(0); + let count = one(&mut ctx, body); + let slot = AllocaOp::new(&mut ctx, i32_ty.into(), count, 0); + llvm_export::ops::set_op_alignment(&mut ctx, slot.get_operation(), 4); + slot.get_operation().insert_at_back(body, &ctx); + let pointer = slot.get_operation().deref(&ctx).get_result(0); + StoreOp::new(&mut ctx, value, pointer) + .get_operation() + .insert_at_back(body, &ctx); + let load = LoadOp::new(&mut ctx, pointer, i32_ty.into()); + let result = load.get_operation().deref(&ctx).get_result(0); + load.get_operation().insert_at_back(body, &ctx); + ReturnOp::new(&mut ctx, Some(result)) + .get_operation() + .insert_at_back(body, &ctx); + func.get_operation().insert_at_back(top, &ctx); + module.get_operation().deref(&ctx).verify(&ctx).unwrap(); + + // Raw export is a printer; even NVVM export must not silently move storage. + let before = module.get_operation().deref(&ctx).disp(&ctx).to_string(); + for config in [ + NvvmExportConfig::new(NvvmIrDialect::Modern), + NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ] { + let ir = export_module_to_string_with_config(&ctx, &module, &config).unwrap(); + let (entry_text, body_text) = ir.split_once("bb0:\n").unwrap(); + assert!(!entry_text.contains(" = alloca "), "{ir}"); + assert!(body_text.contains(" = alloca "), "{ir}"); + } + let ptx_ir = llvm_export::export::export_module_to_string(&ctx, &module).unwrap(); + assert!( + !ptx_ir + .split_once("bb0:\n") + .unwrap() + .0 + .contains(" = alloca ") + ); + assert_eq!( + before, + module.get_operation().deref(&ctx).disp(&ctx).to_string() + ); + for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { + normalize(&mut ctx, &module, dialect); + assert_eq!( + slot.get_operation().deref(&ctx).get_parent_block(), + Some(entry) + ); + let entry_ops: Vec<_> = entry.deref(&ctx).iter(&ctx).collect(); + let count = slot + .get_operation() + .deref(&ctx) + .get_operand(0) + .defining_op() + .unwrap(); + assert!( + entry_ops.iter().position(|op| *op == count).unwrap() + < entry_ops + .iter() + .position(|op| *op == slot.get_operation()) + .unwrap() + ); + assert_eq!(count.deref(&ctx).get_parent_block(), Some(entry)); + assert_eq!(value.get_type(&ctx), i32_ty.into()); + let ir = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .unwrap(); + let (entry, body) = ir.split_once("bb0:\n").unwrap(); + assert_eq!(ir.matches(" = alloca ").count(), 1, "{ir}"); + assert!(entry.contains(" = alloca i32, align 4"), "{ir}"); + assert!( + body.trim_start().starts_with('%') && body.contains(" = phi i32"), + "{ir}" + ); + assert!( + body.contains("store i32") && body.contains("load i32"), + "{ir}" + ); + assert!(!entry.contains("store i32"), "{ir}"); + } +} + +#[test] +fn dynamic_alloca_count_stays_in_its_original_block() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "dynamic_alloca".try_into().unwrap()); + let top = module_top_block(&ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let pointer_ty = PointerType::get(&ctx, 0); + let ty = FuncType::get(&ctx, pointer_ty.into(), vec![i32_ty.into()], false); + let func = FuncOp::new(&mut ctx, "dynamic".try_into().unwrap(), ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let count = entry.deref(&ctx).get_argument(0); + let region = func.get_operation().deref(&ctx).get_region(0); + let body = BasicBlock::new(&mut ctx, None, vec![]); + body.insert_at_back(region, &ctx); + BrOp::new(&mut ctx, body, vec![]) + .get_operation() + .insert_at_back(entry, &ctx); + let slot = AllocaOp::new(&mut ctx, i32_ty.into(), count, 0); + let pointer = slot.get_operation().deref(&ctx).get_result(0); + slot.get_operation().insert_at_back(body, &ctx); + ReturnOp::new(&mut ctx, Some(pointer)) + .get_operation() + .insert_at_back(body, &ctx); + func.get_operation().insert_at_back(top, &ctx); + + normalize(&mut ctx, &module, NvvmIrDialect::Modern); + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .unwrap(); + let (entry, body) = ir.split_once("bb0:\n").unwrap(); + assert!(!entry.contains(" = alloca "), "{ir}"); + assert!(body.contains(" = alloca i32, i32 %v0"), "{ir}"); +} + +#[test] +fn hoisting_preserves_nonunit_count_width_alignment_and_body_uses() { + for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "count_dominance".try_into().unwrap()); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let ty = FuncType::get(&ctx, i64_ty.into(), vec![], false); + let function = FuncOp::new(&mut ctx, "counted".try_into().unwrap(), ty); + let entry = function.get_or_create_entry_block(&mut ctx); + function + .get_operation() + .insert_at_back(module_top_block(&ctx, &module), &ctx); + let region = function.get_operation().deref(&ctx).get_region(0); + let body = BasicBlock::new(&mut ctx, None, vec![]); + body.insert_at_back(region, &ctx); + BrOp::new(&mut ctx, body, vec![]) + .get_operation() + .insert_at_back(entry, &ctx); + let attribute = IntegerAttr::new(i64_ty, APInt::from_u64(3, NonZero::new(64).unwrap())); + let original = ConstantOp::new(&mut ctx, Box::new(attribute)); + original.get_operation().insert_at_back(body, &ctx); + let count = original.get_operation().deref(&ctx).get_result(0); + let slot = AllocaOp::new(&mut ctx, i32_ty.into(), count, 0); + llvm_export::ops::set_op_alignment(&mut ctx, slot.get_operation(), 32); + slot.get_operation().insert_at_back(body, &ctx); + let returned = ReturnOp::new(&mut ctx, Some(count)); + returned.get_operation().insert_at_back(body, &ctx); + + normalize(&mut ctx, &module, dialect); + assert_eq!( + slot.get_operation().deref(&ctx).get_parent_block(), + Some(entry) + ); + let new_count = slot.get_operation().deref(&ctx).get_operand(0); + assert_eq!(new_count.get_type(&ctx), i64_ty.into()); + assert_eq!( + new_count + .defining_op() + .unwrap() + .deref(&ctx) + .get_parent_block(), + Some(entry) + ); + assert_eq!( + original.get_operation().deref(&ctx).get_parent_block(), + Some(body) + ); + assert_eq!(returned.get_operation().deref(&ctx).get_operand(0), count); + let ir = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .unwrap(); + assert!(ir.contains("alloca i32, i64 3, align 32"), "{ir}"); + } +} + +#[test] +fn loop_scratch_hoists_only_when_its_address_does_not_escape() { + for escapes in [false, true] { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "loop_alloca".try_into().unwrap()); + let top = module_top_block(&ctx, &module); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let pointer_ty = PointerType::get(&ctx, 0); + let ty = FuncType::get( + &ctx, + VoidType::get(&ctx).into(), + vec![i1_ty.into(), pointer_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "loop_scratch".try_into().unwrap(), ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let repeat = entry.deref(&ctx).get_argument(0); + let escape_destination = entry.deref(&ctx).get_argument(1); + let region = func.get_operation().deref(&ctx).get_region(0); + let body = BasicBlock::new(&mut ctx, None, vec![]); + let exit = BasicBlock::new(&mut ctx, None, vec![]); + body.insert_at_back(region, &ctx); + exit.insert_at_back(region, &ctx); + BrOp::new(&mut ctx, body, vec![]) + .get_operation() + .insert_at_back(entry, &ctx); + let count = one(&mut ctx, body); + let slot = AllocaOp::new(&mut ctx, i32_ty.into(), count, 0); + let pointer = slot.get_operation().deref(&ctx).get_result(0); + slot.get_operation().insert_at_back(body, &ctx); + let gep = GetElementPtrOp::new( + &mut ctx, + pointer, + vec![GepIndex::Constant(0)], + i32_ty.into(), + ); + let view = gep.get_operation().deref(&ctx).get_result(0); + gep.get_operation().insert_at_back(body, &ctx); + StoreOp::new(&mut ctx, count, view) + .get_operation() + .insert_at_back(body, &ctx); + LoadOp::new(&mut ctx, view, i32_ty.into()) + .get_operation() + .insert_at_back(body, &ctx); + if escapes { + StoreOp::new(&mut ctx, view, escape_destination) + .get_operation() + .insert_at_back(body, &ctx); + } + CondBrOp::new(&mut ctx, repeat, body, vec![], exit, vec![]) + .get_operation() + .insert_at_back(body, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(exit, &ctx); + func.get_operation().insert_at_back(top, &ctx); + module.get_operation().deref(&ctx).verify(&ctx).unwrap(); + + normalize(&mut ctx, &module, NvvmIrDialect::Modern); + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .unwrap(); + let (entry, body) = ir.split_once("bb0:\n").unwrap(); + assert_eq!(entry.contains(" = alloca "), !escapes, "{ir}"); + assert_eq!(body.contains(" = alloca "), escapes, "{ir}"); + assert!( + body.contains("getelementptr") && body.contains("store i32"), + "{ir}" + ); + } +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/Cargo.lock new file mode 100644 index 0000000000..73ea93f3a2 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/Cargo.lock @@ -0,0 +1,767 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[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 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[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.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +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.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[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-device", + "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 2.0.119", +] + +[[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.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[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.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +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.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[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.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[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.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[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 = "nvjitlink-sys" +version = "0.2.1" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "nvvm_entry_allocas" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", + "libnvvm-sys", + "nvjitlink-sys", +] + +[[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 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +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.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[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.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[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.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[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.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[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/nvvm_entry_allocas/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/Cargo.toml new file mode 100644 index 0000000000..11a7513d39 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/Cargo.toml @@ -0,0 +1,21 @@ +cargo-features = ["profile-rustflags"] + +[package] +name = "nvvm_entry_allocas" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } +cuda-core = { path = "../../../../../cuda-core" } +libnvvm-sys = { path = "../../../libnvvm-sys" } +nvjitlink-sys = { path = "../../../nvjitlink-sys" } + +# Keep helper calls visible until LLVM inlining so scratch placement remains +# observable. Scope the flag to this example, as in disjoint_slice_len. +[profile.release.package.nvvm_entry_allocas] +rustflags = ["-Zinline-mir=no"] diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/src/main.rs new file mode 100644 index 0000000000..84f946a62f --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/src/main.rs @@ -0,0 +1,244 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression for entry placement of fixed scratch in NVVM IR. +//! +//! Input enum -> payload scratch -> scalar read -> aggregate call -> output. +//! NVVM normalization should place fixed allocas in each function's true entry +//! block. The companion `verify-code-shape.sh` checks this IR invariant and +//! checks that extended split compilation emits no dynamic-stack instructions +//! while retaining the aggregate call and its 128-bit argument stores. +//! +//! The host verifies every GPU result independently. On CUDA 13.4, numerical +//! results can remain correct when the placement check fails; this example has +//! not demonstrated a misaligned access. +//! +//! Run from `cuda-oxide/`: +//! ```text +//! cargo oxide run nvvm_entry_allocas --emit-nvvm-ir --arch sm_120 +//! bash crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/verify-code-shape.sh +//! ``` + +use cuda_core::simt::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; +use cuda_device::{DisjointSlice, kernel, thread}; +use cuda_host::cuda_module; +use libnvvm_sys::{LibNvvm, Program}; +use nvjitlink_sys::{InputType, LibNvJitLink, Linker}; + +/// Compile this example's emitted NVVM IR with extended split compilation. +/// It also writes PTX for the code-shape check beside the emitted IR. +/// Both links use the same LTOIR and extended split compilation options. +fn compile_split_cubin() -> Vec { + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + let arch = std::fs::read_to_string(root.join("nvvm_entry_allocas.target")) + .expect("read target sidecar; run cargo oxide with --emit-nvvm-ir --arch sm_120"); + let arch = arch + .lines() + .next() + .expect("target sidecar has an architecture") + .trim(); + let compute = arch.replacen("sm_", "compute_", 1); + let ir = std::fs::read(root.join("nvvm_entry_allocas.ll")).expect("read emitted NVVM IR"); + let nvvm = LibNvvm::load().expect("load libNVVM"); + let mut program = Program::new(&nvvm).expect("create NVVM program"); + program + .add_module(&ir, "nvvm_entry_allocas") + .expect("add NVVM IR"); + let ltoir = program + .compile(&[&format!("-arch={compute}"), "-gen-lto"]) + .expect("compile NVVM IR to LTOIR"); + std::fs::write(root.join("nvvm_entry_allocas.ltoir"), <oir).expect("write LTOIR"); + let library = LibNvJitLink::load().expect("load nvJitLink"); + let mut linker = Linker::new( + &library, + &[ + &format!("-arch={arch}"), + "-lto", + "-split-compile=0", + "-split-compile-extended=0", + "-fma=0", + "-no-cache", + ], + ) + .expect("create extended split linker"); + linker + .add(InputType::Ltoir, <oir, "nvvm_entry_allocas.ltoir") + .expect("add LTOIR"); + let cubin = linker.finish().expect("link extended split cubin"); + std::fs::write(root.join("nvvm_entry_allocas.cubin"), &cubin).expect("write cubin"); + { + let mut linker = Linker::new( + &library, + &[ + &format!("-arch={arch}"), + "-lto", + "-ptx", + "-split-compile=0", + "-split-compile-extended=0", + "-fma=0", + "-no-cache", + ], + ) + .unwrap(); + linker + .add(InputType::Ltoir, <oir, "nvvm_entry_allocas.ltoir") + .unwrap(); + std::fs::write( + root.join("nvvm_entry_allocas.ptx"), + linker.finish_ptx().unwrap(), + ) + .unwrap(); + } + cubin +} + +/// Four lanes passed together as an outgoing device argument. +/// This keeps the aggregate call and its 128-bit parameter stores observable. +#[derive(Clone, Copy)] +#[repr(C)] +pub struct Lanes { + /// Four lane values form the aggregate argument consumed after the scratch read. + /// Distinct weights in `consume` make every lane contribute to the oracle. + pub values: [f32; 4], +} + +/// Keep an aggregate call boundary after the scratch slot's last read. +/// Each lane has a distinct weight so corruption cannot cancel in the sum. +#[inline(never)] +fn consume(prefix: u64, value: Lanes) -> f32 { + prefix as f32 + + value.values[0] + + 2.0 * value.values[1] + + 4.0 * value.values[2] + + 8.0 * value.values[3] +} + +/// A tagged aggregate forces payload extraction after a control-flow branch. +/// Loading it by value lets enum lowering introduce a temporary payload slot. +#[derive(Clone, Copy)] +#[repr(C)] +enum Packet { + /// Three lanes occupy 12 bytes and need four-byte alignment. + /// Dynamic indexing keeps the payload read addressable. + Data([f32; 3]), + /// A scalar alternative keeps the match data dependent. + /// It follows a different payload extraction path. + Scalar(f32), +} + +// SAFETY: both variants contain only plain f32 values, with no pointers, +// ownership, or drop glue. Zero bytes are Data([0.0; 3]); the host initializes +// every field before copying, and this program never copies enum bytes back. +unsafe impl cuda_core::DeviceCopy for Packet {} + +/// Read private scratch through a separate device call at a runtime index. +/// The borrowed array keeps its storage addressable while this call executes. +#[inline(never)] +fn read_scratch(values: &[f32; 3], lane: usize) -> f32 { + values[lane] +} + +/// Extract a tagged payload into scratch in a conditional block. +/// Keeping this helper until LLVM inlining exercises NVVM entry normalization. +#[inline(always)] +fn read_packet(packet: Packet, lane: usize) -> f32 { + match packet { + Packet::Data(values) => read_scratch(&values, (lane >> 1) & 1), + Packet::Scalar(value) => value, + } +} + +/// One independent output per thread, including both sides of a branch. +/// Every scratch slot is private; the kernel has no shared-memory operations. +#[cuda_module] +mod kernels { + use super::*; + + /// Invoke the inlined scratch helper after a data-dependent branch. + /// Each active thread owns exactly one output slot through DisjointSlice. + #[kernel] + pub fn branch_scratch(input: &[Packet], mut output: DisjointSlice) { + let index = thread::index_1d(); + let i = index.get(); + if let Some(slot) = output.get_mut(index) { + let packet = input[i]; + let value = i as f32; + let mut sum = 0.0; + let mut round = 0; + while round < (i & 3) + 1 { + let value = value + round as f32; + let selected = read_packet(packet, i + round); + sum += consume( + i as u64, + Lanes { + values: [selected, value + 3.0, value + 4.0, value + 5.0], + }, + ); + round += 1; + } + *slot = sum; + } + } +} + +/// Launch enough threads to cover both variants and both selected scratch indices. +/// Compare every result with scalar host arithmetic before printing SUCCESS. +fn main() { + let context = CudaContext::new(0).expect("create CUDA context"); + let stream = context.default_stream(); + let input: Vec = (0..1024) + .map(|i| { + let value = i as f32; + if i & 1 == 0 { + Packet::Data([value, value + 1.0, value + 2.0]) + } else { + Packet::Scalar(value + 6.0) + } + }) + .collect(); + let device_input = DeviceBuffer::from_host(&stream, &input).expect("copy input"); + let mut output = DeviceBuffer::::zeroed(&stream, input.len()).expect("allocate output"); + let cubin = compile_split_cubin(); + let module = context + .load_module_from_image(&cubin) + .expect("load extended split cubin"); + let module = kernels::from_module(module).expect("load typed nvvm_entry_allocas module"); + // SAFETY: both buffers cover the launch, and each active thread writes + // exactly its own output element and reads the corresponding input. + unsafe { + module.branch_scratch( + &stream, + LaunchConfig::for_num_elems(input.len() as u32), + &device_input, + &mut output, + ) + } + .expect("launch branch_scratch"); + let actual = output + .to_host_vec(&stream) + .expect("copy output and synchronize"); + for (i, (&packet, &got)) in input.iter().zip(&actual).enumerate() { + let expected: f32 = (0..(i & 3) + 1) + .map(|round| { + let value = i as f32 + round as f32; + let lane = i + round; + i as f32 + + (match packet { + Packet::Data(values) => values[(lane >> 1) & 1], + Packet::Scalar(value) => value, + }) + + 2.0 * (value + 3.0) + + 4.0 * (value + 4.0) + + 8.0 * (value + 5.0) + }) + .sum(); + assert_eq!(got, expected, "branch_scratch element {i}"); + } + println!( + "SUCCESS: nvvm_entry_allocas verified {} elements", + actual.len() + ); +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/verify-code-shape.sh new file mode 100755 index 0000000000..7d3e4886c5 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvvm_entry_allocas/verify-code-shape.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Check the entry-normalization contract separately from the numerical oracle. +# Correct GPU results alone do not establish static placement of scratch storage. +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ir="${root}/nvvm_entry_allocas.ll" +test -s "${ir}" + +# This example has only fixed slots. Entry placement is also checkable in a +# headless compile-only lane, before its host runner can produce linked PTX. +python3 - "${ir}" <<'CHECK_IR' +import re +import sys +from pathlib import Path + +function = None +block = None +slots = 0 +for line in Path(sys.argv[1]).read_text().splitlines(): + if line.startswith("define "): + function = line + block = None + elif line == "}": + function = None + elif function and re.fullmatch(r"[A-Za-z0-9_.$]+:", line): + block = line[:-1] + elif function and " = alloca " in line: + slots += 1 + if block != "entry": + sys.exit(f"error: fixed scratch outside true entry block ({block}): {line.strip()}") +if slots == 0: + sys.exit("error: example no longer exercises fixed scratch allocations") +CHECK_IR + +# The runtime runner writes the split compiler's PTX. It includes a trailing +# NUL, hence grep's -a flag. Compile-only lanes have just the LLVM IR above. +ptx="${root}/nvvm_entry_allocas.ptx" +if [[ -s "${ptx}" ]]; then + if grep -aEq '(^|[[:space:]])(alloca|stacksave|stackrestore)\.' "${ptx}"; then + echo "error: fixed scratch introduced dynamic stack instructions in ${ptx}" >&2 + exit 1 + fi + grep -aEq 'st\.param\.v4\.b32' "${ptx}" + grep -aEq 'call[^;]*nvvm_entry_allocas__consume' "${ptx}" +fi + +echo "nvvm_entry_allocas code shape: PASS" diff --git a/cuda-oxide/scripts/smoketest.sh b/cuda-oxide/scripts/smoketest.sh index cb63423e35..a4ea1a92cc 100755 --- a/cuda-oxide/scripts/smoketest.sh +++ b/cuda-oxide/scripts/smoketest.sh @@ -91,7 +91,7 @@ TCGEN05_EXAMPLES=(gemm_sol gemm_sol_final tcgen05 tcgen05_matmul) WGMMA_EXAMPLES=(wgmma) BLACKWELL_MMA_EXAMPLES=(mma_mxf8f6f4) LTOIR_EXAMPLES=(addressof_sharedarray constant_memory cpp_consumes_rust_device device_ffi_test legacy_atomic_fadd legacy_atomic_rmw_cas legacy_nvvm_pointer_shapes manual_launch_libdevice mathdx_ffi_test primitive_stress) -LTOIR_MODERN_EXAMPLES=(small_type_ffi_test) +LTOIR_MODERN_EXAMPLES=(small_type_ffi_test nvvm_entry_allocas) TILELIB_EXAMPLES=(tilelib_relu) AUTO_NVVM_EXAMPLES=(libdevice_math) IKET_EXAMPLES=(iket_trace)