cubecl_ir/dialect/
branch.rs

1use pliron::{
2    attribute::AttrObj,
3    basic_block::BasicBlock,
4    builtin::attributes::{IntegerAttr, VecAttr},
5    irbuild::inserter::OpInsertionPoint,
6    linked_list::ContainsLinkedList,
7    opts::{constants::ConstFoldInterface, dce::SideEffects},
8    region::Region,
9    utils::const_bound_n::I,
10    verify_err,
11};
12use thiserror::Error;
13
14use crate::{
15    CanMaterialize, NoMemoryEffect, Pure,
16    attributes::{BoolAttr, ZeroAttr},
17    prelude::*,
18    types::scalar::BoolType,
19};
20
21/// Marker for terminators that do not return, i.e. `UnreachableOp`. In Rust terms, they return `!`.
22#[op_interface]
23pub trait IsExitTerminator {
24    fn verify(_op: &dyn Op, _ctx: &Context) -> Result<()>
25    where
26        Self: Sized,
27    {
28        Ok(())
29    }
30}
31
32#[derive(Error, Debug)]
33pub enum YieldOpVerifyErr {
34    #[error("YieldOp operand types do not match parent operation result types")]
35    OperandTypeMismatch,
36    #[error("YieldOp must have a parent operation to verify against")]
37    MissingParentOp,
38}
39
40#[pliron_op(name = "branch.yield", format = "`(` operands(CharSpace(`,`)) `)`")]
41#[op_interfaces(IsTerminatorInterface)]
42#[op_traits(CanMaterialize, NoMemoryEffect)]
43pub struct YieldOp;
44
45impl YieldOp {
46    pub fn new(ctx: &mut Context) -> Self {
47        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);
48        Self { op }
49    }
50
51    pub fn yield_values(&self, ctx: &Context) -> Vec<Value> {
52        self.get_operation().operands(ctx)
53    }
54}
55
56impl Verify for YieldOp {
57    fn verify(&self, ctx: &Context) -> pliron::result::Result<()> {
58        let Some(parent_op) = self.get_operation().deref(ctx).get_parent_op(ctx) else {
59            return verify_err!(self.loc(ctx), YieldOpVerifyErr::MissingParentOp);
60        };
61
62        let expected_types: Vec<_> = parent_op
63            .deref(ctx)
64            .results()
65            .map(|r| r.get_type(ctx))
66            .collect();
67        let actual_types: Vec<_> = self
68            .get_operation()
69            .deref(ctx)
70            .operands()
71            .map(|o| o.get_type(ctx))
72            .collect();
73
74        if expected_types != actual_types {
75            return verify_err!(self.loc(ctx), YieldOpVerifyErr::OperandTypeMismatch);
76        }
77
78        Ok(())
79    }
80}
81
82#[pliron_op(name = "branch.condition", format = "`(` operands(CharSpace(`,`)) `)`")]
83#[op_interfaces(IsTerminatorInterface, OperandNOfType<0, BoolType>)]
84#[op_traits(CanMaterialize, NoMemoryEffect)]
85pub struct ConditionOp;
86
87impl ConditionOp {
88    pub fn new(ctx: &mut Context, cond: Value) -> Self {
89        let op = Operation::new(
90            ctx,
91            Self::get_concrete_op_info(),
92            vec![],
93            vec![cond],
94            vec![],
95            0,
96        );
97        Self { op }
98    }
99
100    pub fn condition(&self, ctx: &Context) -> Value {
101        self.get_operation().operand(ctx, 0)
102    }
103
104    pub fn forward_values(&self, ctx: &Context) -> Vec<Value> {
105        self.get_operation().deref(ctx).operands().skip(1).collect()
106    }
107}
108
109impl Verify for ConditionOp {
110    fn verify(&self, ctx: &Context) -> pliron::result::Result<()> {
111        let Some(parent_op) = self.get_operation().deref(ctx).get_parent_op(ctx) else {
112            return verify_err!(self.loc(ctx), YieldOpVerifyErr::MissingParentOp);
113        };
114
115        let expected_types: Vec<_> = parent_op
116            .deref(ctx)
117            .results()
118            .map(|r| r.get_type(ctx))
119            .collect();
120        let actual_types: Vec<_> = self
121            .forward_values(ctx)
122            .into_iter()
123            .map(|o| o.get_type(ctx))
124            .collect();
125
126        if expected_types != actual_types {
127            return verify_err!(self.loc(ctx), YieldOpVerifyErr::OperandTypeMismatch);
128        }
129
130        Ok(())
131    }
132}
133
134#[pliron_op(
135    name = "branch.return",
136    format = "operands(CharSpace(`,`))",
137    verifier = "succ"
138)]
139#[op_interfaces(IsTerminatorInterface, IsExitTerminator)]
140#[op_traits(CanMaterialize, NoMemoryEffect)]
141pub struct ReturnOp;
142
143impl ReturnOp {
144    pub fn new(ctx: &mut Context) -> Self {
145        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);
146        Self { op }
147    }
148
149    pub fn new_with_value(ctx: &mut Context, value: Value) -> Self {
150        let op = Operation::new(
151            ctx,
152            Self::get_concrete_op_info(),
153            vec![],
154            vec![value],
155            vec![],
156            0,
157        );
158        Self { op }
159    }
160
161    pub fn value(&self, ctx: &Context) -> Option<Value> {
162        self.get_operation().deref(ctx).results().next()
163    }
164}
165
166#[pliron_op(name = "branch.unreachable", format = "", verifier = "succ")]
167#[op_interfaces(IsTerminatorInterface, IsExitTerminator)]
168#[op_traits(CanMaterialize, NoMemoryEffect)]
169pub struct UnreachableOp;
170
171impl UnreachableOp {
172    pub fn new(ctx: &mut Context) -> Self {
173        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);
174        Self { op }
175    }
176}
177
178/// Dead region for constant folding, returns a dummy result so it gets eliminated from dead code
179/// elimination. We can't erase the block straight away because it might contain SCCP candidates
180/// that are already tracked and will cause a dangling ptr deref.
181#[pliron_op(name = "branch.dead_region", format = "region($0)", verifier = "succ")]
182#[op_interfaces(NOpdsInterface<0>, OneResultInterface, OneRegionInterface, SingleBlockRegionInterface)]
183#[op_traits(Pure)]
184pub struct DeadRegionOp;
185
186impl DeadRegionOp {
187    pub fn new(ctx: &mut Context) -> Self {
188        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 1);
189
190        let region = op.deref_mut(ctx).get_region(0);
191        let body = BasicBlock::new(ctx, None, vec![]);
192        body.insert_at_front(region, ctx);
193
194        Self { op }
195    }
196
197    pub fn region(&self, ctx: &Context) -> Ptr<Region> {
198        self.get_operation().deref(ctx).get_region(0)
199    }
200}
201
202pub(super) fn block_side_effects(ctx: &Context, block: Ptr<BasicBlock>) -> bool {
203    block.deref(ctx).iter(ctx).any(|op| {
204        // Yield should not count as an effect in a region, but also can't implement
205        // `SideEffects = true` because then it would immediately get eliminated
206        if op.is_op::<YieldOp>(ctx) {
207            return false;
208        }
209        match op_cast::<dyn SideEffects>(&*op.dyn_op(ctx)) {
210            Some(side_effects) => side_effects.has_side_effects(ctx),
211            None => true,
212        }
213    })
214}
215
216#[pliron_op(
217    name = "branch.if",
218    format = "$0 ` then ` region($0) ` else ` region($1)",
219    verifier = "succ"
220)]
221#[op_interfaces(NOpdsInterface<1>, NResultsInterface<0>, NRegionsInterface<2>, SingleBlockRegionInterface, OperandNOfType<0, BoolType>)]
222pub struct IfOp;
223
224impl IfOp {
225    pub fn new(ctx: &mut Context, cond: Value) -> Self {
226        let op = Operation::new(
227            ctx,
228            Self::get_concrete_op_info(),
229            vec![],
230            vec![cond],
231            vec![],
232            2,
233        );
234
235        let then_region = op.deref_mut(ctx).get_region(0);
236        let then_body = BasicBlock::new(ctx, Some("then".try_into().unwrap()), vec![]);
237        then_body.insert_at_front(then_region, ctx);
238
239        let else_region = op.deref_mut(ctx).get_region(1);
240        let else_body = BasicBlock::new(ctx, Some("else".try_into().unwrap()), vec![]);
241        else_body.insert_at_front(else_region, ctx);
242
243        Self { op }
244    }
245
246    pub fn condition(&self, ctx: &Context) -> Value {
247        self.get_operation().deref(ctx).get_operand(0)
248    }
249
250    pub fn then_region(&self, ctx: &Context) -> Ptr<Region> {
251        self.get_operation().deref(ctx).get_region(0)
252    }
253
254    pub fn then_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
255        self.get_body(ctx, 0)
256    }
257
258    pub fn else_region(&self, ctx: &Context) -> Ptr<Region> {
259        self.get_operation().deref(ctx).get_region(1)
260    }
261
262    pub fn else_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
263        self.get_body(ctx, 1)
264    }
265}
266
267#[op_interface_impl]
268impl ConstFoldInterface for IfOp {
269    fn check_fold(
270        &self,
271        _ctx: &Context,
272        operand_attrs: &[Option<AttrObj>],
273    ) -> Vec<Option<AttrObj>> {
274        operand_attrs.to_vec()
275    }
276
277    fn fold_in_place(
278        &self,
279        ctx: &mut Context,
280        operand_attrs: &[Option<AttrObj>],
281        rewriter: &mut dyn Rewriter,
282    ) -> IRStatus {
283        let op = self.get_operation();
284        let Some(attr) = operand_attrs[0].as_ref() else {
285            return IRStatus::Unchanged;
286        };
287        let zero = attr.downcast_ref::<ZeroAttr>().map(|_| false);
288        let bool = attr.downcast_ref::<BoolAttr>().map(|it| it.0);
289        let Some(const_cond) = zero.or(bool) else {
290            return IRStatus::Unchanged;
291        };
292        let (taken, not_taken) = match const_cond {
293            true => (self.then_block(ctx), self.else_block(ctx)),
294            false => (self.else_block(ctx), self.then_block(ctx)),
295        };
296
297        let not_taken_op = DeadRegionOp::new(ctx);
298        let dead_block = not_taken_op.get_body(ctx, 0);
299        rewriter.append_op(ctx, &not_taken_op);
300
301        inline_block(ctx, rewriter, taken, OpInsertionPoint::BeforeOperation(op));
302        inline_block(
303            ctx,
304            rewriter,
305            not_taken,
306            OpInsertionPoint::AtBlockStart(dead_block),
307        );
308
309        IRStatus::Changed
310    }
311}
312
313fn inline_block(
314    ctx: &Context,
315    rewriter: &mut dyn Rewriter,
316    block: Ptr<BasicBlock>,
317    insertion_point: OpInsertionPoint,
318) {
319    let ops = block.deref(ctx).iter(ctx).collect::<Vec<_>>();
320    let mut insertion_pt = insertion_point;
321    for op in ops {
322        if !op.is_terminator(ctx) {
323            rewriter.move_operation(ctx, op, insertion_pt);
324            insertion_pt = OpInsertionPoint::AfterOperation(op);
325        }
326    }
327}
328
329#[op_interface_impl]
330impl SideEffects for IfOp {
331    fn has_side_effects(&self, ctx: &Context) -> bool {
332        block_side_effects(ctx, self.then_block(ctx))
333            || block_side_effects(ctx, self.else_block(ctx))
334    }
335}
336
337#[pliron_op(
338    name = "branch.switch",
339    format,
340    attributes = (branch_switch_cases: VecAttr),
341    verifier = "succ"
342)]
343#[op_interfaces(NOpdsInterface<1>, NResultsInterface<0>, SingleBlockRegionInterface)]
344pub struct SwitchOp;
345
346impl SwitchOp {
347    pub fn new(ctx: &mut Context, value: Value) -> Self {
348        let op = Operation::new(
349            ctx,
350            Self::get_concrete_op_info(),
351            vec![],
352            vec![value],
353            vec![],
354            1,
355        );
356
357        let default_region = op.deref_mut(ctx).get_region(0);
358        let default_body = BasicBlock::new(ctx, Some("default".try_into().unwrap()), vec![]);
359        default_body.insert_at_front(default_region, ctx);
360
361        Self { op }
362    }
363
364    pub fn value(&self, ctx: &Context) -> Value {
365        self.get_operation().deref(ctx).get_operand(0)
366    }
367
368    pub fn default_region(&self, ctx: &Context) -> Ptr<Region> {
369        self.get_operation().deref(ctx).get_region(0)
370    }
371
372    pub fn default_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
373        self.get_body(ctx, 0)
374    }
375
376    pub fn append_case_block(&self, ctx: &mut Context) -> Ptr<BasicBlock> {
377        let region = Operation::add_region(self.get_operation(), ctx);
378        let body = BasicBlock::new(ctx, None, vec![]);
379        body.insert_at_front(region, ctx);
380        region.deref(ctx).get_head().unwrap()
381    }
382
383    pub fn cases(&self, ctx: &Context) -> Vec<(IntegerAttr, Ptr<BasicBlock>)> {
384        let cases = self.get_attr_branch_switch_cases(ctx).unwrap().clone().0;
385        let out = (0..cases.len()).map(|i| {
386            let value = cases[i].downcast_ref::<IntegerAttr>().unwrap().clone();
387            let block = self.get_body(ctx, i + 1);
388            (value, block)
389        });
390        out.collect()
391    }
392
393    pub fn get_case_destinations(&self, ctx: &Context) -> Vec<Ptr<BasicBlock>> {
394        let op = self.get_operation().deref(ctx);
395        (1..op.regions().count())
396            .map(|i| self.get_body(ctx, i))
397            .collect()
398    }
399
400    pub fn set_attr_cases(&self, ctx: &Context, cases: impl IntoIterator<Item = AttrObj>) {
401        self.set_attr_branch_switch_cases(ctx, VecAttr(cases.into_iter().collect()));
402    }
403}
404
405#[pliron_op(name = "branch.range_loop", format, verifier = "succ")]
406#[op_interfaces(NResultsInterface<0>, OneRegionInterface, SingleBlockRegionInterface, SameOperandsType)]
407pub struct RangeLoopOp;
408
409impl RangeLoopOp {
410    pub fn new(ctx: &mut Context, start: Value, end: Value, step: Value) -> Self {
411        let iter_ty = start.get_type(ctx);
412        let op = Operation::new(
413            ctx,
414            Self::get_concrete_op_info(),
415            vec![],
416            vec![start, end, step],
417            vec![],
418            1,
419        );
420
421        let body_region = op.deref_mut(ctx).get_region(0);
422        let body = BasicBlock::new(ctx, Some("body".try_into().unwrap()), vec![iter_ty]);
423        body.insert_at_front(body_region, ctx);
424
425        Self { op }
426    }
427
428    pub fn iter_var(&self, ctx: &Context) -> Value {
429        self.loop_body(ctx).deref(ctx).get_argument(0)
430    }
431
432    pub fn start(&self, ctx: &Context) -> Value {
433        self.get_operation().deref(ctx).get_operand(0)
434    }
435
436    pub fn end(&self, ctx: &Context) -> Value {
437        self.get_operation().deref(ctx).get_operand(1)
438    }
439
440    pub fn step(&self, ctx: &Context) -> Value {
441        self.get_operation().deref(ctx).get_operand(2)
442    }
443
444    pub fn loop_region(&self, ctx: &Context) -> Ptr<Region> {
445        self.get_operation().deref(ctx).get_region(0)
446    }
447
448    pub fn loop_body(&self, ctx: &Context) -> Ptr<BasicBlock> {
449        self.get_body(ctx, 0)
450    }
451}
452
453#[pliron_op(
454    name = "branch.while",
455    format = "`while ` region($0) ` do ` region($1)",
456    verifier = "succ"
457)]
458#[op_interfaces(
459    NResultsInterface<0>,
460    NRegionsInterface<2>,
461    SingleBlockRegionInterface
462)]
463pub struct WhileOp;
464
465impl WhileOp {
466    pub fn new(ctx: &mut Context) -> Self {
467        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 2);
468
469        let before_region = op.deref_mut(ctx).get_region(0);
470        let before = BasicBlock::new(ctx, Some("before".try_into().unwrap()), vec![]);
471        before.insert_at_front(before_region, ctx);
472
473        let after_region = op.deref_mut(ctx).get_region(1);
474        let after = BasicBlock::new(ctx, Some("after".try_into().unwrap()), vec![]);
475        after.insert_at_front(after_region, ctx);
476
477        Self { op }
478    }
479
480    pub fn before_region(&self, ctx: &Context) -> Ptr<Region> {
481        self.get_region_i(ctx, I::<0>.into())
482    }
483
484    pub fn before_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
485        self.get_body(ctx, 0)
486    }
487
488    pub fn after_region(&self, ctx: &Context) -> Ptr<Region> {
489        self.get_region_i(ctx, I::<1>.into())
490    }
491
492    pub fn after_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
493        self.get_body(ctx, 1)
494    }
495}
cubecl_ir/dialect/branch.rs

cubecl_ir/dialect/
branch.rs
