[mlgo] Support exposing more features than those supported by models

[mtrofin]

Summary

This allows the compiler to support more features than those supported by a
model. The only requirement (development mode only) is that the new features must be appended
at the end of the list of features requested from the model. The support
is transparent to compiler code: for unsupported features, we provide a
valid buffer to copy their values; it's just that this buffer is
disconnected from the model, so insofar as the model is concerned (AOT
or development mode), these features don't exist. The buffers are
allocated at setup - meaning, at steady state, there is no extra allocation
(maintaining the current invariant). These buffers has 2 roles: one,
keep the compiler code simple. Second, allow logging their values in
development mode. The latter allows retraining a model supporting the
larger feature set starting from traces produced with the old model.

For release mode (AOT-ed models), this decouples compiler evolution from
model evolution, which we want in scenarios where the toolchain is
frequently rebuilt and redeployed: we can first deploy the new features,
and continue working with the older model, until a new model is made
available, which can then be picked up the next time the compiler is built.

[yundiqian]

what restricts that "the new features must be appended at the end of the list of features requested from the model"?

Wondering whether it is possible to relax this constraint cause it may cause trouble if there are multiple developers for the same pass

[mtrofin]

clarify "new features at the end" is a requirement for dev mode only

[mtrofin]

In D124565#3481794, @yundiqian wrote:

what restricts that "the new features must be appended at the end of the list of features requested from the model"?

It's a requirement for development mode only (updated to clarify). To your question, see TFUtils.cpp, lines 303 and then the added lines 309 and below.

Wondering whether it is possible to relax this constraint cause it may cause trouble if there are multiple developers for the same pass

How?

[yundiqian]

In D124565#3486601, @mtrofin wrote:

In D124565#3481794, @yundiqian wrote:

what restricts that "the new features must be appended at the end of the list of features requested from the model"?

It's a requirement for development mode only (updated to clarify). To your question, see TFUtils.cpp, lines 303 and then the added lines 309 and below.

Wondering whether it is possible to relax this constraint cause it may cause trouble if there are multiple developers for the same pass

How?

As per offline discussion, will keep this constraint for now, and remove this constraint after migrating to tflite.

[vitalybuka]

It breaks several bots
https://lab.llvm.org/buildbot/#/builders/169
https://lab.llvm.org/buildbot/#/builders/85
https://lab.llvm.org/buildbot/#/builders/74
https://lab.llvm.org/buildbot/#/builders/168

[mtrofin]

Hmm... there may be something else here. I did this, by looking at
https://lab.llvm.org/buildbot/#/builders/169/builds/8182/steps/7/logs/stdio

git pull
mkdir build && cd build
cmake -GNinja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF
'-DLLVM_ENABLE_PROJECTS='\''clang;compiler-rt;lld'\'''
-DLLVM_CCACHE_BUILD=ON ../llvmninja
This completed just fine.
Could it be an issue with the caches?

[vitalybuka]

In D124565#3512815, @mtrofin wrote:

Hmm... there may be something else here. I did this, by looking at
https://lab.llvm.org/buildbot/#/builders/169/builds/8182/steps/7/logs/stdio

git pull
mkdir build && cd build
cmake -GNinja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF
'-DLLVM_ENABLE_PROJECTS='\''clang;compiler-rt;lld'\'''
-DLLVM_CCACHE_BUILD=ON ../llvmninja
This completed just fine.
Could it be an issue with the caches?

No, we can reproduce that on clean VM

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 "llvm/Analysis/TensorSpec.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

 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 std::array<TensorSpec, NumberOfFeatures> FeatureMap;
 
 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/include/llvm/Analysis/MLModelRunner.h

 //===- MLModelRunner.h ---- ML model runner interface -----------*- 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_MLMODELRUNNER_H
 #define LLVM_ANALYSIS_MLMODELRUNNER_H
 
+#include "llvm/Analysis/TensorSpec.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 class LLVMContext;
 
 /// MLModelRunner interface: abstraction of a mechanism for evaluating a
 /// tensorflow "saved model".
 /// NOTE: feature indices are expected to be consistent all accross
return reinterpret_cast<T *>(
         getTensorUntyped(static_cast<size_t>(FeatureID)));
   }
 
   template <typename T, typename I> const T *getTensor(I FeatureID) const {
     return reinterpret_cast<const T *>(
         getTensorUntyped(static_cast<size_t>(FeatureID)));
   }
 
