Unpack the CostEstimate feature in ML inlining models.

[jacobhegna]

Summary

This change yields an additional 2% size reduction on an internal search
binary and an additional 0.5% size reduction on fuchsia's Oz targets.

[jacobhegna]

Fix misc build issues.

[mtrofin]

llvm/include/llvm/Analysis/InlineCost.h

57

I think it'd be preferable if this were defined on the ML side, like in InlineModelFeatureMaps.h.

The reason is that it brings in this concept of feature name (e.g. "sroa_savings"), which is a ML side concept.

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

17

this wouldn't be needed once you move in the changes in InlineCost.h here.

llvm/lib/Analysis/InlineCost.cpp

469

I don't think the ML aspects should intermingle with the non-ml ones.

[jacobhegna]

Move ML-specific code out of InlineCost.[h, cpp]

[jacobhegna]

Switch to visitor design pattern

Additionally, we re-use the heuristic based cost estimate feature in the model,
which is then supplemented by the expanded cost features. This increases
stability in training.

In a later patch, we will revisit completely removing the heuristics from the ML.

[jacobhegna]

Updated the config for mock model generation.

Removed unnecessary edits to heuristic inliner.

[mtrofin]

llvm/include/llvm/Analysis/InlineAdvisor.h

40 ↗

(On Diff #355007)

unrelated change - if the linter did this, just group it (and others like it) upfront, and commit that as a NFC, then rebase this one

llvm/include/llvm/Analysis/InlineCost.h

314

I'd drop the name 'array', suggests too much impl detail.

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

51

Aah! I see why InlineCostFeaturesArray.

How about this enum were InlineCostFeatureIndex, and then the collection is InlineCostFeatures

llvm/lib/Analysis/CMakeLists.txt

8 ↗

(On Diff #355007)

more clear to say "<TO BE UPDATED>"

llvm/lib/Analysis/InlineCost.cpp

132

remove spurious change

416–420

formatting change - see previous comment

427

formatting change

1057

Where is this used, other than in onDisableLoadElimination, where it's set to 0?

1058

this is counting the opportunities for SROACostSavings, right? Maybe rename to SROACostSavingOpportunities?

1061

Is this really a threshold, or can it have a name that better communicates what it accumulates? (apologies, can't right now figure what that may be)

same for the 2 bonuses above.

1089–1091

formatting

1196–1197

formatting

1351

formatting

2093–2095

formatting

2563

formatting

2656

Not sure. what's the value of CFA.features() if !R?

may be more clear to handle the error here and return {} or something like that. It creates less coupling between deep implementation details and their consumers; it also simplifies the state space of the values this API can produce.

Better yet: return an Optional<InlineCostFeaturesArray>, because a 0-valued feature vector isn't necessarily equivalent to CFA.analyze() saying "no". Returning None in that case is a more clear API

2916–2917

formatting

llvm/lib/Analysis/MLInlineAdvisor.cpp

46

why

254

do we still want this here?

[jacobhegna]

llvm/lib/Analysis/InlineCost.cpp

1057

it is meant so that InlineCostFeatures::LoadElimination is only ever set once, but there is a better way to express that. It also was never set to a nonzero value in the first place, which is just a bug.

1061

these are copied from the heuristic inline visitor and used in almost the same way - there is some complexity in the other visitor that isn't replicated here. I'll keep the names the same to reflect that duplication. In the future, there some be some abstraction so that the two visitors could share these values/logic for updating them.

2656

I'm wasn't returning an optional for a very specific reason: std::array cannot be efficiently moved-from. so, if I return an optional, we couldn't explicitly unpack it later without eating another copy of the array. I can swap it back now, then in MLInlineAdvisor.cpp avoid the copy by taking a reference to the array inside the optional. In a later patch, I want to switch the array to a vector, so that we can efficiently std::move it, and then returning an optional here wouldn't come with a performance hit

llvm/lib/Analysis/MLInlineAdvisor.cpp

46

clang-format doesn't seem to understand these back-to-back macros and tries to add another level of indentation to each line, it makes it hard to read.

254

ya - it increases stability in training. still looking in how to completely remove the heuristics from the ML without sacrificing training stability.

[jacobhegna]

Address comments.

[jacobhegna]

Rebase on to clang-format changes.

[mtrofin]

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

51

why not leave this where it was? because you can now populate the remainder in a loop?

64–65

unrelated change?

llvm/lib/Analysis/InlineCost.cpp

442

spurious change?

1061

OK, can you add a FIXME explaining this, above their name or something like that.

1064

I think increment would be more clear. Then Delta can have default value 1, too (For the few cases that'd be the case)

1097

can this '3' be a const with some name (so it's easier to chase down any magic factors)

1162

can we replace this 2 and the 4 above with const values with a good name

same below - 3, 2 etc.

[jacobhegna]

Address mtrofin's comments; add simple unit test.

[jacobhegna]

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

51

it's an artifact from when I was fiddling around with things - will move it back

64–65

are you referring to the newline that is inserted here? I think it just makes things easier to read in the enum

llvm/lib/Analysis/InlineCost.cpp

1097

discussing offline

1162

ok

[jacobhegna]

Forgot to name the constants - now added.

One thing to note is the getExpectedNumberOfCompare function. Because the same
math is used in both visitors, I abstracted it and moved the relevant comment.
Just want to double check that you think that's ok.

[jacobhegna]

(One more constant)

[mtrofin]

lgtm

[jacobhegna]

Silence clang tidy.

llvm/include/llvm/Analysis/InlineCost.h

 /// Do not inline functions which allocate this many bytes on the stack
 /// when the caller is recursive.
 const unsigned TotalAllocaSizeRecursiveCaller = 1024;
 /// Do not inline dynamic allocas that have been constant propagated to be
 /// static allocas above this amount in bytes.
 const uint64_t MaxSimplifiedDynamicAllocaToInline = 65536;
 } // namespace InlineConstants
 
+// The features which determine the computation of InlineCost.
+#define INLINE_COST_FEATURE_ITERATOR(M)                                        \
+  M(SROASavings, "sroa_savings")                                               \
+  M(SROALosses, "sroa_losses")                                                 \
+  M(LoadElimination, "load_elimination")                                       \
+  M(CallPenalty, "call_penalty")                                               \
+  M(CallArgumentSetup, "call_argument_setup")                                  \
+  M(LoadRelativeIntrinsic, "load_relative_intrinsic")                          \
+  M(LoweredCallArgSetup, "lowered_call_arg_setup")                             \
+  M(IndirectCallPenalty, "indirect_call_penalty")                              \
+  M(JumpTablePenalty, "jump_table_penalty")                                    \
+  M(CaseClusterPenalty, "case_cluster_penalty")                                \
+  M(SwitchPenalty, "switch_penalty")                                           \
+  M(UnsimplifiedCommonInstructions, "unsimplified_common_instructions")        \
+  M(NumLoops, "num_loops")                                                     \
+  M(DeadBlocks, "dead_blocks")                                                 \
+  M(SimplifiedInstructions, "simplified_instructions")                         \
+  M(ConstantArgs, "constant_args")                                             \
+  M(ConstantOffsetPtrArgs, "constant_offset_ptr_args")                         \
+  M(CallSiteCost, "callsite_cost")                                             \
+  M(ColdCcPenalty, "cold_cc_penalty")                                          \
+  M(LastCallToStaticBonus, "last_call_to_static_bonus")                        \
+  M(IsMultipleBlocks, "is_multiple_blocks")
+
+// clang-format off
+enum class InlineCostFeatures : size_t {
+#define POPULATE_INDICES(INDEX_NAME, NAME) INDEX_NAME,
+  INLINE_COST_FEATURE_ITERATOR(POPULATE_INDICES)
+#undef POPULATE_INDICES
+
+  NumberOfFeatures
+};
+// clang-format on
+
+using InlineCostFeaturesArray =
+    std::array<int, static_cast<size_t>(InlineCostFeatures::NumberOfFeatures)>;
+
+int reduceCostFeatures(const InlineCostFeaturesArray &Features);
+
 /// Represents the cost of inlining a function.
 ///
 /// This supports special values for functions which should "always" or
 /// "never" be inlined. Otherwise, the cost represents a unitless amount;
 /// smaller values increase the likelihood of the function being inlined.
 ///
 /// Objects of this type also provide the adjusted threshold for inlining
 /// based on the information available for a particular callsite. They can be

 /// Get the cost estimate ignoring thresholds. This is similar to getInlineCost
 /// when passed InlineParams::ComputeFullInlineCost, or a non-null ORE. It
 /// uses default InlineParams otherwise.
 /// Contrary to getInlineCost, which makes a threshold-based final evaluation of
 /// should/shouldn't inline, captured in InlineResult, getInliningCostEstimate
 /// returns:
 /// - None, if the inlining cannot happen (is illegal)
 /// - an integer, representing the cost.
-Optional<int> getInliningCostEstimate(
+Optional<InlineCostFeaturesArray> getInliningCostFeatures(
     CallBase &Call, TargetTransformInfo &CalleeTTI,
     function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
     function_ref<BlockFrequencyInfo &(Function &)> GetBFI = nullptr,
     ProfileSummaryInfo *PSI = nullptr,
     OptimizationRemarkEmitter *ORE = nullptr);
 
 /// Minimal filter to detect invalid constructs for inlining.
 InlineResult isInlineViable(Function &Callee);
 
 // This pass is used to annotate instructions during the inline process for
 // debugging and analysis. The main purpose of the pass is to see and test
 // inliner's decisions when creating new optimizations to InlineCost.
 struct InlineCostAnnotationPrinterPass
     : PassInfoMixin<InlineCostAnnotationPrinterPass> {
   raw_ostream &OS;
 
 public:
   explicit InlineCostAnnotationPrinterPass(raw_ostream &OS) : OS(OS) {}
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
 };
 } // namespace llvm
 
 #endif

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

 //===- InlineModelFeatureMaps.h - common model runner defs ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 
 #ifndef LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H
 #define LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H
 
 #include <array>
 #include <string>
 #include <vector>
 
+#include "llvm/Analysis/InlineCost.h"
+
 namespace llvm {
 
 // List of features. Each feature is defined through a triple:
 // - the name of an enum member, which will be the feature index
 // - a textual name, used for Tensorflow model binding (so it needs to match the
 // names used by the Tensorflow model)
 // - a documentation description. Currently, that is not used anywhere
 // programmatically, and serves as workaround to inability of inserting comments
 // in macros.
 #define INLINE_FEATURE_ITERATOR(M)                                             \
   M(CalleeBasicBlockCount, "callee_basic_block_count",                         \
     "number of basic blocks of the callee")                                    \
   M(CallSiteHeight, "callsite_height",                                         \
     "position of the call site in the original call graph - measured from "    \
     "the farthest SCC")                                                        \
   M(NodeCount, "node_count",                                                   \
     "total current number of defined functions in the module")                 \
   M(NrCtantParams, "nr_ctant_params",                                          \
     "number of parameters in the call site that are constants")                \
-  M(CostEstimate, "cost_estimate", "total cost estimate (threshold - free)")   \
   M(EdgeCount, "edge_count", "total number of calls in the module")            \
   M(CallerUsers, "caller_users",                                               \
     "number of module-internal users of the caller, +1 if the caller is "      \
     "exposed externally")                                                      \
   M(CallerConditionallyExecutedBlocks, "caller_conditionally_executed_blocks", \
     "number of blocks reached from a conditional instruction, in the caller")  \
   M(CallerBasicBlockCount, "caller_basic_block_count",                         \
     "number of basic blocks in the caller")                                    \
   M(CalleeConditionallyExecutedBlocks, "callee_conditionally_executed_blocks", \
     "number of blocks reached from a conditional instruction, in the callee")  \
   M(CalleeUsers, "callee_users",                                               \
     "number of module-internal users of the callee, +1 if the callee is "      \
-    "exposed externally")
+    "exposed externally")                                                      \
+  M(CostEstimate, "cost_estimate", "total cost estimate (threshold - free)")
 
+// clang-format off
 enum class FeatureIndex : size_t {
+// InlineCost features - these must come first
+#define POPULATE_INDICES(INDEX_NAME, NAME) INDEX_NAME,
+  INLINE_COST_FEATURE_ITERATOR(POPULATE_INDICES)
+#undef POPULATE_INDICES
+
+// Non-cost features
 #define POPULATE_INDICES(INDEX_NAME, NAME, COMMENT) INDEX_NAME,
   INLINE_FEATURE_ITERATOR(POPULATE_INDICES)
 #undef POPULATE_INDICES
-      NumberOfFeatures
+
 };
+// clang-format on
+
+constexpr FeatureIndex
+inlineCostFeatureToMlFeature(InlineCostFeatures Feature) {
+  return static_cast<FeatureIndex>(static_cast<size_t>(Feature));
+}
 
 constexpr size_t NumberOfFeatures =
     static_cast<size_t>(FeatureIndex::NumberOfFeatures);
 
 extern const std::array<std::string, NumberOfFeatures> FeatureNameMap;
 
 extern const char *const DecisionName;
 extern const char *const DefaultDecisionName;
 extern const char *const RewardName;
 
 using InlineFeatures = std::vector<int64_t>;
 
 } // namespace llvm
 #endif // LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H

llvm/lib/Analysis/InlineCost.cpp

static cl::opt<bool> InlineCallerSupersetNoBuiltin(
     cl::desc("Allow inlining when caller has a superset of callee's nobuiltin "
              "attributes."));
 
 static cl::opt<bool> DisableGEPConstOperand(
     "disable-gep-const-evaluation", cl::Hidden, cl::init(false),
     cl::desc("Disables evaluation of GetElementPtr with constant operands"));
 
 namespace {
-class InlineCostCallAnalyzer;
+
+template <class CostUpdater> class InlineCostCallAnalyzer;
 
 // This struct is used to store information about inline cost of a
 // particular instruction
 struct InstructionCostDetail {
   int CostBefore = 0;
   int CostAfter = 0;
   int ThresholdBefore = 0;
   int ThresholdAfter = 0;
 
   int getThresholdDelta() const { return ThresholdAfter - ThresholdBefore; }
 
   int getCostDelta() const { return CostAfter - CostBefore; }
 
   bool hasThresholdChanged() const { return ThresholdAfter != ThresholdBefore; }
 };
 
+template <class CostUpdater>
 class InlineCostAnnotationWriter : public AssemblyAnnotationWriter {
 private:
-  InlineCostCallAnalyzer *const ICCA;
+  InlineCostCallAnalyzer<CostUpdater> *const ICCA;
 
 public:
-  InlineCostAnnotationWriter(InlineCostCallAnalyzer *ICCA) : ICCA(ICCA) {}
+  InlineCostAnnotationWriter(InlineCostCallAnalyzer<CostUpdater> *ICCA)
+      : ICCA(ICCA) {}
   virtual void emitInstructionAnnot(const Instruction *I,
                                     formatted_raw_ostream &OS) override;
 };
 
 /// Carry out call site analysis, in order to evaluate inlinability.
 /// NOTE: the type is currently used as implementation detail of functions such
 /// as llvm::getInlineCost. Note the function_ref constructor parameters - the
 /// expectation is that they come from the outer scope, from the wrapper
protected:
   bool visitSwitchInst(SwitchInst &SI);
   bool visitIndirectBrInst(IndirectBrInst &IBI);
   bool visitResumeInst(ResumeInst &RI);
   bool visitCleanupReturnInst(CleanupReturnInst &RI);
   bool visitCatchReturnInst(CatchReturnInst &RI);
   bool visitUnreachableInst(UnreachableInst &I);
 
 public:
-  CallAnalyzer(
-      Function &Callee, CallBase &Call, const TargetTransformInfo &TTI,
+  CallAnalyzer(Function &Callee, CallBase &Call, const TargetTransformInfo &TTI,
                function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
                function_ref<BlockFrequencyInfo &(Function &)> GetBFI = nullptr,
                ProfileSummaryInfo *PSI = nullptr,
                OptimizationRemarkEmitter *ORE = nullptr)
       : TTI(TTI), GetAssumptionCache(GetAssumptionCache), GetBFI(GetBFI),
         PSI(PSI), F(Callee), DL(F.getParent()->getDataLayout()), ORE(ORE),
         CandidateCall(Call), EnableLoadElimination(true) {}
 
   InlineResult analyze();
 
-  Optional<Constant*> getSimplifiedValue(Instruction *I) {
+  Optional<Constant *> getSimplifiedValue(Instruction *I) {
     if (SimplifiedValues.find(I) != SimplifiedValues.end())
       return SimplifiedValues[I];
     return None;
   }
 
   // Keep a bunch of stats about the cost savings found so we can print them
   // out when debugging.
   unsigned NumConstantArgs = 0;
   unsigned NumConstantOffsetPtrArgs = 0;
   unsigned NumAllocaArgs = 0;
   unsigned NumConstantPtrCmps = 0;
   unsigned NumConstantPtrDiffs = 0;
   unsigned NumInstructionsSimplified = 0;
 
   void dump();
 };
 
+bool isExclusiveMLFeature(InlineCostFeatures Feature) {
+  return Feature == InlineCostFeatures::SROASavings ||
+         Feature == InlineCostFeatures::IsMultipleBlocks ||
+         Feature == InlineCostFeatures::DeadBlocks ||
+         Feature == InlineCostFeatures::SimplifiedInstructions ||
+         Feature == InlineCostFeatures::ConstantArgs ||
+         Feature == InlineCostFeatures::ConstantOffsetPtrArgs;
+}
+
+class InlineCostValueUpdater {
+public:
+  using CostType = int;
+
+  bool shouldStop(int64_t Threshold) const { return Cost >= Threshold; }
+
+  void addUnchecked(InlineCostFeatures Feature, int CostDelta) {
+    Cost += CostDelta;
+  }
+
+  void add(InlineCostFeatures Feature, int64_t CostDelta,
+           int64_t UpperBound = std::numeric_limits<int>::max()) {
+    assert(UpperBound > 0 && UpperBound <= std::numeric_limits<int>::max() &&
+           "invalid upper bound");
+
+    if (isExclusiveMLFeature(Feature)) {
+      return;
+    }
+
+    Cost = std::min<int64_t>(UpperBound, Cost + CostDelta);
+  }
+
+  void accumulateBonus(int OtherCost, int Threshold) {
+    Cost -= std::max(0, Threshold - OtherCost);
+  }
+
+  // These all are equivalent in this case, but differ for the
+  // InlineCostFeaturesUpdater
+  int reduce() const { return Cost; }
+  int get() const { return Cost; }
+
+private:
+  int Cost = 0;
+};
+
+class InlineCostFeaturesUpdater {
+public:
+  using CostType = InlineCostFeaturesArray;
+
+  // Never stop, because we want to collect all features
+  bool shouldStop(int64_t Threshold) const {
+    (void)Threshold;
+    return false;
+  }
+
+  void addUnchecked(InlineCostFeatures Feature, int CostDelta) {
+    auto FeatureIdx = static_cast<int>(Feature);
+    Cost[FeatureIdx] += CostDelta;
+  }
+
+  void add(InlineCostFeatures Feature, int64_t CostDelta,
+           int64_t UpperBound = std::numeric_limits<int>::max()) {
+    assert(UpperBound > 0 && UpperBound <= std::numeric_limits<int>::max() &&
+           "invalid upper bound");
+    auto FeatureIdx = static_cast<int>(Feature);
+    Cost[FeatureIdx] =
+        std::min<int64_t>(UpperBound, Cost[FeatureIdx] + CostDelta);
+  }
+
+  void accumulateBonus(const CostType &OtherFeatures, int Threshold) {
+    for (int I = 0; I < static_cast<int>(InlineCostFeatures::NumberOfFeatures);
+         ++I) {
+      Cost[I] -= std::max(0, Threshold - OtherFeatures[I]);
+    }
+  }
+
+  int reduce() const { return reduceCostFeatures(Cost); }
+
+  const InlineCostFeaturesArray &get() const { return Cost; }
+
+private:
+  CostType Cost = {};
+};
+
 /// FIXME: if it is necessary to derive from InlineCostCallAnalyzer, note
 /// the FIXME in onLoweredCall, when instantiating an InlineCostCallAnalyzer
+template <class CostUpdater>
 class InlineCostCallAnalyzer final : public CallAnalyzer {
   const int CostUpperBound = INT_MAX - InlineConstants::InstrCost - 1;
   const bool ComputeFullInlineCost;
   int LoadEliminationCost = 0;
   /// Bonus to be applied when percentage of vector instructions in callee is
   /// high (see more details in updateThreshold).
   int VectorBonus = 0;
   /// Bonus to be applied when the callee has only one reachable basic block.
class InlineCostCallAnalyzer final : public CallAnalyzer {
 
   // True if the cost-benefit-analysis-based inliner is enabled.
   const bool CostBenefitAnalysisEnabled;
 
   /// Inlining cost measured in abstract units, accounts for all the
   /// instructions expected to be executed for a given function invocation.
   /// Instructions that are statically proven to be dead based on call-site
   /// arguments are not counted here.
-  int Cost = 0;
+  CostUpdater Cost = {};
 
   // The cumulative cost at the beginning of the basic block being analyzed.  At
   // the end of analyzing each basic block, "Cost - CostAtBBStart" represents
   // the size of that basic block.
   int CostAtBBStart = 0;
 
   // The static size of live but cold basic blocks.  This is "static" in the
   // sense that it's not weighted by profile counts at all.
void addCost(int64_t Inc, int64_t UpperBound = INT_MAX) {
     assert(UpperBound > 0 && UpperBound <= INT_MAX && "invalid upper bound");
     Cost = (int)std::min(UpperBound, Cost + Inc);
   }
 
   void onDisableSROA(AllocaInst *Arg) override {
     auto CostIt = SROAArgCosts.find(Arg);
     if (CostIt == SROAArgCosts.end())
       return;
-    addCost(CostIt->second);
+    Cost.add(InlineCostFeatures::SROALosses, CostIt->second);
     SROACostSavings -= CostIt->second;
     SROACostSavingsLost += CostIt->second;
     SROAArgCosts.erase(CostIt);
   }
 
   void onDisableLoadElimination() override {
-    addCost(LoadEliminationCost);
+    Cost.add(InlineCostFeatures::LoadElimination, LoadEliminationCost);
     LoadEliminationCost = 0;
   }
-  void onCallPenalty() override { addCost(InlineConstants::CallPenalty); }
+  void onCallPenalty() override {
+    Cost.add(InlineCostFeatures::CallPenalty, InlineConstants::CallPenalty);
+  }
   void onCallArgumentSetup(const CallBase &Call) override {
     // Pay the price of the argument setup. We account for the average 1
     // instruction per call argument setup here.
-    addCost(Call.arg_size() * InlineConstants::InstrCost);
+    Cost.add(InlineCostFeatures::CallArgumentSetup,
+             Call.arg_size() * InlineConstants::InstrCost);
   }
   void onLoadRelativeIntrinsic() override {
     // This is normally lowered to 4 LLVM instructions.
-    addCost(3 * InlineConstants::InstrCost);
+    Cost.add(InlineCostFeatures::LoadRelativeIntrinsic,
+             3 * InlineConstants::InstrCost);
   }
   void onLoweredCall(Function *F, CallBase &Call,
                      bool IsIndirectCall) override {
     // We account for the average 1 instruction per call argument setup here.
-    addCost(Call.arg_size() * InlineConstants::InstrCost);
+    Cost.add(InlineCostFeatures::LoweredCallArgSetup,
+             Call.arg_size() * InlineConstants::InstrCost);
 
     // If we have a constant that we are calling as a function, we can peer
     // through it and see the function target. This happens not infrequently
     // during devirtualization and so we want to give it a hefty bonus for
     // inlining, but cap that bonus in the event that inlining wouldn't pan out.
     // Pretend to inline the function, with a custom threshold.
     if (IsIndirectCall && BoostIndirectCalls) {
       auto IndirectCallParams = Params;
       IndirectCallParams.DefaultThreshold =
           InlineConstants::IndirectCallThreshold;
       /// FIXME: if InlineCostCallAnalyzer is derived from, this may need
       /// to instantiate the derived class.
       InlineCostCallAnalyzer CA(*F, Call, IndirectCallParams, TTI,
                                 GetAssumptionCache, GetBFI, PSI, ORE, false);
       if (CA.analyze().isSuccess()) {
         // We were able to inline the indirect call! Subtract the cost from the
         // threshold to get the bonus we want to apply, but don't go below zero.
-        Cost -= std::max(0, CA.getThreshold() - CA.getCost());
+        Cost.accumulateBonus(CA.getCost(), CA.getThreshold());
       }
     } else
       // Otherwise simply add the cost for merely making the call.
-      addCost(InlineConstants::CallPenalty);
+      Cost.add(InlineCostFeatures::IndirectCallPenalty,
+               InlineConstants::CallPenalty);
   }
 
   void onFinalizeSwitch(unsigned JumpTableSize,
                         unsigned NumCaseCluster) override {
     // If suitable for a jump table, consider the cost for the table size and
     // branch to destination.
     // Maximum valid cost increased in this function.
     if (JumpTableSize) {
       int64_t JTCost = (int64_t)JumpTableSize * InlineConstants::InstrCost +
                        4 * InlineConstants::InstrCost;
 
-      addCost(JTCost, (int64_t)CostUpperBound);
+      Cost.add(InlineCostFeatures::JumpTablePenalty, JTCost,
+               (int64_t)CostUpperBound);
       return;
     }
     // Considering forming a binary search, we should find the number of nodes
     // which is same as the number of comparisons when lowered. For a given
     // number of clusters, n, we can define a recursive function, f(n), to find
     // the number of nodes in the tree. The recursion is :
     // f(n) = 1 + f(n/2) + f (n - n/2), when n > 3,
     // and f(n) = n, when n <= 3.
     // This will lead a binary tree where the leaf should be either f(2) or f(3)
     // when n > 3.  So, the number of comparisons from leaves should be n, while
     // the number of non-leaf should be :
     //   2^(log2(n) - 1) - 1
     //   = 2^log2(n) * 2^-1 - 1
     //   = n / 2 - 1.
     // Considering comparisons from leaf and non-leaf nodes, we can estimate the
     // number of comparisons in a simple closed form :
     //   n + n / 2 - 1 = n * 3 / 2 - 1
     if (NumCaseCluster <= 3) {
       // Suppose a comparison includes one compare and one conditional branch.
-      addCost(NumCaseCluster * 2 * InlineConstants::InstrCost);
+      Cost.add(InlineCostFeatures::CaseClusterPenalty,
+               NumCaseCluster * 2 * InlineConstants::InstrCost);
       return;
     }
 
     int64_t ExpectedNumberOfCompare = 3 * (int64_t)NumCaseCluster / 2 - 1;
     int64_t SwitchCost =
         ExpectedNumberOfCompare * 2 * InlineConstants::InstrCost;
 
-    addCost(SwitchCost, (int64_t)CostUpperBound);
+    Cost.add(InlineCostFeatures::SwitchPenalty, SwitchCost,
+             static_cast<int64_t>(CostUpperBound));
   }
   void onMissedSimplification() override {
-    addCost(InlineConstants::InstrCost);
+    Cost.add(InlineCostFeatures::UnsimplifiedCommonInstructions,
+             InlineConstants::InstrCost);
   }
 
   void onInitializeSROAArg(AllocaInst *Arg) override {
     assert(Arg != nullptr &&
            "Should not initialize SROA costs for null value.");
     SROAArgCosts[Arg] = 0;
   }
 
   void onAggregateSROAUse(AllocaInst *SROAArg) override {
     auto CostIt = SROAArgCosts.find(SROAArg);
     assert(CostIt != SROAArgCosts.end() &&
            "expected this argument to have a cost");
     CostIt->second += InlineConstants::InstrCost;
     SROACostSavings += InlineConstants::InstrCost;
+    Cost.add(InlineCostFeatures::SROASavings, InlineConstants::InstrCost);
   }
 
-  void onBlockStart(const BasicBlock *BB) override { CostAtBBStart = Cost; }
+  void onBlockStart(const BasicBlock *BB) override {
+    CostAtBBStart = Cost.reduce();
+  }
 
   void onBlockAnalyzed(const BasicBlock *BB) override {
     if (CostBenefitAnalysisEnabled) {
       // Keep track of the static size of live but cold basic blocks.  For now,
       // we define a cold basic block to be one that's never executed.
       assert(GetBFI && "GetBFI must be available");
       BlockFrequencyInfo *BFI = &(GetBFI(F));
       assert(BFI && "BFI must be available");
       auto ProfileCount = BFI->getBlockProfileCount(BB);
       assert(ProfileCount.hasValue());
       if (ProfileCount.getValue() == 0)
-        ColdSize += Cost - CostAtBBStart;
+        ColdSize += Cost.reduce() - CostAtBBStart;
     }
 
     auto *TI = BB->getTerminator();
     // If we had any successors at this point, than post-inlining is likely to
     // have them as well. Note that we assume any basic blocks which existed
     // due to branches or switches which folded above will also fold after
     // inlining.
     if (SingleBB && TI->getNumSuccessors() > 1) {
       // Take off the bonus we applied to the threshold.
       Threshold -= SingleBBBonus;
       SingleBB = false;
+
+      Cost.add(InlineCostFeatures::IsMultipleBlocks, 1);
     }
   }
 
   void onInstructionAnalysisStart(const Instruction *I) override {
     // This function is called to store the initial cost of inlining before
     // the given instruction was assessed.
     if (!PrintInstructionComments)
       return;
-    InstructionCostDetailMap[I].CostBefore = Cost;
+    InstructionCostDetailMap[I].CostBefore = Cost.reduce();
     InstructionCostDetailMap[I].ThresholdBefore = Threshold;
   }
 
   void onInstructionAnalysisFinish(const Instruction *I) override {
     // This function is called to find new values of cost and threshold after
     // the instruction has been assessed.
     if (!PrintInstructionComments)
       return;
-    InstructionCostDetailMap[I].CostAfter = Cost;
+    InstructionCostDetailMap[I].CostAfter = Cost.reduce();
     InstructionCostDetailMap[I].ThresholdAfter = Threshold;
   }
 
   bool isCostBenefitAnalysisEnabled() {
     if (!PSI || !PSI->hasProfileSummary())
       return false;
 
     if (!GetBFI)
Optional<bool> costBenefitAnalysis() {
 
     // Compute the total savings for the call site.
     auto *CallerBB = CandidateCall.getParent();
     BlockFrequencyInfo *CallerBFI = &(GetBFI(*(CallerBB->getParent())));
     CycleSavings += getCallsiteCost(this->CandidateCall, DL);
     CycleSavings *= CallerBFI->getBlockProfileCount(CallerBB).getValue();
 
     // Remove the cost of the cold basic blocks.
-    int Size = Cost - ColdSize;
+    int Size = Cost.reduce() - ColdSize;
 
     // Allow tiny callees to be inlined regardless of whether they meet the
     // savings threshold.
     Size = Size > InlineSizeAllowance ? Size - InlineSizeAllowance : 1;
 
     // Return true if the savings justify the cost of inlining.  Specifically,
     // we evaluate the following inequality:
     //
if (Caller->hasMinSize()) {
       LoopInfo LI(DT);
       int NumLoops = 0;
       for (Loop *L : LI) {
         // Ignore loops that will not be executed
         if (DeadBlocks.count(L->getHeader()))
           continue;
         NumLoops++;
       }
-      addCost(NumLoops * InlineConstants::CallPenalty);
+      Cost.add(InlineCostFeatures::NumLoops,
+               NumLoops * InlineConstants::CallPenalty);
     }
 
+    Cost.add(InlineCostFeatures::DeadBlocks, DeadBlocks.size());
+    Cost.add(InlineCostFeatures::SimplifiedInstructions,
+             NumInstructionsSimplified);
+    Cost.add(InlineCostFeatures::ConstantArgs, NumConstantArgs);
+    Cost.add(InlineCostFeatures::ConstantOffsetPtrArgs,
+             NumConstantOffsetPtrArgs);
+
     // We applied the maximum possible vector bonus at the beginning. Now,
     // subtract the excess bonus, if any, from the Threshold before
     // comparing against Cost.
-    if (NumVectorInstructions <= NumInstructions / 10)
+    if (NumVectorInstructions <= NumInstructions / 10) {
       Threshold -= VectorBonus;
-    else if (NumVectorInstructions <= NumInstructions / 2)
+    } else if (NumVectorInstructions <= NumInstructions / 2) {
       Threshold -= VectorBonus / 2;
+    }
 
     if (auto Result = costBenefitAnalysis()) {
       DecidedByCostBenefit = true;
       if (Result.getValue())
         return InlineResult::success();
       else
         return InlineResult::failure("Cost over threshold.");
     }
 
-    if (IgnoreThreshold || Cost < std::max(1, Threshold))
+    if (IgnoreThreshold || Cost.reduce() < std::max(1, Threshold))
       return InlineResult::success();
     return InlineResult::failure("Cost over threshold.");
   }
   bool shouldStop() override {
     // Bail out the moment we cross the threshold. This means we'll under-count
     // the cost, but only when undercounting doesn't matter.
-    return !IgnoreThreshold && Cost >= Threshold && !ComputeFullInlineCost;
+    return !IgnoreThreshold && Cost.reduce() >= Threshold &&
+           !ComputeFullInlineCost;
   }
 
   void onLoadEliminationOpportunity() override {
     LoadEliminationCost += InlineConstants::InstrCost;
   }
 
   InlineResult onAnalysisStart() override {
     // Perform some tweaks to the cost and threshold based on the direct
InlineResult onAnalysisStart() override {
 
     // Speculatively apply all possible bonuses to Threshold. If cost exceeds
     // this Threshold any time, and cost cannot decrease, we can stop processing
     // the rest of the function body.
     Threshold += (SingleBBBonus + VectorBonus);
 
     // Give out bonuses for the callsite, as the instructions setting them up
     // will be gone after inlining.
-    addCost(-getCallsiteCost(this->CandidateCall, DL));
+    Cost.add(InlineCostFeatures::CallSiteCost,
+             -1 * getCallsiteCost(this->CandidateCall, DL));
 
     // If this function uses the coldcc calling convention, prefer not to inline
     // it.
     if (F.getCallingConv() == CallingConv::Cold)
-      Cost += InlineConstants::ColdccPenalty;
+      Cost.addUnchecked(InlineCostFeatures::ColdCcPenalty,
+                        InlineConstants::ColdccPenalty);
 
     // Check if we're done. This can happen due to bonuses and penalties.
-    if (Cost >= Threshold && !ComputeFullInlineCost)
+    if (Cost.shouldStop(Threshold) && !ComputeFullInlineCost)
       return InlineResult::failure("high cost");
 
     return InlineResult::success();
   }
 
 public:
   InlineCostCallAnalyzer(
       Function &Callee, CallBase &Call, const InlineParams &Params,
       const TargetTransformInfo &TTI,
       function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
       function_ref<BlockFrequencyInfo &(Function &)> GetBFI = nullptr,
       ProfileSummaryInfo *PSI = nullptr,
       OptimizationRemarkEmitter *ORE = nullptr, bool BoostIndirect = true,
       bool IgnoreThreshold = false)
       : CallAnalyzer(Callee, Call, TTI, GetAssumptionCache, GetBFI, PSI, ORE),
         ComputeFullInlineCost(OptComputeFullInlineCost ||
                               Params.ComputeFullInlineCost || ORE ||
                               isCostBenefitAnalysisEnabled()),
         Params(Params), Threshold(Params.DefaultThreshold),
         BoostIndirectCalls(BoostIndirect), IgnoreThreshold(IgnoreThreshold),
         CostBenefitAnalysisEnabled(isCostBenefitAnalysisEnabled()),
         Writer(this) {}
 
   /// Annotation Writer for instruction details
-  InlineCostAnnotationWriter Writer;
+  InlineCostAnnotationWriter<CostUpdater> Writer;
 
   void dump();
 
   // Prints the same analysis as dump(), but its definition is not dependent
   // on the build.
   void print();
 
   Optional<InstructionCostDetail> getCostDetails(const Instruction *I) {
     if (InstructionCostDetailMap.find(I) != InstructionCostDetailMap.end())
       return InstructionCostDetailMap[I];
     return None;
   }
 
   virtual ~InlineCostCallAnalyzer() {}
-  int getThreshold() { return Threshold; }
-  int getCost() { return Cost; }
-  bool wasDecidedByCostBenefit() { return DecidedByCostBenefit; }
+  bool wasDecidedByCostBenefit() const { return DecidedByCostBenefit; }
+  int getThreshold() const { return Threshold; }
+  typename CostUpdater::CostType getCost() const { return Cost.get(); }
 };
 } // namespace
 
+llvm::raw_ostream &operator<<(llvm::raw_ostream &O,
+                              const InlineCostValueUpdater &Cost) {
+  O << Cost.reduce();
+  return O;
+}
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &O,
+                              const InlineCostFeaturesUpdater &Cost) {
+  O << Cost.reduce();
+  return O;
+}
+
+int llvm::reduceCostFeatures(const InlineCostFeaturesArray &Features) {
+  int Sum = 0;
+  for (size_t I = 0;
+       I < static_cast<size_t>(InlineCostFeatures::NumberOfFeatures); ++I) {
+    if (!isExclusiveMLFeature(static_cast<InlineCostFeatures>(I))) {
+      Sum += Features[I];
+    }
+  }
+  return Sum;
+}
+
 /// Test whether the given value is an Alloca-derived function argument.
 bool CallAnalyzer::isAllocaDerivedArg(Value *V) {
   return SROAArgValues.count(V);
 }
 
 void CallAnalyzer::disableSROAForArg(AllocaInst *SROAArg) {
   onDisableSROA(SROAArg);
   EnabledSROAAllocas.erase(SROAArg);
   disableLoadElimination();
 }
 
-void InlineCostAnnotationWriter::emitInstructionAnnot(const Instruction *I,
-                                                formatted_raw_ostream &OS) {
+template <class CostUpdater>
+void InlineCostAnnotationWriter<CostUpdater>::emitInstructionAnnot(
+    const Instruction *I, formatted_raw_ostream &OS) {
   // The cost of inlining of the given instruction is printed always.
   // The threshold delta is printed only when it is non-zero. It happens
   // when we decided to give a bonus at a particular instruction.
   Optional<InstructionCostDetail> Record = ICCA->getCostDetails(I);
   if (!Record)
     OS << "; No analysis for the instruction";
   else {
     OS << "; cost before = " << Record->CostBefore
        << ", cost after = " << Record->CostAfter
        << ", threshold before = " << Record->ThresholdBefore
        << ", threshold after = " << Record->ThresholdAfter << ", ";
     OS << "cost delta = " << Record->getCostDelta();
     if (Record->hasThresholdChanged())
       OS << ", threshold delta = " << Record->getThresholdDelta();
for (gep_type_iterator GTI = gep_type_begin(GEP), GTE = gep_type_end(GEP);
     }
 
     APInt TypeSize(IntPtrWidth, DL.getTypeAllocSize(GTI.getIndexedType()));
     Offset += OpC->getValue().sextOrTrunc(IntPtrWidth) * TypeSize;
   }
   return true;
 }
 
 /// Use TTI to check whether a GEP is free.
 ///
 /// Respects any simplified values known during the analysis of this callsite.
 bool CallAnalyzer::isGEPFree(GetElementPtrInst &GEP) {
   SmallVector<Value *, 4> Operands;
   Operands.push_back(GEP.getOperand(0));
   for (const Use &Op : GEP.indices())
     if (Constant *SimpleOp = SimplifiedValues.lookup(Op))
       Operands.push_back(SimpleOp);
if (auto *AllocSize = dyn_cast_or_null<ConstantInt>(Size)) {
       // unconditional CFG path. Avoid inlining if this is going to happen above
       // a threshold.
       // FIXME: If the threshold is removed or lowered too much, we could end up
       // being too pessimistic and prevent inlining non-problematic code. This
       // could result in unintended perf regressions. A better overall strategy
       // is needed to track stack usage during inlining.
       Type *Ty = I.getAllocatedType();
       AllocatedSize = SaturatingMultiplyAdd(
-          AllocSize->getLimitedValue(), DL.getTypeAllocSize(Ty).getKnownMinSize(),
-          AllocatedSize);
+          AllocSize->getLimitedValue(),
+          DL.getTypeAllocSize(Ty).getKnownMinSize(), AllocatedSize);
       if (AllocatedSize > InlineConstants::MaxSimplifiedDynamicAllocaToInline)
         HasDynamicAlloca = true;
       return false;
     }
   }
 
   // Accumulate the allocated size.
   if (I.isStaticAlloca()) {
auto IsGEPOffsetConstant = [&](GetElementPtrInst &GEP) {
     for (const Use &Op : GEP.indices())
       if (!isa<Constant>(Op) && !SimplifiedValues.lookup(Op))
         return false;
     return true;
   };
 
   if (!DisableGEPConstOperand)
     if (simplifyInstruction(I, [&](SmallVectorImpl<Constant *> &COps) {
           SmallVector<Constant *, 2> Indices;
-        for (unsigned int Index = 1 ; Index < COps.size() ; ++Index)
+          for (unsigned int Index = 1; Index < COps.size(); ++Index)
             Indices.push_back(COps[Index]);
-        return ConstantExpr::getGetElementPtr(I.getSourceElementType(), COps[0],
-                                              Indices, I.isInBounds());
+          return ConstantExpr::getGetElementPtr(
+              I.getSourceElementType(), COps[0], Indices, I.isInBounds());
         }))
       return true;
 
   if ((I.isInBounds() && canFoldInboundsGEP(I)) || IsGEPOffsetConstant(I)) {
     if (SROAArg)
       SROAArgValues[&I] = SROAArg;
 
     // Constant GEPs are modeled as free.
if (InvokeInst *II = dyn_cast<InvokeInst>(&Call)) {
     if (isa<UnreachableInst>(II->getNormalDest()->getTerminator()))
       return false;
   } else if (isa<UnreachableInst>(Call.getParent()->getTerminator()))
     return false;
 
   return true;
 }
 
-bool InlineCostCallAnalyzer::isColdCallSite(CallBase &Call,
-                                            BlockFrequencyInfo *CallerBFI) {
+template <class CostUpdater>
+bool InlineCostCallAnalyzer<CostUpdater>::isColdCallSite(
+    CallBase &Call, BlockFrequencyInfo *CallerBFI) {
   // If global profile summary is available, then callsite's coldness is
   // determined based on that.
   if (PSI && PSI->hasProfileSummary())
     return PSI->isColdCallSite(Call, CallerBFI);
 
   // Otherwise we need BFI to be available.
   if (!CallerBFI)
     return false;
 
   // Determine if the callsite is cold relative to caller's entry. We could
   // potentially cache the computation of scaled entry frequency, but the added
   // complexity is not worth it unless this scaling shows up high in the
   // profiles.
   const BranchProbability ColdProb(ColdCallSiteRelFreq, 100);
   auto CallSiteBB = Call.getParent();
   auto CallSiteFreq = CallerBFI->getBlockFreq(CallSiteBB);
   auto CallerEntryFreq =
       CallerBFI->getBlockFreq(&(Call.getCaller()->getEntryBlock()));
   return CallSiteFreq < CallerEntryFreq * ColdProb;
 }
 
-Optional<int>
-InlineCostCallAnalyzer::getHotCallSiteThreshold(CallBase &Call,
-                                                BlockFrequencyInfo *CallerBFI) {
+template <class CostUpdater>
+Optional<int> InlineCostCallAnalyzer<CostUpdater>::getHotCallSiteThreshold(
+    CallBase &Call, BlockFrequencyInfo *CallerBFI) {
 
   // If global profile summary is available, then callsite's hotness is
   // determined based on that.
   if (PSI && PSI->hasProfileSummary() && PSI->isHotCallSite(Call, CallerBFI))
     return Params.HotCallSiteThreshold;
 
   // Otherwise we need BFI to be available and to have a locally hot callsite
   // threshold.
Optional<int> InlineCostCallAnalyzer<CostUpdater>::getHotCallSiteThreshold(
   auto CallerEntryFreq = CallerBFI->getEntryFreq();
   if (CallSiteFreq >= CallerEntryFreq * HotCallSiteRelFreq)
     return Params.LocallyHotCallSiteThreshold;
 
   // Otherwise treat it normally.
   return None;
 }
 
-void InlineCostCallAnalyzer::updateThreshold(CallBase &Call, Function &Callee) {
+template <class CostUpdater>
+void InlineCostCallAnalyzer<CostUpdater>::updateThreshold(CallBase &Call,
+                                                          Function &Callee) {
   // If no size growth is allowed for this inlining, set Threshold to 0.
   if (!allowSizeGrowth(Call)) {
     Threshold = 0;
     return;
   }
 
   Function *Caller = Call.getCaller();
 
void InlineCostCallAnalyzer<CostUpdater>::updateThreshold(CallBase &Call,
   VectorBonus = Threshold * VectorBonusPercent / 100;
 
   bool OnlyOneCallAndLocalLinkage =
       F.hasLocalLinkage() && F.hasOneUse() && &F == Call.getCalledFunction();
   // If there is only one call of the function, and it has internal linkage,
   // the cost of inlining it drops dramatically. It may seem odd to update
   // Cost in updateThreshold, but the bonus depends on the logic in this method.
   if (OnlyOneCallAndLocalLinkage)
-    Cost -= LastCallToStaticBonus;
+    Cost.addUnchecked(InlineCostFeatures::LastCallToStaticBonus,
+                      -1 * LastCallToStaticBonus);
 }
 
 bool CallAnalyzer::visitCmpInst(CmpInst &I) {
   Value *LHS = I.getOperand(0), *RHS = I.getOperand(1);
   // First try to handle simplified comparisons.
   if (simplifyInstruction(I, [&](SmallVectorImpl<Constant *> &COps) {
         return ConstantExpr::getCompare(I.getPredicate(), COps[0], COps[1]);
       }))
if (TrueBaseAndOffset == FalseBaseAndOffset && TrueBaseAndOffset.first) {
         SROAArgValues[&SI] = SROAArg;
       return true;
     }
 
     return Base::visitSelectInst(SI);
   }
 
   // Select condition is a constant.
-  Value *SelectedV = CondC->isAllOnesValue()
-                         ? TrueVal
-                         : (CondC->isNullValue()) ? FalseVal : nullptr;
+  Value *SelectedV = CondC->isAllOnesValue()  ? TrueVal
+                     : (CondC->isNullValue()) ? FalseVal
+                                              : nullptr;
   if (!SelectedV) {
     // Condition is a vector constant that is not all 1s or all 0s.  If all
     // operands are constants, ConstantExpr::getSelect() can handle the cases
     // such as select vectors.
     if (TrueC && FalseC) {
       if (auto *C = ConstantExpr::getSelect(CondC, TrueC, FalseC)) {
         SimplifiedValues[&SI] = C;
         return true;
InlineResult CallAnalyzer::analyze() {
   // inlining this would cause the removal of the caller (so the instruction
   // is not actually duplicated, just moved).
   if (!OnlyOneCallAndLocalLinkage && ContainsNoDuplicateCall)
     return InlineResult::failure("noduplicate");
 
   return finalizeAnalysis();
 }
 
-void InlineCostCallAnalyzer::print() {
+template <class CostUpdater> void InlineCostCallAnalyzer<CostUpdater>::print() {
 #define DEBUG_PRINT_STAT(x) dbgs() << "      " #x ": " << x << "\n"
   if (PrintInstructionComments)
     F.print(dbgs(), &Writer);
   DEBUG_PRINT_STAT(NumConstantArgs);
   DEBUG_PRINT_STAT(NumConstantOffsetPtrArgs);
   DEBUG_PRINT_STAT(NumAllocaArgs);
   DEBUG_PRINT_STAT(NumConstantPtrCmps);
   DEBUG_PRINT_STAT(NumConstantPtrDiffs);
   DEBUG_PRINT_STAT(NumInstructionsSimplified);
   DEBUG_PRINT_STAT(NumInstructions);
   DEBUG_PRINT_STAT(SROACostSavings);
   DEBUG_PRINT_STAT(SROACostSavingsLost);
   DEBUG_PRINT_STAT(LoadEliminationCost);
   DEBUG_PRINT_STAT(ContainsNoDuplicateCall);
   DEBUG_PRINT_STAT(Cost);
   DEBUG_PRINT_STAT(Threshold);
 #undef DEBUG_PRINT_STAT
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 /// Dump stats about this call's analysis.
-LLVM_DUMP_METHOD void InlineCostCallAnalyzer::dump() {
+template <class CostUpdater>
+LLVM_DUMP_METHOD void InlineCostCallAnalyzer<CostUpdater>::dump() {
   print();
 }
 #endif
 
 /// Test that there are no attribute conflicts between Caller and Callee
 ///        that prevent inlining.
 static bool functionsHaveCompatibleAttributes(
     Function *Caller, Function *Callee, TargetTransformInfo &TTI,
     function_ref<const TargetLibraryInfo &(Function &)> &GetTLI) {
   // Note that CalleeTLI must be a copy not a reference. The legacy pass manager
InlineCost llvm::getInlineCost(
     function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
     function_ref<const TargetLibraryInfo &(Function &)> GetTLI,
     function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
     ProfileSummaryInfo *PSI, OptimizationRemarkEmitter *ORE) {
   return getInlineCost(Call, Call.getCalledFunction(), Params, CalleeTTI,
                        GetAssumptionCache, GetTLI, GetBFI, PSI, ORE);
 }
 
-Optional<int> llvm::getInliningCostEstimate(
+Optional<InlineCostFeaturesArray> llvm::getInliningCostFeatures(
     CallBase &Call, TargetTransformInfo &CalleeTTI,
     function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
     function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
     ProfileSummaryInfo *PSI, OptimizationRemarkEmitter *ORE) {
   const InlineParams Params = {/* DefaultThreshold*/ 0,
                                /*HintThreshold*/ {},
                                /*ColdThreshold*/ {},
                                /*OptSizeThreshold*/ {},
                                /*OptMinSizeThreshold*/ {},
                                /*HotCallSiteThreshold*/ {},
                                /*LocallyHotCallSiteThreshold*/ {},
                                /*ColdCallSiteThreshold*/ {},
                                /*ComputeFullInlineCost*/ true,
                                /*EnableDeferral*/ true};
 
-  InlineCostCallAnalyzer CA(*Call.getCalledFunction(), Call, Params, CalleeTTI,
-                            GetAssumptionCache, GetBFI, PSI, ORE, true,
-                            /*IgnoreThreshold*/ true);
+  InlineCostCallAnalyzer<InlineCostFeaturesUpdater> CA(
+      *Call.getCalledFunction(), Call, Params, CalleeTTI, GetAssumptionCache,
+      GetBFI, PSI, ORE, true,
   auto R = CA.analyze();
-  if (!R.isSuccess())
-    return None;
-  return CA.getCost();
+  return R.isSuccess() ? CA.getCost() : Optional<InlineCostFeaturesArray>(None);
 }
 
 Optional<InlineResult> llvm::getAttributeBasedInliningDecision(
     CallBase &Call, Function *Callee, TargetTransformInfo &CalleeTTI,
     function_ref<const TargetLibraryInfo &(Function &)> GetTLI) {
 
   // Cannot inline indirect calls.
   if (!Callee)
     return InlineResult::failure("indirect call");
 
   // When callee coroutine function is inlined into caller coroutine function
   // before coro-split pass,
   // coro-early pass can not handle this quiet well.
   // So we won't inline the coroutine function if it have not been unsplited
   if (Callee->isPresplitCoroutine())
     return InlineResult::failure("unsplited coroutine call");
 
   // Never inline calls with byval arguments that does not have the alloca
if (UserDecision->isSuccess())
       return llvm::InlineCost::getAlways("always inline attribute");
     return llvm::InlineCost::getNever(UserDecision->getFailureReason());
   }
 
   LLVM_DEBUG(llvm::dbgs() << "      Analyzing call of " << Callee->getName()
                           << "... (caller:" << Call.getCaller()->getName()
                           << ")\n");
 
-  InlineCostCallAnalyzer CA(*Callee, Call, Params, CalleeTTI,
-                            GetAssumptionCache, GetBFI, PSI, ORE);
+  InlineCostCallAnalyzer<InlineCostValueUpdater> CA(
+      *Callee, Call, Params, CalleeTTI, GetAssumptionCache, GetBFI, PSI, ORE);
   InlineResult ShouldInline = CA.analyze();
 
   LLVM_DEBUG(CA.dump());
 
   // Always make cost benefit based decision explicit.
   // We use always/never here since threshold is not meaningful,
   // as it's not what drives cost-benefit analysis.
   if (CA.wasDecidedByCostBenefit()) {
if (OptLevel > 2)
     Params.LocallyHotCallSiteThreshold = LocallyHotCallSiteThreshold;
   return Params;
 }
 
 PreservedAnalyses
 InlineCostAnnotationPrinterPass::run(Function &F,
                                      FunctionAnalysisManager &FAM) {
   PrintInstructionComments = true;
-  std::function<AssumptionCache &(Function &)> GetAssumptionCache = [&](
-      Function &F) -> AssumptionCache & {
+  std::function<AssumptionCache &(Function &)> GetAssumptionCache =
+      [&](Function &F) -> AssumptionCache & {
     return FAM.getResult<AssumptionAnalysis>(F);
   };
   Module *M = F.getParent();
   ProfileSummaryInfo PSI(*M);
   DataLayout DL(M);
   TargetTransformInfo TTI(DL);
   // FIXME: Redesign the usage of InlineParams to expand the scope of this pass.
   // In the current implementation, the type of InlineParams doesn't matter as
   // the pass serves only for verification of inliner's decisions.
   // We can add a flag which determines InlineParams for this run. Right now,
   // the default InlineParams are used.
   const InlineParams Params = llvm::getInlineParams();
   for (BasicBlock &BB : F) {
     for (Instruction &I : BB) {
       if (CallInst *CI = dyn_cast<CallInst>(&I)) {
         Function *CalledFunction = CI->getCalledFunction();
         if (!CalledFunction || CalledFunction->isDeclaration())
           continue;
         OptimizationRemarkEmitter ORE(CalledFunction);
-        InlineCostCallAnalyzer ICCA(*CalledFunction, *CI, Params, TTI,
-                                    GetAssumptionCache, nullptr, &PSI, &ORE);
+        InlineCostCallAnalyzer<InlineCostValueUpdater> ICCA(
+            *CalledFunction, *CI, Params, TTI, GetAssumptionCache, nullptr,
+            &PSI, &ORE);
         ICCA.analyze();
         OS << "      Analyzing call of " << CalledFunction->getName()
            << "... (caller:" << CI->getCaller()->getName() << ")\n";
         ICCA.print();
       }
     }
   }
   return PreservedAnalyses::all();
 }

llvm/lib/Analysis/MLInlineAdvisor.cpp

 #define DEBUG_TYPE "inline-ml"
 
 static cl::opt<float> SizeIncreaseThreshold(
     "ml-advisor-size-increase-threshold", cl::Hidden,
     cl::desc("Maximum factor by which expected native size may increase before "
              "blocking any further inlining."),
     cl::init(2.0));
 
+// clang-format off
 const std::array<std::string, NumberOfFeatures> llvm::FeatureNameMap{
+// InlineCost features - these must come first
+#define POPULATE_NAMES(INDEX_NAME, NAME) NAME,
+  INLINE_COST_FEATURE_ITERATOR(POPULATE_NAMES)
+#undef POPULATE_NAMES
+
+// Non-cost features
 #define POPULATE_NAMES(INDEX_NAME, NAME, COMMENT) NAME,
-    INLINE_FEATURE_ITERATOR(POPULATE_NAMES)
 #undef POPULATE_NAMES
 };
+// clang-format on
 
 const char *const llvm::DecisionName = "inlining_decision";
 const char *const llvm::DefaultDecisionName = "inlining_default";
 const char *const llvm::RewardName = "delta_size";
 
 CallBase *getInlinableCS(Instruction &I) {
   if (auto *CS = dyn_cast<CallBase>(&I))
     if (Function *Callee = CS->getCalledFunction()) {
std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   if (ForceStop) {
     ORE.emit([&] {
       return OptimizationRemarkMissed(DEBUG_TYPE, "ForceStop", &CB)
              << "Won't attempt inlining because module size grew too much.";
     });
     return std::make_unique<InlineAdvice>(this, CB, ORE, Mandatory);
   }
 
-  int CostEstimate = 0;
+  InlineCostFeaturesArray CostFeatures;
   if (!Mandatory) {
     auto IsCallSiteInlinable =
-        llvm::getInliningCostEstimate(CB, TIR, GetAssumptionCache);
+        llvm::getInliningCostFeatures(CB, TIR, GetAssumptionCache);
     if (!IsCallSiteInlinable) {
       // We can't inline this for correctness reasons, so return the base
       // InlineAdvice, as we don't care about tracking any state changes (which
       // won't happen).
       return std::make_unique<InlineAdvice>(this, CB, ORE, false);
     }
-    CostEstimate = *IsCallSiteInlinable;
+    CostFeatures = std::move(*IsCallSiteInlinable);
   }
 
   if (Mandatory)
     return getMandatoryAdvice(CB, true);
 
   auto NrCtantParams = 0;
   for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I) {
     NrCtantParams += (isa<Constant>(*I));
   }
 
   auto &CallerBefore = FAM.getResult<FunctionPropertiesAnalysis>(Caller);
   auto &CalleeBefore = FAM.getResult<FunctionPropertiesAnalysis>(Callee);
 
   ModelRunner->setFeature(FeatureIndex::CalleeBasicBlockCount,
                           CalleeBefore.BasicBlockCount);
   ModelRunner->setFeature(FeatureIndex::CallSiteHeight,
                           FunctionLevels[&Caller]);
   ModelRunner->setFeature(FeatureIndex::NodeCount, NodeCount);
   ModelRunner->setFeature(FeatureIndex::NrCtantParams, NrCtantParams);
-  ModelRunner->setFeature(FeatureIndex::CostEstimate, CostEstimate);
   ModelRunner->setFeature(FeatureIndex::EdgeCount, EdgeCount);
   ModelRunner->setFeature(FeatureIndex::CallerUsers, CallerBefore.Uses);
   ModelRunner->setFeature(FeatureIndex::CallerConditionallyExecutedBlocks,
                           CallerBefore.BlocksReachedFromConditionalInstruction);
   ModelRunner->setFeature(FeatureIndex::CallerBasicBlockCount,
                           CallerBefore.BasicBlockCount);
   ModelRunner->setFeature(FeatureIndex::CalleeConditionallyExecutedBlocks,
                           CalleeBefore.BlocksReachedFromConditionalInstruction);
   ModelRunner->setFeature(FeatureIndex::CalleeUsers, CalleeBefore.Uses);
+  ModelRunner->setFeature(FeatureIndex::CostEstimate,
+                          reduceCostFeatures(CostFeatures));
+
+  // Add the cost features
+  for (size_t I = 0;
+       I < static_cast<size_t>(InlineCostFeatures::NumberOfFeatures); ++I) {
+    ModelRunner->setFeature(
+        inlineCostFeatureToMlFeature(static_cast<InlineCostFeatures>(I)),
+        CostFeatures[I]);
+  }
+
   return getAdviceFromModel(CB, ORE);
 }
 
 std::unique_ptr<MLInlineAdvice>
 MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
                                     OptimizationRemarkEmitter &ORE) {
   return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run());
 }