[llvm][ModuleInliner] Refactor InlineSizePriority and PriorityInlineOrder

This patch introduces the abstract base class InlinePriority to serve as
the comparison function for the priority queue.  A derived class, such
as SizePriority, may choose to cache the priorities for different
functions for performance reasons.

This design shields the type used for the priority away from classes
outside InlinePriority and classes derived from it.  In turn,
PriorityInlineOrder no longer needs to be a template class.

Reviewed By: kazu

Differential Revision: https://reviews.llvm.org/D126300

llvm/include/llvm/Analysis/InlineOrder.h

@@ -70,34 +70,51 @@ private:
   size_t FirstIndex = 0;
 };
 
-class InlineSizePriority {
+class InlinePriority {
 public:
-  InlineSizePriority(int Size) : Size(Size) {}
-
-  static bool isMoreDesirable(const InlineSizePriority &S1,
-                              const InlineSizePriority &S2) {
-    return S1.Size < S2.Size;
-  }
-
-  static InlineSizePriority evaluate(CallBase *CB) {
-    Function *Callee = CB->getCalledFunction();
-    return InlineSizePriority(Callee->getInstructionCount());
-  }
-
-  int Size;
+  virtual bool hasLowerPriority(const CallBase *L, const CallBase *R) const = 0;
+  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) {
+    Function *Callee = CB->getCalledFunction();
+    return Callee->getInstructionCount();
+  }
+
+  static bool isMoreDesirable(const PriorityT &P1, const PriorityT &P2) {
+    return P1 < P2;
+  }
+
+  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);
+  }
+
+public:
+  // 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);
+  }
 };
 
-template <typename PriorityT>
 class PriorityInlineOrder : public InlineOrder<std::pair<CallBase *, int>> {
   using T = std::pair<CallBase *, int>;
-  using HeapT = std::pair<CallBase *, PriorityT>;
   using reference = T &;
   using const_reference = const T &;
 
-  static bool cmp(const HeapT &P1, const HeapT &P2) {
-    return PriorityT::isMoreDesirable(P2.second, P1.second);
-  }
-
   // A call site could become less desirable for inlining because of the size
   // growth from prior inlining into the callee. This method is used to lazily
   // update the desirability of a call site if it's decreasing. It is only
@@ -106,31 +123,29 @@ class PriorityInlineOrder : public InlineOrder<std::pair<CallBase *, int>> {
   // pushed right back into the heap. For simplicity, those cases where
   // the desirability of a call site increases are ignored here.
   void adjust() {
-    bool Changed = false;
-    do {
-      CallBase *CB = Heap.front().first;
-      const PriorityT PreviousGoodness = Heap.front().second;
-      const PriorityT CurrentGoodness = PriorityT::evaluate(CB);
-      Changed = PriorityT::isMoreDesirable(PreviousGoodness, CurrentGoodness);
-      if (Changed) {
-        std::pop_heap(Heap.begin(), Heap.end(), cmp);
-        Heap.pop_back();
-        Heap.push_back({CB, CurrentGoodness});
-        std::push_heap(Heap.begin(), Heap.end(), cmp);
-      }
-    } while (Changed);
+    while (PriorityPtr->updateAndCheckDecreased(Heap.front())) {
+      std::pop_heap(Heap.begin(), Heap.end(), isLess);
+      std::push_heap(Heap.begin(), Heap.end(), isLess);
+    }
   }
 
 public:
+  PriorityInlineOrder(std::unique_ptr<InlinePriority> PriorityPtr)
+      : PriorityPtr(std::move(PriorityPtr)) {
+    isLess = [this](const CallBase *L, const CallBase *R) {
+      return this->PriorityPtr->hasLowerPriority(L, R);
+    };
+  }
+
   size_t size() override { return Heap.size(); }
 
   void push(const T &Elt) override {
     CallBase *CB = Elt.first;
     const int InlineHistoryID = Elt.second;
-    const PriorityT Goodness = PriorityT::evaluate(CB);
 
-    Heap.push_back({CB, Goodness});
-    std::push_heap(Heap.begin(), Heap.end(), cmp);
+    Heap.push_back(CB);
+    PriorityPtr->update(CB);
+    std::push_heap(Heap.begin(), Heap.end(), isLess);
     InlineHistoryMap[CB] = InlineHistoryID;
   }
 
@@ -138,10 +153,10 @@ public:
     assert(size() > 0);
     adjust();
 
-    CallBase *CB = Heap.front().first;
+    CallBase *CB = Heap.front();
     T Result = std::make_pair(CB, InlineHistoryMap[CB]);
     InlineHistoryMap.erase(CB);
-    std::pop_heap(Heap.begin(), Heap.end(), cmp);
+    std::pop_heap(Heap.begin(), Heap.end(), isLess);
     Heap.pop_back();
     return Result;
   }
@@ -150,21 +165,23 @@ public:
     assert(size() > 0);
     adjust();
 
-    CallBase *CB = Heap.front().first;
+    CallBase *CB = Heap.front();
     return *InlineHistoryMap.find(CB);
   }
 
   void erase_if(function_ref<bool(T)> Pred) override {
-    auto PredWrapper = [=](HeapT P) -> bool {
-      return Pred(std::make_pair(P.first, 0));
+    auto PredWrapper = [=](CallBase *CB) -> bool {
+      return Pred(std::make_pair(CB, 0));
     };
     llvm::erase_if(Heap, PredWrapper);
-    std::make_heap(Heap.begin(), Heap.end(), cmp);
+    std::make_heap(Heap.begin(), Heap.end(), isLess);
   }
 
 private:
-  SmallVector<HeapT, 16> Heap;
+  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/Transforms/IPO/ModuleInliner.cpp

@@ -143,7 +143,8 @@ PreservedAnalyses ModuleInlinerPass::run(Module &M,
   // the SCC inliner, which need some refactoring.
   std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>> Calls;
   if (InlineEnablePriorityOrder)
-    Calls = std::make_unique<PriorityInlineOrder<InlineSizePriority>>();
+    Calls = std::make_unique<PriorityInlineOrder>(
+              std::make_unique<SizePriority>());
   else
     Calls = std::make_unique<DefaultInlineOrder<std::pair<CallBase *, int>>>();
   assert(Calls != nullptr && "Expected an initialized InlineOrder");