[llvm][ModuleInliner] Refactor InlineSizePriority and PriorityInlineOrder

[taolq]

Summary

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.

[kazu]

May I suggest expanding the summary? You could say something like:

ModuleInliner] Refactor InlineSizePriority and PriorityInlineOrder

This patch introduces 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.

Note that you need to click "Edit Revision" to update the summary as Phabricator does not automatically update the summary on arc diff.

llvm/include/llvm/Analysis/InlineOrder.h

106–108

I'm a bit worried about the fact that we are hashing CB three times here. Could we do something like this?

auto It = Priorities.find(CB);
const auto OldPriority = *It;
*It = evaluate(CB);
const auto NewPriority = *It;

[taolq]

• use iterator to decrease rehashing

[tschuett]

The name of the function is compare, but the comment reads more like isLessThan.

For compare, I would expect -1, 0, or +1.

[taolq]

• rename compare to hasLowerPriority

[kazu]

LGTM. Thank you for updating the patch!

[taolq]

• rebase to main

[tianshilei1992]

Compiler now complains:

warning: 'llvm::InlinePriority' has virtual functions but non-virtual destructor [-Wnon-virtual-dtor]
class InlinePriority {
      ^
llvm-project/llvm/include/llvm/Analysis/InlineOrder.h:80:7: warning: 'llvm::SizePriority' has virtual functions but non-virtual destructor [-Wnon-virtual-dtor]
class SizePriority : public InlinePriority {
      ^

llvm/include/llvm/Analysis/InlineOrder.h

Calls.erase(std::remove_if(Calls.begin() + FirstIndex, Calls.end(), Pred),
                 Calls.end());
   }
 
 private:
   Container Calls;
   size_t FirstIndex = 0;
 };
 
-class InlineSizePriority {
+class InlinePriority {
 public:
-  InlineSizePriority(int Size) : Size(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;
+};
 
-  static bool isMoreDesirable(const InlineSizePriority &S1,
-                              const InlineSizePriority &S2) {
-    return S1.Size < S2.Size;
-  }
+class SizePriority : public InlinePriority {
+  using PriorityT = unsigned;
+  DenseMap<const CallBase *, PriorityT> Priorities;
 
-  static InlineSizePriority evaluate(CallBase *CB) {
+  static PriorityT evaluate(const CallBase *CB) {
     Function *Callee = CB->getCalledFunction();
-    return InlineSizePriority(Callee->getInstructionCount());
+    return Callee->getInstructionCount();
+  }
+
+  static bool isMoreDesirable(const PriorityT &P1, const PriorityT &P2) {
+    return P1 < P2;
   }
 
-  int Size;
+  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
   // called on pop() or front(), not every time the desirability changes. When
   // the desirability of the front call site decreases, an updated one would be
   // 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;
   }
 
   T pop() override {
     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;
   }
 
   const_reference front() override {
     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

PreservedAnalyses ModuleInlinerPass::run(Module &M,
   // 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<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");
 
   // 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