Add support for standard Rust atomic compareexchange intrinsic.

Author: FractalFir

crates/rustc_codegen_nvvm/src/builder.rs

@@ -1134,16 +1134,64 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     // Atomic Operations
     fn atomic_cmpxchg(
         &mut self,
-        _dst: &'ll Value,
-        _cmp: &'ll Value,
-        _src: &'ll Value,
-        _order: AtomicOrdering,
-        _failure_order: AtomicOrdering,
-        _weak: bool,
+        dst: &'ll Value,
+        cmp: &'ll Value,
+        src: &'ll Value,
+        order: AtomicOrdering,
+        failure_order: AtomicOrdering,
+        weak: bool,
     ) -> (&'ll Value, &'ll Value) {
-        // allowed but only for some things and with restrictions
-        // https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#cmpxchg-instruction
-        self.fatal("atomic cmpxchg is not supported")
+        // LLVM verifier rejects cases where the `failure_order` is stronger than `order`
+        match (order,failure_order){
+            (AtomicOrdering::SeqCst, _)=>(),
+            (_, AtomicOrdering::Relaxed)=>(),
+            (AtomicOrdering::Release, AtomicOrdering::Release) |  (AtomicOrdering::Release, AtomicOrdering::Acquire) | (AtomicOrdering::Acquire, AtomicOrdering::Acquire)=>(),
+            (AtomicOrdering::AcqRel,AtomicOrdering::Acquire)  => (),
+            (AtomicOrdering::Relaxed, _) | (_, AtomicOrdering::Release | AtomicOrdering::AcqRel | AtomicOrdering::SeqCst)=>{
+                // Invalid cmpxchg - `failure_order` is stronger than `order`! So, we abort.
+                self.abort();
+                return (self.const_undef(self.val_ty(cmp)),self.const_undef(self.type_i1()));
+            }
+        };
+        let res = self.atomic_op(
+            dst,
+             |builder, dst| {
+                // We are in a supported address space - just use ordinary atomics
+                unsafe {
+                    llvm::LLVMRustBuildAtomicCmpXchg(
+                        builder.llbuilder,
+                        dst,
+                        cmp,
+                        src,
+                       crate::llvm::AtomicOrdering::from_generic( order),
+                        crate::llvm::AtomicOrdering::from_generic(failure_order),
+                        weak as u32,
+                    )
+                }
+            },
+         |builder, dst| {
+                // Local space is only accessible to the current thread.
+                // So, there are no synchronization issues, and we can emulate it using a simple load / compare / store. 
+                let load:&'ll Value = unsafe{ llvm::LLVMBuildLoad(builder.llbuilder, dst, UNNAMED) };
+                let compare = builder.icmp(IntPredicate::IntEQ, load, cmp);
+                // We can do something smart & branchless here:
+                // We select either the current value(if the comparison fails), or a new value. 
+                // We then *undconditionally* write that back to local memory(which is very, very cheap).
+                // TODO: measure if this has a positive impact, or if we should just use more blocks, and conditional writes.
+                let value = builder.select(compare, src, load);
+                unsafe { llvm::LLVMBuildStore(builder.llbuilder, value, dst)};
+                let res_type = builder.type_struct(&[builder.val_ty(cmp),builder.type_ix(1)], false);
+                // We pack the result, to match the behaviour of proper atomics / emulated thread-local atomics. 
+                let res = builder.const_undef(res_type);
+                let res = builder.insert_value(res, load, 0);
+                let res = builder.insert_value(res, compare, 1);
+                res
+            },
+        );
+        // Unpack the result
+        let val = self.extract_value(res, 0);
+        let success = self.extract_value(res, 1);
+        (val, success)
     }
     fn atomic_rmw(
         &mut self,
@@ -1609,3 +1657,98 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         }
     }
 }
