diff --git a/xdsl/transforms/__init__.py b/xdsl/transforms/__init__.py index c598e7aa11..f560e17ba9 100644 --- a/xdsl/transforms/__init__.py +++ b/xdsl/transforms/__init__.py @@ -340,6 +340,11 @@ def get_function_persist_arg_names(): return function_transformations.FunctionPersistArgNamesPass + def get_func_inline(): + from xdsl.transforms import func_inline + + return func_inline.FuncInlinePass + def get_func_to_pdl_rewrite(): from xdsl.transforms.experimental import func_to_pdl_rewrite @@ -746,6 +751,7 @@ def get_verify_register_allocation(): "frontend-desymrefy": get_frontend_desymrefy, "function-constant-pinning": get_function_constant_pinning, "function-persist-arg-names": get_function_persist_arg_names, + "func-inline": get_func_inline, "func-to-pdl-rewrite": get_func_to_pdl_rewrite, "gpu-map-parallel-loops": get_gpu_map_parallel_loops, "hls-convert-stencil-to-ll-mlir": get_hls_convert_stencil_to_ll_mlir, diff --git a/xdsl/transforms/func_inline.py b/xdsl/transforms/func_inline.py new file mode 100644 index 0000000000..d347f4a273 --- /dev/null +++ b/xdsl/transforms/func_inline.py @@ -0,0 +1,57 @@ +from xdsl.context import Context +from xdsl.dialects import builtin, cf, func +from xdsl.passes import ModulePass +from xdsl.pattern_rewriter import ( + PatternRewriter, + PatternRewriteWalker, + RewritePattern, + op_type_rewrite_pattern, +) +from xdsl.rewriter import BlockInsertPoint, InsertPoint +from xdsl.traits import SymbolTable + + +class FuncInlinePattern(RewritePattern): + @op_type_rewrite_pattern + def match_and_rewrite(self, op: func.CallOp, rewriter: PatternRewriter, /): + if op.parent is None: + return + + f = SymbolTable.lookup_symbol(op, op.callee) + + if not isinstance(f, func.FuncOp): + return + + if f.is_declaration: + return + + body = f.body.clone() + first_block = body.blocks.first + last_block = body.blocks.last + assert first_block is not None + assert last_block is not None + term = last_block.ops.last + assert isinstance(term, func.ReturnOp) + + before = op.parent + after = before.split_before(op, arg_types=term.arguments.types) + + rewriter.insert_op( + cf.BranchOp(first_block, *op.arguments), InsertPoint.at_end(before) + ) + + rewriter.inline_region(body, BlockInsertPoint.before(after)) + + rewriter.replace_op(term, cf.BranchOp(after, *term.arguments)) + + rewriter.replace_op(op, (), after.args) + + +class FuncInlinePass(ModulePass): + name = "func-inline" + + def apply(self, ctx: Context, op: builtin.ModuleOp) -> None: + PatternRewriteWalker( + FuncInlinePattern(), + apply_recursively=False, + ).rewrite_module(op)