[NFC][TFUtils] Extract out the output spec loader

[mtrofin]

Summary

It's generic for the 'development mode', not specific to the inliner
case.

[yundiqian]

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

204–212

Do we still need this function?

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

486

why not passing in OutputSpecs?

[mtrofin]

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

204–212

We can keep it as a lower primitive.

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

486

because those are LoggedFeatureSpecs, so they include the extra name for logging. The evaluator doesn't care about that, and we shouldn't muddy the abstraction layers.

[yundiqian]

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

204–212

Is it a higher primitive? This function calls TFModelEvalutor(..., function_ref<TensorSpec(size_t)> GetOutputSpecs,, ...)

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

486

How about making a copy into std::vector<TensorSpec&> outputspec? It would make the code much easier

[mtrofin]

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

204–212

sorry, I misread what you were referring to - the ctor. Tests use it. You are correct, it's a higher primitive.

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

486

because I'd prefer avoiding copies, if possible.

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

 ///   "port": <int>,
 ///   "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
 ///   "shape": <array of ints> }
 /// For the "type" field, see the C++ primitive types used in
 /// TFUTILS_SUPPORTED_TYPES.
 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
                                            const json::Value &Value);
 
+struct LoggedFeatureSpec {
+  TensorSpec Spec;
+  Optional<std::string> LoggingName;
+};
+
+bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName,
+                     StringRef ExpectedDecisionName,
+                     std::vector<LoggedFeatureSpec> &Ret);
+
 /// Logging utility - given an ordered specification of features, and assuming
 /// a scalar reward, allow logging feature values and rewards, and then print
 /// as tf.train.SequenceExample text protobuf.
 /// The assumption is that, for an event to be logged (i.e. a set of feature
 /// values and a reward), the user calls the log* API for each feature exactly
 /// once, providing the index matching the position in the feature spec list
 /// provided at construction:
 /// event 0:
 ///   logTensorValue(0, ...)
 ///   logTensorValue(1, ...)
 ///   ...
 ///   logReward(...)
 /// event 1:
 ///   logTensorValue(0, ...)
 ///   logTensorValue(1, ...)
 ///   ...
 ///   logReward(...)
 ///
 /// At the end, call print to generate the protobuf.
 class Logger final {
 public:
-  struct LoggedFeatureSpec {
-    TensorSpec Spec;
-    Optional<std::string> LoggingName;
-  };
-
   /// Construct a Logger. If IncludeReward is false, then logReward shouldn't
   /// be called, and the reward feature won't be printed out.
   Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
          const TensorSpec &RewardSpec, bool IncludeReward)
       : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
         RawLogData(FeatureSpecs.size() + IncludeReward),
         IncludeReward(IncludeReward) {}
 
public:
     const void *getUntypedTensorValue(size_t Index) const;
 
   private:
     friend class TFModelEvaluator;
     EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl);
     std::unique_ptr<EvaluationResultImpl> Impl;
   };
 
   TFModelEvaluator(StringRef SavedModelPath,
                    const std::vector<TensorSpec> &InputSpecs,
                    const std::vector<TensorSpec> &OutputSpecs,
                    const char *Tags = "serve");
+  TFModelEvaluator(StringRef SavedModelPath,
+                   const std::vector<TensorSpec> &InputSpecs,
+                   function_ref<TensorSpec(size_t)> GetOutputSpecs,
+                   size_t OutputSpecsSize, const char *Tags = "serve");
+
   ~TFModelEvaluator();
   TFModelEvaluator(const TFModelEvaluator &) = delete;
   TFModelEvaluator(TFModelEvaluator &&) = delete;
 
   /// Evaluate the model, assuming it is valid. Returns None if the evaluation
   /// fails or the model is invalid, or an EvaluationResult otherwise. The
   /// inputs are assumed to have been already provided via getInput(). When
   /// returning None, it also invalidates this object.

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

public:
   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<std::string> &outputNames() const { return OutputNames; }
-
-  const std::vector<TensorSpec> &outputSpecs() const { return OutputSpecs; }
+  const std::vector<LoggedFeatureSpec> &outputLoggedFeatureSpecs() const {
+    return OutputSpecs;
+  }
 
   const Optional<TFModelEvaluator::EvaluationResult> &
   lastEvaluationResult() const {
     return LastEvaluationResult;
   }
 
 private:
   std::unique_ptr<TFModelEvaluator> Evaluator;
