Rollup merge of #139340 - beetrees:riscv-float-struct-abi, r=workingjubilee

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

RISC-V passes structs containing only one or two floats (or a float and integer pair) in registers, as long as the individual floats/integers fit in a single corresponding register (see [the ABI specification](https://github.com/riscv-non-isa/riscv-elf-psabi-doc/releases/download/v1.0/riscv-abi.pdf) for details). Before this PR, Rust would not check what offset the second float/integer was at, instead assuming that it was at the standard offset for its default alignment. However, as the offset can be affected by `#[repr(align(N))]` and `#[repr(packed)]`, this caused miscompilations (see #115609). To fix this, this PR introduces a `rest_offset` field to `CastTarget` that can be used to explicitly specify at what offset the `rest` part of the cast is located at.

While fixing this, I discovered another bug: the size of the cast target was being used as the size of the MIR return place (when the function was using a `PassMode::Cast` return type). However, the cast target is allowed to be smaller than the size of the actual type, causing a miscompilation. This PR fixes this issue by using the largest of the size of the type and the size of the cast target as the size of the MIR return place, ensuring all reads/writes will be inbounds.

Fixes the RISC-V part of #115609.

cc target maintainers of `riscv64gc-unknown-linux-gnu`: `@kito-cheng` `@michaelmaitland` `@robin-randhawa-sifive` `@topperc`

r? `@workingjubilee`

Author: Kobzol

compiler/rustc_codegen_cranelift/src/abi/pass_mode.rs

@@ -40,7 +40,18 @@ fn apply_attrs_to_abi_param(param: AbiParam, arg_attrs: ArgAttributes) -> AbiPar
     }
 }
 
-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 {
@@ -55,25 +66,32 @@ 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)));
+
+    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());
+    }
 
     // Append final integer
     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),
-        }));
+        res.push((
+            offset,
+            reg_to_abi_param(Reg { kind: RegKind::Integer, size: Size::from_bytes(rem_bytes) }),
+        ));
     }
 
-    args
+    res
 }
 
 impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
@@ -104,7 +122,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 {
@@ -160,9 +178,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);
                 (
@@ -187,12 +206,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()
@@ -205,7 +226,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`.
@@ -214,16 +235,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

@@ -626,7 +626,7 @@ 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);
+                rustc_codegen_ssa::mir::store_cast(bx, cast, val, llscratch, scratch_align);
 
                 // ... and then memcpy it to the intended destination.
                 bx.memcpy(

compiler/rustc_codegen_llvm/src/abi.rs

@@ -229,7 +229,7 @@ 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);
+                rustc_codegen_ssa::mir::store_cast(bx, cast, 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

@@ -1,6 +1,6 @@
 use std::cmp;
 
-use rustc_abi::{BackendRepr, ExternAbi, HasDataLayout, Reg, Size, WrappingRange};
+use rustc_abi::{Align, BackendRepr, ExternAbi, HasDataLayout, Reg, Size, WrappingRange};
 use rustc_ast as ast;
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_data_structures::packed::Pu128;
@@ -13,7 +13,7 @@ use rustc_middle::{bug, span_bug};
 use rustc_session::config::OptLevel;
 use rustc_span::Span;
 use rustc_span::source_map::Spanned;
-use rustc_target::callconv::{ArgAbi, FnAbi, PassMode};
+use rustc_target::callconv::{ArgAbi, CastTarget, FnAbi, PassMode};
 use tracing::{debug, info};
 
 use super::operand::OperandRef;
@@ -558,8 +558,7 @@ 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)
+                load_cast(bx, cast_ty, llslot, self.fn_abi.ret.layout.align.abi)
             }
         };
         bx.ret(llval);
@@ -1618,8 +1617,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     MemFlags::empty(),
                 );
                 // ...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 = load_cast(bx, cast, llscratch, scratch_align);
                 bx.lifetime_end(llscratch, scratch_size);
             } else {
                 // We can't use `PlaceRef::load` here because the argument
@@ -1969,3 +1967,47 @@ enum ReturnDest<'tcx, V> {
     /// Store a direct return value to an operand local place.
     DirectOperand(mir::Local),
 }
+
+fn load_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    bx: &mut Bx,
+    cast: &CastTarget,
+    ptr: Bx::Value,
+    align: Align,
+) -> Bx::Value {
+    let cast_ty = bx.cast_backend_type(cast);
+    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, ptr, align);
+        let second_ptr = bx.inbounds_ptradd(ptr, bx.const_usize(offset_from_start.bytes()));
+        let second = bx.load(second_ty, second_ptr, 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, ptr, align)
+    }
+}
+
+pub fn store_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    bx: &mut Bx,
+    cast: &CastTarget,
+    value: Bx::Value,
+    ptr: Bx::Value,
+    align: 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(value, 0);
+        let second = bx.extract_value(value, 1);
+        bx.store(first, ptr, align);
+        let second_ptr = bx.inbounds_ptradd(ptr, bx.const_usize(offset_from_start.bytes()));
+        bx.store(second, second_ptr, align.restrict_for_offset(offset_from_start));
+    } else {
+        bx.store(value, ptr, align);
+    };
+}

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -26,6 +26,7 @@ pub mod place;
 mod rvalue;
 mod statement;
 
+pub use self::block::store_cast;
 use self::debuginfo::{FunctionDebugContext, PerLocalVarDebugInfo};
 use self::operand::{OperandRef, OperandValue};
 use self::place::PlaceRef;
@@ -259,7 +260,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>)