cubecl-llvm 0.11.0-pre.4

LLVM compiler for CubeCL
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

//! Atomics with explicit memory ordering.

use crate::prelude::*;
use cubecl_core::{
    ir::dialect::{
        atomic::AtomicLoadOp,
        memory::LoadOp,
        plane::{AtomicUniformLoadOp, UniformLoadOp},
        synchronization::{SyncOp, SyncScope, SyncScopeAttr},
    },
    prelude::*,
};

#[cube_op(name = "cpu.ordered_atomic_load")]
#[result_ty(argument)]
pub struct OrderedAtomicLoadOp {
    pub ptr: Value,
    pub ordering: AtomicOrderingAttr,
}

#[cube_op(name = "cpu.ordered_atomic_store")]
#[result_ty(none)]
pub struct OrderedAtomicStoreOp {
    pub ptr: Value,
    pub value: Value,
    pub ordering: AtomicOrderingAttr,
}

#[cube_op(name = "cpu.ordered_atomic_fetch_add")]
#[result_ty(same_as = value)]
pub struct OrderedAtomicFetchAddOp {
    pub ptr: Value,
    pub value: Value,
    pub ordering: AtomicOrderingAttr,
}

/// Acquire load.
#[cube]
pub fn atomic_load_acquire(atomic: &Atomic<u32>) -> u32 {
    intrinsic!(|scope| {
        let ptr = atomic.value(scope);
        let ty = u32::__expand_as_type(scope);
        let op = OrderedAtomicLoadOp::new(scope.ctx_mut(), ty, ptr, AtomicOrderingAttr::Acquire);
        scope.register_with_result(&op).into()
    })
}

/// Release store.
#[cube]
pub fn atomic_store_release(atomic: &Atomic<u32>, value: u32) {
    intrinsic!(|scope| {
        let ptr = atomic.value(scope);
        let value = value.read_value(scope);
        let op =
            OrderedAtomicStoreOp::new(scope.ctx_mut(), ptr, value, AtomicOrderingAttr::Release);
        scope.register(&op);
    })
}

/// Acquire-release addition. Returns the previous value.
#[cube]
pub fn atomic_fetch_add_acq_rel(atomic: &Atomic<u32>, value: u32) -> u32 {
    intrinsic!(|scope| {
        let ptr = atomic.value(scope);
        let value = value.read_value(scope);
        let op =
            OrderedAtomicFetchAddOp::new(scope.ctx_mut(), ptr, value, AtomicOrderingAttr::AcqRel);
        scope.register_with_result(&op).into()
    })
}

#[op_interface_impl]
impl LowerOp for UniformLoadOp {
    fn lower(&self, scope: &Scope) -> Vec<Value> {
        scope.register(&SyncOp::new(
            scope.ctx_mut(),
            SyncScopeAttr::new(SyncScope::Cube),
        ));
        let ptr = self.ptr(scope.ctx());
        vec![scope.register_with_result(&LoadOp::new(scope.ctx_mut(), ptr))]
    }
}

#[op_interface_impl]
impl LowerOp for AtomicUniformLoadOp {
    fn lower(&self, scope: &Scope) -> Vec<Value> {
        scope.register(&SyncOp::new(
            scope.ctx_mut(),
            SyncScopeAttr::new(SyncScope::Cube),
        ));
        let ptr = self.ptr(scope.ctx());
        vec![scope.register_with_result(&AtomicLoadOp::new(scope.ctx_mut(), ptr))]
    }
}