+impl<'ll, 'tcx, 'a> Builder<'a, 'll, 'tcx> {
+    /// Implements a standard atomic, using LLVM intrinsics(in `atomic_supported`, if `dst` is in a supported address space)
+    /// or emulation(with `emulate_local`, if `dst` points to a thread-local address space).
+    fn atomic_op(
+        &mut self,
+        dst: &'ll Value,
+        atomic_supported: impl FnOnce(&mut Builder<'a,'ll ,'tcx>, &'ll Value) -> &'ll Value,
+        emulate_local: impl FnOnce(&mut Builder<'a,'ll ,'tcx>, &'ll Value) -> &'ll Value,
+    ) -> &'ll Value {
+        // (FractalFir) Atomics in CUDA have some limitations, and we have to work around them.
+        // For example, they are restricted in what address space they operate on.
+        // CUDA has 4 address spaces(and a generic one, which is an union of all of those).
+        // An atomic instruction can soundly operate on:
+        // 1. The global address space
+        // 2. The shared(cluster) address space.
+        // It can't operate on:
+        // 1. The const address space(atomics on consts are UB anyway)
+        // 2. The thread address space(which should be only accessible to 1 thread, anyway?)
+        // So, we do the following:
+        // 1. Check if the pointer is in one of the address spaces atomics support.
+        //  a) if so, we perform an atomic operation
+        // 2. Check if the pointer is in the thread-local address space. If it is, we use non-atomic ops here,
+        // **ASSUMING** only the current thread can access thread-local memory. (FIXME: is this sound?)
+        // 3. If the pointer is not in a supported address space, and is not thread-local, then we bail, and trap.
+
+        // We check if the `dst` pointer is in the `global` address space.
+        let (isspacep_global_ty, isspacep_global_fn) =
+            self.get_intrinsic("llvm.nvvm.isspacep.global");
+        let isspacep_global = self.call(
+            isspacep_global_ty,
+            None,
+            None,
+            isspacep_global_fn,
+            &[dst],
+            None,
+            None,
+        );
+        // We check if the `dst` pointer is in the `shared` address space.
+        let (isspacep_shared_ty, isspacep_shared_fn) =
+            self.get_intrinsic("llvm.nvvm.isspacep.shared");
+        let isspacep_shared = self.call(
+            isspacep_shared_ty,
+            None,
+            None,
+            isspacep_shared_fn,
+            &[dst],
+            None,
+            None,
+        );
+        // Combine those to check if we are in a supported address space.
+        let atomic_supported_addrspace = self.or(isspacep_shared, isspacep_global);
+        // We create 2 blocks here: one we branch to if atomic is in the right address space, and one we branch to otherwise.
+        let supported_bb = self.append_sibling_block("atomic_space_supported");
+        let unsupported_bb = self.append_sibling_block("atomic_space_unsupported");
+        self.cond_br(atomic_supported_addrspace, supported_bb, unsupported_bb);
+        //  We also create a "merge" block we will jump to, after the the atomic ops finish.
+        let merge_bb = self.append_sibling_block("atomic_op_done");
+        // Execute atomic op if supported, then jump to merge
+        self.switch_to_block(supported_bb);
+        let supported_res = atomic_supported(self, dst);
+        self.br(merge_bb);
+        // Check if the pointer is in the thread space. If so, we can emulate it.
+        self.switch_to_block(unsupported_bb);
+        let (isspacep_local_ty, isspacep_local_fn) = self.get_intrinsic("llvm.nvvm.isspacep.local");
+        let isspacep_local = self.call(
+            isspacep_local_ty,
+            None,
+            None,
+            isspacep_local_fn,
+            &[dst],
+            None,
+            None,
+        );
+        let local_bb = self.append_sibling_block("atomic_local_space");
+        let atomic_ub_bb = self.append_sibling_block("atomic_space_ub");
+        self.cond_br(isspacep_local, local_bb, atomic_ub_bb);
+        // The pointer is in the thread(local) space.
+        self.switch_to_block(local_bb);
+        let local_res = emulate_local(self, dst);
+        self.br(merge_bb);
+        // The pointer is neither in the supported address space, nor the local space.
+        // This is very likely UB. So, we trap here.
+        // TODO: should we print some kind of a message here? NVVM supports printf.
+        self.switch_to_block(atomic_ub_bb);
+        self.abort();
+        self.unreachable();
+        // Atomic is impl has finished, and we can now switch to the merge_bb
+        self.switch_to_block(merge_bb);
+        self.phi(
+            self.val_ty(local_res),
+            &[supported_res, local_res],
+            &[supported_bb, local_bb],
+        )
+    }
+}

