mir-importer: read closure upvars from generic suffix

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<Vec<rustc_public::ty::Ty>> {
+    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)?);

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<T>` and `add<T>` with trait bounds
 - Host code calls `scale::<f32>` which triggers monomorphization
+- Defines a captured closure inside a generic kernel body
 - Verifies the monomorphized kernel executes correctly
 
 ## Key Concepts Demonstrated

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<T, F>(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<T>(bias: T, scale: T, input: &[T], mut out: DisjointSlice<T>)
+    where
+        T: Copy + Add<Output = T> + Mul<Output = T>,
+    {
+        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::<i32> produced {errors} errors");
     }
 
-    println!("\n✓ SUCCESS: typed generic launches worked for f32 and i32");
+    println!("Launching closure_capture::<f32> kernel...");
+    {
+        let bias: f32 = 1.25;
+        let scale: f32 = 0.5;
+        let input_data: Vec<f32> = (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::<f32>::zeroed(&stream, N).expect("Failed to alloc f32 output");
+
+        module
+            .closure_capture::<f32>(
+                &stream,
+                LaunchConfig::for_num_elems(N as u32),
+                bias,
+                scale,
+                &input_dev,
+                &mut output_dev,
+            )
+            .expect("closure_capture::<f32> 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::<f32> produced {errors} errors");
+    }
+
+    println!("Launching closure_capture::<i32> kernel...");
+    {
+        let bias: i32 = 7;
+        let scale: i32 = 2;
+        let input_data: Vec<i32> = (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::<i32>::zeroed(&stream, N).expect("Failed to alloc i32 output");
+
+        module
+            .closure_capture::<i32>(
+                &stream,
+                LaunchConfig::for_num_elems(N as u32),
+                bias,
+                scale,
+                &input_dev,
+                &mut output_dev,
+            )
+            .expect("closure_capture::<i32> 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::<i32> produced {errors} errors");
+    }
+
+    println!("\n✓ SUCCESS: typed generic launches and captured closures worked for f32 and i32");
 }