-  virtual void *getTensorUntyped(size_t Index) = 0;
+  void *getTensorUntyped(size_t Index) { return InputBuffers[Index]; }
   const void *getTensorUntyped(size_t Index) const {
     return (const_cast<MLModelRunner *>(this))->getTensorUntyped(Index);
   }
 
   enum class Kind : int { Unknown, Release, Development, NoOp };
   Kind getKind() const { return Type; }
 
 protected:
-  MLModelRunner(LLVMContext &Ctx, Kind Type) : Ctx(Ctx), Type(Type) {
+  MLModelRunner(LLVMContext &Ctx, Kind Type, size_t NrInputs)
+      : Ctx(Ctx), Type(Type), InputBuffers(NrInputs) {
     assert(Type != Kind::Unknown);
   }
   virtual void *evaluateUntyped() = 0;
 
+  void setUpBufferForTensor(size_t Index, const TensorSpec &Spec,
+                            void *Buffer) {
+    if (!Buffer) {
+      OwnedBuffers.emplace_back(Spec.getTotalTensorBufferSize());
+      Buffer = OwnedBuffers.back().data();
+    }
+    InputBuffers[Index] = Buffer;
+  }
+
   LLVMContext &Ctx;
   const Kind Type;
+
+private:
+  std::vector<void *> InputBuffers;
+  std::vector<std::vector<char *>> OwnedBuffers;
 };
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_MLMODELRUNNER_H

llvm/include/llvm/Analysis/ModelUnderTrainingRunner.h

 //===- ModelUnderTrainingRunner.h -- 'development' mode runner --*- 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_MODELUNDERTRAININGRUNNER_H
 #define LLVM_ANALYSIS_MODELUNDERTRAININGRUNNER_H
 
+#include "llvm/Analysis/TensorSpec.h"
 #include "llvm/Config/llvm-config.h"
 
 #ifdef LLVM_HAVE_TF_API
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PassManager.h"
 
static bool classof(const MLModelRunner *R) {
     return R->getKind() == MLModelRunner::Kind::Development;
   }
 
   static std::unique_ptr<ModelUnderTrainingRunner>
   createAndEnsureValid(LLVMContext &Ctx, const std::string &ModelPath,
                        StringRef DecisionName,
                        const std::vector<TensorSpec> &InputSpecs,
                        StringRef OutputSpecsPathOverride = "");
+  static std::unique_ptr<ModelUnderTrainingRunner>
+  createAndEnsureValid(LLVMContext &Ctx, const std::string &ModelPath,
+                       StringRef DecisionName,
+                       const std::vector<TensorSpec> &InputSpecs,
+                       const std::vector<LoggedFeatureSpec> &OutputSpecs);
 
 private:
   ModelUnderTrainingRunner(LLVMContext &Ctx, const std::string &ModelPath,
                            const std::vector<TensorSpec> &InputSpecs,
                            const std::vector<LoggedFeatureSpec> &OutputSpecs);
 
   std::unique_ptr<TFModelEvaluator> Evaluator;
   const std::vector<LoggedFeatureSpec> OutputSpecs;
   Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
   void *evaluateUntyped() override;
-  void *getTensorUntyped(size_t Index) override;
   bool isValid() const { return !!Evaluator; }
 };
 
 } // namespace llvm
 #endif // define(LLVM_HAVE_TF_API)
 #endif // LLVM_ANALYSIS_MODELUNDERTRAININGRUNNER_H

llvm/include/llvm/Analysis/NoInferenceModelRunner.h

 //===- NoInferenceModelRunner.h ---- noop ML model runner  ------*- 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_NOINFERENCEMODELRUNNER_H
 #define LLVM_ANALYSIS_NOINFERENCEMODELRUNNER_H
 
-#include "llvm/Config/llvm-config.h"
-
-/// While not strictly necessary to conditionally compile this, it really
-/// has no usecase outside the 'development' mode.
-#ifdef LLVM_HAVE_TF_API
 #include "llvm/Analysis/MLModelRunner.h"
