Skip to content

Inline

Inline dataclass

Inline(heuristic: Callable[[ir.Statement], bool])

Bases: RewriteRule

heuristic instance-attribute

heuristic: Callable[[Statement], bool]

inline heuristic that determines whether a function should be inlined

inline_call_like

inline_call_like(
    call_like: ir.Statement,
    args: tuple[ir.SSAValue, ...],
    region: ir.Region,
)

Inline a function call-like statement

Parameters:

Name Type Description Default
call_like Statement

the call-like statement

required
args tuple[SSAValue, ...]

the arguments of the call (first one is the callee)

required
region Region

the region of the callee

required
Source code in src/kirin/rewrite/inline.py
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def inline_call_like(
    self,
    call_like: ir.Statement,
    args: tuple[ir.SSAValue, ...],
    region: ir.Region,
):
    """
    Inline a function call-like statement

    Args:
        call_like (ir.Statement): the call-like statement
        args (tuple[ir.SSAValue, ...]): the arguments of the call (first one is the callee)
        region (ir.Region): the region of the callee
    """
    # <stmt>
    # <stmt>
    # <br (a, b, c)>

    # <block (a, b,c)>:
    # <block>:
    # <block>:
    # <br>

    # ^<block>:
    # <stmt>
    # <stmt>

    # 1. we insert the entry block of the callee function
    # 2. we insert the rest of the blocks into the parent region
    # 3.1 if the return is in the entry block, means no control flow,
    #     replace the call results with the return values
    # 3.2 if the return is some of the blocks, means control flow,
    #     split the current block into two, and replace the return with
    #     the branch instruction
    # 4. remove the call
    if not call_like.parent_block:
        return

    if not call_like.parent_region:
        return

    # NOTE: we cannot change region because it may be used elsewhere
    inline_region: ir.Region = region.clone()
    parent_block: ir.Block = call_like.parent_block
    parent_region: ir.Region = call_like.parent_region

    # wrap what's after invoke into a block
    after_block = ir.Block()
    stmt = call_like.next_stmt
    while stmt is not None:
        stmt.detach()
        after_block.stmts.append(stmt)
        stmt = call_like.next_stmt

    for result in call_like.results:
        block_arg = after_block.args.append_from(result.type, result.name)
        result.replace_by(block_arg)

    parent_block_idx = parent_region._block_idx[parent_block]
    entry_block = inline_region.blocks.popfirst()
    idx, block = 0, entry_block
    while block is not None:
        idx += 1

        if block.last_stmt and isinstance(block.last_stmt, func.Return):
            block.last_stmt.replace_by(
                cf.Branch(
                    arguments=(block.last_stmt.value,),
                    successor=after_block,
                )
            )

        parent_region.blocks.insert(parent_block_idx + idx, block)
        block = inline_region.blocks.popfirst()

    parent_region.blocks.append(after_block)

    # NOTE: we expect always to have an entry block
    # but we still check for it cuz we are not gonna
    # error for empty regions here.
    if entry_block:
        cf.Branch(
            arguments=args,
            successor=entry_block,
        ).insert_before(call_like)
    call_like.delete()
    return