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