cubecl-ir 0.11.0-pre.3

Intermediate representation for CubeCL
Documentation 
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

use cubecl_macros_internal::{cube_op, op_traits};

use crate::{CanMaterialize, dialect::ptr_value_ty, prelude::*};

macro_rules! atomic_binop {
    ($name: literal, $ty: ident) => {
        #[cube_op(name = $name)]
        #[result_ty(same_as = value)]
        #[op_traits(CanMaterialize)]
        pub struct $ty {
            #[operand(ptr_read, ptr_write)]
            pub ptr: Value,
            pub value: Value,
        }
    };
}

atomic_binop!("atomic.exchange", AtomicExchangeOp);
atomic_binop!("atomic.i_add", AtomicIAddOp);
atomic_binop!("atomic.f_add", AtomicFAddOp);
atomic_binop!("atomic.i_sub", AtomicISubOp);
atomic_binop!("atomic.f_sub", AtomicFSubOp);
atomic_binop!("atomic.s_min", AtomicSMinOp);
atomic_binop!("atomic.u_min", AtomicUMinOp);
atomic_binop!("atomic.f_min", AtomicFMinOp);
atomic_binop!("atomic.s_max", AtomicSMaxOp);
atomic_binop!("atomic.u_max", AtomicUMaxOp);
atomic_binop!("atomic.f_max", AtomicFMaxOp);
atomic_binop!("atomic.and", AtomicAndOp);
atomic_binop!("atomic.or", AtomicOrOp);
atomic_binop!("atomic.xor", AtomicXorOp);

#[cube_op(name = "atomic.load")]
#[result_ty(from_inputs = ptr_value_ty)]
#[op_traits(CanMaterialize)]
pub struct AtomicLoadOp {
    #[operand(ptr_read)]
    pub ptr: Value,
}

#[cube_op(name = "atomic.store")]
#[result_ty(none)]
#[op_traits(CanMaterialize)]
pub struct AtomicStoreOp {
    #[operand(ptr_write)]
    pub ptr: Value,
    pub value: Value,
}

#[cube_op(name = "atomic.compare_exchange_weak")]
#[result_ty(same_as = value)]
#[op_traits(CanMaterialize)]
pub struct AtomicCompareExchangeWeakOp {
    #[operand(ptr_read, ptr_write)]
    pub ptr: Value,
    pub cmp: Value,
    pub value: Value,
}