atomic_load intrinsic: use const generic parameter for ordering

Author: RalfJung

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -791,11 +791,12 @@ fn codegen_regular_intrinsic_call<'tcx>(
             // FIXME use a compiler fence once Cranelift supports it
             fx.bcx.ins().fence();
         }
-        _ if intrinsic.as_str().starts_with("atomic_load") => {
+        sym::atomic_load => {
             intrinsic_args!(fx, args => (ptr); intrinsic);
             let ptr = ptr.load_scalar(fx);
 
             let ty = generic_args.type_at(0);
+            let _ord = generic_args.const_at(1).to_value(); // FIXME: forward this to cranelift once they support that
             match ty.kind() {
                 ty::Uint(UintTy::U128) | ty::Int(IntTy::I128) => {
                     // FIXME implement 128bit atomics

compiler/rustc_codegen_ssa/src/lib.rs

@@ -6,6 +6,7 @@
 #![doc(rust_logo)]
 #![feature(assert_matches)]
 #![feature(box_patterns)]
+#![feature(core_intrinsics)]
 #![feature(file_buffered)]
 #![feature(if_let_guard)]
 #![feature(negative_impls)]

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -99,6 +99,19 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let llret_ty = bx.backend_type(bx.layout_of(ret_ty));
         let result = PlaceRef::new_sized(llresult, fn_abi.ret.layout);
 
+        let ret_llval = |bx: &mut Bx, llval| {
+            if !fn_abi.ret.is_ignore() {
+                if let PassMode::Cast { .. } = &fn_abi.ret.mode {
+                    bx.store_to_place(llval, result.val);
+                } else {
+                    OperandRef::from_immediate_or_packed_pair(bx, llval, result.layout)
+                        .val
+                        .store(bx, result);
+                }
+            }
+            Ok(())
+        };
+
         let llval = match name {
             sym::abort => {
                 bx.abort();
@@ -331,6 +344,56 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 use crate::common::AtomicOrdering::*;
                 use crate::common::{AtomicRmwBinOp, SynchronizationScope};
 
+                let invalid_monomorphization = |ty| {
+                    bx.tcx().dcx().emit_err(InvalidMonomorphization::BasicIntegerType {
+                        span,
+                        name,
+                        ty,
+                    });
+                };
+
+                let parse_const_generic_ordering = |ord: ty::Value<'tcx>| {
+                    let discr = ord.valtree.unwrap_branch()[0].unwrap_leaf();
+                    let ord = discr.to_atomic_ordering();
+                    // We have to translate from the intrinsic ordering to the backend ordering.
+                    use std::intrinsics::AtomicOrdering;
+                    match ord {
+                        AtomicOrdering::Relaxed => Relaxed,
+                        AtomicOrdering::Release => Release,
+                        AtomicOrdering::Acquire => Acquire,
+                        AtomicOrdering::AcqRel => AcquireRelease,
+                        AtomicOrdering::SeqCst => SequentiallyConsistent,
+                    }
+                };
+
+                // Some intrinsics have the ordering already converted to a const generic parameter, we handle those first.
+                match name {
+                    sym::atomic_load => {
+                        let ty = fn_args.type_at(0);
+                        let ordering = fn_args.const_at(1).to_value();
+                        let llval =
+                            if int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr() {
+                                let layout = bx.layout_of(ty);
+                                let size = layout.size;
+                                let source = args[0].immediate();
+                                bx.atomic_load(
+                                    bx.backend_type(layout),
+                                    source,
+                                    parse_const_generic_ordering(ordering),
+                                    size,
+                                )
+                            } else {
+                                invalid_monomorphization(ty);
+                                return Ok(());
+                            };
+
+                        return ret_llval(bx, llval);
+                    }
+
+                    // The rest falls back to below.
+                    _ => {}
+                }
+
                 let Some((instruction, ordering)) = atomic.split_once('_') else {
                     bx.sess().dcx().emit_fatal(errors::MissingMemoryOrdering);
                 };
@@ -344,14 +407,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     _ => bx.sess().dcx().emit_fatal(errors::UnknownAtomicOrdering),
                 };
 
-                let invalid_monomorphization = |ty| {
-                    bx.tcx().dcx().emit_err(InvalidMonomorphization::BasicIntegerType {
-                        span,
-                        name,
-                        ty,
-                    });
-                };
-
                 match instruction {
                     "cxchg" | "cxchgweak" => {
                         let Some((success, failure)) = ordering.split_once('_') else {
@@ -384,24 +439,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         return Ok(());
                     }
 
-                    "load" => {
-                        let ty = fn_args.type_at(0);
-                        if int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr() {
-                            let layout = bx.layout_of(ty);
-                            let size = layout.size;
-                            let source = args[0].immediate();
-                            bx.atomic_load(
-                                bx.backend_type(layout),
-                                source,
-                                parse_ordering(bx, ordering),
-                                size,
-                            )
-                        } else {
-                            invalid_monomorphization(ty);
-                            return Ok(());
-                        }
-                    }
-
                     "store" => {
                         let ty = fn_args.type_at(0);
                         if int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr() {
@@ -532,16 +569,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
         };
 
-        if !fn_abi.ret.is_ignore() {
-            if let PassMode::Cast { .. } = &fn_abi.ret.mode {
-                bx.store_to_place(llval, result.val);
-            } else {
-                OperandRef::from_immediate_or_packed_pair(bx, llval, result.layout)
-                    .val
-                    .store(bx, result);
-            }
-        }
-        Ok(())
+        ret_llval(bx, llval)
     }
 }
 

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -204,24 +204,25 @@ pub(crate) fn check_intrinsic_type(
 
         // Each atomic op has variants with different suffixes (`_seq_cst`, `_acquire`, etc.). Use
         // string ops to strip the suffixes, because the variants all get the same treatment here.
-        let (n_tps, inputs, output) = match split[1] {
+        let (n_tps, n_cts, inputs, output) = match split[1] {
             "cxchg" | "cxchgweak" => (
                 1,
+                0,
                 vec![Ty::new_mut_ptr(tcx, param(0)), param(0), param(0)],
                 Ty::new_tup(tcx, &[param(0), tcx.types.bool]),
             ),
-            "load" => (1, vec![Ty::new_imm_ptr(tcx, param(0))], param(0)),
-            "store" => (1, vec![Ty::new_mut_ptr(tcx, param(0)), param(0)], tcx.types.unit),
+            "load" => (1, 1, vec![Ty::new_imm_ptr(tcx, param(0))], param(0)),
+            "store" => (1, 0, vec![Ty::new_mut_ptr(tcx, param(0)), param(0)], tcx.types.unit),
 
             "xchg" | "xadd" | "xsub" | "and" | "nand" | "or" | "xor" | "max" | "min" | "umax"
-            | "umin" => (1, vec![Ty::new_mut_ptr(tcx, param(0)), param(0)], param(0)),
-            "fence" | "singlethreadfence" => (0, Vec::new(), tcx.types.unit),
+            | "umin" => (1, 0, vec![Ty::new_mut_ptr(tcx, param(0)), param(0)], param(0)),
+            "fence" | "singlethreadfence" => (0, 0, Vec::new(), tcx.types.unit),
             op => {
                 tcx.dcx().emit_err(UnrecognizedAtomicOperation { span, op });
                 return;
             }
         };
-        (n_tps, 0, 0, inputs, output, hir::Safety::Unsafe)
+        (n_tps, 0, n_cts, inputs, output, hir::Safety::Unsafe)
     } else if intrinsic_name == sym::contract_check_ensures {
         // contract_check_ensures::<Ret, C>(Ret, C) -> Ret
         // where C: for<'a> Fn(&'a Ret) -> bool,

compiler/rustc_middle/src/ty/consts/int.rs

@@ -318,6 +318,27 @@ impl ScalarInt {
         self.to_uint(tcx.data_layout.pointer_size).try_into().unwrap()
     }
 
+    #[inline]
+    pub fn to_atomic_ordering(self) -> std::intrinsics::AtomicOrdering {
+        use std::intrinsics::AtomicOrdering::*;
+        // Here we rely on `AtomicOrdering` having the same values in host and target!
+        // We use `u32` as that's how valtrees encode the enum discriminant.
+        let val = self.to_u32();
+        if val == Relaxed as u32 {
+            Relaxed
+        } else if val == Release as u32 {
+            Release
+        } else if val == Acquire as u32 {
+            Acquire
+        } else if val == AcqRel as u32 {
+            AcqRel
+        } else if val == SeqCst as u32 {
+            SeqCst
+        } else {
+            panic!("not a valid atomic ordering")
+        }
+    }
+
     /// Converts the `ScalarInt` to `bool`.
     /// Panics if the `size` of the `ScalarInt` is not equal to 1 byte.
     /// Errors if it is not a valid `bool`.
@@ -488,7 +509,7 @@ from_scalar_int_for_x_signed!(i8, i16, i32, i64, i128);
 impl From<std::cmp::Ordering> for ScalarInt {
     #[inline]
     fn from(c: std::cmp::Ordering) -> Self {
-        // Here we rely on `Ordering` having the same values in host and target!
+        // Here we rely on `cmp::Ordering` having the same values in host and target!
         ScalarInt::from(c as i8)
     }
 }

compiler/rustc_span/src/symbol.rs

@@ -512,6 +512,7 @@ symbols! {
         async_iterator_poll_next,
         async_trait_bounds,
         atomic,
+        atomic_load,
         atomic_mod,
         atomics,
         att_syntax,

library/core/src/intrinsics/mod.rs

@@ -30,7 +30,7 @@
 //!
 //! The atomic intrinsics provide common atomic operations on machine
 //! words, with multiple possible memory orderings. See the
-//! [atomic types][crate::sync::atomic] docs for details.
+//! [atomic types][atomic] docs for details.
 //!
 //! # Unwinding
 //!
@@ -50,7 +50,7 @@
 )]
 #![allow(missing_docs)]
 
-use crate::marker::{DiscriminantKind, Tuple};
+use crate::marker::{ConstParamTy, DiscriminantKind, Tuple};
 use crate::ptr;
 
 pub mod fallback;
@@ -60,7 +60,21 @@ pub mod simd;
 // These imports are used for simplifying intra-doc links
 #[allow(unused_imports)]
 #[cfg(all(target_has_atomic = "8", target_has_atomic = "32", target_has_atomic = "ptr"))]
-use crate::sync::atomic::{self, AtomicBool, AtomicI32, AtomicIsize, AtomicU32, Ordering};
+use crate::sync::atomic::{AtomicBool, AtomicI32, AtomicIsize, AtomicU32, Ordering};
+
+/// A type for atomic ordering parameters for intrinsics. This is a separate type from
+/// `atomic::Ordering` so that we can make it `ConstParamTy` and fix the values used here without a
+/// risk of leaking that to stable code.
+///
+/// The exact values of these discriminants are mirrored in the compiler so they cannot be changed!
+#[derive(Debug, ConstParamTy, PartialEq, Eq)]
+pub enum AtomicOrdering {
+    Relaxed = 0,
+    Release = 1,
+    Acquire = 2,
+    AcqRel = 3,
+    SeqCst = 4,
+}
 
 // N.B., these intrinsics take raw pointers because they mutate aliased
 // memory, which is not valid for either `&` or `&mut`.
@@ -391,6 +405,15 @@ pub unsafe fn atomic_cxchgweak_seqcst_acquire<T: Copy>(dst: *mut T, old: T, src:
 #[rustc_nounwind]
 pub unsafe fn atomic_cxchgweak_seqcst_seqcst<T: Copy>(dst: *mut T, old: T, src: T) -> (T, bool);
 
+/// Loads the current value of the pointer.
+/// `T` must be an integer or pointer type.
+///
+/// The stabilized version of this intrinsic is available on the
+/// [`atomic`] types via the `load` method. For example, [`AtomicBool::load`].
+#[rustc_intrinsic]
+#[rustc_nounwind]
+#[cfg(not(bootstrap))]
+pub unsafe fn atomic_load<T: Copy, const ORD: AtomicOrdering>(src: *const T) -> T;
 /// Loads the current value of the pointer.
 /// `T` must be an integer or pointer type.
 ///
@@ -399,6 +422,7 @@ pub unsafe fn atomic_cxchgweak_seqcst_seqcst<T: Copy>(dst: *mut T, old: T, src:
 /// [`Ordering::SeqCst`] as the `order`. For example, [`AtomicBool::load`].
 #[rustc_intrinsic]
 #[rustc_nounwind]
+#[cfg(bootstrap)]
 pub unsafe fn atomic_load_seqcst<T: Copy>(src: *const T) -> T;
 /// Loads the current value of the pointer.
 /// `T` must be an integer or pointer type.
@@ -408,6 +432,7 @@ pub unsafe fn atomic_load_seqcst<T: Copy>(src: *const T) -> T;
 /// [`Ordering::Acquire`] as the `order`. For example, [`AtomicBool::load`].
 #[rustc_intrinsic]
 #[rustc_nounwind]
+#[cfg(bootstrap)]
 pub unsafe fn atomic_load_acquire<T: Copy>(src: *const T) -> T;
 /// Loads the current value of the pointer.
 /// `T` must be an integer or pointer type.
@@ -417,6 +442,7 @@ pub unsafe fn atomic_load_acquire<T: Copy>(src: *const T) -> T;
 /// [`Ordering::Relaxed`] as the `order`. For example, [`AtomicBool::load`].
 #[rustc_intrinsic]
 #[rustc_nounwind]
+#[cfg(bootstrap)]
 pub unsafe fn atomic_load_relaxed<T: Copy>(src: *const T) -> T;
 
 /// Stores the value at the specified memory location.

library/core/src/sync/atomic.rs

@@ -3822,6 +3822,7 @@ unsafe fn atomic_store<T: Copy>(dst: *mut T, val: T, order: Ordering) {
 
 #[inline]
 #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
+#[cfg(bootstrap)]
 unsafe fn atomic_load<T: Copy>(dst: *const T, order: Ordering) -> T {
     // SAFETY: the caller must uphold the safety contract for `atomic_load`.
     unsafe {
@@ -3835,6 +3836,23 @@ unsafe fn atomic_load<T: Copy>(dst: *const T, order: Ordering) -> T {
     }
 }
 
+#[inline]
+#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
+#[cfg(not(bootstrap))]
+unsafe fn atomic_load<T: Copy>(dst: *const T, order: Ordering) -> T {
+    use intrinsics::AtomicOrdering;
+    // SAFETY: the caller must uphold the safety contract for `atomic_load`.
+    unsafe {
+        match order {
+            Relaxed => intrinsics::atomic_load::<T, { AtomicOrdering::Relaxed }>(dst),
+            Acquire => intrinsics::atomic_load::<T, { AtomicOrdering::Acquire }>(dst),
+            SeqCst => intrinsics::atomic_load::<T, { AtomicOrdering::SeqCst }>(dst),
+            Release => panic!("there is no such thing as a release load"),
+            AcqRel => panic!("there is no such thing as an acquire-release load"),
+        }
+    }
+}
+
 #[inline]
 #[cfg(target_has_atomic)]
 #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces

src/tools/miri/src/intrinsics/atomic.rs

@@ -19,6 +19,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn emulate_atomic_intrinsic(
         &mut self,
         intrinsic_name: &str,
+        generic_args: ty::GenericArgsRef<'tcx>,
         args: &[OpTy<'tcx>],
         dest: &MPlaceTy<'tcx>,
     ) -> InterpResult<'tcx, EmulateItemResult> {
@@ -35,6 +36,16 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             }
         }
 
+        fn read_ord_const_generic(o: std::intrinsics::AtomicOrdering) -> AtomicReadOrd {
+            use std::intrinsics::AtomicOrdering;
+            match o {
+                AtomicOrdering::SeqCst => AtomicReadOrd::SeqCst,
+                AtomicOrdering::Acquire => AtomicReadOrd::Acquire,
+                AtomicOrdering::Relaxed => AtomicReadOrd::Relaxed,
+                _ => panic!("invalid read ordering `{o:?}`"),
+            }
+        }
+
         fn write_ord(ord: &str) -> AtomicWriteOrd {
             match ord {
                 "seqcst" => AtomicWriteOrd::SeqCst,
@@ -66,7 +77,15 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         }
 
         match &*intrinsic_structure {
-            ["load", ord] => this.atomic_load(args, dest, read_ord(ord))?,
+            // New-style intrinsics that use const generics
+            ["load"] => {
+                let ordering = generic_args.const_at(1).to_value();
+                let ordering =
+                    ordering.valtree.unwrap_branch()[0].unwrap_leaf().to_atomic_ordering();
+                this.atomic_load(args, dest, read_ord_const_generic(ordering))?;
+            }
+
+            // Old-style intrinsics that have the ordering in the intrinsic name
             ["store", ord] => this.atomic_store(args, write_ord(ord))?,
 
             ["fence", ord] => this.atomic_fence_intrinsic(args, fence_ord(ord))?,

src/tools/miri/src/intrinsics/mod.rs

@@ -94,7 +94,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let this = self.eval_context_mut();
 
         if let Some(name) = intrinsic_name.strip_prefix("atomic_") {
-            return this.emulate_atomic_intrinsic(name, args, dest);
+            return this.emulate_atomic_intrinsic(name, generic_args, args, dest);
         }
         if let Some(name) = intrinsic_name.strip_prefix("simd_") {
             return this.emulate_simd_intrinsic(name, generic_args, args, dest);