diff --git a/crates/mir-importer/src/translator/types.rs b/crates/mir-importer/src/translator/types.rs index 4dbbdfcf21..6698db6881 100644 --- a/crates/mir-importer/src/translator/types.rs +++ b/crates/mir-importer/src/translator/types.rs @@ -67,6 +67,23 @@ pub fn get_usize_type( pliron::builtin::types::IntegerType::get(ctx, 64, pliron::builtin::types::Signedness::Unsigned) } +/// Returns the tupled-upvars types for a `RigidTy::Closure`. +/// +/// rustc suffix-encodes closure substitutions as +/// `[parent_args..., closure_kind, closure_sig, tupled_upvars]`, so the upvars +/// tuple is the last generic arg, not a fixed index. +fn closure_upvar_tys(substs: &rustc_public::ty::GenericArgs) -> Option> { + let rustc_public::ty::GenericArgKind::Type(upvar_tuple_ty) = substs.0.last()? else { + return None; + }; + let rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Tuple(upvar_tys)) = + upvar_tuple_ty.kind() + else { + return None; + }; + Some(upvar_tys) +} + /// Returns the 32-bit floating point type. pub fn get_f32_type( ctx: &mut Context, @@ -128,15 +145,9 @@ pub fn is_rust_type_zst(rust_ty: &rustc_public::ty::Ty) -> bool { false } } - // Closures with no captures are ZST, closures with captures are not + // Closures with no captures are ZST, closures with captures are not. rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Closure(_, substs)) => { - // Check substs[2] which is the tuple of upvar types - if substs.0.len() >= 3 - && let rustc_public::ty::GenericArgKind::Type(upvar_tuple_ty) = &substs.0[2] - && let rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Tuple( - upvar_tys, - )) = upvar_tuple_ty.kind() - { + if let Some(upvar_tys) = closure_upvar_tys(&substs) { // ZST if no captures return upvar_tys.is_empty(); } @@ -858,25 +869,21 @@ pub fn translate_type( // Handle Closure types // Closures are represented as structs with fields for each captured variable (upvar). // The substs for a closure contain: - // [0] Internal marker type (usually i8) - // [1] Function signature - // [2] Tuple of upvar types (the captured variables) + // [parent_args..] Captured parent generics, when the closure appears in a generic item + // [n - 3] Closure kind + // [n - 2] Function signature + // [n - 1] Tuple of upvar types (the captured variables) rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Closure( closure_def, substs, )) => { let closure_name = format!("{:?}", closure_def.def_id()); - // Extract upvar types from substs[2] (the tuple of captured types) + // Extract upvar types from the tupled-upvars generic arg. let mut field_names = Vec::new(); let mut field_types = Vec::new(); - if substs.0.len() >= 3 - && let rustc_public::ty::GenericArgKind::Type(upvar_tuple_ty) = &substs.0[2] - && let rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Tuple( - upvar_tys, - )) = upvar_tuple_ty.kind() - { + if let Some(upvar_tys) = closure_upvar_tys(&substs) { for (i, upvar_ty) in upvar_tys.iter().enumerate() { field_names.push(format!("capture_{}", i)); field_types.push(translate_type(ctx, upvar_ty)?); diff --git a/crates/rustc-codegen-cuda/examples/generic/README.md b/crates/rustc-codegen-cuda/examples/generic/README.md index 3427e46bc3..5d84179449 100644 --- a/crates/rustc-codegen-cuda/examples/generic/README.md +++ b/crates/rustc-codegen-cuda/examples/generic/README.md @@ -8,6 +8,7 @@ Tests whether the collector correctly handles monomorphized generic kernels. Thi - Defines generic kernels `scale` and `add` with trait bounds - Host code calls `scale::` which triggers monomorphization +- Defines a captured closure inside a generic kernel body - Verifies the monomorphized kernel executes correctly ## Key Concepts Demonstrated diff --git a/crates/rustc-codegen-cuda/examples/generic/src/main.rs b/crates/rustc-codegen-cuda/examples/generic/src/main.rs index 40706f70aa..5f622a8a71 100644 --- a/crates/rustc-codegen-cuda/examples/generic/src/main.rs +++ b/crates/rustc-codegen-cuda/examples/generic/src/main.rs @@ -57,6 +57,32 @@ mod kernels { *c_elem = a[idx_raw] + b[idx_raw]; } } + + #[inline(never)] + fn apply_unary(f: F, x: T) -> T + where + F: FnOnce(T) -> T, + { + f(x) + } + + /// Generic kernel whose body defines a captured closure. + /// + /// rustc prepends the parent generic args before a closure's fixed + /// `[kind, sig, tupled_upvars]` suffix. The importer must read the upvars + /// tuple from the suffix, not from a hard-coded substitution index. + #[kernel] + pub fn closure_capture(bias: T, scale: T, input: &[T], mut out: DisjointSlice) + where + T: Copy + Add + Mul, + { + let idx = thread::index_1d(); + let idx_raw = idx.get(); + if let Some(out_elem) = out.get_mut(idx) { + let transform = move |x: T| (x + bias) * scale; + *out_elem = apply_unary(transform, input[idx_raw]); + } + } } // ============================================================================= @@ -139,5 +165,65 @@ fn main() { assert_eq!(errors, 0, "scale:: produced {errors} errors"); } - println!("\n✓ SUCCESS: typed generic launches worked for f32 and i32"); + println!("Launching closure_capture:: kernel..."); + { + let bias: f32 = 1.25; + let scale: f32 = 0.5; + let input_data: Vec = (0..N).map(|i| i as f32).collect(); + let input_dev = + DeviceBuffer::from_host(&stream, &input_data).expect("Failed to copy f32 input"); + let mut output_dev = + DeviceBuffer::::zeroed(&stream, N).expect("Failed to alloc f32 output"); + + module + .closure_capture::( + &stream, + LaunchConfig::for_num_elems(N as u32), + bias, + scale, + &input_dev, + &mut output_dev, + ) + .expect("closure_capture:: launch failed"); + + let output_host = output_dev + .to_host_vec(&stream) + .expect("Failed to copy f32 output back"); + let errors = (0..N) + .filter(|&i| (output_host[i] - (input_data[i] + bias) * scale).abs() > 1e-5) + .count(); + assert_eq!(errors, 0, "closure_capture:: produced {errors} errors"); + } + + println!("Launching closure_capture:: kernel..."); + { + let bias: i32 = 7; + let scale: i32 = 2; + let input_data: Vec = (0..N as i32).collect(); + let input_dev = + DeviceBuffer::from_host(&stream, &input_data).expect("Failed to copy i32 input"); + let mut output_dev = + DeviceBuffer::::zeroed(&stream, N).expect("Failed to alloc i32 output"); + + module + .closure_capture::( + &stream, + LaunchConfig::for_num_elems(N as u32), + bias, + scale, + &input_dev, + &mut output_dev, + ) + .expect("closure_capture:: launch failed"); + + let output_host = output_dev + .to_host_vec(&stream) + .expect("Failed to copy i32 output back"); + let errors = (0..N) + .filter(|&i| output_host[i] != (input_data[i] + bias) * scale) + .count(); + assert_eq!(errors, 0, "closure_capture:: produced {errors} errors"); + } + + println!("\n✓ SUCCESS: typed generic launches and captured closures worked for f32 and i32"); }