[llvm] 2d6ec14 - [ModuleInliner] Add MLPriority

[Kazu Hirata via llvm-commits]

Author: Kazu Hirata
Date: 2022-12-02T14:25:13-08:00
New Revision: 2d6ec146dd1f61d5293dd22d5b395f0e580ab0d0

URL: https://github.com/llvm/llvm-project/commit/2d6ec146dd1f61d5293dd22d5b395f0e580ab0d0
DIFF: https://github.com/llvm/llvm-project/commit/2d6ec146dd1f61d5293dd22d5b395f0e580ab0d0.diff

LOG: [ModuleInliner] Add MLPriority

This patch adds MLPriority as the first step toward the ML-based
function inlining with the module inliner.

For now, MLPriority is completely identical to CostPriority.

Once this patch lands, I'm planning to:

- integrate NoInferenceModelRunner,

- memoize the priority computation so that the priority remains the
  same for given values of metrics even with the noise injected during
  training, and

- port/take more features into account.

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

Added: 
    

Modified: 
    llvm/lib/Analysis/InlineOrder.cpp
    llvm/test/Transforms/Inline/module-inliner-basic.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/InlineOrder.cpp b/llvm/lib/Analysis/InlineOrder.cpp
index ae2c4366ef4c..75f8bb69e3ae 100644
--- a/llvm/lib/Analysis/InlineOrder.cpp
+++ b/llvm/lib/Analysis/InlineOrder.cpp
@@ -22,7 +22,7 @@ using namespace llvm;
 
 #define DEBUG_TYPE "inline-order"
 
-enum class InlinePriorityMode : int { Size, Cost, CostBenefit };
+enum class InlinePriorityMode : int { Size, Cost, CostBenefit, ML };
 
 static cl::opt<InlinePriorityMode> UseInlinePriority(
     "inline-priority-mode", cl::init(InlinePriorityMode::Size), cl::Hidden,
@@ -32,7 +32,9 @@ static cl::opt<InlinePriorityMode> UseInlinePriority(
                clEnumValN(InlinePriorityMode::Cost, "cost",
                           "Use inline cost priority."),
                clEnumValN(InlinePriorityMode::CostBenefit, "cost-benefit",
-                          "Use cost-benefit ratio.")));
+                          "Use cost-benefit ratio."),
+               clEnumValN(InlinePriorityMode::ML, "ml",
+                          "Use ML.")));
 
 static cl::opt<int> ModuleInlinerTopPriorityThreshold(
     "moudle-inliner-top-priority-threshold", cl::Hidden, cl::init(0),
@@ -175,6 +177,26 @@ class CostBenefitPriority {
   Optional<CostBenefitPair> CostBenefit;
 };
 
+class MLPriority {
+public:
+  MLPriority() = default;
+  MLPriority(const CallBase *CB, FunctionAnalysisManager &FAM,
+             const InlineParams &Params) {
+    auto IC = getInlineCostWrapper(const_cast<CallBase &>(*CB), FAM, Params);
+    if (IC.isVariable())
+      Cost = IC.getCost();
+    else
+      Cost = IC.isNever() ? INT_MAX : INT_MIN;
+  }
+
+  static bool isMoreDesirable(const MLPriority &P1, const MLPriority &P2) {
+    return P1.Cost < P2.Cost;
+  }
+
+private:
+  int Cost = INT_MAX;
+};
+
 template <typename PriorityT>
 class PriorityInlineOrder : public InlineOrder<std::pair<CallBase *, int>> {
   using T = std::pair<CallBase *, int>;
@@ -274,6 +296,10 @@ llvm::getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params) {
     LLVM_DEBUG(
         dbgs() << "    Current used priority: cost-benefit priority ---- \n");
     return std::make_unique<PriorityInlineOrder<CostBenefitPriority>>(FAM, Params);
+  case InlinePriorityMode::ML:
+    LLVM_DEBUG(
+        dbgs() << "    Current used priority: ML priority ---- \n");
+    return std::make_unique<PriorityInlineOrder<MLPriority>>(FAM, Params);
   }
   return nullptr;
 }

diff  --git a/llvm/test/Transforms/Inline/module-inliner-basic.ll b/llvm/test/Transforms/Inline/module-inliner-basic.ll
index 4fede8a00cbe..f4553914485f 100644
--- a/llvm/test/Transforms/Inline/module-inliner-basic.ll
+++ b/llvm/test/Transforms/Inline/module-inliner-basic.ll
@@ -2,6 +2,7 @@
 ; RUN: opt -passes=module-inline -inline-priority-mode=size         -S < %s | FileCheck %s
 ; RUN: opt -passes=module-inline -inline-priority-mode=cost         -S < %s | FileCheck %s
 ; RUN: opt -passes=module-inline -inline-priority-mode=cost-benefit -S < %s | FileCheck %s
+; RUN: opt -passes=module-inline -inline-priority-mode=ml           -S < %s | FileCheck %s
 
 define i32 @callee(i32 %a) {
 entry: