[llvm][ModuleInliner] Add inline cost priority for module inliner

[taolq]

Summary

This patch introduces the inline cost priority into the
module inliner, which uses the same computation as
InlineCost.

[kazu]

LGTM modulo some comments. Thanks!

llvm/lib/Analysis/InlineOrder.cpp

2

Could you put the standard file header here?

//===- InlineOrder.cpp - Inlining order abstraction -*- C++ ---*-----------===//
//
// 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
//
//===----------------------------------------------------------------------===//

llvm/lib/Transforms/IPO/ModuleInliner.cpp

153–154

I'd suggest:

auto Calls = getInlineOrder(UseInlinePriority, FAM, Params);

[taolq]

• rebase to main
• resolve reviews

llvm/include/llvm/Analysis/InlineOrder.h

 //===- InlineOrder.h - Inlining order abstraction -*- C++ ---*-------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 #ifndef LLVM_ANALYSIS_INLINEORDER_H
 #define LLVM_ANALYSIS_INLINEORDER_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/InlineCost.h"
 #include "llvm/IR/InstrTypes.h"
 #include <algorithm>
 #include <utility>
 
 namespace llvm {
 class CallBase;
 class Function;
 
+enum class InlinePriorityMode : int { NoPriority, Size, Cost, OptRatio };
+
 template <typename T> class InlineOrder {
 public:
   using reference = T &;
   using const_reference = const T &;
 
   virtual ~InlineOrder() = default;
 
   virtual size_t size() = 0;
 
   virtual void push(const T &Elt) = 0;
 
   virtual T pop() = 0;
 
   virtual const_reference front() = 0;
 
   virtual void erase_if(function_ref<bool(T)> Pred) = 0;
 
   bool empty() { return !size(); }
 };
 
+std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
+getInlineOrder(InlinePriorityMode UseInlinePriority,
+               FunctionAnalysisManager &FAM, const InlineParams &Params);
+
 template <typename T, typename Container = SmallVector<T, 16>>
 class DefaultInlineOrder : public InlineOrder<T> {
   using reference = T &;
   using const_reference = const T &;
 
 public:
   size_t size() override { return Calls.size() - FirstIndex; }
 
public:
   virtual void update(const CallBase *CB) = 0;
   virtual bool updateAndCheckDecreased(const CallBase *CB) = 0;
 };
 
 class SizePriority : public InlinePriority {
   using PriorityT = unsigned;
   DenseMap<const CallBase *, PriorityT> Priorities;
 
-  static PriorityT evaluate(const CallBase *CB) {
+  PriorityT evaluate(const CallBase *CB) {
     Function *Callee = CB->getCalledFunction();
     return Callee->getInstructionCount();
   }
 
-  static bool isMoreDesirable(const PriorityT &P1, const PriorityT &P2) {
+  bool isMoreDesirable(const PriorityT &P1, const PriorityT &P2) const {
     return P1 < P2;
   }
 
+public:
   bool hasLowerPriority(const CallBase *L, const CallBase *R) const override {
     const auto I1 = Priorities.find(L);
     const auto I2 = Priorities.find(R);
     assert(I1 != Priorities.end() && I2 != Priorities.end());
     return isMoreDesirable(I2->second, I1->second);
   }
 
+  // Update the priority associated with CB.
+  void update(const CallBase *CB) override { Priorities[CB] = evaluate(CB); };
+
+  bool updateAndCheckDecreased(const CallBase *CB) override {
+    auto It = Priorities.find(CB);
+    const auto OldPriority = It->second;
+    It->second = evaluate(CB);
+    const auto NewPriority = It->second;
+    return isMoreDesirable(OldPriority, NewPriority);
+  }
+};
+
+class CostPriority : public InlinePriority {
+  using PriorityT = int;
+  DenseMap<const CallBase *, PriorityT> Priorities;
+  std::function<InlineCost(const CallBase *)> getInlineCost;
+
+  PriorityT evaluate(const CallBase *CB) {
+    auto IC = getInlineCost(CB);
+    int cost = 0;
+    if (IC.isVariable())
+      cost = IC.getCost();
+    else
+      cost = IC.isNever() ? INT_MAX : INT_MIN;
+    return cost;
+  }
+
+  bool isMoreDesirable(const PriorityT &P1, const PriorityT &P2) const {
+    return P1 < P2;
+  }
+
 public:
+  CostPriority() = delete;
+  CostPriority(std::function<InlineCost(const CallBase *)> getInlineCost)
+      : getInlineCost(getInlineCost){};
+
+  bool hasLowerPriority(const CallBase *L, const CallBase *R) const override {
+    const auto I1 = Priorities.find(L);
+    const auto I2 = Priorities.find(R);
+    assert(I1 != Priorities.end() && I2 != Priorities.end());
+    return isMoreDesirable(I2->second, I1->second);
+  }
+
   // Update the priority associated with CB.
   void update(const CallBase *CB) override { Priorities[CB] = evaluate(CB); };
 
   bool updateAndCheckDecreased(const CallBase *CB) override {
     auto It = Priorities.find(CB);
     const auto OldPriority = It->second;
     It->second = evaluate(CB);
     const auto NewPriority = It->second;
public:
   }
 
 private:
   SmallVector<CallBase *, 16> Heap;
   std::function<bool(const CallBase *L, const CallBase *R)> isLess;
   DenseMap<CallBase *, int> InlineHistoryMap;
   std::unique_ptr<InlinePriority> PriorityPtr;
 };
+
 } // namespace llvm
 #endif // LLVM_ANALYSIS_INLINEORDER_H

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   HeatUtils.cpp
   IRSimilarityIdentifier.cpp
   IVDescriptors.cpp
   IVUsers.cpp
   ImportedFunctionsInliningStatistics.cpp
   IndirectCallPromotionAnalysis.cpp
   InlineCost.cpp
   InlineAdvisor.cpp
+  InlineOrder.cpp
   InlineSizeEstimatorAnalysis.cpp
   InstCount.cpp
   InstructionPrecedenceTracking.cpp
   InstructionSimplify.cpp
   Interval.cpp
   IntervalPartition.cpp
   LazyBranchProbabilityInfo.cpp
   LazyBlockFrequencyInfo.cpp

llvm/lib/Analysis/InlineOrder.cpp

+//===- InlineOrder.cpp - Inlining order abstraction -*- C++ ---*-----------===//
+//
+// 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 "llvm/Analysis/InlineOrder.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/InlineAdvisor.h"
+#include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "inline-order"
+
+static llvm::InlineCost getInlineCostWrapper(CallBase &CB,
+                                             FunctionAnalysisManager &FAM,
+                                             const InlineParams &Params) {
+  Function &Caller = *CB.getCaller();
+  ProfileSummaryInfo *PSI =
+      FAM.getResult<ModuleAnalysisManagerFunctionProxy>(Caller)
+          .getCachedResult<ProfileSummaryAnalysis>(
+              *CB.getParent()->getParent()->getParent());
+
+  auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(Caller);
+  auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
+    return FAM.getResult<AssumptionAnalysis>(F);
+  };
+  auto GetBFI = [&](Function &F) -> BlockFrequencyInfo & {
+    return FAM.getResult<BlockFrequencyAnalysis>(F);
+  };
+  auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
+    return FAM.getResult<TargetLibraryAnalysis>(F);
+  };
+
+  Function &Callee = *CB.getCalledFunction();
+  auto &CalleeTTI = FAM.getResult<TargetIRAnalysis>(Callee);
+  bool RemarksEnabled =
+      Callee.getContext().getDiagHandlerPtr()->isMissedOptRemarkEnabled(
+          DEBUG_TYPE);
+  return getInlineCost(CB, Params, CalleeTTI, GetAssumptionCache, GetTLI,
+                       GetBFI, PSI, RemarksEnabled ? &ORE : nullptr);
+}
+
+std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
+llvm::getInlineOrder(InlinePriorityMode UseInlinePriority,
+                     FunctionAnalysisManager &FAM, const InlineParams &Params) {
+  switch (UseInlinePriority) {
+  case InlinePriorityMode::NoPriority:
+    return std::make_unique<DefaultInlineOrder<std::pair<CallBase *, int>>>();
+
+  case InlinePriorityMode::Size:
+    LLVM_DEBUG(dbgs() << "    Current used priority: Size priority ---- \n");
+    return std::make_unique<PriorityInlineOrder>(
+        std::make_unique<SizePriority>());
+
+  case InlinePriorityMode::Cost:
+    LLVM_DEBUG(dbgs() << "    Current used priority: Cost priority ---- \n");
+    return std::make_unique<PriorityInlineOrder>(
+        std::make_unique<CostPriority>([&](const CallBase *CB) -> InlineCost {
+          return getInlineCostWrapper(const_cast<CallBase &>(*CB), FAM, Params);
+        }));
+
+  default:
+    assert("Unsupported Inline Priority Mode");
+    break;
+  }
+  return nullptr;
+}

