[NFC][MLGO] Factor ModelUnderTrainingRunner for reuse

[mtrofin]

Summary

This is so we may reuse it. It was very non-inliner specific already.

[yundiqian]

llvm/lib/Analysis/ModelUnderTrainingRunner.cpp

43

this is inline & regalloc specific, up to you to decide whether to fix it this time (when checking in regalloc code) or next time when we have different output types

[mtrofin]

llvm/lib/Analysis/ModelUnderTrainingRunner.cpp

43

good catch! We can just return the untyped value, we have an API for that. Fixing now.

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/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"
+
+namespace llvm {
+
+/// 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,
+                           const std::vector<TensorSpec> &InputSpecs,
+                           const std::vector<LoggedFeatureSpec> &OutputSpecs);
+
+  // Disallows copy and assign.
+  ModelUnderTrainingRunner(const ModelUnderTrainingRunner &) = delete;
+  ModelUnderTrainingRunner &
+  operator=(const ModelUnderTrainingRunner &) = delete;
+
+  bool isValid() const { return !!Evaluator; }
+
+  const std::vector<LoggedFeatureSpec> &outputLoggedFeatureSpecs() const {
+    return OutputSpecs;
+  }
+
+  const Optional<TFModelEvaluator::EvaluationResult> &
+  lastEvaluationResult() const {
+    return LastEvaluationResult;
+  }
+
+private:
+  std::unique_ptr<TFModelEvaluator> Evaluator;
+  const std::vector<LoggedFeatureSpec> OutputSpecs;
+  Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
+  void *evaluateUntyped() override;
+  void *getTensorUntyped(size_t Index) override;
+};
+
+} // namespace llvm
+#endif // define(LLVM_HAVE_TF_API)
+#endif // LLVM_ANALYSIS_MODELUNDERTRAININGRUNNER_H

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   MLInlineAdvisor.cpp
   MemDepPrinter.cpp
   MemDerefPrinter.cpp
   MemoryBuiltins.cpp
   MemoryDependenceAnalysis.cpp
   MemoryLocation.cpp
   MemorySSA.cpp
   MemorySSAUpdater.cpp
+  ModelUnderTrainingRunner.cpp
   ModuleDebugInfoPrinter.cpp
   ModuleSummaryAnalysis.cpp
   MustExecute.cpp
   NoInferenceModelRunner.cpp
   ObjCARCAliasAnalysis.cpp
   ObjCARCAnalysisUtils.cpp
   ObjCARCInstKind.cpp
   OptimizationRemarkEmitter.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/ModelUnderTrainingRunner.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>
 

 
 /// Collect data we may use for training a model, and write it as a textual
 /// Tensorflow SequenceExample
 /// (https://www.tensorflow.org/api_docs/python/tf/train/SequenceExample)
 /// protobuf (https://developers.google.com/protocol-buffers).
 /// Because this is a protobuf, we cannot just stream the events as they come.
 /// Internally, TrainingLogger stores data in column-major format, because that
 /// lines up with how TF SequenceExample represents it.
-class ModelUnderTrainingRunner;
 class TrainingLogger final {
 public:
   TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR);
 
   /// Log one inlining event.
   void logInlineEvent(const InlineEvent &Event,
                       const MLModelRunner &ModelRunner);
 
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;
 };
 
-/// 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);
-
-  // Disallows copy and assign.
-  ModelUnderTrainingRunner(const ModelUnderTrainingRunner &) = delete;
-  ModelUnderTrainingRunner &
-  operator=(const ModelUnderTrainingRunner &) = delete;
-
-  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> 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}));
-    append_range(InputSpecs, TrainingOnlyFeatures);
-    return InputSpecs;
-  }
+    InputSpecs.push_back(
+        TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
 
-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 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,
                                const ModelUnderTrainingRunner *MUTR)
     : LogFileName(LogFileName), MUTR(MUTR) {
   // The first output is the inlining decision.
   if (MUTR)
     OutputCount = MUTR->outputLoggedFeatureSpecs().size();
if (F.isDeclaration())
       continue;
     if (isFunctionDeleted(&F))
       continue;
     Ret += *getNativeSizeEstimate(F);
   }
   return Ret;
 }
 
-ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
-                                                   const std::string &ModelPath)
-    : MLModelRunner(Ctx) {
-  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;
-  }
-}
-
-void *ModelUnderTrainingRunner::evaluateUntyped() {
-  LastEvaluationResult = Evaluator->evaluate();
-  if (!LastEvaluationResult.hasValue()) {
-    Ctx.emitError("Error evaluating model.");
-    return nullptr;
-  }
-  return LastEvaluationResult->getTensorValue<int64_t>(0);
-}
-
-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, ModelUnderTrainingRunner::getInputFeatures()));
+    Runner.reset(new NoInferenceModelRunner(Ctx, getInputFeatures()));
   else {
-    auto MUTR = std::make_unique<ModelUnderTrainingRunner>(
-        Ctx, TFModelUnderTrainingPath);
+    std::unique_ptr<ModelUnderTrainingRunner> MUTR;
+    if (auto MaybeOutputSpecs = loadOutputSpecs(
+            Ctx, DecisionName, TFModelUnderTrainingPath, TFOutputSpecOverride))
+      MUTR = std::make_unique<ModelUnderTrainingRunner>(
+          Ctx, TFModelUnderTrainingPath, getInputFeatures(), *MaybeOutputSpecs);
     if (!MUTR || !MUTR->isValid()) {
       Ctx.emitError("Could not load the policy model from the provided path");
       return nullptr;
     }
     IsDoingInference = true;
     MUTRPtr = MUTR.get();
     Runner = std::move(MUTR);
   }

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/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), 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 inliner saved model evaluator");
+    Evaluator.reset();
+    return;
+  }
+}
+
+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);
+}
+
+#endif // defined(LLVM_HAVE_TF_API)