-#include "llvm/Analysis/Utils/TFUtils.h"
+#include "llvm/Analysis/TensorSpec.h"
+#include "llvm/Config/llvm-config.h"
 namespace llvm {
 /// A pseudo model runner. We use it to store feature values when collecting
 /// logs for the default policy, in 'development' mode, but never ask it to
 /// 'run'.
 class NoInferenceModelRunner : public MLModelRunner {
 public:
   NoInferenceModelRunner(LLVMContext &Ctx,
                          const std::vector<TensorSpec> &Inputs);
 
   static bool classof(const MLModelRunner *R) {
     return R->getKind() == MLModelRunner::Kind::NoOp;
   }
 
 private:
   void *evaluateUntyped() override {
     llvm_unreachable("We shouldn't call run on this model runner.");
   }
-  void *getTensorUntyped(size_t Index) override;
-
-  std::vector<std::unique_ptr<char[]>> ValuesBuffer;
 };
 } // namespace llvm
-#endif // defined(LLVM_HAVE_TF_API)
 #endif // LLVM_ANALYSIS_NOINFERENCEMODELRUNNER_H

llvm/include/llvm/Analysis/ReleaseModeModelRunner.h

 // Only inference is supported.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
 #define LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
 
 #include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/Analysis/TensorSpec.h"
 #include "llvm/Support/ErrorHandling.h"
 
 #include <memory>
 #include <vector>
 
 namespace llvm {
 
 /// ReleaseModeModelRunner - production mode implementation of the
 /// MLModelRunner. It uses an AOT-compiled SavedModel for efficient execution.
 template <class TGen>
 class ReleaseModeModelRunner final : public MLModelRunner {
 public:
   /// FeatureNames' type should be an indexed collection of std::string, like
   /// std::array or std::vector, that has a size() method.
   template <class FType>
-  ReleaseModeModelRunner(LLVMContext &Ctx, const FType &FeatureNames,
+  ReleaseModeModelRunner(LLVMContext &Ctx, const FType &InputSpec,
                          StringRef DecisionName, StringRef FeedPrefix = "feed_",
                          StringRef FetchPrefix = "fetch_")
-      : MLModelRunner(Ctx, MLModelRunner::Kind::Release),
+      : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()),
         CompiledModel(std::make_unique<TGen>()) {
     assert(CompiledModel && "The CompiledModel should be valid");
 
-    const size_t FeatureCount = FeatureNames.size();
-    FeatureIndices.resize(FeatureCount);
-
-    for (size_t I = 0; I < FeatureCount; ++I) {
+    for (size_t I = 0; I < InputSpec.size(); ++I) {
       const int Index =
-          CompiledModel->LookupArgIndex(FeedPrefix.str() + FeatureNames[I]);
-      assert(Index >= 0 && "Cannot find Feature in inlining model");
-      FeatureIndices[I] = Index;
+          CompiledModel->LookupArgIndex(FeedPrefix.str() + InputSpec[I].name());
+      void *Buffer = nullptr;
+      if (Index >= 0)
+        Buffer = CompiledModel->arg_data(Index);
+      setUpBufferForTensor(I, InputSpec[I], Buffer);
     }
 
     ResultIndex = CompiledModel->LookupResultIndex(FetchPrefix.str() +
                                                    DecisionName.str());
     assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
   }
 
   virtual ~ReleaseModeModelRunner() = default;
 
   static bool classof(const MLModelRunner *R) {
     return R->getKind() == MLModelRunner::Kind::Release;
   }
 
 private:
   void *evaluateUntyped() override {
     CompiledModel->Run();
     return CompiledModel->result_data(ResultIndex);
   }
 
-  void *getTensorUntyped(size_t Index) override {
-    return reinterpret_cast<char *>(
-        CompiledModel->arg_data(FeatureIndices[Index]));
-  }
-
-  std::vector<int32_t> FeatureIndices;
   int32_t ResultIndex = -1;
   std::unique_ptr<TGen> CompiledModel;
 };
 
 /// A mock class satisfying the interface expected by ReleaseModeModelRunner for
 /// its `TGen` parameter. Useful to avoid conditional compilation complexity, as
 /// a compile-time replacement for a real AOT-ed model.
 class NoopSavedModelImpl final {

llvm/include/llvm/Analysis/TensorSpec.h

public:
   }
 
   bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }
 
   /// Get the number of elements in a tensor with this shape.
   size_t getElementCount() const { return ElementCount; }
   /// Get the size, in bytes, of one element.
   size_t getElementByteSize() const { return ElementSize; }
+  /// Get the total size of a memory buffer needed to store the whole tensor.
+  size_t getTotalTensorBufferSize() const { return ElementCount * ElementSize; }
 
   template <typename T> bool isElementType() const {
     return getDataType<T>() == Type;
   }
 
 private:
   TensorSpec(const std::string &Name, int Port, TensorType Type,
              size_t ElementSize, const std::vector<int64_t> &Shape);

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

static const std::vector<TensorSpec> TrainingOnlyFeatures{
     TensorSpec::createSpec<int64_t>(TFFeedPrefix + "inlining_default", {1}),
     TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),
     TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),
     TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})};
 
 static const std::vector<TensorSpec> getInputFeatures() {
   std::vector<TensorSpec> InputSpecs;
   for (size_t I = 0; I < NumberOfFeatures; ++I)
-    InputSpecs.push_back(
-        TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
+    InputSpecs.push_back(TensorSpec::createSpec<int64_t>(
+        TFFeedPrefix + FeatureMap[I].name(), FeatureMap[I].shape()));
   append_range(InputSpecs, TrainingOnlyFeatures);
   return InputSpecs;
 }
 
 } // namespace
 
 TrainingLogger::TrainingLogger(StringRef LogFileName,
                                const ModelUnderTrainingRunner *MUTR)
     : LogFileName(LogFileName), MUTR(MUTR) {
   // The first output is the inlining decision.
   if (MUTR)
     OutputCount = MUTR->outputLoggedFeatureSpecs().size();
   std::vector<LoggedFeatureSpec> FT;
 
   for (size_t I = 0; I < NumberOfFeatures; ++I)
-    FT.push_back(
-        {TensorSpec::createSpec<int64_t>(FeatureNameMap.at(I), {1}), None});
+    FT.push_back({FeatureMap.at(I), None});
   if (MUTR && MUTR->outputLoggedFeatureSpecs().size() > 1)
     append_range(FT, drop_begin(MUTR->outputLoggedFeatureSpecs()));
 
   DefaultDecisionPos = FT.size();
   FT.push_back(
       {TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}), None});
 
   DecisionPos = FT.size();

