Skip to content

Commit 631a0ef

Browse files
committed
Added support for most atomic operations
1 parent 843b714 commit 631a0ef

File tree

3 files changed

+86
-12
lines changed

3 files changed

+86
-12
lines changed

crates/rustc_codegen_nvvm/src/builder.rs

Lines changed: 59 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use libc::{c_char, c_uint};
66
use rustc_abi as abi;
77
use rustc_abi::{AddressSpace, Align, HasDataLayout, Size, TargetDataLayout, WrappingRange};
88
use rustc_codegen_ssa::MemFlags;
9-
use rustc_codegen_ssa::common::{IntPredicate, RealPredicate, TypeKind};
9+
use rustc_codegen_ssa::common::{IntPredicate, RealPredicate, TypeKind,AtomicRmwBinOp};
1010
use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
1111
use rustc_codegen_ssa::mir::place::PlaceRef;
1212
use rustc_codegen_ssa::traits::*;
@@ -1195,13 +1195,65 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
11951195
}
11961196
fn atomic_rmw(
11971197
&mut self,
1198-
_op: rustc_codegen_ssa::common::AtomicRmwBinOp,
1199-
_dst: &'ll Value,
1200-
_src: &'ll Value,
1201-
_order: AtomicOrdering,
1198+
op: AtomicRmwBinOp,
1199+
dst: &'ll Value,
1200+
src: &'ll Value,
1201+
order: AtomicOrdering,
12021202
) -> &'ll Value {
1203-
// see cmpxchg comment
1204-
self.fatal("atomic rmw is not supported")
1203+
if matches!(op,AtomicRmwBinOp::AtomicNand){
1204+
self.fatal("Atomic NAND not supported yet!")
1205+
}
1206+
self.atomic_op(
1207+
dst,
1208+
|builder, dst| {
1209+
// We are in a supported address space - just use ordinary atomics
1210+
unsafe {
1211+
llvm::LLVMBuildAtomicRMW(
1212+
builder.llbuilder,
1213+
op,
1214+
dst,
1215+
src,
1216+
crate::llvm::AtomicOrdering::from_generic( order),
1217+
0,
1218+
)
1219+
}
1220+
},
1221+
|builder, dst| {
1222+
// Local space is only accessible to the current thread.
1223+
// So, there are no synchronization issues, and we can emulate it using a simple load / compare / store.
1224+
let load:&'ll Value = unsafe{ llvm::LLVMBuildLoad(builder.llbuilder, dst, UNNAMED) };
1225+
let next_val = match op{
1226+
AtomicRmwBinOp::AtomicXchg => src,
1227+
AtomicRmwBinOp::AtomicAdd => builder.add(load, src),
1228+
AtomicRmwBinOp::AtomicSub => builder.sub(load, src),
1229+
AtomicRmwBinOp::AtomicAnd => builder.and(load, src),
1230+
AtomicRmwBinOp::AtomicNand => {
1231+
let and = builder.and(load, src);
1232+
builder.not(and)
1233+
},
1234+
AtomicRmwBinOp::AtomicOr => builder.or(load, src),
1235+
AtomicRmwBinOp::AtomicXor => builder.xor(load, src),
1236+
AtomicRmwBinOp::AtomicMax => {
1237+
let is_src_bigger = builder.icmp(IntPredicate::IntSGT, src, load);
1238+
builder.select(is_src_bigger,src,load)
1239+
}
1240+
AtomicRmwBinOp::AtomicMin => {
1241+
let is_src_smaller = builder.icmp(IntPredicate::IntSLT, src, load);
1242+
builder.select(is_src_smaller,src,load)
1243+
}
1244+
AtomicRmwBinOp::AtomicUMax => {
1245+
let is_src_bigger = builder.icmp(IntPredicate::IntUGT, src, load);
1246+
builder.select(is_src_bigger,src,load)
1247+
},
1248+
AtomicRmwBinOp::AtomicUMin => {
1249+
let is_src_smaller = builder.icmp(IntPredicate::IntULT, src, load);
1250+
builder.select(is_src_smaller,src,load)
1251+
}
1252+
};
1253+
unsafe { llvm::LLVMBuildStore(builder.llbuilder, next_val, dst)};
1254+
load
1255+
},
1256+
)
12051257
}
12061258

12071259
fn atomic_fence(

crates/rustc_codegen_nvvm/src/nvvm.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,9 +125,10 @@ pub fn codegen_bitcode_modules(
125125
let res = match prog.compile(&args.nvvm_options) {
126126
Ok(b) => b,
127127
Err(error) => {
128+
let log = prog.compiler_log().unwrap().unwrap_or_default();
128129
// this should never happen, if it does, something went really bad or its a bug on libnvvm's end
129130
panic!(
130-
"libnvvm returned an error that was not previously caught by the verifier: {error:?}"
131+
"libnvvm returned an error that was not previously caught by the verifier: {error:?} {log:?}"
131132
);
132133
}
133134
};
Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// Test CUDA atomic operations compile correctly
22
// build-pass
3-
// compile-flags: -Z verify-llvm-ir
3+
// compile-flags: -Z verify-llvm-ir
44
use core::sync::atomic::{AtomicUsize,Ordering};
55

66
use cuda_std::atomic::{
@@ -10,14 +10,35 @@ use cuda_std::kernel;
1010
static GLOBAL:AtomicUsize = AtomicUsize::new(0);
1111
#[kernel]
1212
pub unsafe fn test_cuda_atomic_floats() {
13-
let local = AtomicUsize::new(0);
13+
let local = AtomicUsize::new(0);
1414
// `compare_exchange` should succeed
1515
local.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
1616
// `compare_exchange` should fail
1717
local.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
18-
// `compare_exchange` should succeed
18+
// `compare_exchange` should succeed
1919
GLOBAL.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
2020
// `compare_exchange` should fail
2121
GLOBAL.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed);
22-
22+
// Ops
23+
local.swap(1, Ordering::Relaxed);
24+
GLOBAL.swap(1, Ordering::Relaxed);
25+
local.fetch_add(1, Ordering::Relaxed);
26+
GLOBAL.fetch_add(1, Ordering::Relaxed);
27+
local.fetch_sub(1, Ordering::Relaxed);
28+
GLOBAL.fetch_sub(1, Ordering::Relaxed);
29+
local.fetch_and(1, Ordering::Relaxed);
30+
GLOBAL.fetch_and(1, Ordering::Relaxed);
31+
local.fetch_and(1, Ordering::Relaxed);
32+
GLOBAL.fetch_and(1, Ordering::Relaxed);
33+
local.fetch_or(1, Ordering::Relaxed);
34+
GLOBAL.fetch_or(1, Ordering::Relaxed);
35+
local.fetch_xor(1, Ordering::Relaxed);
36+
GLOBAL.fetch_xor(1, Ordering::Relaxed);
37+
local.fetch_max(1, Ordering::Relaxed);
38+
GLOBAL.fetch_max(1, Ordering::Relaxed);
39+
local.fetch_min(1, Ordering::Relaxed);
40+
GLOBAL.fetch_min(1, Ordering::Relaxed);
41+
// Atomic NAND is not supported quite yet
42+
//local.fetch_nand(1, Ordering::Relaxed);
43+
//GLOBAL.fetch_nand(1, Ordering::Relaxed);
2344
}

0 commit comments

Comments
 (0)