make atomic operations const

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -2,6 +2,7 @@
 //! looking at their MIR. Intrinsics/functions supported here are shared by CTFE
 //! and miri.
 
+mod atomic;
 mod simd;
 
 use std::assert_matches;
@@ -20,9 +21,9 @@ use tracing::trace;
 use super::memory::MemoryKind;
 use super::util::ensure_monomorphic_enough;
 use super::{
-    AllocId, CheckInAllocMsg, ImmTy, InterpCx, InterpResult, Machine, OpTy, PlaceTy, Pointer,
-    PointerArithmetic, Projectable, Provenance, Scalar, err_ub_format, err_unsup_format, interp_ok,
-    throw_inval, throw_ub, throw_ub_format,
+    AllocId, AtomicRmwOp, CheckInAllocMsg, ImmTy, Immediate, InterpCx, InterpResult, Machine, OpTy,
+    PlaceTy, Pointer, PointerArithmetic, Projectable, Provenance, Scalar, err_ub_format,
+    err_unsup_format, interp_ok, throw_inval, throw_ub, throw_ub_format,
 };
 use crate::interpret::{MPlaceTy, Writeable};
 
@@ -177,6 +178,9 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         let instance_args = instance.args;
         let intrinsic_name = self.tcx.item_name(instance.def_id());
 
+        if intrinsic_name.as_str().starts_with("atomic_") {
+            return self.eval_atomic_intrinsic(intrinsic_name, instance_args, args, dest, ret);
+        }
         if intrinsic_name.as_str().starts_with("simd_") {
             return self.eval_simd_intrinsic(intrinsic_name, instance_args, args, dest, ret);
         }

compiler/rustc_const_eval/src/interpret/intrinsics/atomic.rs