crates/rustc_codegen_nvvm/src/ctx_intrinsics.rs

@@ -449,5 +449,11 @@ impl<'ll> CodegenCx<'ll, '_> {
             "__nv_ynf",
             fn(t_i32, t_f32) -> t_f32
         );
+        // Address space checks
+        ifn!(map, "llvm.nvvm.isspacep.const", fn(i8p) -> i1);
+        ifn!(map, "llvm.nvvm.isspacep.global", fn(i8p) -> i1);
+        ifn!(map, "llvm.nvvm.isspacep.local", fn(i8p) -> i1);
+        ifn!(map, "llvm.nvvm.isspacep.shared", fn(i8p) -> i1);
+  
     }
 }

crates/rustc_codegen_nvvm/src/llvm.rs

@@ -16,7 +16,7 @@
 // but likely will use in the future, so we ignore any unused functions
 // in case we need them in the future for things like debug info or LTO.
 #![allow(dead_code)]
-
+use rustc_codegen_ssa::common::AtomicRmwBinOp;
 use libc::{c_char, c_uint, c_void, size_t};
 use libc::{c_int, c_ulonglong};
 use std::ffi::{CStr, CString};
@@ -1947,4 +1947,50 @@ unsafe extern "C" {
     pub(crate) fn LLVMRustAddDereferenceableOrNullAttr(Fn: &Value, index: c_uint, bytes: u64);
 
     pub(crate) fn LLVMRustPositionBuilderAtStart<'a>(B: &Builder<'a>, BB: &'a BasicBlock);
+    // Atomics
+    pub fn LLVMRustBuildAtomicCmpXchg<'a>(
+        B: &Builder<'a>,
+        LHS: &Value,
+        CMP: &Value,
+        RHS: &Value,
+        Order: AtomicOrdering,
+        FailureOrder: AtomicOrdering,
+        Weak: Bool,
+    ) -> &'a Value;
+
+    pub fn LLVMBuildAtomicRMW<'a>(
+        B: &Builder<'a>,
+        Op: AtomicRmwBinOp,
+        LHS: &Value,
+        RHS: &Value,
+        Order: AtomicOrdering,
+        SingleThreaded: Bool,
+    ) -> &'a Value;
 }
+/// LLVMAtomicOrdering
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub(crate) enum AtomicOrdering {
+    #[allow(dead_code)]
+    NotAtomic = 0,
+    #[allow(dead_code)]
+    Unordered = 1,
+    Monotonic = 2,
+    // Consume = 3,  // Not specified yet.
+    Acquire = 4,
+    Release = 5,
+    AcquireRelease = 6,
+    SequentiallyConsistent = 7,
+}
+impl AtomicOrdering {
+    pub(crate) fn from_generic(ao: rustc_middle::ty::AtomicOrdering) -> Self {
+        use rustc_middle::ty::AtomicOrdering as Common;
+        match ao {
+            Common::Relaxed => Self::Monotonic,
+            Common::Acquire => Self::Acquire,
+            Common::Release => Self::Release,
+            Common::AcqRel => Self::AcquireRelease,
+            Common::SeqCst => Self::SequentiallyConsistent,
+        }
+    }
+}

tests/compiletests/ui/atomic/std_atomic_ops.rs

@@ -0,0 +1,23 @@
+// Test CUDA atomic operations compile correctly
+// build-pass
+// compile-flags: -Z verify-llvm-ir
+use core::sync::atomic::{AtomicUsize,Ordering};
+
+use cuda_std::atomic::{
+    AtomicF32, AtomicF64, BlockAtomicF32, BlockAtomicF64, SystemAtomicF32, SystemAtomicF64,
+};
+use cuda_std::kernel;
+static GLOBAL:AtomicUsize = AtomicUsize::new(0);
+#[kernel]
+pub unsafe fn test_cuda_atomic_floats() {
+      let local = AtomicUsize::new(0);
+    // `compare_exchange` should succeed
+    local.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
+    // `compare_exchange` should fail
+    local.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
+     // `compare_exchange` should succeed
+    GLOBAL.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
+    // `compare_exchange` should fail
+    GLOBAL.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
+
+}