cubecl_ir/
reflect.rs

1use alloc::collections::VecDeque;
2
3use alloc::vec;
4use alloc::vec::Vec;
5
6use crate::{Builtin, Instruction, Memory, Scope, Type, Value};
7
8/// An operation that can be reflected on
9pub trait OperationReflect: Sized {
10    /// Type of the op codes for this operation
11    type OpCode;
12
13    /// Get the opcode for this operation
14    fn op_code(&self) -> Self::OpCode;
15    /// Get the list of arguments for this operation. If not all arguments are [`Value`], returns
16    /// `None` instead.
17    fn args(&self) -> Option<Vec<Value>> {
18        None
19    }
20    fn args_mut(&mut self) -> Option<Vec<&mut Value>> {
21        None
22    }
23    /// Create typed operation from an opcode and a list of arguments. Returns `None` if not all
24    /// arguments are [`Value`].
25    #[allow(unused)]
26    fn from_code_and_args(op_code: Self::OpCode, args: &[Value]) -> Option<Self> {
27        None
28    }
29    /// Sanitize args, i.e. loading inputs from pointers for ops that take values
30    fn sanitize_args(&mut self, scope: &Scope);
31    /// Whether this operation is commutative (arguments can be freely reordered). Ignored for
32    /// single argument operations.
33    fn is_commutative(&self) -> bool {
34        false
35    }
36    /// Whether this operation is pure (has no side effects). Things like uniform/plane operations
37    /// are considered impure, because they affect other units.
38    fn is_pure(&self) -> bool {
39        false
40    }
41
42    fn read_pointers(&self) -> Vec<Value> {
43        Vec::new()
44    }
45
46    fn write_pointers(&self) -> Vec<Value> {
47        Vec::new()
48    }
49}
50
51/// A type that represents an operation's arguments
52pub trait OperationArgs: Sized {
53    /// Sanitize args for the `ptr` constraint, loading inputs from pointers for ops that take values.
54    /// This is needed because Rust "helpfully" auto-derefs references, skipping the explicit deref
55    /// code that inserts a [`Memory::Load`].
56    fn sanitize_args_ptr(&mut self, scope: &Scope);
57
58    /// Construct this type from a list of arguments. If not all arguments are [`Value`], returns
59    /// `None`
60    #[allow(unused)]
61    fn from_args(args: &[Value]) -> Option<Self> {
62        None
63    }
64
65    /// Turns this type into a flat list of arguments. If not all arguments are [`Value`],
66    /// returns `None`
67    fn as_args(&self) -> Option<Vec<Value>> {
68        None
69    }
70
71    /// Turns this type into a flat list of arguments. If not all arguments are [`Value`],
72    /// returns `None`
73    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
74        None
75    }
76
77    fn read_pointers(&self) -> Vec<Value> {
78        Vec::new()
79    }
80
81    fn write_pointers(&self) -> Vec<Value> {
82        Vec::new()
83    }
84}
85
86impl OperationArgs for Value {
87    fn sanitize_args_ptr(&mut self, scope: &Scope) {
88        *self = read_value(scope, *self)
89    }
90
91    fn from_args(args: &[Value]) -> Option<Self> {
92        Some(args[0])
93    }
94
95    fn as_args(&self) -> Option<Vec<Value>> {
96        Some(vec![*self])
97    }
98
99    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
100        Some(vec![self])
101    }
102}
103
104impl<T: OperationArgs> OperationArgs for Vec<T> {
105    fn sanitize_args_ptr(&mut self, scope: &Scope) {
106        self.iter_mut().for_each(|it| it.sanitize_args_ptr(scope));
107    }
108
109    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
110        let inner = self.iter_mut().map(|it| it.as_args_mut());
111        let inner = inner.collect::<Option<Vec<_>>>()?;
112        Some(inner.into_iter().flatten().collect())
113    }
114}
115
116impl<T: OperationArgs> OperationArgs for Option<T> {
117    fn sanitize_args_ptr(&mut self, scope: &Scope) {
118        if let Some(it) = self.as_mut() {
119            it.sanitize_args_ptr(scope)
120        }
121    }
122
123    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
124        self.as_mut().and_then(|inner| inner.as_args_mut())
125    }
126}
127
128impl OperationArgs for usize {
129    fn sanitize_args_ptr(&mut self, _: &Scope) {}
130
131    fn from_args(args: &[Value]) -> Option<Self> {
132        Some(args[0].as_const().unwrap().as_usize())
133    }
134    fn as_args(&self) -> Option<Vec<Value>> {
135        Some(vec![(*self).into()])
136    }
137    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
138        Some(vec![])
139    }
140}
141impl OperationArgs for u32 {
142    fn sanitize_args_ptr(&mut self, _: &Scope) {}
143
144    fn from_args(args: &[Value]) -> Option<Self> {
145        Some(args[0].as_const().unwrap().as_u32())
146    }
147    fn as_args(&self) -> Option<Vec<Value>> {
148        Some(vec![(*self).into()])
149    }
150    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
151        Some(vec![])
152    }
153}
154impl OperationArgs for bool {
155    fn sanitize_args_ptr(&mut self, _: &Scope) {}
156
157    fn from_args(args: &[Value]) -> Option<Self> {
158        Some(args[0].as_const().unwrap().as_bool())
159    }
160    fn as_args(&self) -> Option<Vec<Value>> {
161        Some(vec![(*self).into()])
162    }
163    fn as_args_mut(&mut self) -> Option<Vec<&mut Value>> {
164        Some(vec![])
165    }
166}
167
168/// Types that can be destructured into and created from a list of [`Value`]s.
169pub trait FromArgList: Sized {
170    /// Creates this type from a list of values. This works like a parse stream, where consumed
171    /// values are popped from the front.
172    fn from_arg_list(args: &mut VecDeque<Value>) -> Self;
173    /// Turns this type into a list of [`Value`]s.
174    fn as_arg_list(&self) -> impl IntoIterator<Item = Value>;
175    fn as_arg_list_mut(&mut self) -> impl IntoIterator<Item = &mut Value>;
176}
177
178impl FromArgList for Value {
179    fn from_arg_list(args: &mut VecDeque<Value>) -> Self {
180        args.pop_front().expect("Missing value from arg list")
181    }
182
183    fn as_arg_list(&self) -> impl IntoIterator<Item = Value> {
184        [*self]
185    }
186
187    fn as_arg_list_mut(&mut self) -> impl IntoIterator<Item = &mut Value> {
188        [self]
189    }
190}
191
192impl FromArgList for Vec<Value> {
193    fn from_arg_list(args: &mut VecDeque<Value>) -> Self {
194        core::mem::take(args).into_iter().collect()
195    }
196
197    fn as_arg_list(&self) -> impl IntoIterator<Item = Value> {
198        self.iter().cloned()
199    }
200
201    fn as_arg_list_mut(&mut self) -> impl IntoIterator<Item = &mut Value> {
202        self.iter_mut()
203    }
204}
205
206impl FromArgList for bool {
207    fn from_arg_list(args: &mut VecDeque<Value>) -> Self {
208        args.pop_front()
209            .expect("Missing value from arg list")
210            .as_const()
211            .unwrap()
212            .as_bool()
213    }
214
215    fn as_arg_list(&self) -> impl IntoIterator<Item = Value> {
216        [(*self).into()]
217    }
218
219    fn as_arg_list_mut(&mut self) -> impl IntoIterator<Item = &mut Value> {
220        []
221    }
222}
223
224impl FromArgList for u32 {
225    fn from_arg_list(args: &mut VecDeque<Value>) -> Self {
226        args.pop_front()
227            .expect("Missing value from arg list")
228            .as_const()
229            .unwrap()
230            .as_u32()
231    }
232
233    fn as_arg_list(&self) -> impl IntoIterator<Item = Value> {
234        [(*self).into()]
235    }
236
237    fn as_arg_list_mut(&mut self) -> impl IntoIterator<Item = &mut Value> {
238        []
239    }
240}
241
242impl OperationArgs for Builtin {
243    fn sanitize_args_ptr(&mut self, _: &crate::Scope) {}
244}
245
246impl FromArgList for usize {
247    fn from_arg_list(args: &mut VecDeque<Value>) -> Self {
248        args.pop_front()
249            .expect("Missing value from arg list")
250            .as_const()
251            .unwrap()
252            .as_usize()
253    }
254
255    fn as_arg_list(&self) -> impl IntoIterator<Item = Value> {
256        [(*self).into()]
257    }
258
259    fn as_arg_list_mut(&mut self) -> impl IntoIterator<Item = &mut Value> {
260        []
261    }
262}
263
264fn read_value(scope: &Scope, val: Value) -> Value {
265    if let Type::Pointer(inner, _) = val.ty {
266        let out = scope.create_value(*inner);
267        scope.register(Instruction::new(Memory::Load(val), out));
268        out
269    } else {
270        val
271    }
272}
cubecl_ir/reflect.rs

cubecl_ir/
reflect.rs