llvm/lib/Analysis/MLInlineAdvisor.cpp

 #include "llvm/Analysis/ReleaseModeModelRunner.h"
 // codegen-ed file
 #include "InlinerSizeModel.h" // NOLINT
 
 std::unique_ptr<InlineAdvisor>
 llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
   auto AOTRunner =
       std::make_unique<ReleaseModeModelRunner<llvm::InlinerSizeModel>>(
-          M.getContext(), FeatureNameMap, DecisionName);
+          M.getContext(), FeatureMap, DecisionName);
   return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
 }
 #endif
 
 #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{
+const std::array<TensorSpec, NumberOfFeatures> llvm::FeatureMap{
+#define POPULATE_NAMES(_, NAME) TensorSpec::createSpec<int64_t>(NAME, {1} ),
 // 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,
+#define POPULATE_NAMES(_, NAME, __) TensorSpec::createSpec<int64_t>(NAME, {1} ),
   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";
MLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) {
   return std::make_unique<MLInlineAdvice>(this, CB, getCallerORE(CB), true);
 }
 
 void MLInlineAdvice::reportContextForRemark(
     DiagnosticInfoOptimizationBase &OR) {
   using namespace ore;
   OR << NV("Callee", Callee->getName());
   for (size_t I = 0; I < NumberOfFeatures; ++I)
-    OR << NV(FeatureNameMap[I],
+    OR << NV(FeatureMap[I].name(),
              *getAdvisor()->getModelRunner().getTensor<int64_t>(I));
   OR << NV("ShouldInline", isInliningRecommended());
 }
 
 void MLInlineAdvice::recordInliningImpl() {
   ORE.emit([&]() {
     OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block);
     reportContextForRemark(R);

llvm/lib/Analysis/ModelUnderTrainingRunner.cpp

 //===- ModelUnderTrainingRunner.cpp - 'development' mode runner -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of a MLModelRunner for 'development' mode, i.e. evaluation
 // happens off a model that's provided from the command line and is interpreted.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Analysis/TensorSpec.h"
 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
 
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 
 using namespace llvm;
 
 ModelUnderTrainingRunner::ModelUnderTrainingRunner(
     LLVMContext &Ctx, const std::string &ModelPath,
     const std::vector<TensorSpec> &InputSpecs,
     const std::vector<LoggedFeatureSpec> &OutputSpecs)
-    : MLModelRunner(Ctx, MLModelRunner::Kind::Development),
+    : MLModelRunner(Ctx, MLModelRunner::Kind::Development, InputSpecs.size()),
       OutputSpecs(OutputSpecs) {
   Evaluator = std::make_unique<TFModelEvaluator>(
       ModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I].Spec; },
       OutputSpecs.size());
   if (!Evaluator || !Evaluator->isValid()) {
     Ctx.emitError("Failed to create saved model evaluator");
     Evaluator.reset();
     return;
   }
+
+  for (size_t I = 0, E = InputSpecs.size(); I < E; ++I) {
+    setUpBufferForTensor(I, InputSpecs[I], Evaluator->getUntypedInput(I));
+  }
 }
 
 void *ModelUnderTrainingRunner::evaluateUntyped() {
   LastEvaluationResult = Evaluator->evaluate();
   if (!LastEvaluationResult.hasValue()) {
     Ctx.emitError("Error evaluating model.");
     return nullptr;
   }
   return LastEvaluationResult->getUntypedTensorValue(0);
 }
 
-void *ModelUnderTrainingRunner::getTensorUntyped(size_t Index) {
-  return Evaluator->getUntypedInput(Index);
-}
-
 std::unique_ptr<ModelUnderTrainingRunner>
 ModelUnderTrainingRunner::createAndEnsureValid(
     LLVMContext &Ctx, const std::string &ModelPath, StringRef DecisionName,
     const std::vector<TensorSpec> &InputSpecs,
     StringRef OutputSpecsPathOverride) {
-  std::unique_ptr<ModelUnderTrainingRunner> MUTR;
   if (auto MaybeOutputSpecs = loadOutputSpecs(Ctx, DecisionName, ModelPath,
                                               OutputSpecsPathOverride))
-    MUTR.reset(new ModelUnderTrainingRunner(Ctx, ModelPath, InputSpecs,
-                                            *MaybeOutputSpecs));
+    return createAndEnsureValid(Ctx, ModelPath, DecisionName, InputSpecs,
+                                *MaybeOutputSpecs);
+  Ctx.emitError("Could not load the policy model from the provided path");
+  return nullptr;
+}
+
+std::unique_ptr<ModelUnderTrainingRunner>
+ModelUnderTrainingRunner::createAndEnsureValid(
+    LLVMContext &Ctx, const std::string &ModelPath, StringRef DecisionName,
+    const std::vector<TensorSpec> &InputSpecs,
+    const std::vector<LoggedFeatureSpec> &OutputSpecs) {
+  std::unique_ptr<ModelUnderTrainingRunner> MUTR;
+  MUTR.reset(
+      new ModelUnderTrainingRunner(Ctx, ModelPath, InputSpecs, OutputSpecs));
   if (MUTR && MUTR->isValid())
     return MUTR;
 
-  Ctx.emitError("Could not load the policy model from the provided path");
+  Ctx.emitError("Could not load or create model evaluator.");
   return nullptr;
 }
 
 #endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/NoInferenceModelRunner.cpp

 //===- NoInferenceModelRunner.cpp - noop ML model runner   ----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // A pseudo model runner. We use it to store feature values when collecting
 // logs for the default policy, in 'development' mode, but never ask it to
 // 'run'.
 //===----------------------------------------------------------------------===//
-#include "llvm/Config/config.h"
-#if defined(LLVM_HAVE_TF_API)
-
 #include "llvm/Analysis/NoInferenceModelRunner.h"
-#include "llvm/Analysis/Utils/TFUtils.h"
 
 using namespace llvm;
 
 NoInferenceModelRunner::NoInferenceModelRunner(
     LLVMContext &Ctx, const std::vector<TensorSpec> &Inputs)
-    : MLModelRunner(Ctx, MLModelRunner::Kind::NoOp) {
-  ValuesBuffer.reserve(Inputs.size());
+    : MLModelRunner(Ctx, MLModelRunner::Kind::NoOp, Inputs.size()) {
+  size_t Index = 0;
   for (const auto &TS : Inputs)
-    ValuesBuffer.push_back(std::make_unique<char[]>(TS.getElementCount() *
-                                                    TS.getElementByteSize()));
-}
-
-void *NoInferenceModelRunner::getTensorUntyped(size_t Index) {
-  return ValuesBuffer[Index].get();
+    setUpBufferForTensor(Index++, TS, nullptr);
 }
-#endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/TFUtils.cpp

TFModelEvaluatorImpl::TFModelEvaluatorImpl(
 
   Session = TF_LoadSessionFromSavedModel(Options.get(), nullptr,
                                          SavedModelPath.str().c_str(), &Tags, 1,
                                          Graph.get(), nullptr, Status.get());
   if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
     errs() << TF_Message(Status.get());
     invalidate();
   }
+  size_t NrSupported = 0;
   for (size_t I = 0; I < InputSpecs.size(); ++I) {
     auto &InputSpec = InputSpecs[I];
     InputFeed[I] = {
         TF_GraphOperationByName(Graph.get(), (InputSpec.name()).c_str()),
         InputSpec.port()};
+    if (!InputFeed[I].oper) {
+      continue;
+    }
+    if (NrSupported++ != I) {
+      errs()
+          << "Unsupported features must be placed at the end of the InputSpecs";
+      invalidate();
+      return;
+    }
     if (!checkReportAndInvalidate(InputFeed[I], InputSpec))
       return;
     initInput(I, static_cast<TF_DataType>(getTFTypeIndex(InputSpec.type())),
               InputSpec.shape());
   }
+  InputFeed.resize(NrSupported);
+  Input.resize(NrSupported);
+
   for (size_t I = 0; I < OutputSpecsSize; ++I) {
     auto OutputSpec = GetOutputSpecs(I);
     OutputFeed[I] = {
         TF_GraphOperationByName(Graph.get(), (OutputSpec.name()).c_str()),
         OutputSpec.port()};
     if (!checkReportAndInvalidate(OutputFeed[I], OutputSpec))
       return;
   }
for (auto &D : Dimensions)
     TotalSize *= D;
 
   Input[Index] =
       TF_AllocateTensor(Type, Dimensions.data(), Dimensions.size(), TotalSize);
   std::memset(TF_TensorData(Input[Index]), 0, TotalSize);
 }
 
 void *TFModelEvaluator::getUntypedInput(size_t Index) {
+  if (Index < Impl->getInput().size())
     return TF_TensorData(Impl->getInput()[Index]);
+  return nullptr;
 }
 
 TFModelEvaluator::EvaluationResult::EvaluationResult(
     std::unique_ptr<EvaluationResultImpl> Impl)
     : Impl(std::move(Impl)) {}
 
 TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
     : Impl(std::move(Other.Impl)) {}

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
 #include "RegAllocEvictionAdvisor.h"
 #include "RegAllocGreedy.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/Analysis/TensorSpec.h"
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TF_API) 
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
 #endif
 #include "llvm/Analysis/ReleaseModeModelRunner.h"
 #include "llvm/CodeGen/CalcSpillWeights.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
