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};
13use thiserror::Error;
14
15use crate::{
16    Builtin, CanMaterialize, Commutative, ConstantValue, HasSideEffects, NoMemoryEffect,
17    PropagatesUniformity, Pure,
18    attributes::{BoolAttr, IndexAttr},
19    dialect::{
20        math::{index_attr, int_attr},
21        pure_binop, pure_unop,
22    },
23    interfaces::{
24        ScalarType, TriviallyUnrollable, TypedExt,
25        aliasing::AliasingOp,
26        uniformity::{UniformOpInterface, Uniformity},
27    },
28    prelude::*,
29    types::scalar::IndexType,
30};
31
32#[derive(Error, Debug)]
33pub enum SymbolUserOpVerifyErr {
34    #[error("Symbol {0} not found")]
35    SymbolNotFound(String),
36    #[error("Function {0} should have been builtin.func type")]
37    NotFunc(String),
38    #[error("Function call has incorrect type: {0}")]
39    FuncTypeErr(String),
40}
41
42#[cube_op(name = "cube.copy")]
43#[result_ty(same_as = value)]
44#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
45pub struct CopyOp {
46    pub value: Value,
47}
48
49simplify!(CopyOp, {
50    |_| Some(self.value(ctx)),
51});
52
53#[op_interface_impl]
54impl AliasingOp for CopyOp {
55    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
56        Some(self.value(ctx))
57    }
58}
59
60#[cube_op(name = "cube.poison")]
61#[result_ty(argument)]
62#[op_traits(Pure, CanMaterialize)]
63pub struct PoisonOp {}
64
65pure_binop!("cube.bool_and", BoolAndOp);
66Commutative!(BoolAndOp);
67const_eval!(BoolAndOp, {
68    BoolAttr: |lhs, rhs| lhs && rhs,
69    // false && x -> false
70    custom: |lhs, _| match lhs?.as_const_val(ctx) {
71        ConstantValue::Bool(false) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
72        _ => None
73    },
74    // x && false -> false
75    custom: |_, rhs| match rhs?.as_const_val(ctx) {
76        ConstantValue::Bool(false) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
77        _ => None
78    }
79});
80simplify!(BoolAndOp, {
81    // true && x -> x
82    |lhs, _| match lhs?.as_const_val(ctx) {
83        ConstantValue::Bool(true) => Some(self.rhs(ctx)),
84        _ => None,
85    },
86    // x && true -> x
87    |_, rhs| match rhs?.as_const_val(ctx) {
88        ConstantValue::Bool(true) => Some(self.lhs(ctx)),
89        _ => None,
90    },
91    // x && x -> x
92    |_, _| match self.lhs(ctx) == self.rhs(ctx) {
93        true => Some(self.lhs(ctx)),
94        false => None
95    }
96});
97
98pure_binop!("cube.bool_or", BoolOrOp);
99Commutative!(BoolOrOp);
100const_eval!(BoolOrOp, {
101    BoolAttr: |lhs, rhs| lhs || rhs,
102    // true || x -> true
103    custom: |lhs, _| match lhs?.as_const_val(ctx) {
104        ConstantValue::Bool(true) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
105        _ => None
106    },
107    // x || true -> true
108    custom: |_, rhs| match rhs?.as_const_val(ctx) {
109        ConstantValue::Bool(true) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
110        _ => None
111    }
112});
113simplify!(BoolOrOp, {
114    // false || x -> x
115    |lhs, _| match lhs?.as_const_val(ctx) {
116        ConstantValue::Bool(false) => Some(self.rhs(ctx)),
117        _ => None,
118    },
119    // false || x -> x
120    |_, rhs| match rhs?.as_const_val(ctx) {
121        ConstantValue::Bool(false) => Some(self.lhs(ctx)),
122        _ => None,
123    },
124    // x || x -> x
125    |_, _| match self.lhs(ctx) == self.rhs(ctx) {
126        true => Some(self.lhs(ctx)),
127        false => None
128    }
129});
130
131pure_unop!("cube.bool_not", BoolNotOp);
132const_eval!(BoolNotOp, {
133    BoolAttr: |inp| !inp
134});
135
136#[cube_op(name = "cube.cast")]
137#[result_ty(argument)]
138#[op_interfaces(TriviallyUnrollable)]
139#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
140pub struct CastOp {
141    pub input: Value,
142}
143const_eval!(CastOp, {
144    custom: |inp| {
145        // TF32 uses FP32 storage, but casts require backend-specific rounding.
146        // Keep the cast until lowering instead of folding it as an FP32 alias.
147        if self.result_type(ctx).scalar_ty(ctx).is_tfloat32(ctx) {
148            return None;
149        }
150        let val = inp?.as_const_val(ctx);
151        let out_ty = self.get_result(ctx).get_type(ctx).deref(ctx);
152        let elem = type_cast::<dyn ScalarType>(&*out_ty)?.elem_type(ctx);
153        Some(val.cast_to(elem).as_attribute(ctx, elem))
154    }
155});
156simplify!(CastOp, {
157    |_| {
158        if self.input(ctx).get_type(ctx) == self.result_type(ctx) {
159            Some(self.input(ctx))
160        } else {
161            None
162        }
163    }
164});
165
166#[cube_op(name = "cube.reinterpret_cast")]
167#[result_ty(argument)]
168#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
169pub struct ReinterpretCastOp {
170    pub input: Value,
171}
172const_eval!(ReinterpretCastOp, {
173    custom: |inp| {
174        // Too much weirdness around floats, don't bother dealing with it
175        let val = match inp?.as_const_val(ctx) {
176            ConstantValue::Int(val) => val as u64,
177            ConstantValue::UInt(val) => val,
178            _ => None?,
179        };
180        let out_ty = self.get_result(ctx).get_type(ctx);
181        if out_ty.is_int(ctx) {
182            Some(int_attr(ctx, out_ty, val as i128))
183        } else if out_ty.is_index(ctx) {
184            Some(index_attr(val as usize))
185        } else {
186            None
187        }
188    }
189});
190simplify!(ReinterpretCastOp, {
191    |_| {
192        if self.input(ctx).get_type(ctx) == self.result_type(ctx) {
193            Some(self.input(ctx))
194        } else {
195            None
196        }
197    }
198});
199
200#[op_interface_impl]
201impl AliasingOp for ReinterpretCastOp {
202    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
203        Some(self.input(ctx))
204    }
205}
206
207#[cube_op(name = "cube.select")]
208#[result_ty(same_as = true_value)]
209#[op_interfaces(TriviallyUnrollable)]
210#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
211pub struct SelectOp {
212    pub condition: Value,
213    pub true_value: Value,
214    pub false_value: Value,
215}
216simplify!(SelectOp, {
217    |cond, _, _| match cond?.as_const_val(ctx) {
218        ConstantValue::Bool(true) => Some(self.true_value(ctx)),
219        ConstantValue::Bool(false) => Some(self.false_value(ctx)),
220        _ => None,
221    },
222    // select(cond, x, x) -> x
223    |_, _, _| match self.true_value(ctx) == self.false_value(ctx) {
224        true => Some(self.true_value(ctx)),
225        false => None
226    }
227});
228
229#[pliron_attr(name = "cube.builtin", format, verifier = "succ")]
230#[derive(new, From, PartialEq, Clone, Debug, Hash)]
231pub struct BuiltinAttr(pub Builtin);
232
233#[cube_op(
234    name = "cube.read_builtin",
235    format = "attr($builtin, $BuiltinAttr) ` : ` type($0)"
236)]
237#[result_ty(argument)]
238#[op_traits(Pure, CanMaterialize)]
239pub struct ReadBuiltinOp {
240    pub builtin: BuiltinAttr,
241}
242
243#[op_interface_impl]
244impl UniformOpInterface for ReadBuiltinOp {
245    fn uniformity(&self, ctx: &Context, _operands: &[Uniformity]) -> Uniformity {
246        match self.builtin(ctx).0 {
247            Builtin::CubeDim
248            | Builtin::CubeDimX
249            | Builtin::CubeDimY
250            | Builtin::CubeDimZ
251            | Builtin::CubeClusterDim
252            | Builtin::CubeClusterDimX
253            | Builtin::CubeClusterDimY
254            | Builtin::CubeClusterDimZ
255            | Builtin::CubeCount
256            | Builtin::CubeCountX
257            | Builtin::CubeCountY
258            | Builtin::CubeCountZ
259            | Builtin::PlaneDim => Uniformity::Device,
260            Builtin::CubePosCluster
261            | Builtin::CubePosClusterX
262            | Builtin::CubePosClusterY
263            | Builtin::CubePosClusterZ
264            | Builtin::CubePos
265            | Builtin::CubePosX
266            | Builtin::CubePosY
267            | Builtin::CubePosZ => Uniformity::Cube,
268            Builtin::PlanePos => Uniformity::Plane,
269            Builtin::UnitPos
270            | Builtin::UnitPosX
271            | Builtin::UnitPosY
272            | Builtin::UnitPosZ
273            | Builtin::UnitPosPlane
274            | Builtin::AbsolutePos
275            | Builtin::AbsolutePosX
276            | Builtin::AbsolutePosY
277            | Builtin::AbsolutePosZ => Uniformity::None,
278        }
279    }
280}
281
282#[cube_op(
283    name = "cube.read_scalar",
284    format = "`<` attr($ty, $TypeAttr) `>[` attr($id, $IndexAttr) `] : ` type($0)"
285)]
286#[result_ty(from_inputs = |ctx, ty: &TypeAttr, _| ty.get_type(ctx))]
287#[op_traits(Pure, CanMaterialize)]
288pub struct ReadScalarOp {
289    pub ty: TypeAttr,
290    pub id: IndexAttr,
291}
292
293#[op_interface_impl]
294impl UniformOpInterface for ReadScalarOp {
295    fn uniformity(&self, _ctx: &Context, _operands: &[Uniformity]) -> Uniformity {
296        Uniformity::Device
297    }
298}
299
300#[cube_op(name = "cube.free")]
301#[result_ty(none)]
302pub struct FreeOp {
303    pub memory: Value,
304}
305
306#[cube_op(
307    name = "cube.buffer_len",
308    format = "`buffer_` attr($buffer_idx, $IndexAttr) `: ` type($0)"
309)]
310#[result_ty(fixed = IndexType::get(ctx).into())]
311#[op_traits(Pure, CanMaterialize)]
312pub struct BufferLenOp {
313    pub buffer_idx: IndexAttr,
314}
315
316#[op_interface_impl]
317impl UniformOpInterface for BufferLenOp {
318    fn uniformity(&self, _ctx: &Context, _operands: &[Uniformity]) -> Uniformity {
319        Uniformity::Device
320    }
321}
322
323#[cube_op(name = "cube.shape")]
324#[result_ty(fixed = IndexType::get(ctx).into())]
325#[op_traits(Pure, CanMaterialize)]
326pub struct ShapeOp {
327    pub dim: Value,
328    pub buffer_idx: IndexAttr,
329}
330
331#[op_interface_impl]
332impl UniformOpInterface for ShapeOp {
333    fn uniformity(&self, _ctx: &Context, _operands: &[Uniformity]) -> Uniformity {
334        Uniformity::Device
335    }
336}
337
338#[cube_op(name = "cube.stride")]
339#[result_ty(fixed = IndexType::get(ctx).into())]
340#[op_traits(Pure, CanMaterialize)]
341pub struct StrideOp {
342    pub dim: Value,
343    pub buffer_idx: IndexAttr,
344}
345
346#[op_interface_impl]
347impl UniformOpInterface for StrideOp {
348    fn uniformity(&self, _ctx: &Context, _operands: &[Uniformity]) -> Uniformity {
349        Uniformity::Device
350    }
351}
352
353#[cube_op(name = "cube.comment")]
354#[result_ty(none)]
355#[op_traits(NoMemoryEffect, HasSideEffects)]
356pub struct CommentOp {
357    pub comment: StringAttr,
358}
359
360#[pliron_op(name = "cube.printf", format, attributes = (cube_printf_format_string: StringAttr), verifier = "succ")]
361#[op_traits(NoMemoryEffect, HasSideEffects)]
362pub struct PrintfOp;
363
364impl PrintfOp {
365    pub fn new(ctx: &mut Context, format_string: String, values: Vec<Value>) -> Self {
366        let op = Self {
367            op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], values, vec![], 0),
368        };
369        op.set_attr_cube_printf_format_string(ctx, StringAttr::new(format_string));
370        op
371    }
372
373    pub fn format_string<'a>(&self, ctx: &'a Context) -> Ref<'a, StringAttr> {
374        self.get_attr_cube_printf_format_string(ctx).unwrap()
375    }
376
377    pub fn args(&self, ctx: &Context) -> Vec<Value> {
378        self.get_operation().deref(ctx).operands().collect()
379    }
380}
cubecl_ir/dialect/general.rs

cubecl_ir/dialect/
general.rs
