Fix RISC-V C function ABI when passing/returning structs containing floats

Author: beetrees

compiler/rustc_codegen_cranelift/src/abi/pass_mode.rs

@@ -39,7 +39,18 @@ fn apply_arg_attrs_to_abi_param(mut param: AbiParam, arg_attrs: ArgAttributes) -
     param
 }
 
-fn cast_target_to_abi_params(cast: &CastTarget) -> SmallVec<[AbiParam; 2]> {
+fn cast_target_to_abi_params(cast: &CastTarget) -> SmallVec<[(Size, AbiParam); 2]> {
+    if let Some(offset_from_start) = cast.rest_offset {
+        assert!(cast.prefix[1..].iter().all(|p| p.is_none()));
+        assert_eq!(cast.rest.unit.size, cast.rest.total);
+        let first = cast.prefix[0].unwrap();
+        let second = cast.rest.unit;
+        return smallvec![
+            (Size::ZERO, reg_to_abi_param(first)),
+            (offset_from_start, reg_to_abi_param(second))
+        ];
+    }
+
     let (rest_count, rem_bytes) = if cast.rest.unit.size.bytes() == 0 {
         (0, 0)
     } else {
@@ -54,25 +65,31 @@ fn cast_target_to_abi_params(cast: &CastTarget) -> SmallVec<[AbiParam; 2]> {
     // different types in Cranelift IR. Instead a single array of primitive types is used.
 
     // Create list of fields in the main structure
-    let mut args = cast
+    let args = cast
         .prefix
         .iter()
         .flatten()
         .map(|&reg| reg_to_abi_param(reg))
-        .chain((0..rest_count).map(|_| reg_to_abi_param(cast.rest.unit)))
-        .collect::<SmallVec<_>>();
+        .chain((0..rest_count).map(|_| reg_to_abi_param(cast.rest.unit)));
 
     // Append final integer
-    if rem_bytes != 0 {
+    let args = args.chain(if rem_bytes != 0 {
         // Only integers can be really split further.
         assert_eq!(cast.rest.unit.kind, RegKind::Integer);
-        args.push(reg_to_abi_param(Reg {
-            kind: RegKind::Integer,
-            size: Size::from_bytes(rem_bytes),
-        }));
+        Some(reg_to_abi_param(Reg { kind: RegKind::Integer, size: Size::from_bytes(rem_bytes) }))
+    } else {
+        None
+    });
+
+    let mut res = SmallVec::new();
+    let mut offset = Size::ZERO;
+
+    for arg in args {
+        res.push((offset, arg));
+        offset += Size::from_bytes(arg.value_type.bytes());
     }
 
-    args
+    res
 }
 
 impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
@@ -103,7 +120,7 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
             },
             PassMode::Cast { ref cast, pad_i32 } => {
                 assert!(!pad_i32, "padding support not yet implemented");
-                cast_target_to_abi_params(cast)
+                cast_target_to_abi_params(cast).into_iter().map(|(_, param)| param).collect()
             }
             PassMode::Indirect { attrs, meta_attrs: None, on_stack } => {
                 if on_stack {
@@ -149,9 +166,10 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                 }
                 _ => unreachable!("{:?}", self.layout.backend_repr),
             },
-            PassMode::Cast { ref cast, .. } => {
-                (None, cast_target_to_abi_params(cast).into_iter().collect())
-            }
+            PassMode::Cast { ref cast, .. } => (
+                None,
+                cast_target_to_abi_params(cast).into_iter().map(|(_, param)| param).collect(),
+            ),
             PassMode::Indirect { attrs: _, meta_attrs: None, on_stack } => {
                 assert!(!on_stack);
                 (Some(AbiParam::special(pointer_ty(tcx), ArgumentPurpose::StructReturn)), vec![])
@@ -170,12 +188,14 @@ pub(super) fn to_casted_value<'tcx>(
 ) -> SmallVec<[Value; 2]> {
     let (ptr, meta) = arg.force_stack(fx);
     assert!(meta.is_none());
-    let mut offset = 0;
     cast_target_to_abi_params(cast)
         .into_iter()
-        .map(|param| {
-            let val = ptr.offset_i64(fx, offset).load(fx, param.value_type, MemFlags::new());
-            offset += i64::from(param.value_type.bytes());
+        .map(|(offset, param)| {
+            let val = ptr.offset_i64(fx, offset.bytes() as i64).load(
+                fx,
+                param.value_type,
+                MemFlags::new(),
+            );
             val
         })
         .collect()
@@ -188,7 +208,7 @@ pub(super) fn from_casted_value<'tcx>(
     cast: &CastTarget,
 ) -> CValue<'tcx> {
     let abi_params = cast_target_to_abi_params(cast);
-    let abi_param_size: u32 = abi_params.iter().map(|param| param.value_type.bytes()).sum();
+    let abi_param_size: u32 = abi_params.iter().map(|(_, param)| param.value_type.bytes()).sum();
     let layout_size = u32::try_from(layout.size.bytes()).unwrap();
     let ptr = fx.create_stack_slot(
         // Stack slot size may be bigger for example `[u8; 3]` which is packed into an `i32`.
@@ -197,16 +217,13 @@ pub(super) fn from_casted_value<'tcx>(
         std::cmp::max(abi_param_size, layout_size),
         u32::try_from(layout.align.abi.bytes()).unwrap(),
     );
-    let mut offset = 0;
     let mut block_params_iter = block_params.iter().copied();
-    for param in abi_params {
-        let val = ptr.offset_i64(fx, offset).store(
+    for (offset, _) in abi_params {
+        ptr.offset_i64(fx, offset.bytes() as i64).store(
             fx,
             block_params_iter.next().unwrap(),
             MemFlags::new(),
-        );
-        offset += i64::from(param.value_type.bytes());
-        val
+        )
     }
     assert_eq!(block_params_iter.next(), None, "Leftover block param");
     CValue::by_ref(ptr, layout)

compiler/rustc_codegen_gcc/src/intrinsic/mod.rs

@@ -551,7 +551,23 @@ impl<'gcc, 'tcx> ArgAbiExt<'gcc, 'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                 bx.lifetime_start(llscratch, scratch_size);
 
                 // ... where we first store the value...
-                bx.store(val, llscratch, scratch_align);
+                if let Some(offset_from_start) = cast.rest_offset {
+                    assert!(cast.prefix[1..].iter().all(|p| p.is_none()));
+                    assert_eq!(cast.rest.unit.size, cast.rest.total);
+                    assert!(cast.prefix[0].is_some());
+                    let first = bx.extract_value(val, 0);
+                    let second = bx.extract_value(val, 1);
+                    bx.store(first, llscratch, scratch_align);
+                    let second_ptr =
+                        bx.inbounds_ptradd(llscratch, bx.const_usize(offset_from_start.bytes()));
+                    bx.store(
+                        second,
+                        second_ptr,
+                        scratch_align.restrict_for_offset(offset_from_start),
+                    );
+                } else {
+                    bx.store(val, llscratch, scratch_align);
+                };
 
                 // ... and then memcpy it to the intended destination.
                 bx.memcpy(

compiler/rustc_codegen_llvm/src/abi.rs

@@ -234,7 +234,23 @@ impl<'ll, 'tcx> ArgAbiExt<'ll, 'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                 let llscratch = bx.alloca(scratch_size, scratch_align);
                 bx.lifetime_start(llscratch, scratch_size);
                 // ...store the value...
-                bx.store(val, llscratch, scratch_align);
+                if let Some(offset_from_start) = cast.rest_offset {
+                    assert!(cast.prefix[1..].iter().all(|p| p.is_none()));
+                    assert_eq!(cast.rest.unit.size, cast.rest.total);
+                    assert!(cast.prefix[0].is_some());
+                    let first = bx.extract_value(val, 0);
+                    let second = bx.extract_value(val, 1);
+                    bx.store(first, llscratch, scratch_align);
+                    let second_ptr =
+                        bx.inbounds_ptradd(llscratch, bx.const_usize(offset_from_start.bytes()));
+                    bx.store(
+                        second,
+                        second_ptr,
+                        scratch_align.restrict_for_offset(offset_from_start),
+                    );
+                } else {
+                    bx.store(val, llscratch, scratch_align);
+                };
                 // ... and then memcpy it to the intended destination.
                 bx.memcpy(
                     dst.val.llval,

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -559,7 +559,25 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     ZeroSized => bug!("ZST return value shouldn't be in PassMode::Cast"),
                 };
                 let ty = bx.cast_backend_type(cast_ty);
-                bx.load(ty, llslot, self.fn_abi.ret.layout.align.abi)
+                if let Some(offset_from_start) = cast_ty.rest_offset {
+                    assert!(cast_ty.prefix[1..].iter().all(|p| p.is_none()));
+                    assert_eq!(cast_ty.rest.unit.size, cast_ty.rest.total);
+                    let first_ty = bx.reg_backend_type(&cast_ty.prefix[0].unwrap());
+                    let second_ty = bx.reg_backend_type(&cast_ty.rest.unit);
+                    let first = bx.load(first_ty, llslot, self.fn_abi.ret.layout.align.abi);
+                    let second_ptr =
+                        bx.inbounds_ptradd(llslot, bx.const_usize(offset_from_start.bytes()));
+                    let second = bx.load(
+                        second_ty,
+                        second_ptr,
+                        self.fn_abi.ret.layout.align.abi.restrict_for_offset(offset_from_start),
+                    );
+                    let res = bx.cx().const_poison(ty);
+                    let res = bx.insert_value(res, first, 0);
+                    bx.insert_value(res, second, 1)
+                } else {
+                    bx.load(ty, llslot, self.fn_abi.ret.layout.align.abi)
+                }
             }
         };
         bx.ret(llval);
@@ -1588,7 +1606,25 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 );
                 // ...and then load it with the ABI type.
                 let cast_ty = bx.cast_backend_type(cast);
-                llval = bx.load(cast_ty, llscratch, scratch_align);
+                llval = if let Some(offset_from_start) = cast.rest_offset {
+                    assert!(cast.prefix[1..].iter().all(|p| p.is_none()));
+                    assert_eq!(cast.rest.unit.size, cast.rest.total);
+                    let first_ty = bx.reg_backend_type(&cast.prefix[0].unwrap());
+                    let second_ty = bx.reg_backend_type(&cast.rest.unit);
+                    let first = bx.load(first_ty, llscratch, self.fn_abi.ret.layout.align.abi);
+                    let second_ptr =
+                        bx.inbounds_ptradd(llscratch, bx.const_usize(offset_from_start.bytes()));
+                    let second = bx.load(
+                        second_ty,
+                        second_ptr,
+                        scratch_align.restrict_for_offset(offset_from_start),
+                    );
+                    let res = bx.cx().const_poison(cast_ty);
+                    let res = bx.insert_value(res, first, 0);
+                    bx.insert_value(res, second, 1)
+                } else {
+                    bx.load(cast_ty, llscratch, scratch_align)
+                };
                 bx.lifetime_end(llscratch, scratch_size);
             } else {
                 // We can't use `PlaceRef::load` here because the argument

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -254,7 +254,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
                     }
                     PassMode::Cast { ref cast, .. } => {
                         debug!("alloc: {:?} (return place) -> place", local);
-                        let size = cast.size(&start_bx);
+                        let size = cast.size(&start_bx).max(layout.size);
                         return LocalRef::Place(PlaceRef::alloca_size(&mut start_bx, size, layout));
                     }
                     _ => {}

compiler/rustc_target/src/callconv/mips64.rs

@@ -2,9 +2,7 @@ use rustc_abi::{
     BackendRepr, FieldsShape, Float, HasDataLayout, Primitive, Reg, Size, TyAbiInterface,
 };
 
-use crate::callconv::{
-    ArgAbi, ArgAttribute, ArgAttributes, ArgExtension, CastTarget, FnAbi, PassMode, Uniform,
-};
+use crate::callconv::{ArgAbi, ArgExtension, CastTarget, FnAbi, PassMode, Uniform};
 
 fn extend_integer_width_mips<Ty>(arg: &mut ArgAbi<'_, Ty>, bits: u64) {
     // Always sign extend u32 values on 64-bit mips
@@ -140,16 +138,7 @@ where
 
     // Extract first 8 chunks as the prefix
     let rest_size = size - Size::from_bytes(8) * prefix_index as u64;
-    arg.cast_to(CastTarget {
-        prefix,
-        rest: Uniform::new(Reg::i64(), rest_size),
-        attrs: ArgAttributes {
-            regular: ArgAttribute::default(),
-            arg_ext: ArgExtension::None,
-            pointee_size: Size::ZERO,
-            pointee_align: None,
-        },
-    });
+    arg.cast_to(CastTarget::prefixed(prefix, Uniform::new(Reg::i64(), rest_size)));
 }
 
 pub(crate) fn compute_abi_info<'a, Ty, C>(cx: &C, fn_abi: &mut FnAbi<'a, Ty>)

compiler/rustc_target/src/callconv/mod.rs

@@ -196,6 +196,17 @@ impl ArgAttributes {
     }
 }
 
+impl From<ArgAttribute> for ArgAttributes {
+    fn from(value: ArgAttribute) -> Self {
+        Self {
+            regular: value,
+            arg_ext: ArgExtension::None,
+            pointee_size: Size::ZERO,
+            pointee_align: None,
+        }
+    }
+}
+
 /// An argument passed entirely registers with the
 /// same kind (e.g., HFA / HVA on PPC64 and AArch64).
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, HashStable_Generic)]
@@ -250,6 +261,9 @@ impl Uniform {
 #[derive(Clone, PartialEq, Eq, Hash, Debug, HashStable_Generic)]
 pub struct CastTarget {
     pub prefix: [Option<Reg>; 8],
+    /// The offset of `rest` from the start of the value. Currently only implemented for a `Reg`
+    /// pair created by the `offset_pair` method.
+    pub rest_offset: Option<Size>,
     pub rest: Uniform,
     pub attrs: ArgAttributes,
 }
@@ -262,42 +276,45 @@ impl From<Reg> for CastTarget {
 
 impl From<Uniform> for CastTarget {
     fn from(uniform: Uniform) -> CastTarget {
-        CastTarget {
-            prefix: [None; 8],
-            rest: uniform,
-            attrs: ArgAttributes {
-                regular: ArgAttribute::default(),
-                arg_ext: ArgExtension::None,
-                pointee_size: Size::ZERO,
-                pointee_align: None,
-            },
-        }
+        Self::prefixed([None; 8], uniform)
     }
 }
 
 impl CastTarget {
-    pub fn pair(a: Reg, b: Reg) -> CastTarget {
-        CastTarget {
+    pub fn prefixed(prefix: [Option<Reg>; 8], rest: Uniform) -> Self {
+        Self { prefix, rest_offset: None, rest, attrs: ArgAttributes::new() }
+    }
+
+    pub fn offset_pair(a: Reg, offset_from_start: Size, b: Reg) -> Self {
+        Self {
             prefix: [Some(a), None, None, None, None, None, None, None],
-            rest: Uniform::from(b),
-            attrs: ArgAttributes {
-                regular: ArgAttribute::default(),
-                arg_ext: ArgExtension::None,
-                pointee_size: Size::ZERO,
-                pointee_align: None,
-            },
+            rest_offset: Some(offset_from_start),
+            rest: b.into(),
+            attrs: ArgAttributes::new(),
         }
     }
 
+    pub fn with_attrs(mut self, attrs: ArgAttributes) -> Self {
+        self.attrs = attrs;
+        self
+    }
+
+    pub fn pair(a: Reg, b: Reg) -> CastTarget {
+        Self::prefixed([Some(a), None, None, None, None, None, None, None], Uniform::from(b))
+    }
+
     /// When you only access the range containing valid data, you can use this unaligned size;
     /// otherwise, use the safer `size` method.
     pub fn unaligned_size<C: HasDataLayout>(&self, _cx: &C) -> Size {
         // Prefix arguments are passed in specific designated registers
-        let prefix_size = self
-            .prefix
-            .iter()
-            .filter_map(|x| x.map(|reg| reg.size))
-            .fold(Size::ZERO, |acc, size| acc + size);
+        let prefix_size = if let Some(offset_from_start) = self.rest_offset {
+            offset_from_start
+        } else {
+            self.prefix
+                .iter()
+                .filter_map(|x| x.map(|reg| reg.size))
+                .fold(Size::ZERO, |acc, size| acc + size)
+        };
         // Remaining arguments are passed in chunks of the unit size
         let rest_size =
             self.rest.unit.size * self.rest.total.bytes().div_ceil(self.rest.unit.size.bytes());
@@ -321,9 +338,22 @@ impl CastTarget {
     /// Checks if these two `CastTarget` are equal enough to be considered "the same for all
     /// function call ABIs".
     pub fn eq_abi(&self, other: &Self) -> bool {
-        let CastTarget { prefix: prefix_l, rest: rest_l, attrs: attrs_l } = self;
-        let CastTarget { prefix: prefix_r, rest: rest_r, attrs: attrs_r } = other;
-        prefix_l == prefix_r && rest_l == rest_r && attrs_l.eq_abi(attrs_r)
+        let CastTarget {
+            prefix: prefix_l,
+            rest_offset: rest_offset_l,
+            rest: rest_l,
+            attrs: attrs_l,
+        } = self;
+        let CastTarget {
+            prefix: prefix_r,
+            rest_offset: rest_offset_r,
+            rest: rest_r,
+            attrs: attrs_r,
+        } = other;
+        prefix_l == prefix_r
+            && rest_offset_l == rest_offset_r
+            && rest_l == rest_r
+            && attrs_l.eq_abi(attrs_r)
     }
 }