private:
   // This could be static and shared, but its initialization is non-trivial.
   std::bitset<FeatureIDs::FeatureCount> DoNotNormalize;
   const float InitialQSize;
 
   using RegID = unsigned;
   mutable DenseMap<RegID, LIFeatureComponents> CachedFeatures;
 };
 
+#define _DECL_FEATURES(type, name, shape, _)                                   \
+  TensorSpec::createSpec<type>(#name, shape),
+
+static const std::vector<TensorSpec> InputFeatures{
+    {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)},
+};
+#undef _DECL_FEATURES
 // ===================================
 // Release (AOT) - specifics
 // ===================================
-const std::array<std::string, FeatureIDs::FeatureCount> FeatureNames{
-#define _GETNAME(_, NAME, __, ___) #NAME,
-    RA_EVICT_FEATURES_LIST(_GETNAME)
-#undef _GETNAME
-};
 class ReleaseModeEvictionAdvisorAnalysis final
     : public RegAllocEvictionAdvisorAnalysis {
 public:
   ReleaseModeEvictionAdvisorAnalysis()
       : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Release) {}
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Release;
   }
 
 private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineLoopInfo>();
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     if (!Runner)
       Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
-          MF.getFunction().getContext(), FeatureNames, DecisionName);
+          MF.getFunction().getContext(), InputFeatures, DecisionName);
     return std::make_unique<MLEvictAdvisor>(
         MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
         getAnalysis<MachineLoopInfo>());
   }
   std::unique_ptr<ReleaseModeModelRunner<CompiledModelType>> Runner;
 };
 
 // ===================================
 // Development mode-specifics
 // ===================================
 //
 // Features we log
 #ifdef LLVM_HAVE_TF_API
