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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use crate::prelude::*;
use cubecl_core::ir::dialect::atomic::*;

macro_rules! lower_atomic_rmw {
    ($cube_op:ty => $pred:expr) => {
        #[op_interface_impl]
        impl ToLLVMDialect for $cube_op {
            fn rewrite(
                &self,
                ctx: &mut Context,
                rewriter: &mut DialectConversionRewriter,
                _operands_info: &OperandsInfo,
            ) -> Result<()> {
                let ptr = self.ptr(ctx);
                let val = self.value(ctx);
                let kind = $pred;
                let ordering = AtomicOrderingAttr::Monotonic;
                let sync_scope = SyncScopeAttr::System;

                let op = llvm::AtomicRmwOp::new(ctx, ptr, val, kind, ordering, sync_scope);
                rewriter.insert_op(ctx, &op);
                rewriter.replace_operation_with_values(
                    ctx,
                    self.get_operation(),
                    vec![op.get_result(ctx)],
                );
                Ok(())
            }
        }
    };
}

lower_atomic_rmw!(AtomicExchangeOp => AtomicRmwKindAttr::Xchg);
lower_atomic_rmw!(AtomicIAddOp => AtomicRmwKindAttr::Add);
lower_atomic_rmw!(AtomicFAddOp => AtomicRmwKindAttr::FAdd);
lower_atomic_rmw!(AtomicISubOp => AtomicRmwKindAttr::Sub);
lower_atomic_rmw!(AtomicFSubOp => AtomicRmwKindAttr::FSub);
lower_atomic_rmw!(AtomicSMinOp => AtomicRmwKindAttr::Min);
lower_atomic_rmw!(AtomicUMinOp => AtomicRmwKindAttr::UMin);
lower_atomic_rmw!(AtomicFMinOp => AtomicRmwKindAttr::FMin);
lower_atomic_rmw!(AtomicSMaxOp => AtomicRmwKindAttr::Max);
lower_atomic_rmw!(AtomicUMaxOp => AtomicRmwKindAttr::UMax);
lower_atomic_rmw!(AtomicFMaxOp => AtomicRmwKindAttr::FMax);
lower_atomic_rmw!(AtomicAndOp => AtomicRmwKindAttr::And);
lower_atomic_rmw!(AtomicOrOp => AtomicRmwKindAttr::Or);
lower_atomic_rmw!(AtomicXorOp => AtomicRmwKindAttr::Xor);

#[op_interface_impl]
impl ToLLVMDialect for AtomicLoadOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        operands_info: &OperandsInfo,
    ) -> Result<()> {
        let ptr = self.ptr(ctx);
        let ordering = AtomicOrderingAttr::Monotonic;
        let result = self.get_result(ctx);
        let res_cube_ty = operands_info
            .lookup_most_recent_type(result)
            .unwrap_or_else(|| result.get_type(ctx));
        let align = scalar_alignment(ctx, res_cube_ty);
        let res_ty = cube_type_to_llvm(ctx, res_cube_ty);

        let sync_scope = SyncScopeAttr::System;
        let op = llvm::AtomicLoadOp::new(ctx, ptr, res_ty, ordering, sync_scope);
        op.set_alignment(ctx, align);
        rewriter.insert_op(ctx, &op);
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![op.get_result(ctx)]);
        Ok(())
    }
}

#[op_interface_impl]
impl ToLLVMDialect for AtomicStoreOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        operands_info: &OperandsInfo,
    ) -> Result<()> {
        let value = self.value(ctx);
        let ptr = self.ptr(ctx);
        let value_cube_ty = operands_info
            .lookup_most_recent_type(value)
            .unwrap_or_else(|| value.get_type(ctx));
        let align = scalar_alignment(ctx, value_cube_ty);
        let ordering = AtomicOrderingAttr::Monotonic;
        let sync_scope = SyncScopeAttr::System;

        let store = llvm::AtomicStoreOp::new(ctx, value, ptr, ordering, sync_scope);
        store.set_alignment(ctx, align);
        rewriter.insert_op(ctx, &store);
        rewriter.replace_operation(ctx, self.get_operation(), store.get_operation());
        Ok(())
    }
}

#[op_interface_impl]
impl ToLLVMDialect for AtomicCompareExchangeWeakOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let ptr = self.ptr(ctx);
        let cmp = self.cmp(ctx);
        let new_val = self.value(ctx);
        let before = AtomicOrderingAttr::Monotonic;
        let after = AtomicOrderingAttr::Monotonic;

        let sync_scope = SyncScopeAttr::System;
        let op = llvm::AtomicCmpxchgOp::new(ctx, ptr, cmp, new_val, before, after, sync_scope);
        rewriter.insert_op(ctx, &op);
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![op.get_result(ctx)]);
        Ok(())
    }
}