[NFC][mlgo] Generalize model runner interface

This prepares it for the regalloc work. Part of it is making model
evaluation accross 'development' and 'release' scenarios more reusable.
This patch:
- extends support to tensors of any shape (not just scalars, like we had
in the inliner -Oz case). While the tensor shape can be anything, we
assume row-major layout and expose the tensor as a buffer.
- exposes the NoInferenceModelRunner, which we use in the 'development'
mode to keep the evaluation code path consistent and simplify logging,
as we'll want to reuse it in the regalloc case.

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

llvm/include/llvm/Analysis/MLModelRunner.h

@@ -10,7 +10,6 @@
 #ifndef LLVM_ANALYSIS_MLMODELRUNNER_H
 #define LLVM_ANALYSIS_MLMODELRUNNER_H
 
-#include "llvm/Analysis/InlineModelFeatureMaps.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PassManager.h"
 
@@ -25,12 +24,27 @@ public:
   MLModelRunner &operator=(const MLModelRunner &) = delete;
   virtual ~MLModelRunner() = default;
 
-  virtual bool run() = 0;
-  virtual void setFeature(FeatureIndex Index, int64_t Value) = 0;
-  virtual int64_t getFeature(int Index) const = 0;
+  template <typename T> T evaluate() {
+    return *reinterpret_cast<T *>(evaluateUntyped());
+  }
+
+  template <typename T, typename I> T *getTensor(I FeatureID) {
+    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)));
+  }
 
 protected:
   MLModelRunner(LLVMContext &Ctx) : Ctx(Ctx) {}
+  virtual void *evaluateUntyped() = 0;
+  virtual void *getTensorUntyped(size_t Index) = 0;
+  const void *getTensorUntyped(size_t Index) const {
+    return (const_cast<MLModelRunner *>(this))->getTensorUntyped(Index);
+  }
 
   LLVMContext &Ctx;
 };

llvm/include/llvm/Analysis/NoInferenceModelRunner.h

@@ -0,0 +1,39 @@
+//===- 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"
+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);
+
+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 // defined(LLVM_ANALYSIS_NOINFERENCEMODELRUNNER_H)

llvm/include/llvm/Analysis/Utils/TFUtils.h

@@ -246,8 +246,10 @@ public:
   /// otherwise.
   bool isValid() const { return !!Impl; }
 
-private:
+  /// Untyped access to input.
   void *getUntypedInput(size_t Index);
+
+private:
   std::unique_ptr<TFModelEvaluatorImpl> Impl;
 };
 

llvm/lib/Analysis/CMakeLists.txt

@@ -105,6 +105,7 @@ add_llvm_component_library(LLVMAnalysis
   ModuleDebugInfoPrinter.cpp
   ModuleSummaryAnalysis.cpp
   MustExecute.cpp
+  NoInferenceModelRunner.cpp
   ObjCARCAliasAnalysis.cpp
   ObjCARCAnalysisUtils.cpp
   ObjCARCInstKind.cpp

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

@@ -16,6 +16,7 @@
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/InlineSizeEstimatorAnalysis.h"
 #include "llvm/Analysis/MLInlineAdvisor.h"
+#include "llvm/Analysis/NoInferenceModelRunner.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/CommandLine.h"
@@ -261,25 +262,6 @@ private:
   const int64_t Mandatory;
 };
 
-/// A pseudo model runner. We use it to store feature values when collecting
-/// logs for the default policy, but never ask it to 'run'.
-class NoInferenceModelRunner : public MLModelRunner {
-public:
-  NoInferenceModelRunner(LLVMContext &Ctx)
-      : MLModelRunner(Ctx), Features(NumberOfFeatures) {}
-  void setFeature(FeatureIndex Index, int64_t Value) override {
-    Features[static_cast<int>(Index)] = Value;
-  }
-
-  int64_t getFeature(int Index) const override { return Features[Index]; }
-  bool run() override {
-    llvm_unreachable("We shouldn't call run on this model runner.");
-  }
-
-private:
-  InlineFeatures Features;
-};
-
 /// ModelUnderTrainingRunner - training mode implementation. It uses TF C APIs
 /// to dynamically load and evaluate a TF SavedModel
 /// (https://www.tensorflow.org/guide/saved_model). Runtime performance is
