cubecl_cpp/shared/
branch.rs

1use cubecl_core::ir::{
2    dialect::{branch::*, general::SelectOp},
3    prelude::*,
4};
5use pliron::{basic_block::BasicBlock, linked_list::ContainsLinkedList};
6
7use crate::{
8    error::EmissionErrors,
9    shared::{
10        CppValue, OpExtCPP, scoped_block, shared_op, shared_op_with_out, ty::TypeExtCPP,
11        unroll::unrolling,
12    },
13};
14
15pub fn block_to_cpp(ctx: &Context, block: Ptr<BasicBlock>) -> String {
16    let mut out = String::new();
17    let ops = block.deref(ctx).iter(ctx);
18    for op in ops {
19        // `Display` can't fail, so record the error and let `compile_ir` fail the compilation.
20        match op.to_cpp(ctx) {
21            Ok(cpp) => out.push_str(&cpp),
22            Err(err) => ctx.aux_ty::<EmissionErrors>().record(err),
23        }
24    }
25    out
26}
27
28shared_op!(IfOp, |op, ctx| {
29    let cond = op.condition(ctx).name(ctx);
30    let else_block = op.else_block(ctx);
31    let mut out = format!("if({cond}) {{\n");
32    out.push_str(&block_to_cpp(ctx, op.then_block(ctx)));
33    if else_block.deref(ctx).iter(ctx).count() > 1 {
34        out.push_str("}\n else {\n");
35        out.push_str(&block_to_cpp(ctx, else_block));
36    }
37    out.push_str("}\n");
38    out
39});
40
41shared_op!(SwitchOp, |op, ctx| {
42    let value = op.value(ctx).name(ctx);
43    let mut out = format!("switch({value}) {{\n");
44    for (value, block) in op.cases(ctx) {
45        let block = block_to_cpp(ctx, block);
46        let case = format!("case {}: {{ {block} break; }}\n", value.value().to_i128());
47        out.push_str(&case);
48    }
49    let block = block_to_cpp(ctx, op.default_block(ctx));
50    out.push_str(&format!("default: {{ {block} break; }}\n"));
51    out.push_str("}\n");
52    out
53});
54
55// Only relevant for IR structure
56shared_op!(YieldOp, |_, _| String::new());
57shared_op!(ConditionOp, |op, ctx| {
58    format!("return {};", op.condition(ctx).name(ctx))
59});
60
61shared_op!(ReturnOp, |op, ctx| {
62    if let Some(value) = op.value(ctx) {
63        format!("return {};", value.name(ctx))
64    } else {
65        "return;".into()
66    }
67});
68
69shared_op!(UnreachableOp, |_, _| "__builtin_unreachable();".into());
70
71shared_op!(RangeLoopOp, |op, ctx| {
72    let i = op.iter_var(ctx).name(ctx);
73    let i_ty = op.iter_var(ctx).get_type(ctx).to_cpp(ctx);
74    let start = op.start(ctx).name(ctx);
75    let end = op.end(ctx).name(ctx);
76    let step = op.step(ctx).name(ctx);
77    let mut out = format!("for({i_ty} {i} = {start}; {i} < {end}; {i} += {step}) {{\n");
78    out.push_str(&block_to_cpp(ctx, op.loop_body(ctx)));
79    out.push_str("}\n");
80    out
81});
82
83shared_op!(WhileOp, |op, ctx| {
84    let cond = scoped_block! {
85        block_to_cpp(ctx, op.before_block(ctx))
86    };
87    let mut out = format!("while({cond}) {{\n");
88    out.push_str(&block_to_cpp(ctx, op.after_block(ctx)));
89    out.push_str("}\n");
90    out
91});
92
93shared_op_with_out!(SelectOp, |op, ctx| {
94    let cond = op.condition(ctx).name(ctx);
95    let then = op.true_value(ctx).name(ctx);
96    let or_else = op.false_value(ctx).name(ctx);
97    format!("{} ? {} : {}", cond, then, or_else)
98});
99unrolling!(SelectOp);
cubecl_cpp/shared/branch.rs

cubecl_cpp/shared/
branch.rs
