[ModuleInliner] Add MLPriority

[kazu]

Summary

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.

[mtrofin]

can you also add a test? thanks!

llvm/lib/Analysis/InlineOrder.cpp

197

can you initialize it here - lowers the risk of undefined values later (like if somehow the ctor forgets to initialize it)

[kazu]

In D139140#3965039, @mtrofin wrote:

can you also add a test? thanks!

Will do.

llvm/lib/Analysis/InlineOrder.cpp

197

Would int Cost = INT_MAX; be OK here? The actual initialization is a bit complicated.

[mtrofin]

llvm/lib/Analysis/InlineOrder.cpp

197

sgtm

[kazu]

Add a test. Also, initialize Cost.

[kazu]

PTAL. Thanks!

llvm/lib/Analysis/InlineOrder.cpp

 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Support/CommandLine.h"
 
 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,
     cl::desc("Choose the priority mode to use in module inline"),
     cl::values(clEnumValN(InlinePriorityMode::Size, "size",
                           "Use callee size priority."),
                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),
     cl::desc("The cost threshold for call sites that get inlined without the "
              "cost-benefit analysis"));
 
 namespace {
 
public:
   }
 
 private:
   int Cost = INT_MAX;
   int StaticBonusApplied = 0;
   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>;
 
   bool hasLowerPriority(const CallBase *L, const CallBase *R) const {
     const auto I1 = Priorities.find(L);
     const auto I2 = Priorities.find(R);
     assert(I1 != Priorities.end() && I2 != Priorities.end());
llvm::getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params) {
   case InlinePriorityMode::Cost:
     LLVM_DEBUG(dbgs() << "    Current used priority: Cost priority ---- \n");
     return std::make_unique<PriorityInlineOrder<CostPriority>>(FAM, Params);
 
   case InlinePriorityMode::CostBenefit:
     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;
 }

llvm/test/Transforms/Inline/module-inliner-basic.ll

 ; RUN: opt -passes=module-inline                                    -S < %s | FileCheck %s
 ; 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:
   %add = add nsw i32 %a, 1
   ret i32 %add
 }
 
 define i32 @caller(i32 %a) {
 entry:
   %call = call i32 @callee(i32 %a)
   ret i32 %call
 }
 
 ; CHECK-LABEL: @caller
 ; CHECK-NOT:     call
 ; CHECK:         ret