Unpack the CostEstimate feature in ML inlining models.

This change yields an additional 2% size reduction on an internal search
binary, and an additional 0.5% size reduction on fuchsia.

Differential Revision: https://reviews.llvm.org/D104751

llvm/include/llvm/Analysis/InlineCost.h

@@ -15,6 +15,7 @@
 
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/Analysis/InlineModelFeatureMaps.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include <cassert>
 #include <climits>
@@ -270,6 +271,15 @@ Optional<int> getInliningCostEstimate(
     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);
 

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

@@ -16,6 +16,61 @@
 
 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(NestedInlineCostEstimate, "nested_inline_cost_estimate")                   \
+  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::NestedInlineCostEstimate &&
+         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
@@ -48,12 +103,26 @@ namespace llvm {
     "number of module-internal users of the callee, +1 if the callee is "      \
     "exposed externally")
 
+// 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
+
+  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);

llvm/lib/Analysis/CMakeLists.txt

@@ -5,7 +5,7 @@ if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
   # 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