-  std::vector<std::string> OutputNames;
-  std::vector<TensorSpec> OutputSpecs;
+  std::vector<LoggedFeatureSpec> OutputSpecs;
   Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
 
-  bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName);
-
   // 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})};
 };
 } // namespace
 
 TrainingLogger::TrainingLogger(StringRef LogFileName,
                                const ModelUnderTrainingRunner *MUTR)
     : LogFileName(LogFileName), MUTR(MUTR) {
   // The first output is the inlining decision.
   if (MUTR)
-    OutputCount = MUTR->outputSpecs().size();
-  std::vector<Logger::LoggedFeatureSpec> FT;
+    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});
-  for (size_t I = 1; I < OutputCount; ++I)
-    FT.push_back({MUTR->outputSpecs()[I], MUTR->outputNames()[I]});
+  if (MUTR && MUTR->outputLoggedFeatureSpecs().size() > 1)
+    FT.insert(FT.end(), MUTR->outputLoggedFeatureSpecs().begin() + 1,
+              MUTR->outputLoggedFeatureSpecs().end());
 
   DefaultDecisionPos = FT.size();
   FT.push_back(
       {TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}), None});
 
   DecisionPos = FT.size();
   FT.push_back({TensorSpec::createSpec<int64_t>(DecisionName, {1}), None});
 
   L = std::make_unique<Logger>(
       FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
       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);
     L->logTensorValue(CurrentFeature, &F);
   }
 
   for (size_t I = 1; I < OutputCount; ++I) {
     const auto &Result = *MUTR->lastEvaluationResult();
-    auto &Spec = MUTR->outputSpecs()[I];
+    auto &Spec = MUTR->outputLoggedFeatureSpecs()[I].Spec;
     const char *RawData =
         reinterpret_cast<const char *>(Result.getUntypedTensorValue(I));
     L->logTensorValue(CurrentFeature, RawData,
                       Spec.getElementCount() * Spec.getElementByteSize());
     ++CurrentFeature;
   }
 
   assert(CurrentFeature == DefaultDecisionPos);
ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
   InputSpecs.insert(InputSpecs.end(), TrainingOnlyFeatures.begin(),
                     TrainingOnlyFeatures.end());
   SmallVector<char, 128> OutputSpecsPath;
   StringRef OutputSpecPath = TFOutputSpecOverride;
   if (OutputSpecPath.empty()) {
     llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
     OutputSpecPath = {OutputSpecsPath.data(), OutputSpecsPath.size()};
   }
-  if (!loadOutputSpecs(Ctx, OutputSpecPath))
+
+  if (!loadOutputSpecs(Ctx, OutputSpecPath, DecisionName, OutputSpecs))
     return;
 
-  Evaluator =
-      std::make_unique<TFModelEvaluator>(ModelPath, InputSpecs, 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;
   }
 }
 
-bool ModelUnderTrainingRunner::loadOutputSpecs(LLVMContext &Ctx,
-                                               StringRef FileName) {
-  auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
-  if (!BufferOrError) {
-    Ctx.emitError("Error opening output specs file: " + FileName + " : " +
-                  BufferOrError.getError().message());
-    return false;
-  }
-  auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
-  if (!ParsedJSONValues) {
-    Ctx.emitError("Could not parse specs file: " + FileName);
-    return false;
-  }
-  auto ValuesArray = ParsedJSONValues->getAsArray();
-  if (!ValuesArray) {
-    Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
-                  "logging_name:<name>} dictionaries");
-    return false;
-  }
-
-  for (const auto &Value : *ValuesArray)
-    if (const auto *Obj = Value.getAsObject())
-      if (const auto *SpecPart = Obj->get("tensor_spec"))
-        if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
-          if (auto LoggingName = Obj->getString("logging_name")) {
-            if (!TensorSpec->isElementType<int64_t>() &&
-                !TensorSpec->isElementType<int32_t>() &&
-                !TensorSpec->isElementType<float>()) {
-              Ctx.emitError(
-                  "Only int64, int32, and float tensors are supported. "
-                  "Found unsupported type for tensor named " +
-                  TensorSpec->name());
-              return false;
-            }
-            OutputNames.push_back(LoggingName->str());
-            OutputSpecs.push_back(*TensorSpec);
-          }
-
-  if (ValuesArray->size() != OutputNames.size()) {
-    Ctx.emitError(
-        "Unable to parse output spec. It should be a json file containing an "
-        "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
-        "with a json object describing a TensorSpec; and a 'logging_name' key, "
-        "which is a string to use as name when logging this tensor in the "
-        "training log.");
-    return false;
-  }
-  assert(OutputNames.size() == OutputSpecs.size());
-  if (OutputNames.empty() || OutputNames[0] != DecisionName) {
-    Ctx.emitError("The first output spec must describe the decision tensor, "
-                  "and must have the logging_name " +
-                  StringRef(DecisionName));
-    return false;
-  }
-  return true;
-}
-
 bool ModelUnderTrainingRunner::run() {
   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);

