[ModuleInliner] Add memorization to MLPriority

[kazu]

Summary

This patch adds memoization to MLPriority.

The trainer will inject noise, but we don't want the priority function
to produce different values for a given metric input. The reason is
that we could potentially evaluate the priority of a call site over
and over to maintain the priority queue, and it's possible that no
metric has changed since the last evaluation.

Without this patch, Priority is the same as Cost, so we've never
bothered to rename Cost to Priority until now. In future, Priority
will be computed by an ML model, so the mental model is:

Priority = MLModel(Cost);

However, since we don't want to produce different outputs for a given
input, we add memorization as:

Priority = MemoisedMLModel(Cost);

where MemoisedMLModel calls MLModel whenever we encounter a given
value of Cost for the first time.

[mtrofin]

What do you mean by "trainer introduces noise"?

llvm/lib/Analysis/InlineOrder.cpp

243

please initialize it at decl.

272

this could be a non-static member with the user of the priority, to avoid the data hanging out for the duration of the compilation, after it's no longer needed. I assume it's static right now as a placeholder?

llvm/lib/Analysis/InlineOrder.cpp

 
 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 {
 
+struct MLInlineOrderMetrics {
+  enum class ObjectState { Empty, Tombstone, Valid } State = ObjectState::Empty;
+  int Cost = 0;
+
+  static MLInlineOrderMetrics makeEmptyKey() {
+    MLInlineOrderMetrics M;
+    M.State = ObjectState::Empty;
+    return M;
+  }
+
+  static MLInlineOrderMetrics makeTombstone() {
+    MLInlineOrderMetrics M;
+    M.State = ObjectState::Tombstone;
+    return M;
+  }
+
+  static MLInlineOrderMetrics makeValid(int Cost) {
+    MLInlineOrderMetrics M;
+    M.State = ObjectState::Valid;
+    M.Cost = Cost;
+    return M;
+  }
+
+  bool isValid() const {
+    return State == ObjectState::Valid;
+  }
+  bool operator==(const MLInlineOrderMetrics &R) const {
+    const MLInlineOrderMetrics &L = *this;
+    return L.State == R.State && L.Cost == R.Cost;
+  }
+};
+
+}  // namespace
+
+namespace llvm {
+template <> struct DenseMapInfo<MLInlineOrderMetrics> {
+  static inline MLInlineOrderMetrics getEmptyKey() {
+    return MLInlineOrderMetrics::makeEmptyKey();
+  }
+
+  static inline MLInlineOrderMetrics getTombstoneKey() {
+    return MLInlineOrderMetrics::makeTombstone();
+  }
+
+  static unsigned getHashValue(const MLInlineOrderMetrics &M) {
+    assert(M.isValid() && "Must be in a valid state for hasing");
+    return static_cast<unsigned>(M.Cost);
+  }
+
+  static bool isEqual(const MLInlineOrderMetrics &A, const MLInlineOrderMetrics &B) {
+    return A == B;
+  }
+};
+}
+
+namespace {
+
 llvm::InlineCost getInlineCostWrapper(CallBase &CB,
                                       FunctionAnalysisManager &FAM,
                                       const InlineParams &Params) {
   Function &Caller = *CB.getCaller();
   ProfileSummaryInfo *PSI =
       FAM.getResult<ModuleAnalysisManagerFunctionProxy>(Caller)
           .getCachedResult<ProfileSummaryAnalysis>(
               *CB.getParent()->getParent()->getParent());

 };
 
 class MLPriority {
 public:
   MLPriority() = default;
   MLPriority(const CallBase *CB, FunctionAnalysisManager &FAM,
              const InlineParams &Params) {
     auto IC = getInlineCostWrapper(const_cast<CallBase &>(*CB), FAM, Params);
+    int Cost;
     if (IC.isVariable())
       Cost = IC.getCost();
     else
       Cost = IC.isNever() ? INT_MAX : INT_MIN;
+    Priority = computePriority(MLInlineOrderMetrics::makeValid(Cost));
   }
 
   static bool isMoreDesirable(const MLPriority &P1, const MLPriority &P2) {
-    return P1.Cost < P2.Cost;
+    return P1.Priority < P2.Priority;
   }
 
 private:
-  int Cost = INT_MAX;
+  int computePriority(const MLInlineOrderMetrics &M) {
+    auto It = MetricsToPriority->find(M);
+    if (It != MetricsToPriority->end()) {
+      return It->second;
+    }
+
+    // Compute the order.  This is to be replaced with ML.  Use the identify
+    // function for now.
+    int Priority = M.Cost;
+
+    // Cache the order.
+    (*MetricsToPriority)[M] = Priority;
+
+    return Priority;
+  }
+
+  static ManagedStatic<DenseMap<MLInlineOrderMetrics, int>> MetricsToPriority;
+
+  int Priority = INT_MAX;
 };
 
+ManagedStatic<DenseMap<MLInlineOrderMetrics, int>> MLPriority::MetricsToPriority;
+
 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());