Removing NumberOfFeatures from enum FeatureIndex

Author: svkeerthy

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

@@ -132,6 +132,11 @@ constexpr bool isHeuristicInlineCostFeature(InlineCostFeatureIndex Feature) {
     "not "                                                                     \
     "fully inlined by ElimAvailExtern)")
 
+// Not all features listed in FeatureIndex are used by the ML model.
+// Specifically, callee_embedding and caller_embedding are used only when the
+// usage of IR2Vec embeddings is explicitly enabled. Meaning, the size/number of
+// features is not static. So, we cannot determine number of features based on
+// the number of elements in this enum.
 // clang-format off
 enum class FeatureIndex : size_t {
 #define POPULATE_INDICES(DTYPE, SHAPE, NAME, COMMENT) NAME,
@@ -146,9 +151,7 @@ enum class FeatureIndex : size_t {
 // Dimensions of embeddings are not known in the compile time (until vocab is 
 // read). Hence macros cannot be used here.
   callee_embedding,
-  caller_embedding,
-
-  NumberOfFeatures
+  caller_embedding
 };
 // clang-format on
 
@@ -157,9 +160,6 @@ inlineCostFeatureToMlFeature(InlineCostFeatureIndex Feature) {
   return static_cast<FeatureIndex>(static_cast<size_t>(Feature));
 }
 
-constexpr size_t NumberOfFeatures =
-    static_cast<size_t>(FeatureIndex::NumberOfFeatures);
-
 LLVM_ABI extern std::vector<TensorSpec> FeatureMap;
 
 LLVM_ABI extern const char *const DecisionName;

llvm/lib/Analysis/MLInlineAdvisor.cpp

@@ -473,7 +473,7 @@ std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   }
   // This one would have been set up to be right at the end.
   if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault)
-    *ModelRunner->getTensor<int64_t>(FeatureIndex::NumberOfFeatures) =
+    *ModelRunner->getTensor<int64_t>((int)FeatureMap.size()) =
         GetDefaultAdvice(CB);
   return getAdviceFromModel(CB, ORE);
 }
@@ -552,7 +552,7 @@ void MLInlineAdvice::reportContextForRemark(
     DiagnosticInfoOptimizationBase &OR) {
   using namespace ore;
   OR << NV("Callee", Callee->getName());
-  for (size_t I = 0; I < NumberOfFeatures; ++I)
+  for (size_t I = 0; I < FeatureMap.size(); ++I)
     OR << NV(FeatureMap[I].name(),
              *getAdvisor()->getModelRunner().getTensor<int64_t>(I));
   OR << NV("ShouldInline", isInliningRecommended());