cubecl-wgpu 0.11.0-pre.4

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
68
69
70
71

use cubecl_ir::{
    AddressSpace, CanMaterialize, GlobalState, NoSideEffects, Scope, dialect::memory::*, ident,
    prelude::*,
};
use pliron::attribute::AttrObj;

use crate::compiler::wgsl::{
    AddressOfOp, GlobalVariableOp,
    lower::LowerOp,
    ops::general::attr_to_wgsl,
    to_wgsl::{TypeExtWgsl, wgsl_op, wgsl_op_with_out},
    value::WgslValue,
};

#[cube_op(name = "wgsl.declare_local")]
#[result_ty(argument)]
#[op_traits(NoSideEffects, CanMaterialize)]
pub struct DeclareLocalOp {
    pub value_ty: TypeAttr,
    #[attribute(optional, untyped)]
    pub initializer: AttrObj,
}

wgsl_op!(DeclareLocalOp, |op, ctx| {
    let value_ty = op.value_ty(ctx).get_type(ctx).to_wgsl(ctx);
    let name = op.get_result(ctx).name(ctx);
    let init = op.initializer(ctx).map(|init| attr_to_wgsl(ctx, &**init));
    let init = init.map(|init| format!("= {init}")).unwrap_or_default();
    format!("var<function> {name}_store: {value_ty}{init};\nlet {name} = &{name}_store;\n")
});

#[op_interface_impl]
impl LowerOp for DeclareVariableOp {
    fn lower(&self, scope: &Scope) -> Vec<Value> {
        let ctx = scope.ctx_mut();
        let addr_space = self.addr_space(ctx).0;
        let ty = self.result_type(ctx);
        let value_ty = self.value_ty(ctx).get_type(ctx);
        match addr_space {
            AddressSpace::Global(_) => panic!("Global vars not supported"),
            AddressSpace::Shared => {
                let name = ident(format!("shared_{}", self.get_result(ctx).name(ctx)));
                let var = GlobalVariableOp::new(ctx, value_ty, addr_space, None, None);
                let entry = ctx.aux_ty::<GlobalState>().entry_func.get_operation();
                var.set_symbol_name(ctx, name.clone());
                var.get_operation().insert_before(ctx, entry);
                vec![scope.register_with_result(&AddressOfOp::new(ctx, ty, name))]
            }
            AddressSpace::Local => {
                let init = self.initializer(ctx).map(|it| it.clone());
                let local = DeclareLocalOp::new(ctx, ty, value_ty, init);
                vec![scope.register_with_result(&local)]
            }
        }
    }
}

wgsl_op_with_out!(IndexOp; |op, ctx| {
    format!("&(*{})[{}]", op.base(ctx).name(ctx), op.index(ctx).name(ctx))
});

wgsl_op_with_out!(LoadOp; |op, ctx| format!("*{}", op.ptr(ctx).name(ctx)));
wgsl_op!(StoreOp, |op, ctx| {
    let ptr = op.ptr(ctx).name(ctx);
    format!("*{ptr} = {};\n", op.value(ctx).name(ctx))
});
wgsl_op!(CopyOp, |op, ctx| {
    assert_eq!(op.len(ctx).0, 1, "WGSL doesn't support bulk copy");
    let dest = op.destination(ctx).name(ctx);
    format!("*{dest} = *{};\n", op.source(ctx).name(ctx))
});