@@ -0,0 +1,131 @@
+use rustc_middle::mir::BinOp;
+use rustc_middle::{mir, span_bug, ty};
+use rustc_span::{Symbol, sym};
+use tracing::trace;
+
+use super::{
+    AtomicRmwOp, Immediate, InterpCx, InterpResult, Machine, OpTy, PlaceTy, Scalar, interp_ok,
+};
+
+impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
+    /// Returns `true` if emulation happened.
+    /// Here we implement the intrinsics that are common to all CTFE instances; individual machines can add their own
+    /// intrinsic handling.
+    pub fn eval_atomic_intrinsic(
+        &mut self,
+        intrinsic_name: Symbol,
+        generic_args: ty::GenericArgsRef<'tcx>,
+        args: &[OpTy<'tcx, M::Provenance>],
+        dest: &PlaceTy<'tcx, M::Provenance>,
+        ret: Option<mir::BasicBlock>,
+    ) -> InterpResult<'tcx, bool> {
+        let get_ord_at = |i: usize| {
+            let ordering = generic_args.const_at(i).to_value();
+            ordering.to_branch()[0].to_value().to_leaf().to_atomic_ordering()
+        };
+
+        match intrinsic_name {
+            sym::atomic_load => {
+                let ord = get_ord_at(1);
+                let [ptr] = args else { span_bug!(self.cur_span(), "invalid `atomic_load` call") };
+
+                let place = self.deref_pointer(ptr)?;
+                let val = M::atomic_load(self, &place, ord)?;
+                self.write_scalar(val, dest)?;
+            }
+            sym::atomic_store => {
+                let ord = get_ord_at(1);
+                let [ptr, val] = args else {
+                    span_bug!(self.cur_span(), "invalid `atomic_store` call")
+                };
+
+                let place = self.deref_pointer(ptr)?;
+                let val = self.read_immediate(val)?;
+                M::atomic_store(self, &place, &val, ord)?;
+            }
+            sym::atomic_or
+            | sym::atomic_xor
+            | sym::atomic_and
+            | sym::atomic_nand
+            | sym::atomic_xadd
+            | sym::atomic_xsub
+            | sym::atomic_min
+            | sym::atomic_umin
+            | sym::atomic_max
+            | sym::atomic_umax
+            | sym::atomic_xchg => {
+                let num_ty_generics = match intrinsic_name {
+                    sym::atomic_min
+                    | sym::atomic_umin
+                    | sym::atomic_max
+                    | sym::atomic_umax
+                    | sym::atomic_xchg => 1,
+                    _ => 2,
+                };
+                let ord = get_ord_at(num_ty_generics);
+                let [ptr, operand] = args else {
+                    span_bug!(self.cur_span(), "invalid `{intrinsic_name}` call")
+                };
+
+                let place = self.deref_pointer(ptr)?;
+                let operand = self.read_immediate(operand)?;
+
+                let op = match intrinsic_name {
+                    sym::atomic_or => AtomicRmwOp::MirOp { op: BinOp::BitOr, neg: false },
+                    sym::atomic_xor => AtomicRmwOp::MirOp { op: BinOp::BitXor, neg: false },
+                    sym::atomic_and => AtomicRmwOp::MirOp { op: BinOp::BitAnd, neg: false },
+                    sym::atomic_nand => AtomicRmwOp::MirOp { op: BinOp::BitAnd, neg: true },
+                    sym::atomic_xadd => AtomicRmwOp::MirOp { op: BinOp::Add, neg: false },
+                    sym::atomic_xsub => AtomicRmwOp::MirOp { op: BinOp::Sub, neg: false },
+                    sym::atomic_min => AtomicRmwOp::Min,
+                    sym::atomic_umin => AtomicRmwOp::Min,
+                    sym::atomic_max => AtomicRmwOp::Max,
+                    sym::atomic_umax => AtomicRmwOp::Max,
+                    sym::atomic_xchg => AtomicRmwOp::Swap,
+                    _ => unreachable!(),
+                };
+
+                let res = M::atomic_rmw(self, &place, op, &operand, ord)?;
+                self.write_scalar(res, dest)?;
+            }
+            sym::atomic_cxchg | sym::atomic_cxchgweak => {
+                let success_ord = get_ord_at(1);
+                let failure_ord = get_ord_at(2);
+                let [ptr, expected_old, new] = args else {
+                    span_bug!(self.cur_span(), "invalid `{intrinsic_name}` call")
+                };
+
+                let place = self.deref_pointer(ptr)?;
+                let expected_old = self.read_immediate(expected_old)?;
+                let new = self.read_immediate(new)?;
+
+                let (actual_old, success) = M::atomic_compare_exchange(
+                    self,
+                    &place,
+                    &expected_old,
+                    &new,
+                    /* can_fail_spuriously */ intrinsic_name == sym::atomic_cxchgweak,
+                    success_ord,
+                    failure_ord,
+                )?;
+                let res = Immediate::ScalarPair(actual_old, Scalar::from_bool(success));
+                self.write_immediate(res, dest)?;
+            }
+            sym::atomic_fence | sym::atomic_singlethreadfence => {
+                let ord = get_ord_at(0);
+                let [] = args else {
+                    span_bug!(self.cur_span(), "invalid `{intrinsic_name}` call")
+                };
+
+                M::atomic_fence(self, ord, intrinsic_name == sym::atomic_singlethreadfence)?;
+            }
+
+            // Unsupported intrinsic: skip the return_to_block below.
+            _ => return interp_ok(false),
+        }
+
+        trace!("{:?}", self.dump_place(&dest.clone().into()));
+        self.return_to_block(ret)?;
+        interp_ok(true)
+    }
+}

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -9,16 +9,17 @@ use std::hash::Hash;
 use rustc_abi::{Align, Size};
 use rustc_apfloat::{Float, FloatConvert};
 use rustc_middle::query::TyCtxtAt;
-use rustc_middle::ty::Ty;
 use rustc_middle::ty::layout::TyAndLayout;
+use rustc_middle::ty::{AtomicOrdering, Ty};
 use rustc_middle::{mir, ty};
 use rustc_span::def_id::DefId;
 use rustc_target::callconv::FnAbi;
 
 use super::{
-    AllocBytes, AllocId, AllocKind, AllocRange, Allocation, CTFE_ALLOC_SALT, ConstAllocation,
-    CtfeProvenance, EnteredTraceSpan, FnArg, Frame, ImmTy, InterpCx, InterpResult, MPlaceTy,
-    MemoryKind, Misalignment, OpTy, PlaceTy, Pointer, Provenance, RangeSet, interp_ok, throw_unsup,
+    AllocBytes, AllocId, AllocKind, AllocRange, Allocation, AtomicRmwOp, CTFE_ALLOC_SALT,
+    ConstAllocation, CtfeProvenance, EnteredTraceSpan, FnArg, Frame, ImmTy, InterpCx, InterpResult,
+    MPlaceTy, MemoryKind, Misalignment, OpTy, PlaceTy, Pointer, Provenance, RangeSet, Scalar,
+    interp_ok, throw_unsup,
 };
 
 /// Data returned by [`Machine::after_stack_pop`], and consumed by
