From 6ebeb7f8baf3feb6657c4e9c523fe61eca2dc43d Mon Sep 17 00:00:00 2001
From: Liqiang Tao <taolq@outlook.com>
Date: Mon, 12 Jul 2021 10:28:17 +0800
Subject: [PATCH] [llvm][Inliner] Templatize PriorityInlineOrder

The patch templatize PriorityInlinerOrder so that it can accept any type priority metric.

Reviewed By: kazu

Differential Revision: https://reviews.llvm.org/D104972
---
 llvm/lib/Transforms/IPO/Inliner.cpp | 48 ++++++++++++++++++++++++-------------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/Inliner.cpp b/llvm/lib/Transforms/IPO/Inliner.cpp
index 5d6a01d..d0d4b9c 100644
--- a/llvm/lib/Transforms/IPO/Inliner.cpp
+++ b/llvm/lib/Transforms/IPO/Inliner.cpp
@@ -724,21 +724,31 @@ private:
   size_t FirstIndex = 0;
 };
 
+class Priority {
+public:
+  Priority(int Size) : Size(Size) {}
+
+  static bool isMoreDesirable(const Priority &S1, const Priority &S2) {
+    return S1.Size < S2.Size;
+  }
+
+  static Priority evaluate(CallBase *CB) {
+    Function *Callee = CB->getCalledFunction();
+    return Priority(Callee->getInstructionCount());
+  }
+
+  int Size;
+};
+
+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 &;
 
-  // Return true if S1 is more desirable than S2.
-  static bool isMoreDesirable(int S1, int S2) { return S1 < S2; }
-
-  static bool cmp(const T &P1, const T &P2) {
-    return isMoreDesirable(P2.second, P1.second);
-  }
-
-  int evaluate(CallBase *CB) {
-    Function *Callee = CB->getCalledFunction();
-    return (int)Callee->getInstructionCount();
+  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
@@ -752,9 +762,9 @@ class PriorityInlineOrder : public InlineOrder<std::pair<CallBase *, int>> {
     bool Changed = false;
     do {
       CallBase *CB = Heap.front().first;
-      const int PreviousGoodness = Heap.front().second;
-      const int CurrentGoodness = evaluate(CB);
-      Changed = isMoreDesirable(PreviousGoodness, CurrentGoodness);
+      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();
@@ -770,7 +780,7 @@ public:
   void push(const T &Elt) override {
     CallBase *CB = Elt.first;
     const int InlineHistoryID = Elt.second;
-    const int Goodness = evaluate(CB);
+    const PriorityT Goodness = PriorityT::evaluate(CB);
 
     Heap.push_back({CB, Goodness});
     std::push_heap(Heap.begin(), Heap.end(), cmp);
@@ -798,12 +808,16 @@ public:
   }
 
   void erase_if(function_ref<bool(T)> Pred) override {
-    Heap.erase(std::remove_if(Heap.begin(), Heap.end(), Pred), Heap.end());
+    auto PredWrapper = [=](HeapT P) -> bool {
+      return Pred(std::make_pair(P.first, 0));
+    };
+    Heap.erase(std::remove_if(Heap.begin(), Heap.end(), PredWrapper),
+               Heap.end());
     std::make_heap(Heap.begin(), Heap.end(), cmp);
   }
 
 private:
-  SmallVector<T, 16> Heap;
+  SmallVector<HeapT, 16> Heap;
   DenseMap<CallBase *, int> InlineHistoryMap;
 };
 
@@ -854,7 +868,7 @@ PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   // incrementally maknig a single function grow in a super linear fashion.
   std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>> Calls;
   if (InlineEnablePriorityOrder)
-    Calls = std::make_unique<PriorityInlineOrder>();
+    Calls = std::make_unique<PriorityInlineOrder<Priority>>();
   else
     Calls = std::make_unique<DefaultInlineOrder<std::pair<CallBase *, int>>>();
   assert(Calls != nullptr && "Expected an initialized InlineOrder");
-- 
2.7.4