[NFC] Expose the ReleaseModeModelRunner

[mtrofin]

Summary

The type was pretty much generic, just needed a bit of parameterization.

[yundiqian]

llvm/include/llvm/Analysis/ReleaseModeModelRunner.h

30

what's the concern for just use std::vector?

[mtrofin]

llvm/include/llvm/Analysis/ReleaseModeModelRunner.h

30

Feature lists are compile-time known. We use std::array for example for the list of the inliner, it's overall cheaper. we can't though use it as a parameter here because we don't know the count (plus it's not fixed, it's per instantiation of this template)

we could do

`template <std::size_t N>
ReleaseModeModelRunner(... const std::array<std::string, N>&...), `

but I prefer the current choice, it's more flexible and doesn't make any significant tradeoffs.

llvm/include/llvm/Analysis/MLModelRunner.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".
+/// NOTE: feature indices are expected to be consistent all accross
+/// MLModelRunners (pertaining to the same model), and also Loggers (see
+/// TFUtils.h)
 class MLModelRunner {
 public:
   // Disallows copy and assign.
   MLModelRunner(const MLModelRunner &) = delete;
   MLModelRunner &operator=(const MLModelRunner &) = delete;
   virtual ~MLModelRunner() = default;
 
   template <typename T> T evaluate() {

llvm/include/llvm/Analysis/ReleaseModeModelRunner.h

+//===- ReleaseModeModelRunner.h - Fast, precompiled 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a model runner wrapping an AOT compiled ML model.
+// Only inference is supported.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/Analysis/MLModelRunner.h"
+
+#include <memory>
+#include <vector>
+
+using namespace llvm;
+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,
+                         StringRef DecisionName, StringRef FeedPrefix = "feed_",
+                         StringRef FetchPrefix = "fetch_")
+      : MLModelRunner(Ctx), 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) {
+      const int Index =
+          CompiledModel->LookupArgIndex(FeedPrefix.str() + FeatureNames[I]);
+      assert(Index >= 0 && "Cannot find Feature in inlining model");
+      FeatureIndices[I] = Index;
+    }
+
+    ResultIndex = CompiledModel->LookupResultIndex(FetchPrefix.str() +
+                                                   DecisionName.str());
+    assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
+  }
+
+  virtual ~ReleaseModeModelRunner() = default;
+
+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;
+};
+} // namespace llvm

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

 /// Alternatively, don't call logReward at the end of each event, just
 /// log{Float|Int32|Int64}FinalReward at the end.
 class LoggerDataImpl;
 class Logger final {
 public:
   /// Construct a Logger. If IncludeReward is false, then logReward or
   /// logFinalReward shouldn't be called, and the reward feature won't be
   /// printed out.
+  /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have
+  /// corresponding indices) with any MLModelRunner implementations
+  /// corresponding to the model being trained/logged.
   Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
          const TensorSpec &RewardSpec, bool IncludeReward);
 
   ~Logger();
 
   void logFloatReward(float Value);
   void logInt32Reward(int32_t Value);
   void logInt64Reward(int64_t Value);

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   PHITransAddr.cpp
   PhiValues.cpp
   PostDominators.cpp
   ProfileSummaryInfo.cpp
   PtrUseVisitor.cpp
   RegionInfo.cpp
   RegionPass.cpp
   RegionPrinter.cpp
-  ReleaseModeModelRunner.cpp
   ReplayInlineAdvisor.cpp
   ScalarEvolution.cpp
   ScalarEvolutionAliasAnalysis.cpp
   ScalarEvolutionDivision.cpp
   ScalarEvolutionNormalization.cpp
   StackLifetime.cpp
   StackSafetyAnalysis.cpp
   SyncDependenceAnalysis.cpp

llvm/lib/Analysis/MLInlineAdvisor.cpp

 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Path.h"
 
 using namespace llvm;
 
+#ifdef LLVM_HAVE_TF_AOT
+#include "llvm/Analysis/ReleaseModeModelRunner.h"
+// codegen-ed file
+#include "InlinerSizeModel.h" // NOLINT
+#include "llvm/Analysis/InlineModelFeatureMaps.h"
+
+std::unique_ptr<InlineAdvisor>
+llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
+  auto AOTRunner =
+      std::make_unique<ReleaseModeModelRunner<llvm::InlinerSizeModel>>(
+          M.getContext(), FeatureNameMap, 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));
 

llvm/lib/Analysis/ReleaseModeModelRunner.cpp

-//===- ReleaseModeModelRunner.cpp - Fast, precompiled 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
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements a model runner wrapping an AOT compiled ML model.
-// Only inference is supported.
-//
-//===----------------------------------------------------------------------===//
-#include "llvm/Config/config.h"
-#if defined(LLVM_HAVE_TF_AOT)
-
-#include "llvm/Analysis/InlineModelFeatureMaps.h"
-#include "llvm/Analysis/MLInlineAdvisor.h"
-
-// codegen-ed file
-#include "InlinerSizeModel.h" // NOLINT
-
-#include <memory>
-#include <vector>
-
-using namespace llvm;
-namespace {
-
-const char FeedPrefix[] = "feed_";
-const char FetchPrefix[] = "fetch_";
-
-/// 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;
-
-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),
-      CompiledModel(std::make_unique<llvm::InlinerSizeModel>()) {
-  assert(CompiledModel && "The CompiledModel should be valid");
-
-  FeatureIndices.resize(NumberOfFeatures);
-
-  for (size_t I = 0; I < NumberOfFeatures; ++I) {
-    const int Index =
-        CompiledModel->LookupArgIndex(FeedPrefix + FeatureNameMap[I]);
-    assert(Index >= 0 && "Cannot find Feature in inlining model");
-    FeatureIndices[I] = Index;
-  }
-
-  ResultIndex =
-      CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName);
-  assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
-}
-
-void *ReleaseModeModelRunner::getTensorUntyped(size_t Index) {
-  return reinterpret_cast<char *>(
-      CompiledModel->arg_data(FeatureIndices[Index]));
-}
-
-void *ReleaseModeModelRunner::evaluateUntyped() {
-  CompiledModel->Run();
-  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)