1 //===- TestInlining.cpp - Pass to inline calls in the test dialect --------===// 2 // 3 // Part of the MLIR 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(riverriddle) 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/Ops.h" 17 #include "mlir/IR/Function.h" 18 #include "mlir/Pass/Pass.h" 19 #include "mlir/Transforms/InliningUtils.h" 20 #include "mlir/Transforms/Passes.h" 21 #include "llvm/ADT/StringSet.h" 22 23 using namespace mlir; 24 25 namespace { 26 struct Inliner : public FunctionPass<Inliner> { 27 void runOnFunction() override { 28 auto function = getFunction(); 29 30 // Collect each of the direct function calls within the module. 31 SmallVector<CallIndirectOp, 16> callers; 32 function.walk([&](CallIndirectOp caller) { callers.push_back(caller); }); 33 34 // Build the inliner interface. 35 InlinerInterface interface(&getContext()); 36 37 // Try to inline each of the call operations. 38 for (auto caller : callers) { 39 auto callee = dyn_cast_or_null<FunctionalRegionOp>( 40 caller.getCallee().getDefiningOp()); 41 if (!callee) 42 continue; 43 44 // Inline the functional region operation, but only clone the internal 45 // region if there is more than one use. 46 if (failed(inlineRegion( 47 interface, &callee.body(), caller, 48 llvm::to_vector<8>(caller.getArgOperands()), 49 SmallVector<Value, 8>(caller.getResults()), caller.getLoc(), 50 /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse()))) 51 continue; 52 53 // If the inlining was successful then erase the call and callee if 54 // possible. 55 caller.erase(); 56 if (callee.use_empty()) 57 callee.erase(); 58 } 59 } 60 }; 61 } // end anonymous namespace 62 63 static PassRegistration<Inliner> pass("test-inline", 64 "Test inlining region calls"); 65