cubecl_core/frontend/
asm.rs

1use alloc::{string::String, vec::Vec};
2
3use cubecl_ir::{
4    AddressSpace, Scope,
5    dialect::{InlineAsmOp, MemoryClobbers, OperationPtrExt},
6    interfaces::TypeExt,
7};
8use pliron::{op::Op, r#type::Typed, value::Value};
9
10use crate::frontend::{HasValue, assign};
11
12pub use cubecl_ir::dialect::{InputKind, InputSpec, RegSpec};
13
14#[derive(Default)]
15pub struct BuildAsmExpand {
16    asm: String,
17    out_values: Vec<Value>,
18    out_specs: Vec<RegSpec>,
19    in_values: Vec<Value>,
20    in_specs: Vec<InputSpec>,
21    pure: bool,
22    nomem: bool,
23    explicit_mem: bool,
24    readonly: bool,
25    reads_spaces: Vec<AddressSpace>,
26    writes_spaces: Vec<AddressSpace>,
27}
28
29impl BuildAsmExpand {
30    pub fn new(asm: String) -> Self {
31        BuildAsmExpand {
32            asm,
33            ..Default::default()
34        }
35    }
36
37    // Takes by reference because of syntax reasons, since normal Rust allows immutables that are
38    // uninitialized and only assigned once (i.e. as the output for an assembly macro).
39    // We also only assign once, so reference gives the correct semantics.
40    pub fn push_output<T: HasValue>(mut self, scope: &Scope, output: &T, spec: RegSpec) -> Self {
41        let value = output.value(scope);
42        self.out_values.push(value);
43        self.out_specs.push(spec);
44        self
45    }
46
47    pub fn push_input<T: HasValue>(
48        mut self,
49        scope: &Scope,
50        input: T,
51        kind: InputKind,
52        spec: RegSpec,
53    ) -> Self {
54        let value = input.value(scope);
55        self.in_values.push(value);
56        self.in_specs.push(InputSpec::new(kind, spec));
57        self
58    }
59
60    pub fn pure(mut self) -> Self {
61        self.pure = true;
62        self
63    }
64
65    pub fn nomem(mut self) -> Self {
66        self.nomem = true;
67        self
68    }
69
70    pub fn explicit_mem(mut self) -> Self {
71        self.explicit_mem = true;
72        self
73    }
74
75    pub fn readonly(mut self) -> Self {
76        self.readonly = true;
77        self
78    }
79
80    pub fn reads_local(mut self) -> Self {
81        self.reads_spaces.push(AddressSpace::Local);
82        self
83    }
84
85    pub fn reads_shared(mut self) -> Self {
86        self.reads_spaces.push(AddressSpace::Shared);
87        self
88    }
89
90    pub fn reads_global(mut self) -> Self {
91        self.reads_spaces.push(AddressSpace::Global(0));
92        self
93    }
94
95    pub fn writes_local(mut self) -> Self {
96        self.writes_spaces.push(AddressSpace::Local);
97        self
98    }
99
100    pub fn writes_shared(mut self) -> Self {
101        self.writes_spaces.push(AddressSpace::Shared);
102        self
103    }
104
105    pub fn writes_global(mut self) -> Self {
106        self.writes_spaces.push(AddressSpace::Global(0));
107        self
108    }
109
110    pub fn register(self, scope: &Scope) {
111        let ctx = scope.ctx_mut();
112        let result_types = self
113            .out_values
114            .iter()
115            .map(|it| it.get_type(ctx).as_ptr(ctx).inner)
116            .collect();
117        let memory_clobbers = if self.nomem {
118            MemoryClobbers::Nomem
119        } else if self.readonly {
120            MemoryClobbers::Readonly
121        } else if self.explicit_mem {
122            MemoryClobbers::Explicit {
123                reads_spaces: self.reads_spaces.into(),
124                writes_spaces: self.writes_spaces.into(),
125            }
126        } else {
127            MemoryClobbers::ReadWrite
128        };
129        let op = InlineAsmOp::new(
130            ctx,
131            result_types,
132            self.out_specs,
133            self.asm,
134            memory_clobbers,
135            self.in_values,
136            self.in_specs,
137        );
138        if self.pure {
139            op.set_pure(ctx);
140        }
141        scope.register(&op);
142        // Store results back to out expand values
143        for (&out_ptr, result) in self.out_values.iter().zip(op.get_operation().results(ctx)) {
144            assign::expand_element(scope, result.into(), out_ptr.into());
145        }
146    }
147}
cubecl_core/frontend/asm.rs

cubecl_core/frontend/
asm.rs