llvm/lib/Analysis/TFUtils.cpp

 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
 
 #include <cassert>
 #include <numeric>
 

 }
 
 /// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
 /// The tensors are assumed to be stored contiguously, in row-major format,
 /// in the TensorData buffer. Each tensor has the shape given by Spec. The
 /// feature name in the output is either the provided LoggingName, if
 /// specified, otherwise it's the name of the tensor (as given by Spec).
 void writeRawTensorsAsFeatureLists(raw_ostream &OutFile,
-                                   const Logger::LoggedFeatureSpec &LoggedSpec,
+                                   const LoggedFeatureSpec &LoggedSpec,
                                    const char *TensorData, size_t TensorCount,
                                    bool FinalReward = false) {
   const char *FieldName = "<invalid>";
   std::function<void(const char *)> ValueWriter;
   const auto &Spec = LoggedSpec.Spec;
   // The 'Feature' protobuf only has 3 possible fields: float_list,
   // int64_list, or bytes_list, so we capture int32 values as int64. We don't
   // support any other types.

 #define PARSE_TYPE(T, E)                                                       \
   if (TensorType == #T)                                                        \
     return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
   TFUTILS_SUPPORTED_TYPES(PARSE_TYPE)
 #undef PARSE_TYPE
   return None;
 }
 
+bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName,
+                     StringRef ExpectedDecisionName,
+                     std::vector<LoggedFeatureSpec> &Ret) {
+  auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
+  if (!BufferOrError) {
+    Ctx.emitError("Error opening output specs file: " + FileName + " : " +
+                  BufferOrError.getError().message());
+    return false;
+  }
+  auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
+  if (!ParsedJSONValues) {
+    Ctx.emitError("Could not parse specs file: " + FileName);
+    return false;
+  }
+  auto ValuesArray = ParsedJSONValues->getAsArray();
+  if (!ValuesArray) {
+    Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
+                  "logging_name:<name>} dictionaries");
+    return false;
+  }
+
+  for (const auto &Value : *ValuesArray)
+    if (const auto *Obj = Value.getAsObject())
+      if (const auto *SpecPart = Obj->get("tensor_spec"))
+        if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
+          if (auto LoggingName = Obj->getString("logging_name")) {
+            if (!TensorSpec->isElementType<int64_t>() &&
+                !TensorSpec->isElementType<int32_t>() &&
+                !TensorSpec->isElementType<float>()) {
+              Ctx.emitError(
+                  "Only int64, int32, and float tensors are supported. "
+                  "Found unsupported type for tensor named " +
+                  TensorSpec->name());
+              return false;
+            }
+            Ret.push_back({*TensorSpec, LoggingName->str()});
+          }
+
+  if (ValuesArray->size() != Ret.size()) {
+    Ctx.emitError(
+        "Unable to parse output spec. It should be a json file containing an "
+        "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
+        "with a json object describing a TensorSpec; and a 'logging_name' key, "
+        "which is a string to use as name when logging this tensor in the "
+        "training log.");
+    return false;
+  }
+  if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) {
+    Ctx.emitError("The first output spec must describe the decision tensor, "
+                  "and must have the logging_name " +
+                  StringRef(ExpectedDecisionName));
+    return false;
+  }
+  return true;
+}
+
 class TFModelEvaluatorImpl {
 public:
   TFModelEvaluatorImpl(StringRef SavedModelPath,
                        const std::vector<TensorSpec> &InputSpecs,
-                       const std::vector<TensorSpec> &OutputSpecs,
-                       const char *Tags);
+                       function_ref<TensorSpec(size_t)> GetOutputSpecs,
+                       size_t OutputSpecsSize, const char *Tags);
 
   bool isValid() const { return IsValid; }
   size_t OutputSize() const { return OutputFeed.size(); }
 
   void evaluate(TF_Tensor **Output, TF_Status *Status) {
     TF_SessionRun(Session, nullptr, InputFeed.data(), Input.data(),
                   Input.size(), OutputFeed.data(), Output, OutputFeed.size(),
                   nullptr, 0, nullptr, Status);
private:
   /// the SavedModel Graph.
   bool checkReportAndInvalidate(const TF_Output &Output,
                                 const TensorSpec &OutputSpec);
 };
 } // namespace llvm
 
 TFModelEvaluatorImpl::TFModelEvaluatorImpl(
     StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
-    const std::vector<TensorSpec> &OutputSpecs, const char *Tags)
+    function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
+    const char *Tags = "serve")
     : Graph(createTFGraph()), Options(createTFSessionOptions()),
       InputFeed(InputSpecs.size()), Input(InputSpecs.size()),