-#define _DECL_FEATURES(type, name, shape, _)                                   \
-  TensorSpec::createSpec<type>(#name, shape),
-
-static const std::vector<TensorSpec> InputFeatures{
-    {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)},
-};
-#undef _DECL_FEATURES
 static const TensorSpec Output =
     TensorSpec::createSpec<int64_t>(DecisionName, {1});
 static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
 
 // Features we bind on the model. The tensor names have a prefix, and we also
 // need to include some tensors that are expected to be present by the training
 // algo.
 // TODO: can we just get rid of these?

llvm/unittests/Analysis/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   Support
   TransformUtils
   )
 
-set(MLGO_TESTS TFUtilsTest.cpp MLModelRunnerTest.cpp)
+set(MLGO_TESTS TFUtilsTest.cpp)
 if (DEFINED LLVM_HAVE_TF_API)
   LIST(APPEND EXTRA_TESTS ${MLGO_TESTS})
 else()
   LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS})
 endif()
 
 add_llvm_unittest_with_input_files(AnalysisTests
   AliasAnalysisTest.cpp
add_llvm_unittest_with_input_files(AnalysisTests
   IRSimilarityIdentifierTest.cpp
   IVDescriptorsTest.cpp
   LazyCallGraphTest.cpp
   LoadsTest.cpp
   LoopInfoTest.cpp
   LoopNestTest.cpp
   MemoryBuiltinsTest.cpp
   MemorySSATest.cpp
+  MLModelRunnerTest.cpp
   PhiValuesTest.cpp
   ProfileSummaryInfoTest.cpp
   ScalarEvolutionTest.cpp
   VectorFunctionABITest.cpp
   SparsePropagation.cpp
   TargetLibraryInfoTest.cpp
   TensorSpecTest.cpp
   TBAATest.cpp
   UnrollAnalyzerTest.cpp
   ValueLatticeTest.cpp
   ValueTrackingTest.cpp
   VectorUtilsTest.cpp
   ${EXTRA_TESTS}
   )
 
  target_link_libraries(AnalysisTests PRIVATE LLVMTestingSupport)

llvm/unittests/Analysis/MLModelRunnerTest.cpp

 //===- MLModelRunnerTest.cpp - test for MLModelRunner ---------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
+#include "llvm/Analysis/ReleaseModeModelRunner.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
+namespace llvm {
+// This is a mock of the kind of AOT-generated model evaluator. It has 2 tensors
+// of shape {1}, and 'evaluation' adds them.
+// The interface is the one expected by ReleaseModelRunner.
+class MockAOTModel final {
+  int64_t A = 0;
+  int64_t B = 0;
+  int64_t R = 0;
+
+public:
+  MockAOTModel() = default;
+  int LookupArgIndex(const std::string &Name) {
+    if (Name == "prefix_a")
+      return 0;
+    if (Name == "prefix_b")
+      return 1;
+    return -1;
+  }
+  int LookupResultIndex(const std::string &) { return 0; }
+  void Run() { R = A + B; }
+  void *result_data(int RIndex) {
+    if (RIndex == 0)
+      return &R;
+    return nullptr;
+  }
+  void *arg_data(int Index) {
+    switch (Index) {
+    case 0:
+      return &A;
+    case 1:
+      return &B;
+    default:
+      return nullptr;
+    }
+  }
+};
+} // namespace llvm
+
 TEST(NoInferenceModelRunner, AccessTensors) {
   const std::vector<TensorSpec> Inputs{
       TensorSpec::createSpec<int64_t>("F1", {1}),
       TensorSpec::createSpec<int64_t>("F2", {10}),
       TensorSpec::createSpec<float>("F2", {5}),
   };
   LLVMContext Ctx;
   NoInferenceModelRunner NIMR(Ctx, Inputs);
   NIMR.getTensor<int64_t>(0)[0] = 1;
   std::memcpy(NIMR.getTensor<int64_t>(1),
               std::vector<int64_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(),
               10 * sizeof(int64_t));
   std::memcpy(NIMR.getTensor<float>(2),
               std::vector<float>{0.1, 0.2, 0.3, 0.4, 0.5}.data(),
               5 * sizeof(float));
   ASSERT_EQ(NIMR.getTensor<int64_t>(0)[0], 1);
   ASSERT_EQ(NIMR.getTensor<int64_t>(1)[8], 9);
   ASSERT_EQ(NIMR.getTensor<float>(2)[1], 0.2f);
+}
+
+TEST(ReleaseModeRunner, NormalUse) {
+  LLVMContext Ctx;
+  std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+                                 TensorSpec::createSpec<int64_t>("b", {1})};
+  auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
+      Ctx, Inputs, "", "prefix_");
+  *Evaluator->getTensor<int64_t>(0) = 1;
+  *Evaluator->getTensor<int64_t>(1) = 2;
+  EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
+}
+
+TEST(ReleaseModeRunner, ExtraFeatures) {
+  LLVMContext Ctx;
+  std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+                                 TensorSpec::createSpec<int64_t>("b", {1}),
+                                 TensorSpec::createSpec<int64_t>("c", {1})};
+  auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
+      Ctx, Inputs, "", "prefix_");
+  *Evaluator->getTensor<int64_t>(0) = 1;
+  *Evaluator->getTensor<int64_t>(1) = 2;
+  *Evaluator->getTensor<int64_t>(2) = -3;
+  EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
+}
+
+TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {
+  LLVMContext Ctx;
+  std::vector<TensorSpec> Inputs{
+      TensorSpec::createSpec<int64_t>("a", {1}),
+      TensorSpec::createSpec<int64_t>("c", {1}),
+      TensorSpec::createSpec<int64_t>("b", {1}),
+  };
+  auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
+      Ctx, Inputs, "", "prefix_");
+  *Evaluator->getTensor<int64_t>(0) = 1;         // a
+  *Evaluator->getTensor<int64_t>(1) = 2;         // c
+  *Evaluator->getTensor<int64_t>(2) = -3;        // b
+  EXPECT_EQ(Evaluator->evaluate<int64_t>(), -2); // a + b
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
+  EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
 }
 No newline at end of file

