cubecl_cpp/cuda/
atomic.rs

1//! CUDA's C++ atomic APIs are a mess of inconsistency. Old-style APIs use weird types, newer
2//! `std::atomic` APIs don't support many types. So for ops where type support is complex, we lower
3//! the atomics directly to PTX. This allows a consistent API across all the different type and
4//! vectorization options, and significantly extends the interface that can be accessed from C++.
5
6use cubecl_core::{
7    frontend::reinterpret_value,
8    ir::{Scope, dialect::atomic::*, interfaces::TypedExt, prelude::*, types::VectorType},
9};
10use pliron::{
11    builtin::types::{IntegerType, Signedness},
12    value::Value,
13};
14
15use crate::{
16    cuda::{
17        packed_ops::packable,
18        ptx::InlinePtxOp,
19        ty::{BFloat16x2Type, Float16x2Type},
20    },
21    shared::{lowering::LowerOpAfterUnroll, ty::TypedExtCPP},
22    target::Cuda,
23};
24
25fn atom_vec(ctx: &Context, val: impl Typed) -> &'static str {
26    match val.vector_size(ctx) {
27        1 => "",
28        2 => ".v2",
29        4 => ".v4",
30        8 => ".v8",
31        _ => unreachable!(),
32    }
33}
34
35fn atom_ftz(ctx: &Context, val: impl Typed) -> &'static str {
36    if val.is_half(ctx) || val.is_half2(ctx) {
37        ".noftz"
38    } else {
39        ""
40    }
41}
42
43// Signed only matters for cmp, addition is signless. And it's not supported for `s64`, only `u64`
44fn atom_ty(ctx: &Context, val: impl Typed) -> &'static str {
45    let scalar_ty = val.scalar_ty(ctx);
46    if scalar_ty.is_float64(ctx) {
47        "f64"
48    } else if scalar_ty.is_float32(ctx) {
49        "f32"
50    } else if scalar_ty.is_float16(ctx) {
51        "f16"
52    } else if scalar_ty.deref(ctx).is::<Float16x2Type>() {
53        "f16x2"
54    } else if scalar_ty.is_bfloat16(ctx) {
55        "bf16"
56    } else if scalar_ty.deref(ctx).is::<BFloat16x2Type>() {
57        "bf16x2"
58    } else if scalar_ty.is_int_of_width(ctx, 64) {
59        "u64"
60    } else if scalar_ty.is_int_of_width(ctx, 32) {
61        "u32"
62    } else {
63        panic!("Unsupported type")
64    }
65}
66
67fn atom_ty_cmp(ctx: &Context, val: impl Typed) -> &'static str {
68    let scalar_ty = val.scalar_ty(ctx);
69    if scalar_ty.is_int_of_width(ctx, 64) && scalar_ty.is_signed_int(ctx) {
70        "s64"
71    } else if scalar_ty.is_int_of_width(ctx, 32) && scalar_ty.is_signed_int(ctx) {
72        "s32"
73    } else {
74        atom_ty(ctx, val)
75    }
76}
77
78// Reinterpet f16 etc
79fn as_registers(scope: &Scope, val: Value) -> Value {
80    let vec = val.vector_size(scope.ctx());
81    let u16 = IntegerType::get(scope.ctx(), 16, Signedness::Unsigned).to_handle();
82    let u32 = IntegerType::get(scope.ctx(), 32, Signedness::Unsigned).to_handle();
83    if vec > 1 && val.is_half(scope.ctx()) {
84        let vec_ty = VectorType::get(scope.ctx(), u16, vec);
85        reinterpret_value(scope, val, vec_ty.to_handle())
86    } else if vec > 1 && val.is_half2(scope.ctx()) {
87        let vec_ty = VectorType::get(scope.ctx(), u32, vec);
88        reinterpret_value(scope, val, vec_ty.to_handle())
89    } else if val.is_half(scope.ctx()) {
90        reinterpret_value(scope, val, u16)
91    } else if val.is_half2(scope.ctx()) {
92        reinterpret_value(scope, val, u32)
93    } else {
94        val
95    }
96}
97
98macro_rules! atomic_binop {
99    ($ty: ty, $op: literal, $atom_ty: ident) => {
100        #[op_interface_impl]
101        impl LowerOpAfterUnroll<Cuda> for $ty {
102            fn lower(&self, scope: &Scope) -> Vec<Value> {
103                let ctx = scope.ctx_mut();
104                let ptr = self.ptr(ctx);
105                let value = self.value(ctx);
106                let out_ty = self.get_result(ctx).get_type(ctx);
107
108                let vec = atom_vec(ctx, value);
109                let ftz = atom_ftz(ctx, value);
110                let ty = $atom_ty(ctx, value);
111                let value = as_registers(scope, value);
112
113                let ptx = format!("atom.relaxed.{}{ftz}{vec}.{ty} $0, [$1], $2;", $op);
114                let op = InlinePtxOp::new_volatile(
115                    ctx,
116                    Some(value.get_type(ctx)),
117                    ptx,
118                    vec![ptr, value],
119                );
120                scope.register(&op);
121                vec![reinterpret_value(scope, op.result(ctx).unwrap(), out_ty)]
122            }
123        }
124    };
125}
126
127packable!(AtomicFAddOp);
128packable!(AtomicFMinOp);
129packable!(AtomicFMaxOp);
130
131atomic_binop!(AtomicIAddOp, "add", atom_ty);
132atomic_binop!(AtomicFAddOp, "add", atom_ty);
133atomic_binop!(AtomicSMinOp, "min", atom_ty_cmp);
134atomic_binop!(AtomicUMinOp, "min", atom_ty_cmp);
135atomic_binop!(AtomicFMinOp, "min", atom_ty_cmp);
136atomic_binop!(AtomicSMaxOp, "max", atom_ty_cmp);
137atomic_binop!(AtomicUMaxOp, "max", atom_ty_cmp);
138atomic_binop!(AtomicFMaxOp, "max", atom_ty_cmp);
cubecl_cpp/cuda/atomic.rs

cubecl_cpp/cuda/
atomic.rs
