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

cubecl_core/frontend/
branch.rs
