1//===- InlinerExtension.cpp - Func Inliner Extension ----------------------===//
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#include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
10#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
11#include "mlir/Dialect/Func/IR/FuncOps.h"
12#include "mlir/IR/DialectInterface.h"
13#include "mlir/Transforms/InliningUtils.h"
14
15using namespace mlir;
16using namespace mlir::func;
17
18//===----------------------------------------------------------------------===//
19// FuncDialect Interfaces
20//===----------------------------------------------------------------------===//
21namespace {
22/// This class defines the interface for handling inlining with func operations.
23struct FuncInlinerInterface : public DialectInlinerInterface {
24 using DialectInlinerInterface::DialectInlinerInterface;
25
26 //===--------------------------------------------------------------------===//
27 // Analysis Hooks
28 //===--------------------------------------------------------------------===//
29
30 /// All call operations can be inlined.
31 bool isLegalToInline(Operation *call, Operation *callable,
32 bool wouldBeCloned) const final {
33 return true;
34 }
35
36 /// All operations can be inlined.
37 bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
38 return true;
39 }
40
41 /// All functions can be inlined.
42 bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
43 return true;
44 }
45
46 //===--------------------------------------------------------------------===//
47 // Transformation Hooks
48 //===--------------------------------------------------------------------===//
49
50 /// Handle the given inlined terminator by replacing it with a new operation
51 /// as necessary.
52 void handleTerminator(Operation *op, Block *newDest) const final {
53 // Only return needs to be handled here.
54 auto returnOp = dyn_cast<ReturnOp>(op);
55 if (!returnOp)
56 return;
57
58 // Replace the return with a branch to the dest.
59 OpBuilder builder(op);
60 builder.create<cf::BranchOp>(op->getLoc(), newDest, returnOp.getOperands());
61 op->erase();
62 }
63
64 /// Handle the given inlined terminator by replacing it with a new operation
65 /// as necessary.
66 void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
67 // Only return needs to be handled here.
68 auto returnOp = cast<ReturnOp>(op);
69
70 // Replace the values directly with the return operands.
71 assert(returnOp.getNumOperands() == valuesToRepl.size());
72 for (const auto &it : llvm::enumerate(returnOp.getOperands()))
73 valuesToRepl[it.index()].replaceAllUsesWith(it.value());
74 }
75};
76} // namespace
77
78//===----------------------------------------------------------------------===//
79// Registration
80//===----------------------------------------------------------------------===//
81
82void mlir::func::registerInlinerExtension(DialectRegistry &registry) {
83 registry.addExtension(extensionFn: +[](MLIRContext *ctx, func::FuncDialect *dialect) {
84 dialect->addInterfaces<FuncInlinerInterface>();
85
86 // The inliner extension relies on the ControlFlow dialect.
87 ctx->getOrLoadDialect<cf::ControlFlowDialect>();
88 });
89}
90

source code of mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp