cubecl_ir/
convert.rs

1use alloc::{vec, vec::Vec};
2
3use pliron::{
4    basic_block::BasicBlock,
5    builtin::ops::FuncOp,
6    context::{Context, Ptr},
7    graph::walkers::{IRNode, WALKCONFIG_PREORDER_FORWARD, uninterruptible::immutable::walk_op},
8    op::Op,
9    operation::Operation,
10    r#type::Typed,
11    value::{DefiningEntity, Value},
12};
13
14use crate::{FuncOpExt, dialect::OperationPtrExt};
15
16/// Lift a closure to a standalone function with the captures as extra args. Returns the value that
17/// should be passed to these new args.
18pub fn lift_closure(ctx: &Context, func: &FuncOp) -> Vec<Value> {
19    let func_op = func.get_operation();
20    let entry = func.get_entry_block(ctx);
21    let mut captures = vec![];
22
23    walk_op(
24        ctx,
25        &mut (func, &mut captures),
26        &WALKCONFIG_PREORDER_FORWARD,
27        func.get_operation(),
28        |ctx, (func, captures), node| {
29            if let IRNode::Operation(ptr) = node {
30                for opd in ptr.operands(ctx) {
31                    if !is_in_func(ctx, func, opd.defining_entity()) {
32                        captures.push(opd);
33                    }
34                }
35            }
36        },
37    );
38
39    for capture in captures.iter() {
40        let id = func.push_argument(ctx, capture.get_type(ctx));
41        let arg_value = entry.deref(ctx).get_argument(id);
42        capture.replace_some_uses_with(
43            ctx,
44            |ctx, r#use| op_is_in_func(ctx, func_op, r#use.user_op()),
45            &arg_value,
46        );
47    }
48
49    captures
50}
51
52fn op_is_in_func(ctx: &Context, func: Ptr<Operation>, mut op: Ptr<Operation>) -> bool {
53    while let Some(parent) = op.deref(ctx).get_parent_op(ctx) {
54        if parent == func {
55            return true;
56        }
57        op = parent;
58    }
59    false
60}
61
62fn block_is_in_func(ctx: &Context, func: Ptr<BasicBlock>, mut block: Ptr<BasicBlock>) -> bool {
63    if func == block {
64        return true;
65    }
66    while let Some(parent) = block.deref(ctx).get_parent_block(ctx) {
67        if parent == func {
68            return true;
69        }
70        block = parent;
71    }
72    false
73}
74
75fn is_in_func(ctx: &Context, func: &FuncOp, entity: DefiningEntity) -> bool {
76    match entity {
77        DefiningEntity::Op(op) => op_is_in_func(ctx, func.get_operation(), op),
78        DefiningEntity::Block(block) => block_is_in_func(ctx, func.get_entry_block(ctx), block),
79    }
80}
cubecl_ir/convert.rs

cubecl_ir/
convert.rs
