1use alloc::{
2 boxed::Box,
3 string::{String, ToString},
4};
5use cubecl_macros_internal::{cube_op, op_traits};
6use pliron::{
7 attribute::AttrObj,
8 printable::Printable,
9 r#type::TypeHandle,
10 utils::table::{HMap, SmallSet},
11 verify_err,
12};
13use thiserror::Error;
14
15use crate::{
16 CanMaterialize, Pure,
17 attributes::IndexAttr,
18 interfaces::{
19 aliasing::AliasingOp,
20 memory_slot::{
21 DeletionKind, DestructurableAccessorOpInterface, DestructurableConstructorOpInterface,
22 DestructurableTypeInterface, DestructurableValueSlot, ValueSlot,
23 },
24 *,
25 },
26 prelude::*,
27 try_cast_ty,
28 types::{VectorType, aggregate::index_attr, scalar::IndexType},
29};
30
31#[pliron_op(
32 name = "composite.construct",
33 format = "operands(CharSpace(`,`)) ` : ` type($0)"
34)]
35#[op_interfaces(NResultsInterface<1>, OneResultInterface, AtLeastNOpdsInterface<1>)]
36#[op_traits(Pure, CanMaterialize)]
37pub struct CompositeConstructOp;
38
39impl CompositeConstructOp {
40 pub fn new(ctx: &mut Context, ty: TypeHandle, values: Vec<Value>) -> Self {
41 let op = Operation::new(
42 ctx,
43 Self::get_concrete_op_info(),
44 vec![ty],
45 values,
46 vec![],
47 0,
48 );
49 Self { op }
50 }
51
52 pub fn values(&self, ctx: &Context) -> Vec<Value> {
53 self.get_operation().deref(ctx).operands().collect()
54 }
55}
56
57#[op_interface_impl]
58impl DestructurableConstructorOpInterface for CompositeConstructOp {
59 fn destructurable_values(&self, ctx: &Context) -> Vec<DestructurableValueSlot> {
60 let ty = self.result_type(ctx).deref(ctx);
61 let Some(destructurable) = type_cast::<dyn DestructurableTypeInterface>(&*ty) else {
62 return vec![];
63 };
64 let Some(subelement_types) = destructurable.subelement_index_map(ctx) else {
65 return vec![];
66 };
67 vec![DestructurableValueSlot {
68 slot: ValueSlot::new(self.get_result(ctx), self.result_type(ctx)),
69 subelement_types,
70 }]
71 }
72
73 fn destructure(
74 &self,
75 ctx: &mut Context,
76 _value: &DestructurableValueSlot,
77 used_indices: &SmallSet<AttrObj, 8>,
78 _rewriter: &mut PassRewriter,
79 _new_constructors: &mut Vec<TraitOp<dyn DestructurableConstructorOpInterface>>,
80 ) -> HMap<AttrObj, ValueSlot> {
81 let op = self.get_operation();
82 let mut slot_map = HMap::new();
83 for used_index in used_indices {
84 let index = used_index.downcast_ref::<IndexAttr>().unwrap().0;
85 let opd = op.operand(ctx, index);
86 slot_map.insert(used_index.clone(), ValueSlot::new(opd, opd.get_type(ctx)));
87 }
88 slot_map
89 }
90
91 fn handle_destructuring_complete(
92 &self,
93 ctx: &mut Context,
94 value: &DestructurableValueSlot,
95 rewriter: &mut PassRewriter,
96 ) -> Option<TraitOp<dyn DestructurableConstructorOpInterface>> {
97 assert_eq!(value.slot.value, self.get_result(ctx));
98 rewriter.erase_operation(ctx, self.get_operation());
99 None
100 }
101}
102
103#[cube_op(
104 name = "composite.extract",
105 format = "$0 `[` attr($index, $IndexAttr) `] : ` type($0)",
106 verifier = "custom"
107)]
108#[result_ty(from_inputs = composite_extract_type)]
109#[op_traits(Pure, CanMaterialize)]
110pub struct CompositeExtractOp {
111 pub composite: Value,
112 pub index: IndexAttr,
113}
114
115fn composite_extract_type(ctx: &Context, aggregate: &Value, field: &IndexAttr) -> TypeHandle {
116 let aggregate_ty = aggregate.get_type(ctx).deref(ctx);
117 let aggregate_ty =
118 type_cast::<dyn DestructurableTypeInterface>(&*aggregate_ty).expect("Should be aggregate");
119 aggregate_ty.type_at_index(ctx, &index_attr(field.0))
120}
121
122#[derive(Error, Debug)]
123pub enum CompositeConstructError {
124 #[error(
125 "[CompositeConstructOp]: Output composite size doesn't match parameter count: Expected {_0} parameters, got {_1}"
126 )]
127 ParameterCountMismatch(usize, usize),
128 #[error(
129 "[CompositeConstructOp]: Output field type doesn't match parameter type: Expected {_0}, got {_1}"
130 )]
131 ParameterTypeMismatch(String, String),
132}
133
134#[op_interface_impl]
135impl AliasingOp for CompositeExtractOp {
136 fn source_ptr(&self, ctx: &Context) -> Option<Value> {
137 let aggregate = self.composite(ctx);
138 let field = self.index(ctx).0;
139 let aggregate_ty = aggregate.get_type(ctx).deref(ctx);
140 let destruct = try_cast_ty!(aggregate_ty, ctx, dyn DestructurableTypeInterface);
141 if destruct.type_at_index(ctx, &index_attr(field)).is_ptr(ctx) {
142 let construct = aggregate.defining_op().expect("Should be construct");
143 Some(construct.operand(ctx, field))
144 } else {
145 None
146 }
147 }
148}
149
150impl Verify for CompositeConstructOp {
151 fn verify(&self, ctx: &Context) -> Result<()> {
152 let ty = self.result_type(ctx).deref(ctx);
153 let opds = self.get_operation().operands(ctx);
154 let destructurable = try_cast_ty!(ty, ctx, dyn DestructurableTypeInterface);
155 let fields = destructurable.subelement_index_map(ctx).unwrap();
156
157 if opds.len() != fields.len() {
158 return verify_err!(
159 self.loc(ctx),
160 CompositeConstructError::ParameterCountMismatch(fields.len(), opds.len())
161 );
162 }
163
164 for (i, &opd) in opds.iter().enumerate() {
165 let field_ty = fields[&index_attr(i)];
166 if opd.get_type(ctx) != field_ty {
167 return verify_err!(
168 self.loc(ctx),
169 CompositeConstructError::ParameterTypeMismatch(
170 field_ty.disp(ctx).to_string(),
171 opd.get_type(ctx).disp(ctx).to_string()
172 )
173 );
174 }
175 }
176
177 Ok(())
178 }
179}
180
181#[derive(Error, Debug)]
182pub enum CompositeOpError {
183 #[error("[CompositeOp]: Index is out of range: index is {_0} but composite size is {_1}.")]
184 IndexOutOfRange(usize, usize),
185 #[error(
186 "[CompositeOp]: Field type doesn't match the inner type of the composite: expected {_0}, got {_1}"
187 )]
188 MismatchedFieldType(String, String),
189}
190
191impl Verify for CompositeExtractOp {
192 fn verify(&self, ctx: &Context) -> Result<()> {
193 let ty = self.composite(ctx).get_type(ctx).deref(ctx);
194 let destructurable = try_cast_ty!(ty, ctx, dyn DestructurableTypeInterface);
195 let fields = destructurable.subelement_index_map(ctx).unwrap();
196
197 let loc = self.loc(ctx);
198 let index = self.index(ctx).0;
199 if index >= fields.len() {
200 return verify_err!(loc, CompositeOpError::IndexOutOfRange(index, fields.len()));
201 }
202 Ok(())
203 }
204}
205
206#[op_interface_impl]
207impl DestructurableAccessorOpInterface for CompositeExtractOp {
208 fn can_rewire(
209 &self,
210 ctx: &Context,
211 value: &DestructurableValueSlot,
212 used_indices: &mut SmallSet<AttrObj, 8>,
213 _must_be_safely_used: &mut Vec<ValueSlot>,
214 ) -> bool {
215 if value.slot.value != self.composite(ctx) {
216 return false;
217 }
218 used_indices.insert(Box::new(*self.index(ctx)));
219 true
220 }
221
222 fn rewire(
223 &self,
224 ctx: &mut Context,
225 _value: &DestructurableValueSlot,
226 subvalues: &HMap<AttrObj, ValueSlot>,
227 rewriter: &mut PassRewriter,
228 ) -> DeletionKind {
229 let index: AttrObj = Box::new(*self.index(ctx));
230 let slot = &subvalues[&index];
231 rewriter.replace_value_uses_with(ctx, self.get_result(ctx), slot.value);
232 DeletionKind::Delete
233 }
234}
235
236#[cube_op(
237 name = "composite.insert",
238 format = "$1 ` -> ` $0 `[` attr($index, $IndexAttr) `] : ` type($0)",
239 verifier = "custom"
240)]
241#[result_ty(same_as = composite)]
242#[op_traits(CanMaterialize, Pure)]
243pub struct CompositeInsertOp {
244 pub composite: Value,
245 pub value: Value,
246 pub index: IndexAttr,
247}
248
249impl Verify for CompositeInsertOp {
250 fn verify(&self, ctx: &Context) -> Result<()> {
251 let ty = self.result_type(ctx).deref(ctx);
252 let destructurable = try_cast_ty!(ty, ctx, dyn DestructurableTypeInterface);
253 let fields = destructurable.subelement_index_map(ctx).unwrap();
254
255 let loc = self.loc(ctx);
256 let index = self.index(ctx).0;
257 if index >= fields.len() {
258 return verify_err!(loc, CompositeOpError::IndexOutOfRange(index, fields.len()));
259 }
260 let field_ty = fields[&index_attr(index)];
261 let value_ty = self.value(ctx).get_type(ctx);
262 if field_ty != value_ty {
263 return verify_err!(
264 loc,
265 CompositeOpError::MismatchedFieldType(
266 field_ty.disp(ctx).to_string(),
267 value_ty.disp(ctx).to_string()
268 )
269 );
270 }
271 Ok(())
272 }
273}
274
275#[op_interface_impl]
276impl DestructurableConstructorOpInterface for CompositeInsertOp {
277 fn destructurable_values(&self, ctx: &Context) -> Vec<DestructurableValueSlot> {
278 let ty = self.result_type(ctx).deref(ctx);
279 let Some(destructurable) = type_cast::<dyn DestructurableTypeInterface>(&*ty) else {
280 return vec![];
281 };
282 let Some(subelement_types) = destructurable.subelement_index_map(ctx) else {
283 return vec![];
284 };
285 vec![DestructurableValueSlot {
286 slot: ValueSlot::new(self.get_result(ctx), self.result_type(ctx)),
287 subelement_types,
288 }]
289 }
290
291 fn destructure(
292 &self,
293 ctx: &mut Context,
294 _value: &DestructurableValueSlot,
295 used_indices: &SmallSet<AttrObj, 8>,
296 rewriter: &mut PassRewriter,
297 _new_constructors: &mut Vec<TraitOp<dyn DestructurableConstructorOpInterface>>,
298 ) -> HMap<AttrObj, ValueSlot> {
299 let inserted_index = self.index(ctx).0;
300 let mut slot_map = HMap::new();
301 for used_index in used_indices {
302 let index = used_index.downcast_ref::<IndexAttr>().unwrap().0;
303 let value = if index == inserted_index {
304 ValueSlot::new(self.value(ctx), self.value(ctx).get_type(ctx))
305 } else {
306 let extract = CompositeExtractOp::new(ctx, self.composite(ctx), index);
307 rewriter.append_op(ctx, &extract);
308 ValueSlot::new(extract.get_result(ctx), extract.result_type(ctx))
309 };
310 slot_map.insert(used_index.clone(), value);
311 }
312 slot_map
313 }
314
315 fn handle_destructuring_complete(
316 &self,
317 ctx: &mut Context,
318 value: &DestructurableValueSlot,
319 rewriter: &mut PassRewriter,
320 ) -> Option<TraitOp<dyn DestructurableConstructorOpInterface>> {
321 assert_eq!(value.slot.value, self.get_result(ctx));
322 rewriter.erase_operation(ctx, self.get_operation());
323 None
324 }
325}
326
327#[cube_op(
328 name = "vector.broadcast",
329 format = "$0 ` : ` type($0)",
330 verifier = "custom"
331)]
332#[result_ty(argument)]
333#[op_interfaces(ResultNOfType<0, VectorType>, TriviallyUnrollable)]
334#[op_traits(CanMaterialize, Pure)]
335pub struct VectorBroadcastOp {
336 pub input: Value,
337}
338
339impl Verify for VectorBroadcastOp {
340 fn verify(&self, ctx: &Context) -> Result<()> {
341 let loc = self.loc(ctx);
342 let value_ty = self.input(ctx).get_type(ctx);
343 let scalar_ty = self.get_result(ctx).scalar_ty(ctx);
344 if scalar_ty != value_ty {
345 return verify_err!(
346 loc,
347 CompositeOpError::MismatchedFieldType(
348 scalar_ty.disp(ctx).to_string(),
349 value_ty.disp(ctx).to_string()
350 )
351 );
352 }
353 Ok(())
354 }
355}
356
357#[op_interface_impl]
358impl DestructurableConstructorOpInterface for VectorBroadcastOp {
359 fn destructurable_values(&self, ctx: &Context) -> Vec<DestructurableValueSlot> {
360 let ty = self.result_type(ctx).deref(ctx);
361 let Some(destructurable) = type_cast::<dyn DestructurableTypeInterface>(&*ty) else {
362 return vec![];
363 };
364 let Some(subelement_types) = destructurable.subelement_index_map(ctx) else {
365 return vec![];
366 };
367 vec![DestructurableValueSlot {
368 slot: ValueSlot::new(self.get_result(ctx), self.result_type(ctx)),
369 subelement_types,
370 }]
371 }
372
373 fn destructure(
374 &self,
375 ctx: &mut Context,
376 _value: &DestructurableValueSlot,
377 used_indices: &SmallSet<AttrObj, 8>,
378 _rewriter: &mut PassRewriter,
379 _new_constructors: &mut Vec<TraitOp<dyn DestructurableConstructorOpInterface>>,
380 ) -> HMap<AttrObj, ValueSlot> {
381 let mut slot_map = HMap::new();
382 for used_index in used_indices {
383 let slot = ValueSlot::new(self.input(ctx), self.input(ctx).get_type(ctx));
384 slot_map.insert(used_index.clone(), slot);
385 }
386 slot_map
387 }
388
389 fn handle_destructuring_complete(
390 &self,
391 ctx: &mut Context,
392 value: &DestructurableValueSlot,
393 rewriter: &mut PassRewriter,
394 ) -> Option<TraitOp<dyn DestructurableConstructorOpInterface>> {
395 assert_eq!(value.slot.value, self.get_result(ctx));
396 rewriter.erase_operation(ctx, self.get_operation());
397 None
398 }
399}
400
401#[cube_op(
402 name = "vector.insert_dynamic",
403 format = "$1 ` -> ` $0 `[` $2 `] : ` type($0)",
404 verifier = "custom"
405)]
406#[result_ty(same_as = vector)]
407#[op_interfaces(OperandNOfType<0, VectorType>, OperandNOfType<2, IndexType>)]
408#[op_traits(CanMaterialize, Pure)]
409pub struct VectorInsertDynamicOp {
410 pub vector: Value,
411 pub value: Value,
412 pub index: Value,
413}
414
415impl Verify for VectorInsertDynamicOp {
416 fn verify(&self, ctx: &Context) -> Result<()> {
417 let loc = self.loc(ctx);
418 let scalar_ty = self.vector(ctx).scalar_ty(ctx);
419 let value_ty = self.value(ctx).get_type(ctx);
420 if scalar_ty != value_ty {
421 verify_err!(
422 loc,
423 CompositeOpError::MismatchedFieldType(
424 scalar_ty.disp(ctx).to_string(),
425 value_ty.disp(ctx).to_string()
426 )
427 )?;
428 }
429 Ok(())
430 }
431}
432
433#[cube_op(name = "vector.extract_dynamic", format = "$0 `[` $1 `] : ` type($0)")]
434#[result_ty(from_inputs = |ctx, vector, _| scalar_ty(ctx, vector))]
435#[op_interfaces(OperandNOfType<0, VectorType>, OperandNOfType<1, IndexType>)]
436#[op_traits(CanMaterialize, Pure)]
437pub struct VectorExtractDynamicOp {
438 pub vector: Value,
439 pub index: Value,
440}
441
442#[cube_op(name = "vector.magnitude")]
443#[result_ty(from_inputs = scalar_ty)]
444#[op_interfaces(OperandNOfType<0, VectorType>)]
445#[op_traits(CanMaterialize, Pure)]
446pub struct MagnitudeOp {
447 pub input: Value,
448}
449
450#[cube_op(name = "vector.normalize")]
451#[result_ty(same_as = input)]
452#[op_interfaces(SameOperandsType, SameOperandsAndResultType, OperandNOfType<0, VectorType>)]
453#[op_traits(CanMaterialize, Pure)]
454pub struct NormalizeOp {
455 pub input: Value,
456}
457
458#[cube_op(name = "vector.i_sum")]
459#[result_ty(from_inputs = scalar_ty)]
460#[op_interfaces(OperandNOfType<0, VectorType>)]
461#[op_traits(CanMaterialize, Pure)]
462pub struct ISumOp {
463 pub input: Value,
464}
465
466#[cube_op(name = "vector.f_sum")]
467#[result_ty(from_inputs = scalar_ty)]
468#[op_interfaces(OperandNOfType<0, VectorType>)]
469#[op_traits(CanMaterialize, Pure)]
470pub struct FSumOp {
471 pub input: Value,
472}
473
474#[cube_op(name = "vector.s_dot")]
475#[result_ty(from_inputs = |ctx, lhs, _| scalar_ty(ctx, lhs))]
476#[op_interfaces(OperandNOfType<0, VectorType>, OperandNOfType<1, VectorType>)]
477#[op_traits(CanMaterialize, Pure)]
478pub struct SDotOp {
479 pub lhs: Value,
480 pub rhs: Value,
481}
482
483#[cube_op(name = "vector.u_dot")]
484#[result_ty(from_inputs = |ctx, lhs, _| scalar_ty(ctx, lhs))]
485#[op_interfaces(OperandNOfType<0, VectorType>, OperandNOfType<1, VectorType>)]
486#[op_traits(CanMaterialize, Pure)]
487pub struct UDotOp {
488 pub lhs: Value,
489 pub rhs: Value,
490}
491
492#[cube_op(name = "vector.f_dot")]
493#[result_ty(from_inputs = |ctx, lhs, _| scalar_ty(ctx, lhs))]
494#[op_interfaces(OperandNOfType<0, VectorType>, OperandNOfType<1, VectorType>)]
495#[op_traits(CanMaterialize, Pure)]
496pub struct FDotOp {
497 pub lhs: Value,
498 pub rhs: Value,
499}
500
501fn scalar_ty(ctx: &Context, input: &Value) -> TypeHandle {
502 input.scalar_ty(ctx)
503}