llvm/lib/Transforms/IPO/ModuleInliner.cpp

 
 using namespace llvm;
 
 #define DEBUG_TYPE "module-inline"
 
 STATISTIC(NumInlined, "Number of functions inlined");
 STATISTIC(NumDeleted, "Number of functions deleted because all callers found");
 
-static cl::opt<bool> InlineEnablePriorityOrder(
-    "module-inline-enable-priority-order", cl::Hidden, cl::init(true),
-    cl::desc("Enable the priority inline order for the module inliner"));
+static cl::opt<InlinePriorityMode> UseInlinePriority(
+    "inline-priority-mode", cl::init(InlinePriorityMode::Size), cl::Hidden,
+    cl::desc("Choose the priority mode to use in module inline"),
+    cl::values(clEnumValN(InlinePriorityMode::NoPriority, "no priority",
+                          "Use no priority, visit callsites in bottom-up."),
+               clEnumValN(InlinePriorityMode::Size, "size",
+                          "Use callee size priority."),
+               clEnumValN(InlinePriorityMode::Cost, "cost",
+                          "Use inline cost priority.")));
 
 /// Return true if the specified inline history ID
 /// indicates an inline history that includes the specified function.
 static bool inlineHistoryIncludes(
     Function *F, int InlineHistoryID,
     const SmallVectorImpl<std::pair<Function *, int>> &InlineHistory) {
   while (InlineHistoryID != -1) {
     assert(unsigned(InlineHistoryID) < InlineHistory.size() &&
PreservedAnalyses ModuleInlinerPass::run(Module &M,
   // In the module inliner, a priority-based worklist is used for calls across
   // the entire Module. With this module inliner, the inline order is not
   // limited to bottom-up order. More globally scope inline order is enabled.
   // Also, the inline deferral logic become unnecessary in this module inliner.
   // It is possible to use other priority heuristics, e.g. profile-based
   // heuristic.
   //
   // TODO: Here is a huge amount duplicate code between the module inliner and
   // the SCC inliner, which need some refactoring.
-  std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>> Calls;
-  if (InlineEnablePriorityOrder)
-    Calls = std::make_unique<PriorityInlineOrder>(
-              std::make_unique<SizePriority>());
-  else
-    Calls = std::make_unique<DefaultInlineOrder<std::pair<CallBase *, int>>>();
+  auto Calls = getInlineOrder(UseInlinePriority, FAM, Params);
   assert(Calls != nullptr && "Expected an initialized InlineOrder");
 
   // Populate the initial list of calls in this module.
   for (Function &F : M) {
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     // We want to generally process call sites top-down in order for
     // simplifications stemming from replacing the call with the returned value
     // after inlining to be visible to subsequent inlining decisions.