llvm/unittests/Analysis/TFUtilsTest.cpp

 //===- TFUtilsTest.cpp - test for TFUtils ---------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "google/protobuf/struct.pb.h"
 #include "tensorflow/core/example/example.pb.h"
 #include "tensorflow/core/example/feature.pb.h"
+#include "llvm/Analysis/ModelUnderTrainingRunner.h"
+#include "llvm/Analysis/TensorSpec.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
TEST(TFUtilsTest, EvalError) {
   for (auto I = 0; I < KnownSize; ++I) {
     V[I] = 1;
   }
   auto ER = Evaluator.evaluate();
   EXPECT_FALSE(ER.hasValue());
   EXPECT_FALSE(Evaluator.isValid());
 }
 
+TEST(TFUtilsTest, UnsupportedFeature) {
+  const static int64_t KnownSize = 214;
+  std::vector<TensorSpec> InputSpecs{
+      TensorSpec::createSpec<int32_t>("serving_default_input_1",
+                                      {1, KnownSize}),
+      TensorSpec::createSpec<float>("this_feature_does_not_exist", {2, 5})};
+
+  LLVMContext Ctx;
+  auto Evaluator = ModelUnderTrainingRunner::createAndEnsureValid(
+      Ctx, getModelPath(), "StatefulPartitionedCall", InputSpecs,
+      {LoggedFeatureSpec{
+          TensorSpec::createSpec<float>("StatefulPartitionedCall", {1}),
+          None}});
+  int32_t *V = Evaluator->getTensor<int32_t>(0);
+  // Fill it up with 1s, we know the output.
+  for (auto I = 0; I < KnownSize; ++I)
+    V[I] = 1;
+
+  float *F = Evaluator->getTensor<float>(1);
+  for (auto I = 0; I < 2 * 5; ++I)
+    F[I] = 3.14 + I;
+  float Ret = Evaluator->evaluate<float>();
+  EXPECT_EQ(static_cast<int64_t>(Ret), 80);
+  // The input vector should be unchanged
+  for (auto I = 0; I < KnownSize; ++I)
+    EXPECT_EQ(V[I], 1);
+  for (auto I = 0; I < 2 * 5; ++I)
+    EXPECT_FLOAT_EQ(F[I], 3.14 + I);
+}
+
 #define PROTO_CHECKER(FNAME, TYPE, INDEX, EXP)                                 \
   do {                                                                         \
     const auto &V = Expected.feature_lists()                                   \
                         .feature_list()                                        \
                         .at(FNAME)                                             \
                         .feature(INDEX)                                        \
                         .TYPE()                                                \
                         .value();                                              \