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

58

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

465

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

276

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

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

49

Aah! I see why InlineCostFeaturesArray.

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

llvm/lib/Analysis/CMakeLists.txt

8

more clear to say "<TO BE UPDATED>"

llvm/lib/Analysis/InlineCost.cpp

132

remove spurious change

413–417

formatting change - see previous comment

424

formatting change

942

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

943

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

946

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.

1156–1157

formatting

1262–1263

formatting

1417

formatting

2155–2157

formatting

2622

formatting

2716

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

2989–2990

formatting

llvm/lib/Analysis/MLInlineAdvisor.cpp

46

why

260

do we still want this here?

[jacobhegna]

llvm/lib/Analysis/InlineCost.cpp

942

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.

946

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.

2716

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.

260

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

104

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

117

unrelated change?

llvm/lib/Analysis/InlineCost.cpp

441

spurious change?

946

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

949

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

982

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

1047

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

104

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

117

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

982

discussing offline

1047

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_INLINECOST_H
 #define LLVM_ANALYSIS_INLINECOST_H
 
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/Analysis/InlineModelFeatureMaps.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include <cassert>
 #include <climits>
 
 namespace llvm {
 class AssumptionCacheTracker;
 class BlockFrequencyInfo;
 class CallBase;

 /// 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
 
 /// 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
 /// directly tested to determine if inlining should occur given the cost and

 /// - an integer, representing the cost.
 Optional<int> getInliningCostEstimate(
     CallBase &Call, TargetTransformInfo &CalleeTTI,
     function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
     function_ref<BlockFrequencyInfo &(Function &)> GetBFI = nullptr,
     ProfileSummaryInfo *PSI = nullptr,
     OptimizationRemarkEmitter *ORE = nullptr);
 
+/// Get the expanded cost features. The features are returned unconditionally,
+/// even if inlining is impossible.
+Optional<InlineCostFeatures> 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> {

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>
 
 namespace llvm {
 
+// List of cost features. A "cost" feature is a summand of the heuristic-based
+// inline cost, and we define them separately to preserve the original heuristic
+// behavior.
+#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")                                    \
+  M(NestedInlines, "nested_inlines")                                           \
+  M(NestedInlineBonus, "nested_inline_bonus")                                  \
+  M(Threshold, "threshold")
+
+// clang-format off
+enum class InlineCostFeatureIndex : size_t {
+#define POPULATE_INDICES(INDEX_NAME, NAME) INDEX_NAME,
+  INLINE_COST_FEATURE_ITERATOR(POPULATE_INDICES)
+#undef POPULATE_INDICES
+
+  NumberOfFeatures
+};
+// clang-format on
+
+using InlineCostFeatures =
+    std::array<int,
+               static_cast<size_t>(InlineCostFeatureIndex::NumberOfFeatures)>;
+
+constexpr bool isHeuristicInlineCostFeature(InlineCostFeatureIndex Feature) {
+  return Feature != InlineCostFeatureIndex::SROASavings &&
+         Feature != InlineCostFeatureIndex::IsMultipleBlocks &&
+         Feature != InlineCostFeatureIndex::DeadBlocks &&
+         Feature != InlineCostFeatureIndex::SimplifiedInstructions &&
+         Feature != InlineCostFeatureIndex::ConstantArgs &&
+         Feature != InlineCostFeatureIndex::ConstantOffsetPtrArgs &&
+         Feature != InlineCostFeatureIndex::NestedInlines &&
+         Feature != InlineCostFeatureIndex::NestedInlineBonus &&
+         Feature != InlineCostFeatureIndex::Threshold;
+}
+
 // 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(InlineCostFeatureIndex 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/CMakeLists.txt

 if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
   include(TensorFlowCompile)
   set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
 
   # This url points to the most recent most which is known to be compatible with
   # LLVM. When better models are published, this url should be updated to aid
   # discoverability.
-  set(LLVM_INLINER_MODEL_CURRENT_URL "https://github.com/google/ml-compiler-opt/releases/download/inlining-Oz-v0.1/inlining-Oz-acabaf6-v0.1.tar.gz")
+  set(LLVM_INLINER_MODEL_CURRENT_URL "TO_BE_UPDATED")
 
   if (DEFINED LLVM_HAVE_TF_AOT)
     # If the path is empty, autogenerate the model
     if (NOT DEFINED LLVM_INLINER_MODEL_PATH OR "${LLVM_INLINER_MODEL_PATH}" STREQUAL "")
       set(LLVM_INLINER_MODEL_PATH "autogenerate")
       message(WARNING "LLVM_INLINER_MODEL_PATH was not set: autogenerating a model to finish the build.")
     endif()
 

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;
 
 // 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;
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,
                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) {
     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();
 };
-
 /// FIXME: if it is necessary to derive from InlineCostCallAnalyzer, note
 /// the FIXME in onLoweredCall, when instantiating an InlineCostCallAnalyzer
 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.
   int SingleBBBonus = 0;
 
   /// Tunable parameters that control the analysis.
   const InlineParams &Params;
 
   // This DenseMap stores the delta change in cost and threshold after
   // accounting for the given instruction. The map is filled only with the
   // flag PrintInstructionComments on.
   DenseMap<const Instruction *, InstructionCostDetail> InstructionCostDetailMap;
 
   /// Upper bound for the inlining cost. Bonuses are being applied to account
   /// for speculative "expected profit" of the inlining decision.
   int Threshold = 0;
 
   /// Attempt to evaluate indirect calls to boost its inline cost.
   const bool BoostIndirectCalls;
 
   /// Ignore the threshold when finalizing analysis.
   const bool IgnoreThreshold;
 
   // True if the cost-benefit-analysis-based inliner is enabled.
   const bool CostBenefitAnalysisEnabled;
 
Optional<InstructionCostDetail> getCostDetails(const Instruction *I) {
     return None;
   }
 
   virtual ~InlineCostCallAnalyzer() {}
   int getThreshold() { return Threshold; }
   int getCost() { return Cost; }
   bool wasDecidedByCostBenefit() { return DecidedByCostBenefit; }
 };
+
+class InlineCostFeaturesAnalyzer final : public CallAnalyzer {
+private:
+  InlineCostFeatures Cost = {};
+  unsigned SROACostSavingOpportunities = 0;
+  int VectorBonus = 0;
+  int SingleBBBonus = 0;
+  int Threshold = 5;
+
+  DenseMap<AllocaInst *, unsigned> SROACosts;
+
+  void update(InlineCostFeatureIndex Feature, int64_t Delta) {
+    Cost[static_cast<size_t>(Feature)] += Delta;
+  }
+
+  void set(InlineCostFeatureIndex Feature, int64_t Value) {
+    Cost[static_cast<size_t>(Feature)] = Value;
+  }
+
+  void onDisableSROA(AllocaInst *Arg) override {
+    auto CostIt = SROACosts.find(Arg);
+    if (CostIt == SROACosts.end())
+      return;
+
+    update(InlineCostFeatureIndex::SROALosses, CostIt->second);
+    SROACostSavingOpportunities -= CostIt->second;
+    SROACosts.erase(CostIt);
+  }
+
+  void onDisableLoadElimination() override {
+    set(InlineCostFeatureIndex::LoadElimination, 1);
+  }
+
+  void onCallPenalty() override {
+    update(InlineCostFeatureIndex::CallPenalty, InlineConstants::CallPenalty);
+  }
+
+  void onCallArgumentSetup(const CallBase &Call) override {
+    update(InlineCostFeatureIndex::CallArgumentSetup,
+           Call.arg_size() * InlineConstants::InstrCost);
+  }
+
+  void onLoadRelativeIntrinsic() override {
+    update(InlineCostFeatureIndex::LoadRelativeIntrinsic,
+           3 * InlineConstants::InstrCost);
+  }
+
+  void onLoweredCall(Function *F, CallBase &Call,
+                     bool IsIndirectCall) override {
+    update(InlineCostFeatureIndex::LoweredCallArgSetup,
+           Call.arg_size() * InlineConstants::InstrCost);
+
+    if (IsIndirectCall) {
+      InlineParams IndirectCallParams = {/* DefaultThreshold*/ 0,
+                                         /*HintThreshold*/ {},
+                                         /*ColdThreshold*/ {},
+                                         /*OptSizeThreshold*/ {},
+                                         /*OptMinSizeThreshold*/ {},
+                                         /*HotCallSiteThreshold*/ {},
+                                         /*LocallyHotCallSiteThreshold*/ {},
+                                         /*ColdCallSiteThreshold*/ {},
+                                         /*ComputeFullInlineCost*/ true,
+                                         /*EnableDeferral*/ true};
+      IndirectCallParams.DefaultThreshold =
+          InlineConstants::IndirectCallThreshold;
+
+      InlineCostCallAnalyzer CA(*F, Call, IndirectCallParams, TTI,
+                                GetAssumptionCache, GetBFI, PSI, ORE, false,
+                                true);
+      if (CA.analyze().isSuccess()) {
+        update(InlineCostFeatureIndex::NestedInlineBonus, CA.getCost());
+        update(InlineCostFeatureIndex::NestedInlines, 1);
+      }
+    } else {
+      onCallPenalty();
+    }
+  }
+
+  void accumulateBonus(const InlineCostFeatures &OtherCost) {
+    size_t NestedInlineIdx =
+        static_cast<size_t>(InlineCostFeatureIndex::NestedInlines);
+    set(InlineCostFeatureIndex::NestedInlines,
+        std::max(Cost[NestedInlineIdx], OtherCost[NestedInlineIdx] + 1));
+
+    int64_t Bonus = 0;
+    for (int I = 0;
+         I < static_cast<int>(InlineCostFeatureIndex::NumberOfFeatures); ++I) {
+      if (isHeuristicInlineCostFeature(
+              static_cast<InlineCostFeatureIndex>(I)) ||
+          static_cast<InlineCostFeatureIndex>(I) ==
+              InlineCostFeatureIndex::NestedInlineBonus) {
+        Bonus += Cost[I];
+      }
+    }
+    update(InlineCostFeatureIndex::NestedInlineBonus, Bonus);
+  }
+
+  void onFinalizeSwitch(unsigned JumpTableSize,
+                        unsigned NumCaseCluster) override {
+
+    if (JumpTableSize) {
+      int64_t JTCost = (int64_t)JumpTableSize * InlineConstants::InstrCost +
+                       4 * InlineConstants::InstrCost;
+      update(InlineCostFeatureIndex::JumpTablePenalty, JTCost);
+      return;
+    }
+
+    if (NumCaseCluster <= 3) {
+      update(InlineCostFeatureIndex::CaseClusterPenalty,
+             NumCaseCluster * 2 * InlineConstants::InstrCost);
+      return;
+    }
+
+    int64_t ExpectedNumberOfCompare = 3 * (int64_t)NumCaseCluster / 2 - 1;
+    int64_t SwitchCost =
+        ExpectedNumberOfCompare * 2 * InlineConstants::InstrCost;
+    update(InlineCostFeatureIndex::SwitchPenalty, SwitchCost);
+  }
+
+  void onMissedSimplification() override {
+    update(InlineCostFeatureIndex::UnsimplifiedCommonInstructions,
+           InlineConstants::InstrCost);
+  }
+
+  void onInitializeSROAArg(AllocaInst *Arg) override { SROACosts[Arg] = 0; }
+  void onAggregateSROAUse(AllocaInst *Arg) override {
+    SROACosts.find(Arg)->second += InlineConstants::InstrCost;
+    SROACostSavingOpportunities += InlineConstants::InstrCost;
+  }
+
+  void onBlockAnalyzed(const BasicBlock *BB) override {
+    if (BB->getTerminator()->getNumSuccessors() > 1)
+      set(InlineCostFeatureIndex::IsMultipleBlocks, 1);
+    Threshold -= SingleBBBonus;
+  }
+
+  InlineResult finalizeAnalysis() override {
+    auto *Caller = CandidateCall.getFunction();
+    if (Caller->hasMinSize()) {
+      DominatorTree DT(F);
+      LoopInfo LI(DT);
+      for (Loop *L : LI) {
+        // Ignore loops that will not be executed
+        if (DeadBlocks.count(L->getHeader()))
+          continue;
+        update(InlineCostFeatureIndex::NumLoops, InlineConstants::CallPenalty);
+      }
+    }
+    set(InlineCostFeatureIndex::DeadBlocks, DeadBlocks.size());
+    set(InlineCostFeatureIndex::SimplifiedInstructions,
+        NumInstructionsSimplified);
+    set(InlineCostFeatureIndex::ConstantArgs, NumConstantArgs);
+    set(InlineCostFeatureIndex::ConstantOffsetPtrArgs,
+        NumConstantOffsetPtrArgs);
+    set(InlineCostFeatureIndex::SROASavings, SROACostSavingOpportunities);
+
+    if (NumVectorInstructions <= NumInstructions / 10)
+      update(InlineCostFeatureIndex::Threshold, -1 * VectorBonus);
+    else if (NumVectorInstructions <= NumInstructions / 2)
+      update(InlineCostFeatureIndex::Threshold, -1 * (VectorBonus / 2));
+
+    set(InlineCostFeatureIndex::Threshold, Threshold);
+
+    return InlineResult::success();
+  }
+
+  bool shouldStop() override { return false; }
+
+  void onLoadEliminationOpportunity() override {
+    update(InlineCostFeatureIndex::LoadElimination, 1);
+  }
+
+  InlineResult onAnalysisStart() override {
+    update(InlineCostFeatureIndex::CallSiteCost,
+           -1 * getCallsiteCost(this->CandidateCall, DL));
+
+    set(InlineCostFeatureIndex::ColdCcPenalty,
+        (F.getCallingConv() == CallingConv::Cold));
+
+    // FIXME: we shouldn't repeat this logic in both the Features and Cost
+    // analyzer - instead, we should abstract it to a common method in the
+    // CallAnalyzer
+    int SingleBBBonusPercent = 50;
+    int VectorBonusPercent = TTI.getInlinerVectorBonusPercent();
+    Threshold += TTI.adjustInliningThreshold(&CandidateCall);
+    Threshold *= TTI.getInliningThresholdMultiplier();
+    SingleBBBonus = Threshold * SingleBBBonusPercent / 100;
+    VectorBonus = Threshold * VectorBonusPercent / 100;
+    Threshold += (SingleBBBonus + VectorBonus);
+
+    return InlineResult::success();
+  }
+
+public:
+  InlineCostFeaturesAnalyzer(
+      const TargetTransformInfo &TTI,
+      function_ref<AssumptionCache &(Function &)> &GetAssumptionCache,
+      function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
+      ProfileSummaryInfo *PSI, OptimizationRemarkEmitter *ORE, Function &Callee,
+      CallBase &Call)
+      : CallAnalyzer(Callee, Call, TTI, GetAssumptionCache, GetBFI, PSI) {}
+
+  const InlineCostFeatures &features() const { return Cost; }
+};
+
 } // namespace
 
 /// 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) {
   // 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
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);
       if (AllocatedSize > InlineConstants::MaxSimplifiedDynamicAllocaToInline)
         HasDynamicAlloca = true;
       return false;
     }
   }
 
   // Accumulate the allocated size.
   if (I.isStaticAlloca()) {
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)
             Indices.push_back(COps[Index]);
           return ConstantExpr::getGetElementPtr(
               I.getSourceElementType(), COps[0], Indices, I.isInBounds());
         }))
       return true;
 
   if ((I.isInBounds() && canFoldInboundsGEP(I)) || IsGEPOffsetConstant(I)) {
     if (SROAArg)
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;
   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;
#define DEBUG_PRINT_STAT(x) dbgs() << "      " #x ": " << x << "\n"
   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() { 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
InlineCostCallAnalyzer CA(*Call.getCalledFunction(), Call, Params, CalleeTTI,
                             GetAssumptionCache, GetBFI, PSI, ORE, true,
                             /*IgnoreThreshold*/ true);
   auto R = CA.analyze();
   if (!R.isSuccess())
     return None;
   return CA.getCost();
 }
 
+Optional<InlineCostFeatures> llvm::getInliningCostFeatures(
+    CallBase &Call, TargetTransformInfo &CalleeTTI,
+    function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
+    function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
+    ProfileSummaryInfo *PSI, OptimizationRemarkEmitter *ORE) {
+  InlineCostFeaturesAnalyzer CFA(CalleeTTI, GetAssumptionCache, GetBFI, PSI,
+                                 ORE, *Call.getCalledFunction(), Call);
+  auto R = CFA.analyze();
+  if (!R.isSuccess())
+    return None;
+  return CFA.features();
+}
+
 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");
 
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 & {
     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

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()) {
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;
   }
 
+  const auto CostFeatures =
+      llvm::getInliningCostFeatures(CB, TIR, GetAssumptionCache);
+  if (!CostFeatures) {
+    return std::make_unique<InlineAdvice>(this, CB, ORE, false);
+  }
+
   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, CostEstimate);
