xref: /llvm-project/mlir/test/lib/Transforms/TestInlining.cpp (revision 7776b19eed44906e9973bfb240b6279d6feaab41)
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 test;
26 
27 namespace {
28 struct Inliner : public PassWrapper<Inliner, FunctionPass> {
29   StringRef getArgument() const final { return "test-inline"; }
30   StringRef getDescription() const final {
31     return "Test inlining region calls";
32   }
33 
34   void runOnFunction() override {
35     auto function = getFunction();
36 
37     // Collect each of the direct function calls within the module.
38     SmallVector<CallIndirectOp, 16> callers;
39     function.walk([&](CallIndirectOp caller) { callers.push_back(caller); });
40 
41     // Build the inliner interface.
42     InlinerInterface interface(&getContext());
43 
44     // Try to inline each of the call operations.
45     for (auto caller : callers) {
46       auto callee = dyn_cast_or_null<FunctionalRegionOp>(
47           caller.getCallee().getDefiningOp());
48       if (!callee)
49         continue;
50 
51       // Inline the functional region operation, but only clone the internal
52       // region if there is more than one use.
53       if (failed(inlineRegion(
54               interface, &callee.body(), caller, caller.getArgOperands(),
55               caller.getResults(), caller.getLoc(),
56               /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
57         continue;
58 
59       // If the inlining was successful then erase the call and callee if
60       // possible.
61       caller.erase();
62       if (callee.use_empty())
63         callee.erase();
64     }
65   }
66 };
67 } // end anonymous namespace
68 
69 namespace mlir {
70 namespace test {
71 void registerInliner() { PassRegistration<Inliner>(); }
72 } // namespace test
73 } // namespace mlir
74