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 not 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/Func/IR/FuncOps.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 "llvm/ADT/StringSet.h"
22
23 using namespace mlir;
24 using namespace test;
25
26 namespace {
27 struct Inliner : public PassWrapper<Inliner, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon3f50e1be0111::Inliner28 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(Inliner)
29
30 StringRef getArgument() const final { return "test-inline"; }
getDescription__anon3f50e1be0111::Inliner31 StringRef getDescription() const final {
32 return "Test inlining region calls";
33 }
34
runOnOperation__anon3f50e1be0111::Inliner35 void runOnOperation() override {
36 auto function = getOperation();
37
38 // Collect each of the direct function calls within the module.
39 SmallVector<func::CallIndirectOp, 16> callers;
40 function.walk(
41 [&](func::CallIndirectOp caller) { callers.push_back(caller); });
42
43 // Build the inliner interface.
44 InlinerInterface interface(&getContext());
45
46 // Try to inline each of the call operations.
47 for (auto caller : callers) {
48 auto callee = dyn_cast_or_null<FunctionalRegionOp>(
49 caller.getCallee().getDefiningOp());
50 if (!callee)
51 continue;
52
53 // Inline the functional region operation, but only clone the internal
54 // region if there is more than one use.
55 if (failed(inlineRegion(
56 interface, &callee.getBody(), caller, caller.getArgOperands(),
57 caller.getResults(), caller.getLoc(),
58 /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
59 continue;
60
61 // If the inlining was successful then erase the call and callee if
62 // possible.
63 caller.erase();
64 if (callee.use_empty())
65 callee.erase();
66 }
67 }
68 };
69 } // namespace
70
71 namespace mlir {
72 namespace test {
registerInliner()73 void registerInliner() { PassRegistration<Inliner>(); }
74 } // namespace test
75 } // namespace mlir
76