+
+  // Add the cost features
+  for (size_t I = 0;
+       I < static_cast<size_t>(InlineCostFeatureIndex::NumberOfFeatures); ++I) {
+    ModelRunner->setFeature(
+        inlineCostFeatureToMlFeature(static_cast<InlineCostFeatureIndex>(I)),
+        CostFeatures->at(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());
 }

llvm/lib/Analysis/models/inlining/config.py

 
 
 # pylint: disable=g-complex-comprehension
 def get_input_signature():
   """Returns the list of features for LLVM inlining."""
   # int64 features
   inputs = [
       tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in [
-          'caller_basic_block_count', 'caller_conditionally_executed_blocks',
-          'caller_users', 'callee_basic_block_count',
-          'callee_conditionally_executed_blocks', 'callee_users',
-          'nr_ctant_params', 'node_count', 'edge_count', 'callsite_height',
-          'cost_estimate', 'inlining_default'
+          'caller_basic_block_count',
+          'caller_conditionally_executed_blocks',
+          'caller_users',
+          'callee_basic_block_count',
+          'callee_conditionally_executed_blocks',
+          'callee_users',
+          'nr_ctant_params',
+          'node_count',
+          'edge_count',
+          'callsite_height',
+          'cost_estimate',
+          'inlining_default',
+          'sroa_savings',
+          'sroa_losses',
+          'load_elimination',
+          'call_penalty',
+          'call_argument_setup',
+          'load_relative_intrinsic',
+          'lowered_call_arg_setup',
+          'indirect_call_penalty',
+          'jump_table_penalty',
+          'case_cluster_penalty',
+          'switch_penalty',
+          'unsimplified_common_instructions',
+          'num_loops',
+          'dead_blocks',
+          'simplified_instructions',
+          'constant_args',
+          'constant_offset_ptr_args',
+          'callsite_cost',
+          'cold_cc_penalty',
+          'last_call_to_static_bonus',
+          'is_multiple_blocks',
+          'nested_inlines',
+          'nested_inline_bonus',
+          'threshold',
       ]
   ]
 
   # float32 features
   inputs.extend([
       tf.TensorSpec(dtype=tf.float32, shape=(), name=key)
       for key in ['discount', 'reward']
   ])