@@ -288,15 +270,11 @@ class ModelUnderTrainingRunner final : public MLModelRunner {
 public:
   ModelUnderTrainingRunner(LLVMContext &Ctx, const std::string &ModelPath);
 
-  bool run() override;
-
   // Disallows copy and assign.
   ModelUnderTrainingRunner(const ModelUnderTrainingRunner &) = delete;
   ModelUnderTrainingRunner &
   operator=(const ModelUnderTrainingRunner &) = delete;
 
-  void setFeature(FeatureIndex Index, int64_t Value) override;
-  int64_t getFeature(int Index) const override;
   bool isValid() const { return !!Evaluator; }
 
   const std::vector<LoggedFeatureSpec> &outputLoggedFeatureSpecs() const {
@@ -308,18 +286,31 @@ public:
     return LastEvaluationResult;
   }
 
+  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}));
+    append_range(InputSpecs, TrainingOnlyFeatures);
+    return InputSpecs;
+  }
+
 private:
   std::unique_ptr<TFModelEvaluator> Evaluator;
   std::vector<LoggedFeatureSpec> OutputSpecs;
   Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
+  void *evaluateUntyped() override;
+  void *getTensorUntyped(size_t Index) override;
 
   // The training framework needs some additional features.
-  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})};
+  const static std::vector<TensorSpec> TrainingOnlyFeatures;
 };
+
+const std::vector<TensorSpec> ModelUnderTrainingRunner::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})};
 } // namespace
 
 TrainingLogger::TrainingLogger(StringRef LogFileName,
@@ -353,7 +344,7 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event,
                                     const MLModelRunner &ModelRunner) {
   size_t CurrentFeature = 0;
   for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) {
-    int64_t F = ModelRunner.getFeature(CurrentFeature);
+    int64_t F = *ModelRunner.getTensor<int64_t>(CurrentFeature);
     L->logInt64Value(CurrentFeature, &F);
   }
 
@@ -433,7 +424,9 @@ DevelopmentModeMLInlineAdvisor::getAdviceFromModel(
     return MLInlineAdvisor::getAdviceFromModel(CB, ORE);
 
   bool DefaultAdvice = GetDefaultAdvice(CB);
-  auto Recommendation = IsDoingInference ? ModelRunner->run() : DefaultAdvice;
+  auto Recommendation =
+      IsDoingInference ? static_cast<bool>(ModelRunner->evaluate<int64_t>())
+                       : DefaultAdvice;
   return std::make_unique<LoggingMLInlineAdvice>(
       /*Advisor=*/this,
       /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation,
@@ -461,11 +454,8 @@ size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() {
 ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
                                                    const std::string &ModelPath)
     : MLModelRunner(Ctx) {
-  std::vector<TensorSpec> InputSpecs;
-  for (size_t I = 0; I < NumberOfFeatures; ++I)
-    InputSpecs.push_back(
-        TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
-  append_range(InputSpecs, TrainingOnlyFeatures);
+  std::vector<TensorSpec> InputSpecs =
+      ModelUnderTrainingRunner::getInputFeatures();
   if (auto MaybeOutSpecs =
           loadOutputSpecs(Ctx, DecisionName, ModelPath, TFOutputSpecOverride))
     OutputSpecs = std::move(*MaybeOutSpecs);
@@ -482,23 +472,17 @@ ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
   }
 }
 
-bool ModelUnderTrainingRunner::run() {
+void *ModelUnderTrainingRunner::evaluateUntyped() {
   LastEvaluationResult = Evaluator->evaluate();
   if (!LastEvaluationResult.hasValue()) {
     Ctx.emitError("Error evaluating model.");
-    return false;
+    return nullptr;
   }
-  int64_t Decision = *LastEvaluationResult->getTensorValue<int64_t>(0);
-  return static_cast<bool>(Decision);
+  return LastEvaluationResult->getTensorValue<int64_t>(0);
 }
 
-int64_t ModelUnderTrainingRunner::getFeature(int Index) const {
-  return *Evaluator->getInput<int64_t>(Index);
-}
-
-void ModelUnderTrainingRunner::setFeature(FeatureIndex Index, int64_t Value) {
-  size_t NumericIndex = static_cast<size_t>(Index);
-  *(Evaluator->getInput<int64_t>(NumericIndex)) = Value;
+void *ModelUnderTrainingRunner::getTensorUntyped(size_t Index) {
+  return Evaluator->getUntypedInput(Index);
 }
 
 std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
@@ -509,7 +493,8 @@ std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
   ModelUnderTrainingRunner *MUTRPtr = nullptr;
   bool IsDoingInference = false;
   if (TFModelUnderTrainingPath.empty())
-    Runner.reset(new NoInferenceModelRunner(Ctx));
+    Runner.reset(new NoInferenceModelRunner(
+        Ctx, ModelUnderTrainingRunner::getInputFeatures()));
   else {
     auto MUTR = std::make_unique<ModelUnderTrainingRunner>(
         Ctx, TFModelUnderTrainingPath);

llvm/lib/Analysis/MLInlineAdvisor.cpp

@@ -245,29 +245,32 @@ std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   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::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);
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::CalleeBasicBlockCount) =
+      CalleeBefore.BasicBlockCount;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::CallSiteHeight) =
+      FunctionLevels[&Caller];
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::NodeCount) = NodeCount;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::NrCtantParams) = NrCtantParams;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::EdgeCount) = EdgeCount;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::CallerUsers) =
+      CallerBefore.Uses;
+  *ModelRunner->getTensor<int64_t>(
+      FeatureIndex::CallerConditionallyExecutedBlocks) =
+      CallerBefore.BlocksReachedFromConditionalInstruction;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::CallerBasicBlockCount) =
+      CallerBefore.BasicBlockCount;
+  *ModelRunner->getTensor<int64_t>(
+      FeatureIndex::CalleeConditionallyExecutedBlocks) =
+      CalleeBefore.BlocksReachedFromConditionalInstruction;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::CalleeUsers) =
+      CalleeBefore.Uses;
+  *ModelRunner->getTensor<int64_t>(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));
+    *ModelRunner->getTensor<int64_t>(inlineCostFeatureToMlFeature(
+        static_cast<InlineCostFeatureIndex>(I))) = CostFeatures->at(I);
   }
 
   return getAdviceFromModel(CB, ORE);
@@ -276,7 +279,8 @@ std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
 std::unique_ptr<MLInlineAdvice>
 MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
                                     OptimizationRemarkEmitter &ORE) {
-  return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run());
+  return std::make_unique<MLInlineAdvice>(
+      this, CB, ORE, static_cast<bool>(ModelRunner->evaluate<int64_t>()));
 }
 
 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getMandatoryAdvice(CallBase &CB,
@@ -302,7 +306,8 @@ void MLInlineAdvice::reportContextForRemark(
   using namespace ore;
   OR << NV("Callee", Callee->getName());
   for (size_t I = 0; I < NumberOfFeatures; ++I)
-    OR << NV(FeatureNameMap[I], getAdvisor()->getModelRunner().getFeature(I));
+    OR << NV(FeatureNameMap[I],
+             *getAdvisor()->getModelRunner().getTensor<int64_t>(I));
   OR << NV("ShouldInline", isInliningRecommended());
 }
 

llvm/lib/Analysis/NoInferenceModelRunner.cpp

@@ -0,0 +1,33 @@
+//===- 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) {
+  ValuesBuffer.reserve(Inputs.size());
+  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();
+}
+#endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/ReleaseModeModelRunner.cpp

@@ -35,12 +35,10 @@ public:
   ReleaseModeModelRunner(LLVMContext &Ctx);
   virtual ~ReleaseModeModelRunner() = default;
 
-  bool run() override;
-
-  void setFeature(FeatureIndex Index, int64_t Value) override;
-  int64_t getFeature(int Index) const override;
-
 private:
+  void *evaluateUntyped() override;
+  void *getTensorUntyped(size_t Index) override;
+
   std::vector<int32_t> FeatureIndices;
   int32_t ResultIndex = -1;
   std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
@@ -66,20 +64,14 @@ ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
   assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
 }
 
-int64_t ReleaseModeModelRunner::getFeature(int Index) const {
-  return *static_cast<int64_t *>(
+void *ReleaseModeModelRunner::getTensorUntyped(size_t Index) {
+  return reinterpret_cast<char *>(
       CompiledModel->arg_data(FeatureIndices[Index]));
 }
 
-void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
-  *static_cast<int64_t *>(CompiledModel->arg_data(
-      FeatureIndices[static_cast<size_t>(Index)])) = Value;
-}
-
-bool ReleaseModeModelRunner::run() {
+void *ReleaseModeModelRunner::evaluateUntyped() {
   CompiledModel->Run();
-  return static_cast<bool>(
-      *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
+  return CompiledModel->result_data(ResultIndex);
 }
 
 std::unique_ptr<InlineAdvisor>

llvm/unittests/Analysis/CMakeLists.txt

@@ -6,10 +6,11 @@ set(LLVM_LINK_COMPONENTS
   TransformUtils
   )
 
+set(MLGO_TESTS TFUtilsTest.cpp MLModelRunnerTest.cpp)
 if (DEFINED LLVM_HAVE_TF_API)
-  LIST(APPEND EXTRA_TESTS TFUtilsTest.cpp)
+  LIST(APPEND EXTRA_TESTS ${MLGO_TESTS})
 else()
-  LIST(APPEND LLVM_OPTIONAL_SOURCES TFUtilsTest.cpp)
+  LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS})
 endif()
 
 add_llvm_unittest_with_input_files(AnalysisTests

llvm/unittests/Analysis/MLModelRunnerTest.cpp

@@ -0,0 +1,33 @@
+//===- 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 "gtest/gtest.h"
+
+using 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);
+}