cubecl-wgpu 0.11.0-pre.3

WGPU runtime for the 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
59
60
61
62
63
64
65
66
67

use cubecl_ir::dialect::atomic::*;

use crate::compiler::wgsl::{
    to_wgsl::{wgsl_op, wgsl_op_with_out},
    value::WgslValue,
};

wgsl_op_with_out!(AtomicExchangeOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicExchange({ptr}, {value})")
});

wgsl_op_with_out!(AtomicIAddOp, AtomicFAddOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicAdd({ptr}, {value})")
});
wgsl_op_with_out!(AtomicISubOp, AtomicFSubOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicSub({ptr}, {value})")
});

wgsl_op_with_out!(AtomicSMaxOp, AtomicUMaxOp, AtomicFMaxOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicMax({ptr}, {value})")
});
wgsl_op_with_out!(AtomicSMinOp, AtomicUMinOp, AtomicFMinOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicMin({ptr}, {value})")
});

wgsl_op_with_out!(AtomicAndOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicAnd({ptr}, {value})")
});
wgsl_op_with_out!(AtomicOrOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicOr({ptr}, {value})")
});
wgsl_op_with_out!(AtomicXorOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicXor({ptr}, {value})")
});

wgsl_op_with_out!(AtomicLoadOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    format!("atomicLoad({ptr})")
});
wgsl_op!(AtomicStoreOp, |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicStore({ptr}, {value});\n")
});

wgsl_op_with_out!(AtomicCompareExchangeWeakOp; |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    let cmp = op.cmp(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    format!("atomicCompareExchangeWeak({ptr}, {cmp}, {value}).old_value")
});