cubecl_core/frontend/
branch.rs

1use alloc::vec::Vec;
2use cubecl_ir::{
3    ExpandState, ExpandValue, OpInserter,
4    dialect::{
5        branch::{ConditionOp, IfOp, RangeLoopOp, ReturnOp, SwitchOp, UnreachableOp, WhileOp},
6        general::BoolAndOp,
7    },
8    pliron::{irbuild::inserter::Inserter, r#type::TypedHandle},
9};
10use pliron::{
11    basic_block::BasicBlock,
12    builtin::{
13        attributes::IntegerAttr,
14        op_interfaces::OneRegionInterface,
15        types::{IntegerType, Signedness},
16    },
17    irbuild::{
18        listener::DummyListener,
19        rewriter::{IRRewriter, Rewriter},
20    },
21    op::Op,
22    region::Region,
23    r#type::Typed,
24    utils::apint::{APInt, bw},
25};
26
27use crate::{
28    IntoRuntime,
29    frontend::{ReadValue, RuntimeAssign, assign, assign_binop_expand, binary_expand},
30    prelude::{CubeEnum, ExpandTypeClone},
31};
32use crate::{ir::Scope, prelude::Assign};
33
34use super::{Int, NativeExpand};
35
36/// Something that can be iterated on by a for loop. Currently only includes `Range`, `StepBy` and
37/// `Sequence`.
38pub trait Iterable: Sized {
39    type Item;
40
41    /// Expand a runtime loop without unrolling
42    ///
43    /// # Arguments
44    /// # Arguments
45    /// * `scope` - the expansion scope
46    /// * `body` - the loop body to be executed repeatedly
47    fn expand(self, scope: &Scope, body: impl FnMut(&Scope, Self::Item));
48    /// Expand an unrolled loop. The body should be invoced `n` times, where `n` is the number of
49    /// iterations.
50    ///
51    /// # Arguments
52    /// * `scope` - the expansion scope
53    /// * `body` - the loop body to be executed repeatedly
54    fn expand_unroll(self, scope: &Scope, body: impl FnMut(&Scope, Self::Item));
55    /// Return the comptime length of this iterable, if possible
56    fn const_len(&self) -> Option<usize> {
57        None
58    }
59}
60
61pub fn for_expand<I: Iterable>(
62    scope: &Scope,
63    range: I,
64    unroll: bool,
65    body: impl FnMut(&Scope, I::Item),
66) {
67    if unroll || range.const_len() == Some(1) {
68        range.expand_unroll(scope, body);
69    } else {
70        range.expand(scope, body);
71    }
72}
73
74pub fn if_expand(scope: &Scope, condition: NativeExpand<bool>, block: impl FnOnce(&Scope)) {
75    let comptime_cond = condition.expand.as_const().map(|it| it.as_bool());
76    match comptime_cond {
77        Some(cond) => {
78            if cond {
79                block(scope);
80            }
81        }
82        None => {
83            let cond = condition.read_value(scope);
84            let if_op = IfOp::new(scope.ctx_mut(), cond);
85
86            let then_block = if_op.then_block(scope.ctx());
87            let then_child = scope.child(OpInserter::new_at_block_end(then_block));
88            block(&then_child);
89            then_child.terminate_yield();
90
91            let else_block = if_op.else_block(scope.ctx());
92            let else_child = scope.child(OpInserter::new_at_block_end(else_block));
93            else_child.terminate_yield();
94
95            scope.register(&if_op);
96            scope.set_break_return(&[then_child, else_child]);
97        }
98    }
99}
100
101#[allow(clippy::large_enum_variant)]
102pub enum IfElseExpand {
103    ComptimeThen,
104    ComptimeElse,
105    Runtime {
106        runtime_cond: NativeExpand<bool>,
107        if_op: IfOp,
108        then_child: Scope,
109    },
110}
111
112impl IfElseExpand {
113    pub fn or_else(self, scope: &Scope, else_block: impl FnOnce(&Scope)) {
114        match self {
115            Self::Runtime {
116                if_op, then_child, ..
117            } => {
118                let else_body = if_op.else_block(scope.ctx());
119                let else_child = scope.child(OpInserter::new_at_block_end(else_body));
120                else_block(&else_child);
121                else_child.terminate_yield();
122
123                scope.register(&if_op);
124                scope.set_break_return(&[then_child, else_child]);
125            }
126            Self::ComptimeElse => else_block(scope),
127            Self::ComptimeThen => (),
128        }
129    }
130}
131
132pub fn if_else_expand(
133    scope: &Scope,
134    condition: NativeExpand<bool>,
135    then_block: impl FnOnce(&Scope),
136) -> IfElseExpand {
137    let comptime_cond = condition.expand.as_const().map(|it| it.as_bool());
138    match comptime_cond {
139        Some(true) => {
140            then_block(scope);
141            IfElseExpand::ComptimeThen
142        }
143        Some(false) => IfElseExpand::ComptimeElse,
144        None => {
145            let cond = condition.read_value(scope);
146            let if_op = IfOp::new(scope.ctx_mut(), cond);
147            let if_block = if_op.then_block(scope.ctx());
148            let then_child = scope.child(OpInserter::new_at_block_end(if_block));
149            then_block(&then_child);
150            then_child.terminate_yield();
151
152            IfElseExpand::Runtime {
153                runtime_cond: condition,
154                if_op,
155                then_child,
156            }
157        }
158    }
159}
160
161#[allow(clippy::large_enum_variant)]
162pub enum IfElseExprExpand<C: Assign> {
163    ComptimeThen(C),
164    ComptimeElse,
165    Runtime {
166        runtime_cond: NativeExpand<bool>,
167        out: C,
168        if_op: IfOp,
169        then_child: Scope,
170    },
171}
172
173impl<C: Assign> IfElseExprExpand<C> {
174    pub fn or_else<R: RuntimeAssign<Expand = C>>(
175        self,
176        scope: &Scope,
177        else_block: impl FnOnce(&Scope) -> R,
178    ) -> C {
179        match self {
180            Self::Runtime {
181                mut out,
182                if_op,
183                then_child,
184                ..
185            } => {
186                let else_body = if_op.else_block(scope.ctx());
187                let else_child = scope.child(OpInserter::new_at_block_end(else_body));
188                let ret = else_block(&else_child);
189                out.__expand_assign_method(&else_child, ret.into_expand(scope));
190                else_child.terminate_yield();
191
192                scope.register(&if_op);
193                scope.set_break_return(&[then_child, else_child]);
194                out
195            }
196            Self::ComptimeElse => else_block(scope).into_expand(scope),
197            Self::ComptimeThen(ret) => ret,
198        }
199    }
200}
201
202pub fn if_else_expr_expand<C: RuntimeAssign>(
203    scope: &Scope,
204    condition: NativeExpand<bool>,
205    then_block: impl FnOnce(&Scope) -> C,
206) -> IfElseExprExpand<C::Expand> {
207    let comptime_cond = condition.expand.as_const().map(|it| it.as_bool());
208    match comptime_cond {
209        Some(true) => {
210            let ret = then_block(scope);
211            IfElseExprExpand::ComptimeThen(ret.into_expand(scope))
212        }
213        Some(false) => IfElseExprExpand::ComptimeElse,
214        None => {
215            let cond = condition.read_value(scope);
216            let if_op = IfOp::new(scope.ctx_mut(), cond);
217            let then_body = if_op.then_block(scope.ctx());
218            let then_child = scope.child(OpInserter::new_at_block_end(then_body));
219            let ret = then_block(&then_child);
220            let mut out = ret.init_mut(scope);
221            out.__expand_assign_method(&then_child, ret.into_expand(scope));
222            then_child.terminate_yield();
223
224            IfElseExprExpand::Runtime {
225                runtime_cond: condition,
226                out,
227                if_op,
228                then_child,
229            }
230        }
231    }
232}
233
234pub struct SwitchExpand<I: Int> {
235    switch_op: SwitchOp,
236    cases: Vec<I>,
237    children: Vec<Scope>,
238}
239
240impl<I: Int> SwitchExpand<I> {
241    pub fn case(mut self, scope: &Scope, value: impl Int, block: impl FnOnce(&Scope)) -> Self {
242        let value = I::from(value).unwrap();
243        self.cases.push(value);
244        let body = self.switch_op.append_case_block(scope.ctx_mut());
245        let case_child = scope.child(OpInserter::new_at_block_end(body));
246        block(&case_child);
247        case_child.terminate_yield();
248        self.children.push(case_child);
249        self
250    }
251
252    pub fn finish(self, scope: &Scope) {
253        let cases = self.cases.into_iter().map(|case| {
254            let ty = I::__expand_as_type(scope);
255            let ty = TypedHandle::<IntegerType>::from_handle(ty, scope.ctx()).unwrap();
256            let width = bw(ty.deref(scope.ctx()).width() as usize);
257            let val = APInt::from_i128(case.to_i128().unwrap(), width);
258            IntegerAttr::new(ty, val).into()
259        });
260        self.switch_op.set_attr_cases(scope.ctx(), cases);
261        scope.register(&self.switch_op);
262        scope.set_break_return(&self.children);
263    }
264}
265
266pub fn switch_expand<I: Int>(
267    scope: &Scope,
268    value: NativeExpand<I>,
269    default_block: impl FnOnce(&Scope),
270) -> SwitchExpand<I> {
271    let value = value.read_value(scope);
272    let switch_op = SwitchOp::new(scope.ctx_mut(), value);
273
274    let default_body = switch_op.default_block(scope.ctx());
275    let default_child = scope.child(OpInserter::new_at_block_end(default_body));
276    default_block(&default_child);
277    default_child.terminate_yield();
278
279    SwitchExpand {
280        switch_op,
281        cases: Vec::new(),
282        children: alloc::vec![default_child],
283    }
284}
285
286pub struct SwitchExpandExpr<I: Int, C: Assign> {
287    switch_op: SwitchOp,
288    cases: Vec<I>,
289    children: Vec<Scope>,
290    out: C,
291}
292
293impl<I: Int, C: Assign> SwitchExpandExpr<I, C> {
294    pub fn case<T: RuntimeAssign<Expand = C>>(
295        mut self,
296        scope: &Scope,
297        value: impl Int,
298        block: impl FnOnce(&Scope) -> T,
299    ) -> Self {
300        let value = I::from(value).unwrap();
301        self.cases.push(value);
302        let body = self.switch_op.append_case_block(scope.ctx_mut());
303        let case_child = scope.child(OpInserter::new_at_block_end(body));
304        let ret = block(&case_child);
305        self.out
306            .__expand_assign_method(&case_child, ret.into_expand(scope));
307        case_child.terminate_yield();
308        self.children.push(case_child);
309        self
310    }
311
312    pub fn finish(self, scope: &Scope) -> C {
313        let cases = self.cases.into_iter().map(|case| {
314            let ty = I::__expand_as_type(scope);
315            let ty = TypedHandle::<IntegerType>::from_handle(ty, scope.ctx()).unwrap();
316            let width = bw(ty.deref(scope.ctx()).width() as usize);
317            let val = APInt::from_i128(case.to_i128().unwrap(), width);
318            IntegerAttr::new(ty, val).into()
319        });
320        self.switch_op.set_attr_cases(scope.ctx(), cases);
321        scope.register(&self.switch_op);
322        scope.set_break_return(&self.children);
323        self.out
324    }
325}
326
327pub fn switch_expand_expr<I: Int, C: RuntimeAssign>(
328    scope: &Scope,
329    value: NativeExpand<I>,
330    default_block: impl FnOnce(&Scope) -> C,
331) -> SwitchExpandExpr<I, C::Expand> {
332    let value = value.read_value(scope);
333    let switch_op = SwitchOp::new(scope.ctx_mut(), value);
334
335    let default_body = switch_op.default_block(scope.ctx());
336    let default_child = scope.child(OpInserter::new_at_block_end(default_body));
337    let default = default_block(&default_child);
338    let mut out = default.init_mut(scope);
339    out.__expand_assign_method(&default_child, default.into_expand(scope));
340    default_child.terminate_yield();
341
342    SwitchExpandExpr {
343        switch_op,
344        cases: Vec::new(),
345        children: alloc::vec![default_child],
346        out,
347    }
348}
349
350#[allow(clippy::large_enum_variant)]
351pub enum MatchExpand<T: CubeEnum> {
352    ComptimeVariant {
353        variant: i32,
354        runtime_value: T::RuntimeValue,
355        matched: bool,
356    },
357    RuntimeVariant {
358        switch_op: SwitchOp,
359        cases: Vec<i32>,
360        children: Vec<Scope>,
361        has_default: bool,
362        runtime_value: T::RuntimeValue,
363    },
364}
365
366impl<T: CubeEnum> MatchExpand<T> {
367    pub fn case(
368        mut self,
369        scope: &Scope,
370        value: i32,
371        block: impl FnOnce(&Scope, T::RuntimeValue),
372    ) -> Self {
373        match &mut self {
374            Self::RuntimeVariant {
375                switch_op,
376                cases,
377                children,
378                runtime_value,
379                ..
380            } => {
381                cases.push(value);
382                let body = switch_op.append_case_block(scope.ctx_mut());
383                let case_child = scope.child(OpInserter::new_at_block_end(body));
384                block(&case_child, (*runtime_value).clone_unchecked());
385                case_child.terminate_yield();
386                children.push(case_child);
387            }
388            Self::ComptimeVariant {
389                variant,
390                runtime_value,
391                matched,
392            } => {
393                if value == *variant {
394                    block(scope, (*runtime_value).clone_unchecked());
395                    *matched = true;
396                }
397            }
398        }
399        self
400    }
401
402    pub fn default(mut self, scope: &Scope, block: impl FnOnce(&Scope, T::RuntimeValue)) -> Self {
403        match &mut self {
404            Self::RuntimeVariant {
405                switch_op,
406                children,
407                runtime_value,
408                has_default,
409                ..
410            } => {
411                let body = switch_op.default_block(scope.ctx());
412                let case_child = scope.child(OpInserter::new_at_block_end(body));
413                block(&case_child, (*runtime_value).clone_unchecked());
414                case_child.terminate_yield();
415                children.push(case_child);
416                *has_default = true;
417            }
418            Self::ComptimeVariant {
419                runtime_value,
420                matched,
421                ..
422            } => {
423                if !*matched {
424                    block(scope, (*runtime_value).clone_unchecked());
425                    *matched = true;
426                }
427            }
428        }
429        self
430    }
431
432    pub fn finish(self, scope: &Scope) {
433        match self {
434            MatchExpand::ComptimeVariant { .. } => {}
435            MatchExpand::RuntimeVariant {
436                switch_op,
437                cases,
438                children,
439                has_default,
440                ..
441            } => {
442                if !has_default {
443                    let default_body = switch_op.default_block(scope.ctx());
444                    let mut inserter = OpInserter::new_at_block_end(default_body);
445                    let unreachable = UnreachableOp::new(scope.ctx_mut());
446                    inserter.append_op(scope.ctx(), &unreachable);
447                }
448
449                let cases = cases.into_iter().map(|case| {
450                    let ty = IntegerType::get(scope.ctx(), 32, Signedness::Unsigned);
451                    IntegerAttr::new(ty, APInt::from_i32(case, bw(32))).into()
452                });
453                switch_op.set_attr_cases(scope.ctx(), cases);
454                scope.register(&switch_op);
455                scope.set_break_return(&children);
456            }
457        }
458    }
459}
460
461pub fn match_expand<T: CubeEnum>(
462    scope: &Scope,
463    value: T,
464    discriminant0: i32,
465    arm0: impl FnOnce(&Scope, T::RuntimeValue),
466) -> MatchExpand<T> {
467    let discriminant = value.discriminant();
468    match discriminant.constant() {
469        Some(const_variant) if const_variant.as_i32() == discriminant0 => {
470            let runtime_value = value.runtime_value();
471            arm0(scope, runtime_value.clone_unchecked());
472            MatchExpand::ComptimeVariant {
473                variant: const_variant.as_i32(),
474                runtime_value,
475                matched: true,
476            }
477        }
478        Some(const_variant) => MatchExpand::ComptimeVariant {
479            variant: const_variant.as_i32(),
480            runtime_value: value.runtime_value(),
481            matched: false,
482        },
483        None => {
484            let discriminant = discriminant.read_value(scope);
485            let runtime_value = value.runtime_value();
486
487            let switch_op = SwitchOp::new(scope.ctx_mut(), discriminant);
488            let body = switch_op.append_case_block(scope.ctx_mut());
489            let case_child = scope.child(OpInserter::new_at_block_end(body));
490            arm0(&case_child, runtime_value.clone_unchecked());
491            case_child.terminate_yield();
492
493            MatchExpand::RuntimeVariant {
494                switch_op,
495                cases: alloc::vec![discriminant0],
496                children: alloc::vec![case_child],
497                has_default: false,
498                runtime_value,
499            }
500        }
501    }
502}
503
504#[allow(clippy::large_enum_variant)]
505pub enum MatchExpandExpr<T: CubeEnum, C: Assign> {
506    ComptimeVariant {
507        variant: i32,
508        runtime_value: T::RuntimeValue,
509        out: Option<C>,
510        matched: bool,
511    },
512    RuntimeVariant {
513        switch_op: SwitchOp,
514        cases: Vec<i32>,
515        children: Vec<Scope>,
516        has_default: bool,
517        out: C,
518        runtime_value: T::RuntimeValue,
519    },
520}
521
522impl<T: CubeEnum, C: Assign> MatchExpandExpr<T, C> {
523    pub fn case<R: RuntimeAssign<Expand = C>>(
524        mut self,
525        scope: &Scope,
526        value: i32,
527        block: impl FnOnce(&Scope, T::RuntimeValue) -> R,
528    ) -> Self {
529        match &mut self {
530            Self::RuntimeVariant {
531                switch_op,
532                cases,
533                children,
534                out,
535                runtime_value,
536                ..
537            } => {
538                cases.push(value);
539                let body = switch_op.append_case_block(scope.ctx_mut());
540                let case_child = scope.child(OpInserter::new_at_block_end(body));
541                let ret_val = block(&case_child, (*runtime_value).clone_unchecked());
542                out.__expand_assign_method(&case_child, ret_val.into_expand(scope));
543                case_child.terminate_yield();
544                children.push(case_child);
545            }
546            Self::ComptimeVariant {
547                variant,
548                runtime_value,
549                out,
550                matched,
551            } => {
552                if value == *variant {
553                    *out =
554                        Some(block(scope, (*runtime_value).clone_unchecked()).into_expand(scope));
555                    *matched = true;
556                }
557            }
558        }
559        self
560    }
561
562    pub fn default<R: RuntimeAssign<Expand = C>>(
563        mut self,
564        scope: &Scope,
565        block: impl FnOnce(&Scope, T::RuntimeValue) -> R,
566    ) -> Self {
567        match &mut self {
568            Self::RuntimeVariant {
569                switch_op,
570                children,
571                runtime_value,
572                out,
573                has_default,
574                ..
575            } => {
576                let body = switch_op.default_block(scope.ctx());
577                let case_child = scope.child(OpInserter::new_at_block_end(body));
578                let ret_val = block(&case_child, (*runtime_value).clone_unchecked());
579                out.__expand_assign_method(&case_child, ret_val.into_expand(scope));
580                case_child.terminate_yield();
581                children.push(case_child);
582                *has_default = true;
583            }
584            Self::ComptimeVariant {
585                runtime_value,
586                out,
587                matched,
588                ..
589            } => {
590                if !*matched {
591                    *out =
592                        Some(block(scope, (*runtime_value).clone_unchecked()).into_expand(scope));
593                    *matched = true;
594                }
595            }
596        }
597        self
598    }
599
600    pub fn finish(self, scope: &Scope) -> C {
601        match self {
602            MatchExpandExpr::ComptimeVariant { out, .. } => {
603                out.expect("At least one variant should be matched")
604            }
605            MatchExpandExpr::RuntimeVariant {
606                switch_op,
607                cases,
608                children,
609                has_default,
610                out,
611                ..
612            } => {
613                if !has_default {
614                    let default_body = switch_op.default_block(scope.ctx());
615                    let mut inserter = OpInserter::new_at_block_end(default_body);
616                    let unreachable = UnreachableOp::new(scope.ctx_mut());
617                    inserter.append_op(scope.ctx(), &unreachable);
618                }
619
620                let cases = cases.into_iter().map(|case| {
621                    let ty = IntegerType::get(scope.ctx(), 32, Signedness::Unsigned);
622                    IntegerAttr::new(ty, APInt::from_i32(case, bw(32))).into()
623                });
624                switch_op.set_attr_cases(scope.ctx(), cases);
625                scope.register(&switch_op);
626                scope.set_break_return(&children);
627
628                out
629            }
630        }
631    }
632}
633
634pub fn match_expand_expr<T: CubeEnum, C: RuntimeAssign>(
635    scope: &Scope,
636    value: T,
637    discriminant0: i32,
638    arm0: impl FnOnce(&Scope, T::RuntimeValue) -> C,
639) -> MatchExpandExpr<T, C::Expand> {
640    let discriminant = value.discriminant();
641    match discriminant.constant() {
642        Some(const_variant) if const_variant.as_i32() == discriminant0 => {
643            let runtime_value = value.runtime_value();
644            let out = arm0(scope, runtime_value.clone_unchecked());
645            MatchExpandExpr::ComptimeVariant {
646                variant: const_variant.as_i32(),
647                out: Some(out.into_expand(scope)),
648                runtime_value,
649                matched: true,
650            }
651        }
652        Some(const_variant) => MatchExpandExpr::ComptimeVariant {
653            variant: const_variant.as_i32(),
654            out: None,
655            runtime_value: value.runtime_value(),
656            matched: false,
657        },
658        None => {
659            let discriminant = discriminant.read_value(scope);
660            let runtime_value = value.runtime_value();
661
662            let switch_op = SwitchOp::new(scope.ctx_mut(), discriminant);
663            let body = switch_op.append_case_block(scope.ctx_mut());
664            let case_child = scope.child(OpInserter::new_at_block_end(body));
665            let ret_val = arm0(&case_child, runtime_value.clone_unchecked());
666
667            let mut out = ret_val.init_mut(scope);
668            out.__expand_assign_method(&case_child, ret_val.into_expand(scope));
669            case_child.terminate_yield();
670
671            MatchExpandExpr::RuntimeVariant {
672                switch_op,
673                out,
674                cases: alloc::vec![discriminant0],
675                children: alloc::vec![case_child],
676                runtime_value,
677                has_default: false,
678            }
679        }
680    }
681}
682
683pub fn break_expand(scope: &Scope) {
684    let inv_break_flag = scope
685        .expand_state()
686        .inv_break_flag
687        .expect("Should be in loop");
688    let false_ = false.__expand_runtime_method(scope).expand;
689    assign::expand_element(scope, false_, inv_break_flag.into());
690    scope.expand_state_mut().may_break = true;
691}
692
693pub fn return_expand(scope: &Scope) {
694    let inv_return_flag = scope.expand_state().inv_return_flag;
695    if let Some(inv_return_flag) = inv_return_flag {
696        let false_ = false.__expand_runtime_method(scope).expand;
697        assign::expand_element(scope, false_, inv_return_flag.into());
698        scope.expand_state_mut().may_return = true;
699    } else {
700        // We're in a rewrite context, so just break the strict hammock model. We can't apply the
701        // single-return conversion after the scope is already built.
702        scope.register(&ReturnOp::new(scope.ctx_mut()));
703    }
704}
705
706pub mod unreachable_unchecked {
707    use super::*;
708
709    pub fn expand(scope: &Scope) {
710        scope.register(&UnreachableOp::new(scope.ctx_mut()));
711    }
712}
713
714/// `WhileOp` builder. Needs to work like this for borrow checker reasons.
715/// Rust requires all closures passed to a function to be live at once, so we use a builder pattern
716/// and separate the cond and body into separate functions to allow the borrows to work.
717pub struct WhileBuilder {
718    while_op: WhileOp,
719    cond: ExpandValue,
720    cond_scope: Scope,
721}
722
723impl WhileBuilder {
724    pub fn new(scope: &Scope, mut cond: impl FnMut(&Scope) -> NativeExpand<bool>) -> Self {
725        let while_op = WhileOp::new(scope.ctx_mut());
726        let cond_scope = scope.child(OpInserter::new_at_block_start(
727            while_op.before_block(scope.ctx()),
728        ));
729
730        WhileBuilder {
731            while_op,
732            cond: cond(&cond_scope).expand,
733            cond_scope,
734        }
735    }
736
737    pub fn with_body(self, scope: &Scope, mut block: impl FnMut(&Scope)) {
738        let Self {
739            while_op,
740            mut cond,
741            cond_scope,
742        } = self;
743
744        let body = while_op.after_block(scope.ctx());
745        let body = scope.loop_child(OpInserter::new_at_block_end(body));
746        block(&body);
747        body.terminate_yield();
748
749        let expand_state = *body.expand_state();
750        let break_flag = expand_state.inv_break_flag.unwrap().into();
751        let return_flag = expand_state.inv_return_flag.map(Into::into);
752
753        if expand_state.may_break {
754            cond = binary_expand(&cond_scope, cond, break_flag, BoolAndOp::new);
755        }
756        if expand_state.may_return {
757            let return_flag = return_flag.unwrap();
758            cond = binary_expand(&cond_scope, cond, return_flag, BoolAndOp::new);
759        }
760
761        cond_scope.register(&ConditionOp::new(scope.ctx_mut(), cond.read_value(scope)));
762
763        scope.register(&while_op);
764    }
765}
766
767/// register a range loop if it contains no break or return, destructure to while if it does
768pub(crate) fn register_range_loop<I: Int>(scope: &Scope, for_op: &RangeLoopOp, body: &Scope) {
769    let ctx = scope.ctx_mut();
770    let ExpandState {
771        may_return,
772        may_break,
773        inv_return_flag,
774        inv_break_flag,
775    } = *body.expand_state();
776    if !may_break && !may_return {
777        body.terminate_yield();
778        scope.register(for_op);
779        return;
780    }
781
782    let mut rewriter = IRRewriter::<DummyListener>::default();
783
784    let start = for_op.start(ctx);
785    let end = for_op.end(ctx);
786    let step = for_op.step(ctx);
787    let iter_var_old = for_op.iter_var(ctx);
788
789    let iter_var = ExpandValue::from(scope.create_local_mut(iter_var_old.get_type(ctx), None));
790    assign::expand_element(scope, start.into(), iter_var);
791
792    assign_binop_expand::<I>(
793        body,
794        &mut iter_var.into(),
795        step.into(),
796        I::__expand_native_add,
797    );
798    body.terminate_yield();
799
800    body.inserter()
801        .set_insertion_point_to_block_start(for_op.loop_body(ctx));
802    let iter_val = iter_var.read_value(body);
803    rewriter.replace_value_uses_with(ctx, iter_var_old, iter_val);
804    BasicBlock::remove_argument(for_op.loop_body(ctx), ctx, 0);
805
806    let while_op = WhileOp::new(ctx);
807
808    let cond_scope = scope.child(OpInserter::new_at_block_start(
809        while_op.before_block(scope.ctx()),
810    ));
811
812    let mut cond = I::__expand_native_lt(&cond_scope, iter_var, end.into());
813    if may_break {
814        let inv_break_flag = inv_break_flag.unwrap().into();
815        cond = binary_expand(&cond_scope, cond, inv_break_flag, BoolAndOp::new);
816    }
817    if may_return {
818        let inv_return_flag = inv_return_flag.unwrap().into();
819        cond = binary_expand(&cond_scope, cond, inv_return_flag, BoolAndOp::new);
820    }
821
822    cond_scope.register(&ConditionOp::new(scope.ctx_mut(), cond.read_value(scope)));
823
824    rewriter.erase_region(ctx, while_op.after_region(ctx));
825    Region::move_to_op(for_op.get_region(ctx), while_op.get_operation(), ctx);
826    rewriter.erase_operation(ctx, for_op.get_operation());
827
828    scope.register(&while_op);
829}
830
831// Don't make this `FnOnce`, it must be executable multiple times
832pub fn loop_expand(scope: &Scope, block: impl FnMut(&Scope)) {
833    WhileBuilder::new(scope, |scope| true.__expand_runtime_method(scope)).with_body(scope, block);
834}
cubecl_core/frontend/branch.rs

cubecl_core/frontend/
branch.rs