@@ -281,8 +282,6 @@ pub trait Machine<'tcx>: Sized {
     ) -> InterpResult<'tcx>;
 
     /// Called for all binary operations where the LHS has pointer type.
-    ///
-    /// Returns a (value, overflowed) pair if the operation succeeded
     fn binary_ptr_op(
         ecx: &InterpCx<'tcx, Self>,
         bin_op: mir::BinOp,
@@ -315,7 +314,46 @@ pub trait Machine<'tcx>: Sized {
     }
 
     /// Determines whether the `fmuladd` intrinsics fuse the multiply-add or use separate operations.
-    fn float_fuse_mul_add(_ecx: &InterpCx<'tcx, Self>) -> bool;
+    fn float_fuse_mul_add(ecx: &InterpCx<'tcx, Self>) -> bool;
+
+    fn atomic_load(
+        ecx: &InterpCx<'tcx, Self>,
+        place: &MPlaceTy<'tcx, Self::Provenance>,
+        ordering: AtomicOrdering,
+    ) -> InterpResult<'tcx, Scalar<Self::Provenance>>;
+
+    fn atomic_store(
+        ecx: &mut InterpCx<'tcx, Self>,
+        place: &MPlaceTy<'tcx, Self::Provenance>,
+        val: &ImmTy<'tcx, Self::Provenance>,
+        ordering: AtomicOrdering,
+    ) -> InterpResult<'tcx>;
+
+    /// Returns the old value.
+    fn atomic_rmw(
+        ecx: &mut InterpCx<'tcx, Self>,
+        place: &MPlaceTy<'tcx, Self::Provenance>,
+        op: AtomicRmwOp,
+        operand: &ImmTy<'tcx, Self::Provenance>,
+        ordering: AtomicOrdering,
+    ) -> InterpResult<'tcx, Scalar<Self::Provenance>>;
+
+    /// Returns a pair of the old value and a boolean indicating whether the update happened.
+    fn atomic_compare_exchange(
+        ecx: &mut InterpCx<'tcx, Self>,
+        place: &MPlaceTy<'tcx, Self::Provenance>,
+        expected_old: &ImmTy<'tcx, Self::Provenance>,
+        new: &ImmTy<'tcx, Self::Provenance>,
+        can_fail_spuriously: bool,
+        success_ordering: AtomicOrdering,
+        failure_ordering: AtomicOrdering,
+    ) -> InterpResult<'tcx, (Scalar<Self::Provenance>, bool)>;
+
+    fn atomic_fence(
+        ecx: &InterpCx<'tcx, Self>,
+        ordering: AtomicOrdering,
+        singlethread: bool,
+    ) -> InterpResult<'tcx>;
 
     /// Called before a basic block terminator is executed.
     #[inline]
@@ -708,6 +746,70 @@ pub macro compile_time_machine(<$tcx: lifetime>) {
         true
     }
 
+    #[inline(always)]
+    fn atomic_load(
+        ecx: &InterpCx<$tcx, Self>,
+        place: &MPlaceTy<$tcx, Self::Provenance>,
+        _ordering: AtomicOrdering,
+    ) -> InterpResult<$tcx, Scalar<Self::Provenance>> {
+        // Compile-time machines are single-threaded so this is like a regular load.
+        ecx.read_scalar(place)
+    }
+
+    #[inline(always)]
+    fn atomic_store(
+        ecx: &mut InterpCx<$tcx, Self>,
+        place: &MPlaceTy<$tcx, Self::Provenance>,
+        val: &ImmTy<$tcx, Self::Provenance>,
+        _ordering: AtomicOrdering,
+    ) -> InterpResult<$tcx> {
+        // Compile-time machines are single-threaded so this is like a regular store.
+        ecx.write_scalar(val.to_scalar(), place)
+    }
+
+    fn atomic_rmw(
+        ecx: &mut InterpCx<$tcx, Self>,
+        place: &MPlaceTy<$tcx, Self::Provenance>,
+        op: AtomicRmwOp,
+        operand: &ImmTy<$tcx, Self::Provenance>,
+        _ordering: AtomicOrdering,
+    ) -> InterpResult<$tcx, Scalar<Self::Provenance>> {
+        // Compile-time machines are single-threaded so we ignore the ordering.
+        let old_val = ecx.read_immediate(place)?;
+        let new_val = ecx.atomic_rmw_op(op, &old_val, operand)?;
+        ecx.write_immediate(*new_val, place)?;
+        interp_ok(old_val.to_scalar())
+    }
+
+    fn atomic_compare_exchange(
+        ecx: &mut InterpCx<$tcx, Self>,
+        place: &MPlaceTy<$tcx, Self::Provenance>,
+        expected_old: &ImmTy<$tcx, Self::Provenance>,
+        new: &ImmTy<$tcx, Self::Provenance>,
+        _can_fail_spuriously: bool,
+        _success_ordering: AtomicOrdering,
+        _failure_ordering: AtomicOrdering,
+    ) -> InterpResult<$tcx, (Scalar<Self::Provenance>, bool)> {
+        // Compile-time machines are single-threaded so we ignore the ordering.
+        // They are also deterministic so we do not fail spuriously.
+        let actual_old = ecx.read_immediate(place)?;
+        let eq = ecx.binary_op(mir::BinOp::Eq, &actual_old, expected_old)?.to_scalar().to_bool()?;
+        if eq {
+            ecx.write_immediate(**new, place)?;
+        }
+        interp_ok((actual_old.to_scalar(), eq))
+    }
+
+    #[inline(always)]
+    fn atomic_fence(
+        _ecx: &InterpCx<$tcx, Self>,
+        _ordering: AtomicOrdering,
+        _singlethread: bool,
+    ) -> InterpResult<$tcx> {
+        // Compile-time machines are single-threaded so this is a NOP.
+        interp_ok(())
+    }
+
     #[inline(always)]
     fn adjust_global_allocation<'b>(
         _ecx: &InterpCx<$tcx, Self>,

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -35,6 +35,7 @@ pub use self::machine::{
 pub use self::memory::{AllocInfo, AllocKind, AllocRef, AllocRefMut, FnVal, Memory, MemoryKind};
 use self::operand::Operand;
 pub use self::operand::{ImmTy, Immediate, OpTy};
+pub use self::operator::AtomicRmwOp;
 pub use self::place::{MPlaceTy, MemPlaceMeta, PlaceTy, Writeable};
 use self::place::{MemPlace, Place};
 pub use self::projection::{OffsetMode, Projectable};

compiler/rustc_const_eval/src/interpret/operator.rs

@@ -10,6 +10,19 @@ use tracing::trace;
 
 use super::{ImmTy, InterpCx, Machine, MemPlaceMeta, interp_ok, throw_ub};
 
+/// Describes an atomic RMW operation.
+pub enum AtomicRmwOp {
+    MirOp {
+        op: mir::BinOp,
+        /// Indicates whether the result of the operation should be negated (`UnOp::Not`, must be a
+        /// boolean/integer-typed operation).
+        neg: bool,
+    },
+    Max,
+    Min,
+    Swap,
+}
+
 impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     fn three_way_compare<T: Ord>(&self, lhs: T, rhs: T) -> ImmTy<'tcx, M::Provenance> {
         let res = Ord::cmp(&lhs, &rhs);
@@ -504,4 +517,27 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             }
         }
     }
+
+    pub fn atomic_rmw_op(
+        &self,
+        op: AtomicRmwOp,
+        left: &ImmTy<'tcx, M::Provenance>,
+        right: &ImmTy<'tcx, M::Provenance>,
+    ) -> InterpResult<'tcx, ImmTy<'tcx, M::Provenance>> {
+        interp_ok(match op {
+            AtomicRmwOp::MirOp { op, neg } => {
+                let val = self.binary_op(op, &left, right)?;
+                if neg { self.unary_op(mir::UnOp::Not, &val)? } else { val }
+            }
+            AtomicRmwOp::Max => {
+                let lt = self.binary_op(mir::BinOp::Lt, &left, right)?.to_scalar().to_bool()?;
+                if lt { right } else { &left }.clone()
+            }
+            AtomicRmwOp::Min => {
+                let lt = self.binary_op(mir::BinOp::Lt, &left, right)?.to_scalar().to_bool()?;
+                if lt { &left } else { right }.clone()
+            }
+            AtomicRmwOp::Swap => right.clone(),
+        })
+    }
 }