81 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// TODO: This pass is only necessary because the main inlining pass10// has not abstracted away the call+callee relationship. When the inlining11// interface has this support, this pass should be removed.12//13//===----------------------------------------------------------------------===//14 15#include "TestDialect.h"16#include "TestOps.h"17#include "mlir/Dialect/Func/IR/FuncOps.h"18#include "mlir/IR/BuiltinOps.h"19#include "mlir/IR/IRMapping.h"20#include "mlir/Pass/Pass.h"21#include "mlir/Transforms/Inliner.h"22#include "mlir/Transforms/InliningUtils.h"23#include "llvm/ADT/StringSet.h"24 25using namespace mlir;26using namespace test;27 28namespace {29struct InlinerTest30 : public PassWrapper<InlinerTest, OperationPass<func::FuncOp>> {31 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InlinerTest)32 33 StringRef getArgument() const final { return "test-inline"; }34 StringRef getDescription() const final {35 return "Test inlining region calls";36 }37 38 void runOnOperation() override {39 InlinerConfig config;40 41 auto function = getOperation();42 43 // Collect each of the direct function calls within the module.44 SmallVector<func::CallIndirectOp, 16> callers;45 function.walk(46 [&](func::CallIndirectOp caller) { callers.push_back(caller); });47 48 // Build the inliner interface.49 InlinerInterface interface(&getContext());50 51 // Try to inline each of the call operations.52 for (auto caller : callers) {53 auto callee = dyn_cast_or_null<FunctionalRegionOp>(54 caller.getCallee().getDefiningOp());55 if (!callee)56 continue;57 58 // Inline the functional region operation, but only clone the internal59 // region if there is more than one use.60 if (failed(inlineRegion(61 interface, config.getCloneCallback(), &callee.getBody(), caller,62 caller.getArgOperands(), caller.getResults(), caller.getLoc(),63 /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))64 continue;65 66 // If the inlining was successful then erase the call and callee if67 // possible.68 caller.erase();69 if (callee.use_empty())70 callee.erase();71 }72 }73};74} // namespace75 76namespace mlir {77namespace test {78void registerInliner() { PassRegistration<InlinerTest>(); }79} // namespace test80} // namespace mlir81