1 //===- TestInlining.cpp - Pass to inline calls in the test dialect --------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // TODO: This pass is only necessary because the main inlining pass 10 // has no abstracted away the call+callee relationship. When the inlining 11 // interface has this support, this pass should be removed. 12 // 13 //===----------------------------------------------------------------------===// 14 15 #include "TestDialect.h" 16 #include "mlir/Dialect/StandardOps/IR/Ops.h" 17 #include "mlir/IR/BlockAndValueMapping.h" 18 #include "mlir/IR/BuiltinOps.h" 19 #include "mlir/Pass/Pass.h" 20 #include "mlir/Transforms/InliningUtils.h" 21 #include "mlir/Transforms/Passes.h" 22 #include "llvm/ADT/StringSet.h" 23 24 using namespace mlir; 25 using namespace mlir::test; 26 27 namespace { 28 struct Inliner : public PassWrapper<Inliner, FunctionPass> { 29 void runOnFunction() override { 30 auto function = getFunction(); 31 32 // Collect each of the direct function calls within the module. 33 SmallVector<CallIndirectOp, 16> callers; 34 function.walk([&](CallIndirectOp caller) { callers.push_back(caller); }); 35 36 // Build the inliner interface. 37 InlinerInterface interface(&getContext()); 38 39 // Try to inline each of the call operations. 40 for (auto caller : callers) { 41 auto callee = dyn_cast_or_null<FunctionalRegionOp>( 42 caller.getCallee().getDefiningOp()); 43 if (!callee) 44 continue; 45 46 // Inline the functional region operation, but only clone the internal 47 // region if there is more than one use. 48 if (failed(inlineRegion( 49 interface, &callee.body(), caller, caller.getArgOperands(), 50 caller.getResults(), caller.getLoc(), 51 /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse()))) 52 continue; 53 54 // If the inlining was successful then erase the call and callee if 55 // possible. 56 caller.erase(); 57 if (callee.use_empty()) 58 callee.erase(); 59 } 60 } 61 }; 62 } // end anonymous namespace 63 64 namespace mlir { 65 namespace test { 66 void registerInliner() { 67 PassRegistration<Inliner>("test-inline", "Test inlining region calls"); 68 } 69 } // namespace test 70 } // namespace mlir 71