-      OutputFeed(OutputSpecs.size()) {
+      OutputFeed(OutputSpecsSize) {
   if (!ensureInitTF()) {
     errs() << "Tensorflow should have been initialized";
     return;
   }
   auto Status = createTFStatus();
 
   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();
   }
   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 (!checkReportAndInvalidate(InputFeed[I], InputSpec))
       return;
     initInput(I, static_cast<TF_DataType>(InputSpec.typeIndex()),
               InputSpec.shape());
   }
-  for (size_t I = 0; I < OutputSpecs.size(); ++I) {
-    auto &OutputSpec = OutputSpecs[I];
+  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;
   }
 }
 
-TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
-                                   const std::vector<TensorSpec> &InputSpecs,
-                                   const std::vector<TensorSpec> &OutputSpecs,
-                                   const char *Tags)
-    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, OutputSpecs,
-                                    Tags)) {
+TFModelEvaluator::TFModelEvaluator(
+    StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
+    function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
+    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, GetOutputSpecs,
+                                    OutputSpecsSize, Tags)) {
   if (!Impl->isValid())
     Impl.reset();
 }
 
+TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
+                                   const std::vector<TensorSpec> &InputSpecs,
+                                   const std::vector<TensorSpec> &OutputSpecs,
+                                   const char *Tags)
+    : TFModelEvaluator(
+          SavedModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I]; },
+          OutputSpecs.size(), Tags) {}
+
 TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {
   for (auto *T : Input) {
     TF_DeleteTensor(T);
   }
   if (Session == nullptr)
     return;
   auto Status = createTFStatus();
   TF_DeleteSession(Session, Status.get());

llvm/unittests/Analysis/TFUtilsTest.cpp

TEST(TFUtilsTest, TensorSpecSizesAndTypes) {
   EXPECT_EQ(Spec2D.getElementCount(), 1U);
   EXPECT_EQ(Spec1DLarge.getElementCount(), 10U);
   EXPECT_EQ(Spec3DLarge.getElementCount(), 80U);
   EXPECT_EQ(Spec3DLarge.getElementByteSize(), sizeof(float));
   EXPECT_EQ(Spec1D.getElementByteSize(), sizeof(int16_t));
 }
 
 TEST(TFUtilsTest, Logger) {
-  std::vector<Logger::LoggedFeatureSpec> Features;
+  std::vector<LoggedFeatureSpec> Features;
   Features.push_back(
       {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
   Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {2}),
                       std::string("alternate_name")});
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, true);
   float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};

 )";
   std::string Result;
   raw_string_ostream OS(Result);
   L.print(OS);
   EXPECT_EQ(Result, Expected);
 }
 
 TEST(TFUtilsTest, LoggerNoReward) {
-  std::vector<Logger::LoggedFeatureSpec> Features;
+  std::vector<LoggedFeatureSpec> Features;
   Features.push_back(
       {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
   Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {2}),
                       std::string("alternate_name")});
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, false);
   float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};

 )";
   std::string Result;
   raw_string_ostream OS(Result);
   L.print(OS);
   EXPECT_EQ(Result, Expected);
 }
 
 TEST(TFUtilsTest, LoggerFinalReward) {
-  std::vector<Logger::LoggedFeatureSpec> Features;
+  std::vector<LoggedFeatureSpec> Features;
   Features.push_back({TensorSpec::createSpec<float>("the_float", {1}), None});
   Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {1}), None});
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, true);
   for (size_t I = 0; I < 3; ++I) {
     float F = static_cast<float>(I);
     L.logTensorValue(0, &F);