diff --git a/crates/mir-importer/src/translator/terminator/mod.rs b/crates/mir-importer/src/translator/terminator/mod.rs index 6e37cce77f..2949f868f9 100644 --- a/crates/mir-importer/src/translator/terminator/mod.rs +++ b/crates/mir-importer/src/translator/terminator/mod.rs @@ -903,15 +903,13 @@ fn translate_call( } } - // Handle closure trait method calls (FnOnce::call_once, FnMut::call_mut, Fn::call) - // These calls pass arguments as a tuple, but the closure body expects unpacked args. - // We need to unpack the tuple before calling the closure. - // - // MIR shows: <{closure} as FnMut<(u32,)>>::call_mut(self_ref, tuple_args) - // But the closure body expects: fn(self_ref, unpacked_arg1, unpacked_arg2, ...) + // Handle genuine closure trait method calls. The receiver test matters: + // wrapper ADTs can carry a closure in their generic substitutions while + // still expecting the rust-call tuple as one argument. if let Some(ref name) = pattern_name && (name.contains("call_once") || name.contains("call_mut") || name.ends_with("::call")) - && substs_contains("Closure") + && !args.is_empty() + && receiver_is_closure(&args[0], body) { return translate_closure_call( ctx, @@ -1451,6 +1449,46 @@ fn translate_closure_call( } } +/// True only when the rust-call receiver is itself a closure. +/// +/// We type the operand's base local (or constant) and ignore any +/// `place.projection`, and we peel at most one reference. Both are safe for +/// the rust-call receiver specifically: rustc lowers a closure-trait call +/// (`Fn::call` / `FnMut::call_mut` / `FnOnce::call_once`) so that the receiver +/// argument is the closure passed by value, or a single `&`/`&mut` borrow of +/// it, materialized into its own temporary local, never an in-place projection +/// of a larger aggregate and never behind multiple references. So a closure +/// reached through a field (`(self.f)(x)`) still arrives here as a base local +/// of type `&{closure}`, and one `Ref` peel plus a base-local type check covers +/// every genuine closure-call shape. A wrapper ADT that merely carries a +/// closure in its generic substitutions (the case this guard exists to reject) +/// has a non-closure receiver type and is correctly left on the ordinary call +/// path with its rust-call tuple intact. +fn receiver_is_closure(receiver: &mir::Operand, body: &mir::Body) -> bool { + let ty = match receiver { + mir::Operand::Copy(place) | mir::Operand::Move(place) => { + let local: usize = place.local; + let local_decls: Vec<_> = body.local_decls().collect(); + match local_decls.get(local).map(|(_, decl)| decl.ty) { + Some(ty) => ty, + None => return false, + } + } + mir::Operand::Constant(const_op) => const_op.const_.ty(), + _ => return false, + }; + + let inner = match ty.kind() { + rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Ref(_, inner, _)) => inner, + _ => ty, + }; + + matches!( + inner.kind(), + rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Closure(_, _)) + ) +} + /// Extracts the closure body's mangled name from a closure operand. /// /// In unified compilation with `std`, when we see a call like: diff --git a/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs b/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs index efa674e646..f34c1fd4c0 100644 --- a/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs +++ b/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs @@ -136,6 +136,21 @@ mod kernels { f(x) } + struct ClosureWrapper { + f: F, + offset: u32, + } + + impl ClosureWrapper + where + F: Fn(u32) -> u32, + { + #[inline(never)] + fn call(&self, args: (u32,)) -> u32 { + (self.f)(args.0) + self.offset + } + } + // Note: apply_twice with Fn trait is commented out due to control flow issues. // The `Fn` trait requires the closure to be called multiple times, which // can create complex control flow that the backend doesn't fully support yet. @@ -201,6 +216,59 @@ mod kernels { *out_elem = apply_closure(scale, val); } } + + /// Wrapper method named `call` whose generic substitutions contain a + /// closure. This is not a closure trait shim: the receiver is the wrapper, + /// so the tuple argument must remain one ordinary method argument. + #[kernel] + pub fn test_wrapper_method_named_call(factor: u32, mut out: DisjointSlice) { + let idx = thread::index_1d(); + let idx_raw = idx.get(); + if let Some(out_elem) = out.get_mut(idx) { + let scale = |x: u32| x * factor; + let wrapper = ClosureWrapper { + f: scale, + offset: 1, + }; + *out_elem = wrapper.call((idx_raw as u32,)); + } + } + + /// Genuine closure call whose receiver is reached through a struct field + /// (`(holder.f)(x)`). The tightened `receiver_is_closure` guard must still + /// unpack the rust-call tuple here: rustc materializes the `&{closure}` + /// receiver into its own temporary local, so the base-local type check + /// sees a closure even though the closure lives in a field. + #[kernel] + pub fn test_closure_field_projection(factor: u32, mut out: DisjointSlice) { + struct Holder { + f: F, + } + let idx = thread::index_1d(); + let idx_raw = idx.get(); + if let Some(out_elem) = out.get_mut(idx) { + let holder = Holder { + f: |x: u32| x * factor, + }; + *out_elem = (holder.f)(idx_raw as u32); + } + } + + /// Genuine closure call through two reference levels (`(**rr)(x)`). The + /// guard peels a single reference; this still classifies as a closure + /// because rustc resolves the receiver to one `&{closure}` borrow at the + /// call site regardless of the source-level double indirection. + #[kernel] + pub fn test_closure_double_ref(factor: u32, mut out: DisjointSlice) { + let idx = thread::index_1d(); + let idx_raw = idx.get(); + if let Some(out_elem) = out.get_mut(idx) { + let closure = |x: u32| x * factor; + let r = &closure; + let rr = &r; + *out_elem = (**rr)(idx_raw as u32); + } + } } // ============================================================================= @@ -456,5 +524,88 @@ fn main() { println!(" ✓ PASSED\n"); } + // ========================================================================= + // Test 9: Wrapper method named call with closure generic + // ========================================================================= + println!("Test 9: Wrapper::call with closure generic"); + { + const N: usize = 8; + let factor: u32 = 6; + let mut out_dev = DeviceBuffer::::zeroed(&stream, N).unwrap(); + + println!(" factor = {}", factor); + + module + .test_wrapper_method_named_call( + (stream).as_ref(), + LaunchConfig::for_num_elems(N as u32), + factor, + &mut out_dev, + ) + .expect("Kernel launch failed"); + + let out: Vec = out_dev.to_host_vec(&stream).unwrap(); + let expected: Vec = (0..N).map(|i| (i as u32) * factor + 1).collect(); + + println!(" Output: {:?}", out); + println!(" Expected: {:?}", expected); + assert_eq!(out, expected); + println!(" ✓ PASSED\n"); + } + + // ========================================================================= + // Test 10: genuine closure called through a struct field projection + // ========================================================================= + println!("Test 10: closure via field projection"); + { + const N: usize = 8; + let factor: u32 = 5; + let mut out_dev = DeviceBuffer::::zeroed(&stream, N).unwrap(); + + module + .test_closure_field_projection( + (stream).as_ref(), + LaunchConfig::for_num_elems(N as u32), + factor, + &mut out_dev, + ) + .expect("Kernel launch failed"); + + let out: Vec = out_dev.to_host_vec(&stream).unwrap(); + let expected: Vec = (0..N).map(|i| (i as u32) * factor).collect(); + + println!(" Output: {:?}", out); + println!(" Expected: {:?}", expected); + assert_eq!(out, expected); + println!(" ✓ PASSED\n"); + } + + // ========================================================================= + // Test 11: genuine closure called through a double reference + // ========================================================================= + println!("Test 11: closure via double reference"); + { + const N: usize = 8; + let factor: u32 = 7; + let mut out_dev = DeviceBuffer::::zeroed(&stream, N).unwrap(); + + module + .test_closure_double_ref( + (stream).as_ref(), + LaunchConfig::for_num_elems(N as u32), + factor, + &mut out_dev, + ) + .expect("Kernel launch failed"); + + let out: Vec = out_dev.to_host_vec(&stream).unwrap(); + let expected: Vec = (0..N).map(|i| (i as u32) * factor).collect(); + + println!(" Output: {:?}", out); + println!(" Expected: {:?}", expected); + assert_eq!(out, expected); + println!(" ✓ PASSED\n"); + } + println!("=== All device closure tests passed! ==="); }