From 2d6ec146dd1f61d5293dd22d5b395f0e580ab0d0 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 2 Dec 2022 14:25:13 -0800
Subject: [PATCH] [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
---
 llvm/lib/Analysis/InlineOrder.cpp                  | 30 ++++++++++++++++++++--
 .../test/Transforms/Inline/module-inliner-basic.ll |  1 +
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Analysis/InlineOrder.cpp b/llvm/lib/Analysis/InlineOrder.cpp
index ae2c436..75f8bb6 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 @@ private:
   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 4fede8a..f455391 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:
-- 
2.7.4