[NFC][mlgo] Generalize model runner interface

[mtrofin]

Summary

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.

[yundiqian]

llvm/include/llvm/Analysis/MLModelRunner.h

41–45

do they need to be public?

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

269

will part of it get re-factored out of inline to common lib?

[mtrofin]

llvm/include/llvm/Analysis/MLModelRunner.h

41–45

I think I may need them in a subsequent patch, but I can move them to public then.

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

269

I'd like to, yes. Same with the log population from the evaluator - it's pretty generic.

[mtrofin]

moved 2 APIs to protected

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/InlineModelFeatureMaps.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// MLModelRunner interface: abstraction of a mechanism for evaluating a
 /// tensorflow "saved model".
 class MLModelRunner {
 public:
   // Disallows copy and assign.
   MLModelRunner(const MLModelRunner &) = delete;
   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;
 };
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_MLMODELRUNNER_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"
+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)
 No newline at end of file

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

public:
   template <typename T> T *getInput(size_t Index) {
     return static_cast<T *>(getUntypedInput(Index));
   }
 
   /// Returns true if the tensorflow model was loaded successfully, false
   /// otherwise.
   bool isValid() const { return !!Impl; }
 
-private:
+  /// Untyped access to input.
   void *getUntypedInput(size_t Index);
+
+private:
   std::unique_ptr<TFModelEvaluatorImpl> Impl;
 };
 
 /// List of supported types, as a pair:
 /// - C++ type
 /// - enum name (implementation-specific)
 #define TFUTILS_SUPPORTED_TYPES(M)                                             \
   M(float, TF_FLOAT)                                                           \

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   MemoryBuiltins.cpp
   MemoryDependenceAnalysis.cpp
   MemoryLocation.cpp
   MemorySSA.cpp
   MemorySSAUpdater.cpp
   ModuleDebugInfoPrinter.cpp
   ModuleSummaryAnalysis.cpp
   MustExecute.cpp
+  NoInferenceModelRunner.cpp
   ObjCARCAliasAnalysis.cpp
   ObjCARCAnalysisUtils.cpp
   ObjCARCInstKind.cpp
   OptimizationRemarkEmitter.cpp
   OverflowInstAnalysis.cpp
   PHITransAddr.cpp
   PhiValues.cpp
   PostDominators.cpp

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

 //
 //===----------------------------------------------------------------------===//
 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
 
 #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"
 #include "llvm/Support/ManagedStatic.h"
 
 #include <vector>
 
 using namespace llvm;
private:
   static const int64_t NoReward = 0;
   TrainingLogger &Logger;
   const Optional<size_t> CallerSizeEstimateBefore;
   const Optional<size_t> CalleeSizeEstimateBefore;
   const int64_t DefaultDecision;
   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
 /// sacrificed for ease of use while training.
 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 {
     return OutputSpecs;
   }
 
   const Optional<TFModelEvaluator::EvaluationResult> &
   lastEvaluationResult() const {
     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{
 } // 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();
L = std::make_unique<Logger>(
       InlineSizeEstimatorAnalysis::isEvaluatorRequested());
 }
 
 /// Log one inlining event.
 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);
   }
 
   for (size_t I = 1; I < OutputCount; ++I) {
     const auto &Result = *MUTR->lastEvaluationResult();
     const char *RawData =
         reinterpret_cast<const char *>(Result.getUntypedTensorValue(I));
     L->logSpecifiedTensorValue(CurrentFeature, RawData);

 
 std::unique_ptr<MLInlineAdvice>
 DevelopmentModeMLInlineAdvisor::getAdviceFromModel(
     CallBase &CB, OptimizationRemarkEmitter &ORE) {
   if (IsDoingInference && !isLogging())
     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,
       /*Logger=*/*Logger,
       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
       /*CalleeSizeEstimateBefore=*/
       getNativeSizeEstimate(*CB.getCalledFunction()),
       /*DefaultDecision=*/DefaultAdvice);
for (auto &F : M) {
     Ret += *getNativeSizeEstimate(F);
   }
   return Ret;
 }
 
 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);
   else
     return;
 
   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 inliner saved model evaluator");
     Evaluator.reset();
     return;
   }
 }
 
-bool ModelUnderTrainingRunner::run() {
+void *ModelUnderTrainingRunner::evaluateUntyped() {
   LastEvaluationResult = Evaluator->evaluate();
   if (!LastEvaluationResult.hasValue()) {
     Ctx.emitError("Error evaluating model.");
-    return false;
-  }
-  int64_t Decision = *LastEvaluationResult->getTensorValue<int64_t>(0);
-  return static_cast<bool>(Decision);
+    return nullptr;
   }
-
-int64_t ModelUnderTrainingRunner::getFeature(int Index) const {
-  return *Evaluator->getInput<int64_t>(Index);
+  return LastEvaluationResult->getTensorValue<int64_t>(0);
 }
 
-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(
     Module &M, ModuleAnalysisManager &MAM,
     std::function<bool(CallBase &)> GetDefaultAdvice) {
   auto &Ctx = M.getContext();
   std::unique_ptr<MLModelRunner> Runner;
   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);
     if (!MUTR || !MUTR->isValid()) {
       Ctx.emitError("Could not load the policy model from the provided path");
       return nullptr;
     }
     IsDoingInference = true;

llvm/lib/Analysis/MLInlineAdvisor.cpp

std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   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::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);
 }
 
 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,
                                                                   bool Advice) {
   // Make sure we track inlinings in all cases - mandatory or not.
   if (Advice && !ForceStop)
     return getMandatoryAdviceImpl(CB);
 
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], getAdvisor()->getModelRunner().getFeature(I));
+    OR << NV(FeatureNameMap[I],
+             *getAdvisor()->getModelRunner().getTensor<int64_t>(I));
   OR << NV("ShouldInline", isInliningRecommended());
 }
 
 void MLInlineAdvice::recordInliningImpl() {
   ORE.emit([&]() {
     OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block);
     reportContextForRemark(R);
     return R;

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) {
+  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)
 No newline at end of file

llvm/lib/Analysis/ReleaseModeModelRunner.cpp

 
 /// MLModelRunner - production mode implementation. It uses a AOT-compiled
 /// SavedModel for efficient execution.
 class ReleaseModeModelRunner final : public MLModelRunner {
 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;
 };
 } // namespace
 
 ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
     : MLModelRunner(Ctx),
for (size_t I = 0; I < NumberOfFeatures; ++I) {
     FeatureIndices[I] = Index;
   }
 
   ResultIndex =
       CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName);
   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::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
   auto AOTRunner = std::make_unique<ReleaseModeModelRunner>(M.getContext());
   return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
 }
 #endif // defined(LLVM_HAVE_TF_AOT)

llvm/unittests/Analysis/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   Support
   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
   AliasAnalysisTest.cpp
   AliasSetTrackerTest.cpp
   AssumeBundleQueriesTest.cpp
   BasicAliasAnalysisTest.cpp
   BlockFrequencyInfoTest.cpp

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 "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);
+}
 No newline at end of file