cubecl_ir/dialect/
general.rs

1use core::cell::Ref;
2
3use alloc::string::String;
4
5use cubecl_macros_internal::{const_eval, cube_op, simplify};
6use derive_more::From;
7use derive_new::new;
8use pliron::{
9    builtin::attributes::{StringAttr, TypeAttr},
10    derive::pliron_attr,
11    r#type::type_cast,
12};
13
14use crate::{
15    Builtin, CanMaterialize, ConstantValue, Pure,
16    attributes::{BoolAttr, IndexAttr},
17    dialect::{
18        math::{index_attr, int_attr},
19        pure_binop, pure_unop,
20    },
21    interfaces::{ScalarType, TriviallyUnrollable, TypedExt, aliasing::AliasingOp},
22    prelude::*,
23    types::scalar::IndexType,
24};
25
26#[cube_op(name = "cube.copy")]
27#[result_ty(same_as = value)]
28#[op_traits(Pure, CanMaterialize)]
29pub struct CopyOp {
30    pub value: Value,
31}
32
33simplify!(CopyOp, {
34    |_| Some(self.value(ctx)),
35});
36
37#[op_interface_impl]
38impl AliasingOp for CopyOp {
39    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
40        Some(self.value(ctx))
41    }
42}
43
44#[cube_op(name = "cube.poison")]
45#[result_ty(argument)]
46#[op_traits(Pure, CanMaterialize)]
47pub struct PoisonOp {}
48
49pure_binop!("cube.bool_and", BoolAndOp);
50const_eval!(BoolAndOp, {
51    BoolAttr: |lhs, rhs| lhs && rhs,
52    // false && x -> false
53    custom: |lhs, _| match lhs?.as_const_val(ctx) {
54        ConstantValue::Bool(false) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
55        _ => None
56    },
57    // x && false -> false
58    custom: |_, rhs| match rhs?.as_const_val(ctx) {
59        ConstantValue::Bool(false) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
60        _ => None
61    }
62});
63simplify!(BoolAndOp, {
64    // true && x -> x
65    |lhs, _| match lhs?.as_const_val(ctx) {
66        ConstantValue::Bool(true) => Some(self.rhs(ctx)),
67        _ => None,
68    },
69    // x && true -> x
70    |_, rhs| match rhs?.as_const_val(ctx) {
71        ConstantValue::Bool(true) => Some(self.lhs(ctx)),
72        _ => None,
73    },
74    // x && x -> x
75    |_, _| match self.lhs(ctx) == self.rhs(ctx) {
76        true => Some(self.lhs(ctx)),
77        false => None
78    }
79});
80
81pure_binop!("cube.bool_or", BoolOrOp);
82const_eval!(BoolOrOp, {
83    BoolAttr: |lhs, rhs| lhs || rhs,
84    // true || x -> true
85    custom: |lhs, _| match lhs?.as_const_val(ctx) {
86        ConstantValue::Bool(true) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
87        _ => None
88    },
89    // x || true -> true
90    custom: |_, rhs| match rhs?.as_const_val(ctx) {
91        ConstantValue::Bool(true) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
92        _ => None
93    }
94});
95simplify!(BoolOrOp, {
96    // false || x -> x
97    |lhs, _| match lhs?.as_const_val(ctx) {
98        ConstantValue::Bool(false) => Some(self.rhs(ctx)),
99        _ => None,
100    },
101    // false || x -> x
102    |_, rhs| match rhs?.as_const_val(ctx) {
103        ConstantValue::Bool(false) => Some(self.lhs(ctx)),
104        _ => None,
105    },
106    // x || x -> x
107    |_, _| match self.lhs(ctx) == self.rhs(ctx) {
108        true => Some(self.lhs(ctx)),
109        false => None
110    }
111});
112
113pure_unop!("cube.bool_not", BoolNotOp);
114const_eval!(BoolNotOp, {
115    BoolAttr: |inp| !inp
116});
117
118#[cube_op(name = "cube.cast")]
119#[result_ty(argument)]
120#[op_interfaces(TriviallyUnrollable)]
121#[op_traits(Pure, CanMaterialize)]
122pub struct CastOp {
123    pub input: Value,
124}
125const_eval!(CastOp, {
126    custom: |inp| {
127        let val = inp?.as_const_val(ctx);
128        let out_ty = self.get_result(ctx).get_type(ctx).deref(ctx);
129        let elem = type_cast::<dyn ScalarType>(&*out_ty)?.elem_type(ctx);
130        Some(val.cast_to(elem).as_attribute(ctx, elem))
131    }
132});
133simplify!(CastOp, {
134    |_| {
135        if self.input(ctx).get_type(ctx) == self.result_type(ctx) {
136            Some(self.input(ctx))
137        } else {
138            None
139        }
140    }
141});
142
143#[cube_op(name = "cube.reinterpret_cast")]
144#[result_ty(argument)]
145#[op_traits(Pure, CanMaterialize)]
146pub struct ReinterpretCastOp {
147    pub input: Value,
148}
149const_eval!(ReinterpretCastOp, {
150    custom: |inp| {
151        // Too much weirdness around floats, don't bother dealing with it
152        let val = match inp?.as_const_val(ctx) {
153            ConstantValue::Int(val) => val as u64,
154            ConstantValue::UInt(val) => val,
155            _ => None?,
156        };
157        let out_ty = self.get_result(ctx).get_type(ctx);
158        if out_ty.is_int(ctx) {
159            Some(int_attr(ctx, out_ty, val as i128))
160        } else if out_ty.is_index(ctx) {
161            Some(index_attr(val as usize))
162        } else {
163            None
164        }
165    }
166});
167simplify!(ReinterpretCastOp, {
168    |_| {
169        if self.input(ctx).get_type(ctx) == self.result_type(ctx) {
170            Some(self.input(ctx))
171        } else {
172            None
173        }
174    }
175});
176
177#[op_interface_impl]
178impl AliasingOp for ReinterpretCastOp {
179    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
180        Some(self.input(ctx))
181    }
182}
183
184#[cube_op(name = "cube.select")]
185#[result_ty(same_as = true_value)]
186#[op_interfaces(TriviallyUnrollable)]
187#[op_traits(Pure, CanMaterialize)]
188pub struct SelectOp {
189    pub condition: Value,
190    pub true_value: Value,
191    pub false_value: Value,
192}
193simplify!(SelectOp, {
194    |cond, _, _| match cond?.as_const_val(ctx) {
195        ConstantValue::Bool(true) => Some(self.true_value(ctx)),
196        ConstantValue::Bool(false) => Some(self.false_value(ctx)),
197        _ => None,
198    },
199    // select(cond, x, x) -> x
200    |_, _, _| match self.true_value(ctx) == self.false_value(ctx) {
201        true => Some(self.true_value(ctx)),
202        false => None
203    }
204});
205
206#[pliron_attr(name = "cube.builtin", format, verifier = "succ")]
207#[derive(new, From, PartialEq, Clone, Debug, Hash)]
208pub struct BuiltinAttr(pub Builtin);
209
210#[cube_op(
211    name = "cube.read_builtin",
212    format = "attr($builtin, $BuiltinAttr) ` : ` type($0)"
213)]
214#[result_ty(argument)]
215#[op_traits(Pure, CanMaterialize)]
216pub struct ReadBuiltinOp {
217    pub builtin: BuiltinAttr,
218}
219
220#[cube_op(name = "cube.read_scalar")]
221#[result_ty(from_inputs = |ctx, ty: &TypeAttr, _| ty.get_type(ctx))]
222#[op_traits(Pure, CanMaterialize)]
223pub struct ReadScalarOp {
224    pub ty: TypeAttr,
225    pub id: IndexAttr,
226}
227
228#[cube_op(name = "cube.free")]
229#[result_ty(none)]
230pub struct FreeOp {
231    pub memory: Value,
232}
233
234#[cube_op(name = "cube.buffer_len")]
235#[result_ty(fixed = IndexType::get(ctx).into())]
236#[op_traits(Pure, CanMaterialize)]
237pub struct BufferLenOp {
238    pub buffer_idx: IndexAttr,
239}
240
241#[cube_op(name = "cube.shape")]
242#[result_ty(fixed = IndexType::get(ctx).into())]
243#[op_traits(Pure, CanMaterialize)]
244pub struct ShapeOp {
245    pub dim: Value,
246    pub buffer_idx: IndexAttr,
247}
248
249#[cube_op(name = "cube.stride")]
250#[result_ty(fixed = IndexType::get(ctx).into())]
251#[op_traits(Pure, CanMaterialize)]
252pub struct StrideOp {
253    pub dim: Value,
254    pub buffer_idx: IndexAttr,
255}
256
257#[cube_op(name = "cube.comment")]
258#[result_ty(none)]
259pub struct CommentOp {
260    pub comment: StringAttr,
261}
262
263#[pliron_op(name = "cube.printf", format, attributes = (cube_printf_format_string: StringAttr), verifier = "succ")]
264pub struct PrintfOp;
265
266impl PrintfOp {
267    pub fn new(ctx: &mut Context, format_string: String, values: Vec<Value>) -> Self {
268        let op = Self {
269            op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], values, vec![], 0),
270        };
271        op.set_attr_cube_printf_format_string(ctx, StringAttr::new(format_string));
272        op
273    }
274
275    pub fn format_string<'a>(&self, ctx: &'a Context) -> Ref<'a, StringAttr> {
276        self.get_attr_cube_printf_format_string(ctx).unwrap()
277    }
278
279    pub fn args(&self, ctx: &Context) -> Vec<Value> {
280        self.get_operation().deref(ctx).operands().collect()
281    }
282}
cubecl_ir/dialect/general.rs

cubecl_ir/dialect/
general.rs
