a5ef51d786
Promised interfaces allow for a dialect to "promise" the implementation of an interface, i.e. declare that it supports an interface, but have the interface defined in an extension in a library separate from the dialect itself. A promised interface is powerful in that it alerts the user when the interface is attempted to be used (e.g. via cast/dyn_cast/etc.) and the implementation has not yet been provided. This makes the system much more robust against misconfiguration, and ensures that we do not lose the benefit we currently have of defining the interface in the dialect library. Differential Revision: https://reviews.llvm.org/D120368
91 lines
3.4 KiB
C++
91 lines
3.4 KiB
C++
//===- InlinerExtension.cpp - Func Inliner Extension ----------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
|
|
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/IR/DialectInterface.h"
|
|
#include "mlir/Transforms/InliningUtils.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::func;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// FuncDialect Interfaces
|
|
//===----------------------------------------------------------------------===//
|
|
namespace {
|
|
/// This class defines the interface for handling inlining with func operations.
|
|
struct FuncInlinerInterface : public DialectInlinerInterface {
|
|
using DialectInlinerInterface::DialectInlinerInterface;
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Analysis Hooks
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
/// All call operations can be inlined.
|
|
bool isLegalToInline(Operation *call, Operation *callable,
|
|
bool wouldBeCloned) const final {
|
|
return true;
|
|
}
|
|
|
|
/// All operations can be inlined.
|
|
bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
|
|
return true;
|
|
}
|
|
|
|
/// All functions can be inlined.
|
|
bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
|
|
return true;
|
|
}
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Transformation Hooks
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
/// Handle the given inlined terminator by replacing it with a new operation
|
|
/// as necessary.
|
|
void handleTerminator(Operation *op, Block *newDest) const final {
|
|
// Only return needs to be handled here.
|
|
auto returnOp = dyn_cast<ReturnOp>(op);
|
|
if (!returnOp)
|
|
return;
|
|
|
|
// Replace the return with a branch to the dest.
|
|
OpBuilder builder(op);
|
|
builder.create<cf::BranchOp>(op->getLoc(), newDest, returnOp.getOperands());
|
|
op->erase();
|
|
}
|
|
|
|
/// Handle the given inlined terminator by replacing it with a new operation
|
|
/// as necessary.
|
|
void handleTerminator(Operation *op,
|
|
ArrayRef<Value> valuesToRepl) const final {
|
|
// Only return needs to be handled here.
|
|
auto returnOp = cast<ReturnOp>(op);
|
|
|
|
// Replace the values directly with the return operands.
|
|
assert(returnOp.getNumOperands() == valuesToRepl.size());
|
|
for (const auto &it : llvm::enumerate(returnOp.getOperands()))
|
|
valuesToRepl[it.index()].replaceAllUsesWith(it.value());
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Registration
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void mlir::func::registerInlinerExtension(DialectRegistry ®istry) {
|
|
registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
|
|
dialect->addInterfaces<FuncInlinerInterface>();
|
|
|
|
// The inliner extension relies on the ControlFlow dialect.
|
|
ctx->getOrLoadDialect<cf::ControlFlowDialect>();
|
|
});
|
|
}
|