Ensure floats are returned losslessly by the C ABI on 32-bit x86

Author: beetrees

compiler/rustc_codegen_llvm/src/abi.rs

@@ -6,6 +6,7 @@ use rustc_abi::{
     RegKind, Size, X86Call,
 };
 use rustc_codegen_ssa::MemFlags;
+use rustc_codegen_ssa::common::RealPredicate;
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
 use rustc_codegen_ssa::mir::place::{PlaceRef, PlaceValue};
 use rustc_codegen_ssa::traits::*;
@@ -175,6 +176,9 @@ impl LlvmType for Reg {
 
 impl LlvmType for CastTarget {
     fn llvm_type<'ll>(&self, cx: &CodegenCx<'ll, '_>) -> &'ll Type {
+        if self.x87_floating_point_stack {
+            return cx.type_x86_fp80();
+        }
         let rest_ll_unit = self.rest.unit.llvm_type(cx);
         let rest_count = if self.rest.total == Size::ZERO {
             0
@@ -324,6 +328,7 @@ impl<'ll, 'tcx> ArgAbiBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
     ) {
         arg_abi.store_fn_arg(self, idx, dst)
     }
+
     fn store_arg(
         &mut self,
         arg_abi: &ArgAbi<'tcx, Ty<'tcx>>,
@@ -332,6 +337,165 @@ impl<'ll, 'tcx> ArgAbiBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
     ) {
         arg_abi.store(self, val, dst)
     }
+
+    fn x87_lossless_float_to_fp_stack(&mut self, value: &'ll Value, no_undef: bool) -> &'ll Value {
+        // If value is (partially) uninitialized (e.g. when returning `MaybeUninit<f64>`) then e.g.
+        // branching on it could lead to undefined behaviour. To ensure that doesn't happen and that
+        // any initialized bytes within a partially uninitialized value survive the round trip,
+        // freeze the value.
+        let value = if no_undef { value } else { self.freeze(value) };
+        // While we only need to manually convert sNaNs, all NaNs can be converted the same way and
+        // checking whether `value` is NaN only takes a single floating-point x86 instruction,
+        // whereas checking if it is a signalling NaNs requires bit operations. LLVM also generally
+        // won't know a non-constant `value` is not a sNaN but could be a qNaN, so being more
+        // general here doesn't prevent the branch from being optimised out in likely scenarios.
+        let is_nan = self.fcmp(RealPredicate::RealUNO, value, value);
+        let is_nan_block = self.append_sibling_block("float_pre_ret.is_nan");
+        let is_not_nan_block = self.append_sibling_block("float_pre_ret.is_not_nan");
+        let after_block = self.append_sibling_block("float_pre_ret.after");
+        let dbg_loc = self.get_dbg_loc();
+        self.cond_br(is_nan, is_nan_block, is_not_nan_block);
+
+        self.switch_to_block(is_nan_block);
+        if let Some(dbg_loc) = dbg_loc {
+            self.set_dbg_loc(dbg_loc);
+        }
+        // This manually converts a NaN to x86_fp80 to avoid setting the quiet NaN bit of
+        // signalling NaNs.
+        let num_bits = self.float_width(self.val_ty(value)) as u64;
+        assert!(
+            num_bits == 32 || num_bits == 64,
+            "attempt to return float on x87 floating point stack with width {num_bits}"
+        );
+        let bits_ty = self.type_ix(num_bits);
+        let bits = self.bitcast(value, bits_ty);
+        // The high 16 bits of an x86_fp80 are the exponent and sign (the sign is the highest
+        // bit) NaNs always have all bits of the exponent set to 1, so the only bit that is
+        // needed from `value` is the sign bit.
+        // Shift out the lower bits.
+        let exp_and_sign = self.lshr(bits, self.const_uint(bits_ty, num_bits - 16));
+        // Set the exponent to all 1s.
+        let exp_and_sign = self.or(exp_and_sign, self.const_uint(bits_ty, 0x7FFF));
+        // Shift the exponent and sign into position.
+        let exp_and_sign = self.zext(exp_and_sign, self.type_ix(80));
+        let exp_and_sign = self.shl(exp_and_sign, self.const_uint_big(self.type_ix(80), 64));
+
+        // The fraction of the input NaN needs to be shifted left to just before x86_fp80's
+        // explicit integer bit. There's no need to manually set the integer bit itself, as it
+        // will be already set to 1 due to the all 1s exponent in the input NaN.
+        let (fraction_bits, fraction) = match num_bits {
+            32 => (f32::MANTISSA_DIGITS - 1, self.zext(bits, self.type_i64())),
+            64 => (f64::MANTISSA_DIGITS - 1, bits),
+            _ => bug!(),
+        };
+        // Shift the fraction into position.
+        let fraction = self.shl(fraction, self.const_u64(63 - u64::from(fraction_bits)));
+        let fraction = self.zext(fraction, self.type_ix(80));
+
+        let is_nan_res = self.or(exp_and_sign, fraction);
+        let is_nan_res = self.bitcast(is_nan_res, self.type_x86_fp80());
+        self.br(after_block);
+
+        self.switch_to_block(is_not_nan_block);
+        if let Some(dbg_loc) = dbg_loc {
+            self.set_dbg_loc(dbg_loc);
+        }
+        let is_not_nan_res = self.fpext(value, self.type_x86_fp80());
+        self.br(after_block);
+
+        self.switch_to_block(after_block);
+        if let Some(dbg_loc) = dbg_loc {
+            self.set_dbg_loc(dbg_loc);
+        }
+        self.phi(
+            self.type_x86_fp80(),
+            &[is_nan_res, is_not_nan_res],
+            &[is_nan_block, is_not_nan_block],
+        )
+    }
+
+    fn x87_lossless_fp_stack_to_float(
+        &mut self,
+        value: &'ll Value,
+        float_type: &'ll Type,
+        no_undef: bool,
+    ) -> &'ll Value {
+        // If value is (partially) uninitialized (e.g. when returning `MaybeUninit<f64>`) then e.g.
+        // branching on it could lead to undefined behaviour. To ensure that doesn't happen and that
+        // any initialized bytes within a partially uninitialized value survive the round trip,
+        // freeze the value.
+        let value = if no_undef { value } else { self.freeze(value) };
+        let num_bits = self.float_width(float_type) as u64;
+        let fraction_bits = u64::from(match num_bits {
+            32 => f32::MANTISSA_DIGITS - 1,
+            64 => f64::MANTISSA_DIGITS - 1,
+            _ => bug!("attempt to return float on x87 floating point stack with width {num_bits}"),
+        });
+        let dest_bits_type = self.type_ix(num_bits);
+        // While we only need to manually convert sNaNs, all NaNs can be converted the same way and
+        // checking whether `value` is NaN only takes a single floating-point x86 instruction,
+        // whereas checking if it is a signalling NaNs requires bit operations. LLVM also generally
+        // won't know a non-constant `value` is not a sNaN but could be a qNaN, so being more
+        // general here doesn't prevent the branch from being optimised out in likely scenarios.
+        let is_nan = self.fcmp(RealPredicate::RealUNO, value, value);
+        let is_nan_block = self.append_sibling_block("float_post_ret.is_nan");
+        let is_not_nan_block = self.append_sibling_block("float_post_ret.is_not_nan");
+        let after_block = self.append_sibling_block("float_post_ret.after");
+        let dbg_loc = self.get_dbg_loc();
+        self.cond_br(is_nan, is_nan_block, is_not_nan_block);
+
+        self.switch_to_block(is_nan_block);
+        if let Some(dbg_loc) = dbg_loc {
+            self.set_dbg_loc(dbg_loc);
+        }
+        // This block converts a NaN to `x86_fp80` manually to avoid setting the quiet NaN bit of
+        // signalling NaNs.
+        // We don't handle the "invalid operand" bitpatterns here (which are treated like NaNs) as
+        // they can't be generated by any post-80387 hardware, and the return value should have been
+        // converted from an actual `f32`/`f64`. Even on non-SSE targets current compilers don't
+        // seem to miscompile code so badly as to allow user-supplied `x86_fp80` "invalid operands"
+        // to be returned as `f32`/`f64`. Similarly, sNaNs are never produced by the hardware so we
+        // don't handle the case where the only fraction bits set are truncated, as that can never
+        // happen with sNaNs converted from `f32`s/`f64`s.
+        let bits = self.bitcast(value, self.type_ix(80));
+        // Mask out the extra 1s in the exponent as `f32`/`f64` have less bits in their
+        // exponents than `x86_fp80`.
+        let exp_and_sign_bits = num_bits - fraction_bits;
+        let exp_and_sign_mask = u16::MAX << (16 - exp_and_sign_bits);
+        let exp_and_sign_mask = u128::from(exp_and_sign_mask) << 64;
+        let exp_and_sign = self.and(bits, self.const_uint_big(self.type_ix(80), exp_and_sign_mask));
+        // Shift the exponent and sign into position
+        let exp_and_sign = self
+            .lshr(exp_and_sign, self.const_uint_big(self.type_ix(80), 80 - u128::from(num_bits)));
+        let exp_and_sign = self.trunc(exp_and_sign, dest_bits_type);
+
+        // Truncate off the exponent and sign
+        let fraction = self.trunc(bits, self.type_i64());
+        // Shift the fraction in to position. There's no need to mask out `x86_fp80`'s
+        // explicit integer bit as the fraction is right next to the exponent which is all
+        // 1s anyway.
+        let fraction = self.lshr(fraction, self.const_u64(63 - fraction_bits));
+        let fraction = if num_bits != 64 { self.trunc(fraction, dest_bits_type) } else { fraction };
+
+        // Combine the parts into the resulting float
+        let is_nan_res = self.or(exp_and_sign, fraction);
+        let is_nan_res = self.bitcast(is_nan_res, float_type);
+        self.br(after_block);
+
+        self.switch_to_block(is_not_nan_block);
+        if let Some(dbg_loc) = dbg_loc {
+            self.set_dbg_loc(dbg_loc);
+        }
+        // Use a regular floating point conversion when the value it not a NaN.
+        let is_not_nan_res = self.fptrunc(value, float_type);
+        self.br(after_block);
+
+        self.switch_to_block(after_block);
+        if let Some(dbg_loc) = dbg_loc {
+            self.set_dbg_loc(dbg_loc);
+        }
+        self.phi(float_type, &[is_nan_res, is_not_nan_res], &[is_nan_block, is_not_nan_block])
+    }
 }
 
 pub(crate) trait FnAbiLlvmExt<'ll, 'tcx> {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1555,6 +1555,12 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     }
 }
 
+impl<'ll> Builder<'_, 'll, '_> {
+    pub(crate) fn freeze(&mut self, value: &'ll Value) -> &'ll Value {
+        unsafe { llvm::LLVMBuildFreeze(self.llbuilder, value, UNNAMED) }
+    }
+}
+
 impl<'ll> StaticBuilderMethods for Builder<'_, 'll, '_> {
     fn get_static(&mut self, def_id: DefId) -> &'ll Value {
         // Forward to the `get_static` method of `CodegenCx`

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -937,6 +937,7 @@ unsafe extern "C" {
     pub(crate) fn LLVMFloatTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMDoubleTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMFP128TypeInContext(C: &Context) -> &Type;
+    pub(crate) fn LLVMX86FP80TypeInContext(C: &Context) -> &Type;
 
     // Operations on non-IEEE real types
     pub(crate) fn LLVMBFloatTypeInContext(C: &Context) -> &Type;
@@ -1596,6 +1597,11 @@ unsafe extern "C" {
         Index: c_uint,
         Name: *const c_char,
     ) -> &'a Value;
+    pub(crate) fn LLVMBuildFreeze<'a>(
+        B: &Builder<'a>,
+        Val: &'a Value,
+        Name: *const c_char,
+    ) -> &'a Value;
 
     // Atomic Operations
     pub(crate) fn LLVMBuildAtomicCmpXchg<'a>(

compiler/rustc_codegen_llvm/src/type_.rs

@@ -278,6 +278,12 @@ impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {
     }
 }
 
+impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
+    pub(crate) fn type_x86_fp80(&self) -> &'ll Type {
+        unsafe { llvm::LLVMX86FP80TypeInContext(self.llcx) }
+    }
+}
+
 pub(crate) fn llvm_type_ptr(llcx: &llvm::Context) -> &Type {
     llvm_type_ptr_in_address_space(llcx, AddressSpace::ZERO)
 }

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -19,7 +19,7 @@ use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::OptLevel;
 use rustc_span::{Span, Spanned};
-use rustc_target::callconv::{ArgAbi, ArgAttributes, CastTarget, FnAbi, PassMode};
+use rustc_target::callconv::{ArgAbi, ArgAttribute, ArgAttributes, CastTarget, FnAbi, PassMode};
 use tracing::{debug, info};
 
 use super::operand::OperandRef;
@@ -281,7 +281,9 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                 // If the return value was retagged as it was stored,
                 // then we might be in a different basic block now.
                 // Update the cached block for `target` to point to this new
-                // block, where codegen will continue.
+                // block, where codegen will continue. Additionally, store_return() may have
+                // required a branch into a new codegen backend basic block (currently this occurs
+                // when `cast_target.x87_floating_point_stack` is set).
                 fx.cached_llbbs[target] = CachedLlbb::Some(bx.llbb());
             }
             MergingSucc::False
@@ -2511,7 +2513,7 @@ fn load_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     align: Align,
 ) -> Bx::Value {
     let cast_ty = bx.cast_backend_type(cast);
-    if let Some(offset_from_start) = cast.rest_offset {
+    let value = if let Some(offset_from_start) = cast.rest_offset {
         assert_eq!(cast.prefix.len(), 1);
         assert_eq!(cast.rest.unit.size, cast.rest.total);
         let first_ty = bx.reg_backend_type(&cast.prefix[0]);
@@ -2523,7 +2525,21 @@ fn load_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         let res = bx.insert_value(res, first, 0);
         bx.insert_value(res, second, 1)
     } else {
-        bx.load(cast_ty, ptr, align)
+        let load_ty = if cast.x87_floating_point_stack {
+            match cast.rest.unit.size.bytes() {
+                4 => bx.type_f32(),
+                8 => bx.type_f64(),
+                _ => bug!(),
+            }
+        } else {
+            cast_ty
+        };
+        bx.load(load_ty, ptr, align)
+    };
+    if cast.x87_floating_point_stack {
+        bx.x87_lossless_float_to_fp_stack(value, cast.attrs.contains(ArgAttribute::NoUndef))
+    } else {
+        value
     }
 }
 
@@ -2534,6 +2550,20 @@ pub fn store_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     ptr: Bx::Value,
     align: Align,
 ) {
+    let value = if cast.x87_floating_point_stack {
+        let float_type = match cast.rest.unit.size.bytes() {
+            4 => bx.type_f32(),
+            8 => bx.type_f64(),
+            _ => bug!(),
+        };
+        bx.x87_lossless_fp_stack_to_float(
+            value,
+            float_type,
+            cast.attrs.contains(ArgAttribute::NoUndef),
+        )
+    } else {
+        value
+    };
     if let Some(offset_from_start) = cast.rest_offset {
         assert_eq!(cast.prefix.len(), 1);
         assert_eq!(cast.rest.unit.size, cast.rest.total);

compiler/rustc_codegen_ssa/src/traits/type_.rs

@@ -157,6 +157,42 @@ pub trait ArgAbiBuilderMethods<'tcx>: BackendTypes {
         val: Self::Value,
         dst: PlaceRef<'tcx, Self::Value>,
     );
+    /// Losslessly convert a `f32` or `f64` to a float to be returned on the x87 floating point
+    /// stack. Because `MaybeUninit` `f32`s/`f64`s are also returned on the floating point stack,
+    /// `value` being uninitialized must not cause undefined behaviour unless `true` is passed in
+    /// the `no_undef` argument. This method is used to avoid an LLVM bug where signalling NaNs get
+    /// quietened when being returned on the x87 stack on 32-bit x86. For more details, see:
+    /// * <https://github.com/rust-lang/rust/issues/115567>
+    /// * <https://github.com/llvm/llvm-project/issues/66803>
+    fn x87_lossless_float_to_fp_stack(
+        &mut self,
+        value: Self::Value,
+        no_undef: bool,
+    ) -> Self::Value {
+        let _ = no_undef;
+        // Default to leaving the value unchanged. The backend can override this method if it needs
+        // extra codegen to avoid quietening signalling NaNs.
+        value
+    }
+    /// Losslessly convert a `f32` or `f64` from a float that was returned on the x87 floating point
+    /// stack. Because `MaybeUninit` `f32`s/`f64`s are also returned on the floating point stack,
+    /// `value` being uninitialized must not cause undefined behaviour unless `true` is passed in
+    /// the `no_undef` argument. This method is used to avoid an LLVM bug where signalling NaNs get
+    /// quietened when being returned on the x87 stack on 32-bit x86. For more details, see:
+    /// * <https://github.com/rust-lang/rust/issues/115567>
+    /// * <https://github.com/llvm/llvm-project/issues/66803>
+    fn x87_lossless_fp_stack_to_float(
+        &mut self,
+        value: Self::Value,
+        float_type: Self::Type,
+        no_undef: bool,
+    ) -> Self::Value {
+        let _ = float_type;
+        let _ = no_undef;
+        // Default to leaving the value unchanged. The backend can override this method if it needs
+        // extra codegen to avoid quietening signalling NaNs.
+        value
+    }
 }
 
 pub trait TypeCodegenMethods<'tcx> = DerivedTypeCodegenMethods<'tcx>