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

cubecl_ir/dialect/
branch.rs
