cubecl_ir/dialect/
branch.rs

1use pliron::{
2    attribute::AttrObj,
3    basic_block::BasicBlock,
4    builtin::attributes::IntegerAttr,
5    irbuild::inserter::OpInsertionPoint,
6    linked_list::ContainsLinkedList,
7    opts::dce::SideEffects,
8    region::Region,
9    utils::{
10        const_bound_n::I,
11        table::{HMap, SmallMap},
12    },
13    verify_err,
14};
15use thiserror::Error;
16
17use crate::{
18    CanMaterialize, NoMemoryEffect, ReturnLike,
19    attributes::{BoolAttr, IntegerVecAttr, ZeroAttr},
20    dialect::scf::block_mem_val,
21    interfaces::{
22        CanonicalizeInterface,
23        control_flow::{
24            InvocationBounds, RegionBranchOpInterface, RegionBranchTerminatorOpInterface,
25            RegionPredecessor, RegionSuccessor,
26        },
27        memory_slot::{
28            MemoryRegionPredecessor, MemorySSAContext, MemorySSARegionOpInterface, MemoryValue,
29            RegionMemoryPhiInputs, RegionMemoryValue,
30        },
31        uniformity::{UniformRegionTerminatorOpInterface, Uniformity},
32    },
33    prelude::*,
34    small_map,
35    types::scalar::BoolType,
36};
37
38/// Marker for terminators that do not return, i.e. `UnreachableOp`. In Rust terms, they return `!`.
39#[op_interface]
40pub trait IsExitTerminator {
41    fn verify(_op: &dyn Op, _ctx: &Context) -> Result<()>
42    where
43        Self: Sized,
44    {
45        Ok(())
46    }
47}
48
49#[derive(Error, Debug)]
50pub enum YieldOpVerifyErr {
51    #[error("YieldOp operand types do not match parent operation result types")]
52    OperandTypeMismatch,
53    #[error("YieldOp must have a parent operation to verify against")]
54    MissingParentOp,
55}
56
57#[pliron_op(name = "branch.yield", format = "`(` operands(CharSpace(`,`)) `)`")]
58#[op_interfaces(IsTerminatorInterface, NResultsInterface<0>)]
59#[op_traits(NoMemoryEffect, ReturnLike, CanMaterialize)]
60pub struct YieldOp;
61
62impl YieldOp {
63    pub fn new(ctx: &mut Context) -> Self {
64        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);
65        Self { op }
66    }
67
68    pub fn yield_values(&self, ctx: &Context) -> Vec<Value> {
69        self.get_operation().operands(ctx)
70    }
71}
72
73impl Verify for YieldOp {
74    fn verify(&self, ctx: &Context) -> pliron::result::Result<()> {
75        let Some(parent_op) = self.get_operation().deref(ctx).get_parent_op(ctx) else {
76            return verify_err!(self.loc(ctx), YieldOpVerifyErr::MissingParentOp);
77        };
78
79        let expected_types: Vec<_> = parent_op
80            .deref(ctx)
81            .results()
82            .map(|r| r.get_type(ctx))
83            .collect();
84        let actual_types: Vec<_> = self
85            .get_operation()
86            .deref(ctx)
87            .operands()
88            .map(|o| o.get_type(ctx))
89            .collect();
90
91        if expected_types != actual_types {
92            return verify_err!(self.loc(ctx), YieldOpVerifyErr::OperandTypeMismatch);
93        }
94
95        Ok(())
96    }
97}
98
99#[pliron_op(name = "branch.condition", format = "`(` operands(CharSpace(`,`)) `)`")]
100#[op_interfaces(IsTerminatorInterface, NResultsInterface<0>, OperandNOfType<0, BoolType>)]
101#[op_traits(CanMaterialize, NoMemoryEffect)]
102pub struct ConditionOp;
103
104impl ConditionOp {
105    pub fn new(ctx: &mut Context, cond: Value) -> Self {
106        let op = Operation::new(
107            ctx,
108            Self::get_concrete_op_info(),
109            vec![],
110            vec![cond],
111            vec![],
112            0,
113        );
114        Self { op }
115    }
116
117    pub fn condition(&self, ctx: &Context) -> Value {
118        self.get_operation().operand(ctx, 0)
119    }
120
121    pub fn forward_values(&self, ctx: &Context) -> Vec<Value> {
122        self.get_operation().deref(ctx).operands().skip(1).collect()
123    }
124}
125
126impl Verify for ConditionOp {
127    fn verify(&self, ctx: &Context) -> pliron::result::Result<()> {
128        let Some(parent_op) = self.get_operation().deref(ctx).get_parent_op(ctx) else {
129            return verify_err!(self.loc(ctx), YieldOpVerifyErr::MissingParentOp);
130        };
131
132        let expected_types: Vec<_> = parent_op
133            .deref(ctx)
134            .results()
135            .map(|r| r.get_type(ctx))
136            .collect();
137        let actual_types: Vec<_> = self
138            .forward_values(ctx)
139            .into_iter()
140            .map(|o| o.get_type(ctx))
141            .collect();
142
143        if expected_types != actual_types {
144            return verify_err!(self.loc(ctx), YieldOpVerifyErr::OperandTypeMismatch);
145        }
146
147        Ok(())
148    }
149}
150
151#[op_interface_impl]
152impl RegionBranchTerminatorOpInterface for ConditionOp {
153    fn successor_operands(&self, ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
154        self.forward_values(ctx)
155    }
156
157    fn successor_regions(
158        &self,
159        ctx: &Context,
160        operands: &[Option<AttrObj>],
161    ) -> Vec<RegionSuccessor> {
162        let while_op = self.get_operation().deref(ctx).get_parent_op(ctx).unwrap();
163        let after_region = while_op.deref(ctx).get_region(1).into();
164        let Some(attr) = operands[0].as_ref() else {
165            return vec![after_region, RegionSuccessor::AfterOp];
166        };
167        let zero = attr.downcast_ref::<ZeroAttr>().map(|_| false);
168        let bool = attr.downcast_ref::<BoolAttr>().map(|it| it.0);
169        let Some(const_cond) = zero.or(bool) else {
170            return vec![after_region, RegionSuccessor::AfterOp];
171        };
172        match const_cond {
173            true => vec![after_region],
174            false => vec![RegionSuccessor::AfterOp],
175        }
176    }
177}
178
179#[op_interface_impl]
180impl UniformRegionTerminatorOpInterface for ConditionOp {
181    fn successor_region_uniformity(
182        &self,
183        ctx: &Context,
184        operands: &[Uniformity],
185    ) -> Vec<Uniformity> {
186        self.all_successor_regions(ctx)
187            .iter()
188            .map(|_| operands[0])
189            .collect()
190    }
191}
192
193#[pliron_op(
194    name = "branch.return",
195    format = "operands(CharSpace(`,`))",
196    verifier = "succ"
197)]
198#[op_interfaces(IsTerminatorInterface, NResultsInterface<0>, IsExitTerminator)]
199#[op_traits(ReturnLike, CanMaterialize, NoMemoryEffect)]
200pub struct ReturnOp;
201
202impl ReturnOp {
203    pub fn new(ctx: &mut Context) -> Self {
204        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);
205        Self { op }
206    }
207
208    pub fn new_with_value(ctx: &mut Context, value: Value) -> Self {
209        let op = Operation::new(
210            ctx,
211            Self::get_concrete_op_info(),
212            vec![],
213            vec![value],
214            vec![],
215            0,
216        );
217        Self { op }
218    }
219
220    pub fn value(&self, ctx: &Context) -> Option<Value> {
221        self.get_operation().deref(ctx).results().next()
222    }
223}
224
225#[pliron_op(name = "branch.unreachable", format = "", verifier = "succ")]
226#[op_interfaces(IsTerminatorInterface, IsExitTerminator)]
227#[op_traits(CanMaterialize, NoMemoryEffect)]
228pub struct UnreachableOp;
229
230impl UnreachableOp {
231    pub fn new(ctx: &mut Context) -> Self {
232        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);
233        Self { op }
234    }
235}
236
237pub(super) fn block_side_effects(ctx: &Context, block: Ptr<BasicBlock>) -> bool {
238    block.deref(ctx).iter(ctx).any(|op| {
239        // Yield should not count as an effect in a region, but also can't implement
240        // `SideEffects = true` because then it would immediately get eliminated
241        if op.is_op::<YieldOp>(ctx) {
242            return false;
243        }
244        match op_cast::<dyn SideEffects>(&*op.dyn_op(ctx)) {
245            Some(side_effects) => side_effects.has_side_effects(ctx),
246            None => true,
247        }
248    })
249}
250
251#[pliron_op(
252    name = "branch.if",
253    format = "$0 ` then ` region($0) ` else ` region($1)",
254    verifier = "succ"
255)]
256#[op_interfaces(NOpdsInterface<1>, NResultsInterface<0>, NRegionsInterface<2>, SingleBlockRegionInterface, OperandNOfType<0, BoolType>)]
257pub struct IfOp;
258
259impl IfOp {
260    pub fn new(ctx: &mut Context, cond: Value) -> Self {
261        let op = Operation::new(
262            ctx,
263            Self::get_concrete_op_info(),
264            vec![],
265            vec![cond],
266            vec![],
267            2,
268        );
269
270        let then_region = op.deref_mut(ctx).get_region(0);
271        let then_body = BasicBlock::new(ctx, Some("then".try_into().unwrap()), vec![]);
272        then_body.insert_at_front(then_region, ctx);
273
274        let else_region = op.deref_mut(ctx).get_region(1);
275        let else_body = BasicBlock::new(ctx, Some("else".try_into().unwrap()), vec![]);
276        else_body.insert_at_front(else_region, ctx);
277
278        Self { op }
279    }
280
281    pub fn condition(&self, ctx: &Context) -> Value {
282        self.get_operation().deref(ctx).get_operand(0)
283    }
284
285    pub fn then_region(&self, ctx: &Context) -> Ptr<Region> {
286        self.get_operation().deref(ctx).get_region(0)
287    }
288
289    pub fn then_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
290        self.get_body(ctx, 0)
291    }
292
293    pub fn else_region(&self, ctx: &Context) -> Ptr<Region> {
294        self.get_operation().deref(ctx).get_region(1)
295    }
296
297    pub fn else_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
298        self.get_body(ctx, 1)
299    }
300}
301
302fn inline_block(
303    ctx: &Context,
304    rewriter: &mut dyn Rewriter,
305    block: Ptr<BasicBlock>,
306    insertion_point: OpInsertionPoint,
307) {
308    let ops = block.deref(ctx).iter(ctx).collect::<Vec<_>>();
309    let mut insertion_pt = insertion_point;
310    for op in ops {
311        if !op.is_terminator(ctx) {
312            rewriter.move_operation(ctx, op, insertion_pt);
313            insertion_pt = OpInsertionPoint::AfterOperation(op);
314        }
315    }
316}
317
318#[op_interface_impl]
319impl SideEffects for IfOp {
320    fn has_side_effects(&self, ctx: &Context) -> bool {
321        block_side_effects(ctx, self.then_block(ctx))
322            || block_side_effects(ctx, self.else_block(ctx))
323    }
324}
325
326#[op_interface_impl]
327impl MemorySSARegionOpInterface for IfOp {
328    fn setup_memory_ssa(
329        &self,
330        ctx: &Context,
331        _state: &mut MemorySSAContext,
332        reaching_def: MemoryValue,
333        _has_memory_defs: bool,
334        regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
335    ) {
336        regions_to_process.insert(self.then_region(ctx), reaching_def);
337        regions_to_process.insert(self.else_region(ctx), reaching_def);
338    }
339
340    fn finalize_memory_ssa(
341        &self,
342        ctx: &Context,
343        _state: &mut MemorySSAContext,
344        entry_reaching_def: MemoryValue,
345        has_memory_defs: bool,
346        _reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
347        reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
348        _region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
349    ) -> RegionMemoryValue {
350        if !has_memory_defs {
351            return RegionMemoryValue::Forward(entry_reaching_def);
352        }
353
354        let (then_pred, reaching_then) = block_mem_val(
355            self.then_block(ctx),
356            entry_reaching_def,
357            reaching_at_block_end,
358        );
359        let (else_pred, reaching_else) = block_mem_val(
360            self.else_block(ctx),
361            entry_reaching_def,
362            reaching_at_block_end,
363        );
364
365        RegionMemoryValue::RegionPhi(small_map! {
366            then_pred => reaching_then,
367            else_pred => reaching_else
368        })
369    }
370}
371
372#[op_interface_impl]
373impl RegionBranchOpInterface for IfOp {
374    fn entry_successor_regions(
375        &self,
376        ctx: &Context,
377        operands: &[Option<AttrObj>],
378    ) -> Vec<RegionSuccessor> {
379        let Some(attr) = operands[0].as_ref() else {
380            return self.successor_regions(ctx, RegionPredecessor::Parent);
381        };
382        let zero = attr.downcast_ref::<ZeroAttr>().map(|_| false);
383        let bool = attr.downcast_ref::<BoolAttr>().map(|it| it.0);
384        let Some(const_cond) = zero.or(bool) else {
385            return self.successor_regions(ctx, RegionPredecessor::Parent);
386        };
387        match const_cond {
388            true => vec![self.then_region(ctx).into()],
389            false => vec![self.else_region(ctx).into()],
390        }
391    }
392
393    fn successor_regions(&self, ctx: &Context, pred: RegionPredecessor) -> Vec<RegionSuccessor> {
394        match pred {
395            RegionPredecessor::Parent => {
396                vec![self.then_region(ctx).into(), self.else_region(ctx).into()]
397            }
398            RegionPredecessor::Terminator(_) => {
399                vec![RegionSuccessor::AfterOp]
400            }
401        }
402    }
403
404    fn successor_inputs(&self, _ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
405        vec![]
406    }
407
408    fn region_invocation_bounds(
409        &self,
410        _ctx: &Context,
411        operands: &[Option<AttrObj>],
412    ) -> Vec<InvocationBounds> {
413        if let Some(cond) = operands[0]
414            .as_ref()
415            .and_then(|it| it.downcast_ref::<BoolAttr>())
416        {
417            match cond.0 {
418                true => vec![InvocationBounds::once(), InvocationBounds::never()],
419                false => vec![InvocationBounds::never(), InvocationBounds::once()],
420            }
421        } else {
422            vec![InvocationBounds::zero_or_one(); 2]
423        }
424    }
425}
426
427impl IfOp {
428    fn fold(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
429        let op = self.get_operation();
430        let operands = const_operands(ctx, op);
431        let valid_branches = self.entry_successor_regions(ctx, &operands);
432        let &[RegionSuccessor::Region(taken)] = valid_branches.as_slice() else {
433            return Ok(());
434        };
435        let taken = taken.deref(ctx).get_entry_block().unwrap();
436
437        inline_block(ctx, rewriter, taken, OpInsertionPoint::BeforeOperation(op));
438        rewriter.erase_operation(ctx, op);
439
440        Ok(())
441    }
442}
443
444#[op_interface_impl]
445impl CanonicalizeInterface for IfOp {
446    fn canonicalize(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
447        self.fold(ctx, rewriter)?;
448        Ok(())
449    }
450}
451
452#[pliron_op(
453    name = "branch.switch",
454    format,
455    attributes = (branch_switch_cases: IntegerVecAttr),
456    verifier = "succ"
457)]
458#[op_interfaces(NOpdsInterface<1>, NResultsInterface<0>, SingleBlockRegionInterface)]
459pub struct SwitchOp;
460
461impl SwitchOp {
462    pub fn new(ctx: &mut Context, value: Value) -> Self {
463        let op = Operation::new(
464            ctx,
465            Self::get_concrete_op_info(),
466            vec![],
467            vec![value],
468            vec![],
469            1,
470        );
471
472        let default_region = op.deref_mut(ctx).get_region(0);
473        let default_body = BasicBlock::new(ctx, Some("default".try_into().unwrap()), vec![]);
474        default_body.insert_at_front(default_region, ctx);
475
476        Self { op }
477    }
478
479    pub fn value(&self, ctx: &Context) -> Value {
480        self.get_operation().deref(ctx).get_operand(0)
481    }
482
483    pub fn default_region(&self, ctx: &Context) -> Ptr<Region> {
484        self.get_operation().deref(ctx).get_region(0)
485    }
486
487    pub fn default_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
488        self.get_body(ctx, 0)
489    }
490
491    pub fn append_case_block(&self, ctx: &mut Context) -> Ptr<BasicBlock> {
492        let region = Operation::add_region(self.get_operation(), ctx);
493        let body = BasicBlock::new(ctx, None, vec![]);
494        body.insert_at_front(region, ctx);
495        region.deref(ctx).get_head().unwrap()
496    }
497
498    pub fn cases(&self, ctx: &Context) -> Vec<(IntegerAttr, Ptr<BasicBlock>)> {
499        let cases = self.get_attr_branch_switch_cases(ctx).unwrap().clone().0;
500        let out = (0..cases.len()).map(|i| {
501            let value = cases[i].clone();
502            let block = self.get_body(ctx, i + 1);
503            (value, block)
504        });
505        out.collect()
506    }
507
508    pub fn cases_regions(&self, ctx: &Context) -> Vec<(IntegerAttr, Ptr<Region>)> {
509        let cases = self.get_attr_branch_switch_cases(ctx).unwrap().clone().0;
510        let out = (0..cases.len()).map(|i| {
511            let value = cases[i].clone();
512            let block = self.get_operation().deref(ctx).get_region(i + 1);
513            (value, block)
514        });
515        out.collect()
516    }
517
518    pub fn cases_values(&self, ctx: &Context) -> Vec<IntegerAttr> {
519        self.get_attr_branch_switch_cases(ctx).unwrap().0.clone()
520    }
521
522    pub fn get_case_destinations(&self, ctx: &Context) -> Vec<Ptr<BasicBlock>> {
523        let op = self.get_operation().deref(ctx);
524        (1..op.regions().count())
525            .map(|i| self.get_body(ctx, i))
526            .collect()
527    }
528
529    pub fn set_attr_cases(&self, ctx: &Context, cases: impl IntoIterator<Item = IntegerAttr>) {
530        self.set_attr_branch_switch_cases(ctx, IntegerVecAttr(cases.into_iter().collect()));
531    }
532}
533
534#[op_interface_impl]
535impl MemorySSARegionOpInterface for SwitchOp {
536    fn setup_memory_ssa(
537        &self,
538        ctx: &Context,
539        _state: &mut MemorySSAContext,
540        reaching_def: MemoryValue,
541        _has_memory_defs: bool,
542        regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
543    ) {
544        regions_to_process.insert(self.default_region(ctx), reaching_def);
545        for (_, case_region) in self.cases_regions(ctx) {
546            regions_to_process.insert(case_region, reaching_def);
547        }
548    }
549
550    fn finalize_memory_ssa(
551        &self,
552        ctx: &Context,
553        _state: &mut MemorySSAContext,
554        entry_reaching_def: MemoryValue,
555        has_memory_defs: bool,
556        _reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
557        reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
558        _region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
559    ) -> RegionMemoryValue {
560        if !has_memory_defs {
561            return RegionMemoryValue::Forward(entry_reaching_def);
562        }
563
564        let mut phi_inputs = SmallMap::new();
565
566        let (default_pred, reaching_default) = block_mem_val(
567            self.default_block(ctx),
568            entry_reaching_def,
569            reaching_at_block_end,
570        );
571        phi_inputs.insert(default_pred, reaching_default);
572
573        for (_, case_block) in self.cases(ctx) {
574            let (case_pred, reaching_case) =
575                block_mem_val(case_block, entry_reaching_def, reaching_at_block_end);
576            phi_inputs.insert(case_pred, reaching_case);
577        }
578
579        RegionMemoryValue::RegionPhi(phi_inputs)
580    }
581}
582
583#[op_interface_impl]
584impl RegionBranchOpInterface for SwitchOp {
585    fn entry_successor_regions(
586        &self,
587        ctx: &Context,
588        operands: &[Option<AttrObj>],
589    ) -> Vec<RegionSuccessor> {
590        let Some(attr) = &operands[0] else {
591            return self.successor_regions(ctx, RegionPredecessor::Parent);
592        };
593        let Some(attr) = attr.downcast_ref::<IntegerAttr>() else {
594            return self.successor_regions(ctx, RegionPredecessor::Parent);
595        };
596        if let Some(&(_, case)) = self.cases_regions(ctx).iter().find(|(val, _)| val == attr) {
597            vec![case.into()]
598        } else {
599            vec![self.default_region(ctx).into()]
600        }
601    }
602
603    fn successor_regions(&self, ctx: &Context, pred: RegionPredecessor) -> Vec<RegionSuccessor> {
604        match pred {
605            RegionPredecessor::Parent => {
606                let op = self.get_operation().deref(ctx);
607                op.regions().map(Into::into).collect()
608            }
609            RegionPredecessor::Terminator(_) => {
610                vec![RegionSuccessor::AfterOp]
611            }
612        }
613    }
614
615    fn successor_inputs(&self, _ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
616        vec![]
617    }
618
619    fn region_invocation_bounds(
620        &self,
621        ctx: &Context,
622        operands: &[Option<AttrObj>],
623    ) -> Vec<InvocationBounds> {
624        let num_regions = self.get_operation().deref(ctx).num_regions();
625        let Some(attr) = operands[0].as_ref() else {
626            return vec![InvocationBounds::zero_or_one(); num_regions];
627        };
628        let Some(attr) = attr.downcast_ref::<IntegerAttr>() else {
629            return vec![InvocationBounds::zero_or_one(); num_regions];
630        };
631
632        let case_idx = self.cases_values(ctx).iter().position(|it| it == attr);
633        let executed_idx = case_idx.map(|i| i + 1).unwrap_or(0);
634        let mut bounds = vec![InvocationBounds::never(); num_regions];
635        bounds[executed_idx] = InvocationBounds::once();
636        bounds
637    }
638}
639
640impl SwitchOp {
641    fn fold(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
642        let op = self.get_operation();
643        let operands = const_operands(ctx, op);
644        let valid_branches = self.entry_successor_regions(ctx, &operands);
645        let &[RegionSuccessor::Region(taken)] = valid_branches.as_slice() else {
646            return Ok(());
647        };
648        let taken = taken.deref(ctx).get_entry_block().unwrap();
649
650        inline_block(ctx, rewriter, taken, OpInsertionPoint::BeforeOperation(op));
651        rewriter.erase_operation(ctx, op);
652
653        Ok(())
654    }
655}
656
657#[op_interface_impl]
658impl CanonicalizeInterface for SwitchOp {
659    fn canonicalize(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
660        self.fold(ctx, rewriter)?;
661        Ok(())
662    }
663}
664
665#[pliron_op(name = "branch.range_loop", format, verifier = "succ")]
666#[op_interfaces(NResultsInterface<0>, OneRegionInterface, SingleBlockRegionInterface, SameOperandsType)]
667pub struct RangeLoopOp;
668
669impl RangeLoopOp {
670    pub fn new(ctx: &mut Context, start: Value, end: Value, step: Value) -> Self {
671        let iter_ty = start.get_type(ctx);
672        let op = Operation::new(
673            ctx,
674            Self::get_concrete_op_info(),
675            vec![],
676            vec![start, end, step],
677            vec![],
678            1,
679        );
680
681        let body_region = op.deref_mut(ctx).get_region(0);
682        let body = BasicBlock::new(ctx, Some("body".try_into().unwrap()), vec![iter_ty]);
683        body.insert_at_front(body_region, ctx);
684
685        Self { op }
686    }
687
688    pub fn iter_var(&self, ctx: &Context) -> Value {
689        self.loop_body(ctx).deref(ctx).get_argument(0)
690    }
691
692    pub fn start(&self, ctx: &Context) -> Value {
693        self.get_operation().deref(ctx).get_operand(0)
694    }
695
696    pub fn end(&self, ctx: &Context) -> Value {
697        self.get_operation().deref(ctx).get_operand(1)
698    }
699
700    pub fn step(&self, ctx: &Context) -> Value {
701        self.get_operation().deref(ctx).get_operand(2)
702    }
703
704    pub fn loop_region(&self, ctx: &Context) -> Ptr<Region> {
705        self.get_operation().deref(ctx).get_region(0)
706    }
707
708    pub fn loop_body(&self, ctx: &Context) -> Ptr<BasicBlock> {
709        self.get_body(ctx, 0)
710    }
711}
712
713#[op_interface_impl]
714impl MemorySSARegionOpInterface for RangeLoopOp {
715    fn setup_memory_ssa(
716        &self,
717        ctx: &Context,
718        state: &mut MemorySSAContext,
719        reaching_def: MemoryValue,
720        has_memory_defs: bool,
721        regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
722    ) {
723        let body_region = self.loop_region(ctx);
724        if !has_memory_defs {
725            regions_to_process.insert(body_region, reaching_def);
726            return;
727        }
728
729        let new_arg = state.new_value_in_block(self.loop_body(ctx));
730        regions_to_process.insert(body_region, new_arg);
731    }
732
733    fn finalize_memory_ssa(
734        &self,
735        ctx: &Context,
736        _state: &mut MemorySSAContext,
737        entry_reaching_def: MemoryValue,
738        has_memory_defs: bool,
739        _reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
740        reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
741        region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
742    ) -> RegionMemoryValue {
743        let body_region = self.loop_region(ctx);
744        if !has_memory_defs {
745            return RegionMemoryValue::Forward(entry_reaching_def);
746        }
747
748        let (body_pred, reaching_body) = block_mem_val(
749            self.loop_body(ctx),
750            entry_reaching_def,
751            reaching_at_block_end,
752        );
753
754        let phi_inputs = small_map! {
755            MemoryRegionPredecessor::Parent => entry_reaching_def,
756            body_pred => reaching_body
757        };
758
759        region_phis.insert(body_region, phi_inputs.clone());
760        RegionMemoryValue::RegionPhi(phi_inputs)
761    }
762}
763
764#[op_interface_impl]
765impl RegionBranchOpInterface for RangeLoopOp {
766    fn entry_successor_operands(&self, _ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
767        vec![]
768    }
769
770    fn successor_regions(&self, ctx: &Context, _pred: RegionPredecessor) -> Vec<RegionSuccessor> {
771        // TODO: Loop interface for constant trip count
772        vec![self.loop_region(ctx).into(), RegionSuccessor::AfterOp]
773    }
774
775    fn successor_inputs(&self, _ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
776        vec![]
777    }
778}
779
780#[pliron_op(
781    name = "branch.while",
782    format = "`while ` region($0) ` do ` region($1)",
783    verifier = "succ"
784)]
785#[op_interfaces(
786    NResultsInterface<0>,
787    NRegionsInterface<2>,
788    SingleBlockRegionInterface
789)]
790pub struct WhileOp;
791
792impl WhileOp {
793    pub fn new(ctx: &mut Context) -> Self {
794        let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 2);
795
796        let before_region = op.deref_mut(ctx).get_region(0);
797        let before = BasicBlock::new(ctx, Some("before".try_into().unwrap()), vec![]);
798        before.insert_at_front(before_region, ctx);
799
800        let after_region = op.deref_mut(ctx).get_region(1);
801        let after = BasicBlock::new(ctx, Some("after".try_into().unwrap()), vec![]);
802        after.insert_at_front(after_region, ctx);
803
804        Self { op }
805    }
806
807    pub fn before_region(&self, ctx: &Context) -> Ptr<Region> {
808        self.get_region_i(ctx, I::<0>.into())
809    }
810
811    pub fn before_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
812        self.get_body(ctx, 0)
813    }
814
815    pub fn after_region(&self, ctx: &Context) -> Ptr<Region> {
816        self.get_region_i(ctx, I::<1>.into())
817    }
818
819    pub fn after_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
820        self.get_body(ctx, 1)
821    }
822}
823
824#[op_interface_impl]
825impl MemorySSARegionOpInterface for WhileOp {
826    fn setup_memory_ssa(
827        &self,
828        ctx: &Context,
829        state: &mut MemorySSAContext,
830        reaching_def: MemoryValue,
831        has_memory_defs: bool,
832        regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
833    ) {
834        let before_region = self.before_region(ctx);
835        let after_region = self.after_region(ctx);
836        if !has_memory_defs {
837            regions_to_process.insert(before_region, reaching_def);
838            regions_to_process.insert(after_region, reaching_def);
839            return;
840        }
841
842        let new_arg = state.new_value_in_block(self.before_block(ctx));
843        regions_to_process.insert(before_region, new_arg);
844
845        let new_arg = state.new_value_in_block(self.after_block(ctx));
846        regions_to_process.insert(after_region, new_arg);
847    }
848
849    fn finalize_memory_ssa(
850        &self,
851        ctx: &Context,
852        _state: &mut MemorySSAContext,
853        entry_reaching_def: MemoryValue,
854        has_memory_defs: bool,
855        reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
856        reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
857        region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
858    ) -> RegionMemoryValue {
859        if !has_memory_defs {
860            return RegionMemoryValue::Forward(entry_reaching_def);
861        }
862
863        let before_region = self.before_region(ctx);
864        let after_region = self.after_region(ctx);
865
866        let arg = reaching_at_region_entry[&before_region];
867        let (before_pred, reaching_before) =
868            block_mem_val(self.before_block(ctx), arg, reaching_at_block_end);
869
870        let arg = reaching_at_region_entry[&after_region];
871        let (after_pred, reaching_after) =
872            block_mem_val(self.after_block(ctx), arg, reaching_at_block_end);
873
874        let inputs_before = small_map! {
875            MemoryRegionPredecessor::Parent => entry_reaching_def,
876            after_pred => reaching_after
877        };
878        let inputs_after = small_map!(before_pred => reaching_before);
879
880        region_phis.insert(before_region, inputs_before);
881        region_phis.insert(after_region, inputs_after);
882
883        RegionMemoryValue::Forward(reaching_before)
884    }
885}
886
887#[op_interface_impl]
888impl RegionBranchOpInterface for WhileOp {
889    fn entry_successor_operands(&self, _ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
890        vec![]
891    }
892
893    fn successor_regions(&self, ctx: &Context, pred: RegionPredecessor) -> Vec<RegionSuccessor> {
894        match pred {
895            RegionPredecessor::Parent => vec![self.before_region(ctx).into()],
896            RegionPredecessor::Terminator(term) => {
897                let op = term.deref(ctx).get_operation();
898                let parent = op.deref(ctx).get_parent_region(ctx).unwrap();
899                if parent == self.after_region(ctx) {
900                    vec![self.before_region(ctx).into()]
901                } else {
902                    vec![RegionSuccessor::AfterOp, self.after_region(ctx).into()]
903                }
904            }
905        }
906    }
907
908    fn successor_inputs(&self, _ctx: &Context, _successor: RegionSuccessor) -> Vec<Value> {
909        vec![]
910    }
911}
cubecl_ir/dialect/branch.rs

cubecl_ir/dialect/
branch.rs
