fix(abi): preserve array alignment and apply 16-byte workaround generically

This fixes an IllegalAddress error when passing 16-byte aligned arrays (e.g. [u128; 4]) or tuples as kernel parameters.

Previously, arrays indiscriminately forced PassMode::Direct(ArgAttributes::new()), discarding any alignment metadata computed by rustc. Furthermore, the 16-byte PassMode::Cast workaround was only applied to ADTs.

This change abstracts the 16-byte workaround to apply to any aggregate type (arrays, ADTs, tuples) with >=16-byte alignment, and properly preserves alignment metadata for arrays when PassMode::Direct is used.

Author: Snehal-Reddy

crates/rustc_codegen_nvvm/src/abi.rs

@@ -53,38 +53,36 @@ pub(crate) fn readjust_fn_abi<'tcx>(
             arg.mode = PassMode::Pair(ptr_attrs, ArgAttributes::new());
         }
 
-        if arg.layout.ty.is_array() && !matches!(arg.mode, PassMode::Direct { .. }) {
-            arg.mode = PassMode::Direct(ArgAttributes::new());
-        }
-
-        // pass all adts directly as values, ptx wants them to be passed all by value, but rustc's
-        // ptx-kernel abi seems to be wrong, and it's unstable.
-        if arg.layout.ty.is_adt() {
-            let align = arg.layout.align.abi.bytes();
-            if align >= 16 && align.is_power_of_two() {
-                let unit = Reg {
-                    kind: RegKind::Integer,
-                    size: Size::from_bytes(16),
-                };
-                let cast = CastTarget {
-                    prefix: Default::default(),
-                    rest: rustc_target::callconv::Uniform {
-                        unit,
-                        total: arg.layout.size,
-                        is_consecutive: false,
-                    },
-                    rest_offset: Some(Size::ZERO),
-                    attrs: ArgAttributes::new(),
-                };
-                arg.mode = PassMode::Cast {
-                    cast: Box::new(cast),
-                    pad_i32: false,
-                };
-            } else if !matches!(arg.mode, PassMode::Direct { .. }) {
-                let mut attrs = ArgAttributes::new();
-                attrs.pointee_align = Some(arg.layout.align.abi);
-                arg.mode = PassMode::Direct(attrs);
-            }
+        let is_aggregate = arg.layout.ty.is_array()
+            || arg.layout.ty.is_adt()
+            || matches!(arg.layout.ty.kind(), TyKind::Tuple(_));
+        let align = arg.layout.align.abi.bytes();
+
+        if is_aggregate && align >= 16 && align.is_power_of_two() {
+            let unit = Reg {
+                kind: RegKind::Integer,
+                size: Size::from_bytes(16),
+            };
+            let cast = CastTarget {
+                prefix: Default::default(),
+                rest: rustc_target::callconv::Uniform {
+                    unit,
+                    total: arg.layout.size,
+                    is_consecutive: false,
+                },
+                rest_offset: Some(Size::ZERO),
+                attrs: ArgAttributes::new(),
+            };
+            arg.mode = PassMode::Cast {
+                cast: Box::new(cast),
+                pad_i32: false,
+            };
+        } else if (arg.layout.ty.is_array() || arg.layout.ty.is_adt())
+            && !matches!(arg.mode, PassMode::Direct { .. })
+        {
+            let mut attrs = ArgAttributes::new();
+            attrs.pointee_align = Some(arg.layout.align.abi);
+            arg.mode = PassMode::Direct(attrs);
         }
         arg
     };