[nfc][mlgo] Separate logger and training-mode model evaluator

[mtrofin]

Summary

This just shuffles implementations and declarations around. Now the
logger and the TF C API-based model evaluator are separate.

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

 ///   setting internal scalars, for all dimensions (tensors are row-major:
 ///   https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/c/c_api.h#L205)
 /// - call evaluate. The input tensors' values are not consumed after this, and
 ///   may still be read.
 /// - use the outputs in the output vector
 class TFModelEvaluatorImpl;
 class EvaluationResultImpl;
 
-/// 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. The example assumes the first feature's element
-/// type is float, the second is int64, and the reward is float:
-///
-/// event 0:
-///   logFloatValue(0, ...)
-///   logInt64Value(1, ...)
-///   ...
-///   logFloatReward(...)
-/// event 1:
-///   logFloatValue(0, ...)
-///   logInt64Value(1, ...)
-///   ...
-///   logFloatReward(...)
-///
-/// At the end, call print to generate the protobuf.
-/// 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);
-
-  void logFloatFinalReward(float Value);
-  void logInt32FinalReward(int32_t Value);
-  void logInt64FinalReward(int64_t Value);
-
-  void logFloatValue(size_t FeatureID, const float *Value);
-  void logInt32Value(size_t FeatureID, const int32_t *Value);
-  void logInt64Value(size_t FeatureID, const int64_t *Value);
-
-  void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);
-
-  // Warning! For int32_t, the return is set up for int64_t, so the caller needs
-  // to piecemeal cast their int32_t values.
-  // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
-  // not supported by the tensorflow::SequenceExample proto. For small values,
-  // we can consider using bytes.
-  char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
-
-  // Flush the content of the log to the stream, clearing the stored data in the
-  // process.
-  void flush(std::string *Str);
-  void flush(raw_ostream &OS);
-
-  // Flush a set of logs that are produced from the same module, e.g.
-  // per-function regalloc traces, as a google::protobuf::Struct message.
-  static void flushLogs(raw_ostream &OS,
-                        const StringMap<std::unique_ptr<Logger>> &Loggers);
-
-private:
-  std::vector<LoggedFeatureSpec> FeatureSpecs;
-  TensorSpec RewardSpec;
-  const bool IncludeReward;
-  std::unique_ptr<LoggerDataImpl> LoggerData;
-};
-
 class TFModelEvaluator final {
 public:
   /// The result of a model evaluation. Handles the lifetime of the output
   /// tensors, which means that their values need to be used before
   /// the EvaluationResult's dtor is called.
   class EvaluationResult {
   public:
     EvaluationResult(const EvaluationResult &) = delete;

llvm/include/llvm/Analysis/Utils/TrainingLogger.h

-//===- TFUtils.h - utilities for tensorflow C API ---------------*- C++ -*-===//
+//===- TrainingLogger.h - mlgo feature/reward logging  ----------*- 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_UTILS_TFUTILS_H
-#define LLVM_ANALYSIS_UTILS_TFUTILS_H
+#ifndef LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H
+#define LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H
 
 #include "llvm/Config/llvm-config.h"
 
 #ifdef LLVM_HAVE_TF_API
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/JSON.h"
 
 #include <memory>
 #include <vector>
 
 namespace llvm {
 
-/// Load a SavedModel, find the given inputs and outputs, and setup storage
-/// for input tensors. The user is responsible for correctly dimensioning the
-/// input tensors and setting their values before calling evaluate().
-/// To initialize:
-/// - construct the object
-/// - initialize the input tensors using initInput. Indices must correspond to
-///   indices in the InputNames used at construction.
-/// To use:
-/// - set input values by using getInput to get each input tensor, and then
-///   setting internal scalars, for all dimensions (tensors are row-major:
-///   https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/c/c_api.h#L205)
-/// - call evaluate. The input tensors' values are not consumed after this, and
-///   may still be read.
-/// - use the outputs in the output vector
-class TFModelEvaluatorImpl;
-class EvaluationResultImpl;
-
 /// 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. The example assumes the first feature's element
 /// type is float, the second is int64, and the reward is float:

 
 private:
   std::vector<LoggedFeatureSpec> FeatureSpecs;
   TensorSpec RewardSpec;
   const bool IncludeReward;
   std::unique_ptr<LoggerDataImpl> LoggerData;
 };
 
-class TFModelEvaluator final {
-public:
-  /// The result of a model evaluation. Handles the lifetime of the output
-  /// tensors, which means that their values need to be used before
-  /// the EvaluationResult's dtor is called.
-  class EvaluationResult {
-  public:
-    EvaluationResult(const EvaluationResult &) = delete;
-    EvaluationResult &operator=(const EvaluationResult &Other) = delete;
-
-    EvaluationResult(EvaluationResult &&Other);
-    EvaluationResult &operator=(EvaluationResult &&Other);
-
-    ~EvaluationResult();
-
-    /// Get a (const) pointer to the first element of the tensor at Index.
-    template <typename T> T *getTensorValue(size_t Index) {
-      return static_cast<T *>(getUntypedTensorValue(Index));
-    }
-
-    template <typename T> const T *getTensorValue(size_t Index) const {
-      return static_cast<T *>(getUntypedTensorValue(Index));
-    }
-
-    /// Get a (const) pointer to the untyped data of the tensor.
-    void *getUntypedTensorValue(size_t Index);
-    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.
-  Optional<EvaluationResult> evaluate();
-
-  /// Provides access to the input vector.
-  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; }
-
-  /// Untyped access to input.
-  void *getUntypedInput(size_t Index);
-
-private:
-  std::unique_ptr<TFModelEvaluatorImpl> Impl;
-};
-
 } // namespace llvm
 
 #endif // LLVM_HAVE_TF_API
-#endif // LLVM_ANALYSIS_UTILS_TFUTILS_H
+#endif // LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   ScalarEvolutionAliasAnalysis.cpp
   ScalarEvolutionDivision.cpp
   ScalarEvolutionNormalization.cpp
   StackLifetime.cpp
   StackSafetyAnalysis.cpp
   SyncDependenceAnalysis.cpp
   SyntheticCountsUtils.cpp
   TFUtils.cpp
-  TensorSpec.cpp
   TargetLibraryInfo.cpp
   TargetTransformInfo.cpp
+  TensorSpec.cpp
   Trace.cpp
+  TrainingLogger.cpp
   TypeBasedAliasAnalysis.cpp
   TypeMetadataUtils.cpp
   ScopedNoAliasAA.cpp
   ValueLattice.cpp
   ValueLatticeUtils.cpp
   ValueTracking.cpp
   VectorUtils.cpp
   VFABIDemangling.cpp

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

 
 #include "llvm/ADT/BitVector.h"
 #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/Analysis/Utils/TrainingLogger.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ManagedStatic.h"
 
 #include <vector>
 
 using namespace llvm;
 

llvm/lib/Analysis/TFUtils.cpp

 #include "llvm/Support/Base64.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
-#include "google/protobuf/struct.pb.h"
-#include "google/protobuf/text_format.h"
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
-#include "tensorflow/core/example/example.pb.h"
 #include <cassert>
 #include <numeric>
 
 using namespace llvm;
 
-using google::protobuf::Message;
-using google::protobuf::TextFormat;
-
-static cl::opt<bool>
-    ProtobufTextMode("tfutils-text-log", cl::init(false), cl::Hidden,
-                     cl::desc("Output textual (human-readable) protobuf."));
-
 namespace {
 
 using TFGraphPtr = std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)>;
 using TFSessionOptionsPtr =
     std::unique_ptr<TF_SessionOptions, decltype(&TF_DeleteSessionOptions)>;
 using TFStatusPtr = std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)>;
 
 struct TFInitializer {

 TFStatusPtr createTFStatus() {
   return TFStatusPtr(TF_NewStatus(), &TF_DeleteStatus);
 }
 
 TFSessionOptionsPtr createTFSessionOptions() {
   return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions);
 }
 
-void serialize(const Message &SE, std::string *OutStr) {
-  if (ProtobufTextMode) {
-    TextFormat::PrintToString(SE, OutStr);
-  } else {
-    *OutStr = SE.SerializeAsString();
-  }
-}
-
 int getTFTypeIndex(TensorType TType) {
   switch (TType) {
   case TensorType::Double:
     return TF_DOUBLE;
   case TensorType::Float:
     return TF_FLOAT;
   case TensorType::Int8:
     return TF_INT8;
private:
   bool IsValid = true;
 
   /// Reusable utility for ensuring we can bind the requested Name to a node in
   /// the SavedModel Graph.
   bool checkReportAndInvalidate(const TF_Output &Output,
                                 const TensorSpec &OutputSpec);
 };
 
-class LoggerDataImpl {
-  const std::vector<LoggedFeatureSpec> LoggedFeatureSpecs;
-  const TensorSpec RewardSpec;
-  const bool IncludeReward;
-
-  std::vector<tensorflow::FeatureList> FeatureLists;
-  tensorflow::FeatureList Reward;
-
-  bool isSelfConsistent(const tensorflow::SequenceExample &SE,
-                        size_t NrRecords) const {
-    bool Ret = true;
-    for (const auto &TSpecs : LoggedFeatureSpecs) {
-      const auto &Name = TSpecs.getLoggingName();
-      const auto &FL = SE.feature_lists().feature_list().at(Name).feature();
-      if (NrRecords != static_cast<size_t>(FL.size())) {
-        dbgs() << "[TF-UTILS]: " << Name << " has missing records. Expected "
-               << NrRecords << " got " << FL.size() << "\n";
-        Ret = false;
-      }
-    }
-    if (IncludeReward && static_cast<size_t>(SE.feature_lists()
-                                                 .feature_list()
-                                                 .at(RewardSpec.name())
-                                                 .feature()
-                                                 .size()) != NrRecords) {
-      dbgs() << "[TF-UTILS]: reward is missing records.\n";
-      Ret = false;
-    }
-    return Ret;
-  }
-
-  void transferLog(tensorflow::SequenceExample &SE) {
-    auto *FL = SE.mutable_feature_lists()->mutable_feature_list();
-    if (IncludeReward)
-      (*FL)[RewardSpec.name()] = std::move(Reward);
-    assert(FeatureLists.size() == LoggedFeatureSpecs.size());
-    for (size_t I = 0; I < FeatureLists.size(); ++I) {
-      const auto &LFS = LoggedFeatureSpecs[I];
-      (*FL)[LFS.getLoggingName()] = std::move(FeatureLists[I]);
-    }
-  }
-
-public:
-  LoggerDataImpl(const std::vector<LoggedFeatureSpec> &LoggedSpecs,
-                 const TensorSpec &RewardSpec, bool IncludeReward)
-      : LoggedFeatureSpecs(LoggedSpecs), RewardSpec(RewardSpec),
-        IncludeReward(IncludeReward), FeatureLists(LoggedFeatureSpecs.size()) {}
-
-  // flush the logged info to a stream and clear the log contents.
-  void flush(std::string *Str) {
-    size_t NrRecords = getNrRecords();
-    (void)NrRecords;
-    tensorflow::SequenceExample SE;
-    transferLog(SE);
-    assert(isSelfConsistent(SE, NrRecords));
-    serialize(SE, Str);
-  }
-
-  char *addNewTensor(size_t FeatureID) {
-    const auto &Spec = LoggedFeatureSpecs[FeatureID].Spec;
-    if (Spec.isElementType<float>()) {
-      auto *RF = FeatureLists[FeatureID]
-                     .add_feature()
-                     ->mutable_float_list()
-                     ->mutable_value();
-      RF->Resize(Spec.getElementCount(), 0.0);
-      return reinterpret_cast<char *>(RF->mutable_data());
-    } else if (Spec.isElementType<int32_t>() || Spec.isElementType<int64_t>()) {
-      auto *RF = FeatureLists[FeatureID]
-                     .add_feature()
-                     ->mutable_int64_list()
-                     ->mutable_value();
-      RF->Resize(Spec.getElementCount(), 0);
-      return reinterpret_cast<char *>(RF->mutable_data());
-    }
-    llvm_unreachable("Unsupported tensor type.");
-  }
-
-  template <typename T> void logReward(T Value) {
-    assert(IncludeReward);
-    if (RewardSpec.isElementType<float>())
-      Reward.add_feature()->mutable_float_list()->add_value(Value);
-    else if (RewardSpec.isElementType<int32_t>() ||
-             RewardSpec.isElementType<int64_t>())
-      Reward.add_feature()->mutable_int64_list()->add_value(Value);
-    else
-      llvm_unreachable("Unsupported tensor type.");
-  }
-
-  size_t getNrRecords() const {
-    return FeatureLists.empty() ? 0 : FeatureLists[0].feature().size();
-  }
-};
 } // namespace llvm
 
 TFModelEvaluatorImpl::TFModelEvaluatorImpl(
     StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
     function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
     const char *Tags = "serve")
     : Graph(createTFGraph()), Options(createTFSessionOptions()),
       InputFeed(InputSpecs.size()), Input(InputSpecs.size()),

 const void *
 TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) const {
   return TF_TensorData(Impl->getOutput()[Index]);
 }
 
 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 
-Logger::Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
-               const TensorSpec &RewardSpec, bool IncludeReward)
-    : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
-      IncludeReward(IncludeReward),
-      LoggerData(std::make_unique<LoggerDataImpl>(FeatureSpecs, RewardSpec,
-                                                  IncludeReward)) {}
-
-Logger::~Logger() {}
-
-#define LOG_REWARD(NAME, TYPE)                                                 \
-  void Logger::log##NAME##Reward(TYPE Value) {                                 \
-    assert(IncludeReward);                                                     \
-    LoggerData->logReward(Value);                                              \
-  }
-
-LOG_REWARD(Float, float)
-LOG_REWARD(Int32, int32_t)
-LOG_REWARD(Int64, int64_t)
-#undef LOG_REWARD
-
-#define LOG_FINAL_REWARD(NAME, TYPE)                                           \
-  void Logger::log##NAME##FinalReward(TYPE Value) {                            \
-    assert(RewardSpec.isElementType<TYPE>());                                  \
-    for (size_t I = 1; I < LoggerData->getNrRecords(); ++I)                    \
-      log##NAME##Reward(0);                                                    \
-    log##NAME##Reward(Value);                                                  \
-  }
-
-LOG_FINAL_REWARD(Float, float)
-LOG_FINAL_REWARD(Int32, int32_t)
-LOG_FINAL_REWARD(Int64, int64_t)
-#undef LOG_FINAL_REWARD
-
-void Logger::logFloatValue(size_t FeatureID, const float *Value) {
-  assert(FeatureSpecs[FeatureID].Spec.isElementType<float>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
-  assert(FeatureSpecs[FeatureID].Spec.isElementType<int64_t>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
-  assert(FeatureSpecs[FeatureID].Spec.isElementType<int32_t>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
-  const auto &Spec = FeatureSpecs[FeatureID].Spec;
-  char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID);
-  if (Spec.isElementType<int32_t>())
-    for (size_t I = 0; I < Spec.getElementCount(); ++I)
-      (reinterpret_cast<int64_t *>(Buff))[I] =
-          static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]);
-  else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>())
-    std::memcpy(Buff, RawData,
-                Spec.getElementCount() * Spec.getElementByteSize());
-  else
-    llvm_unreachable("Unsupported tensor type");
-}
-
-char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) {
-  return reinterpret_cast<char *>(LoggerData->addNewTensor(FeatureID));
-}
-
-void Logger::flush(std::string *Str) { LoggerData->flush(Str); }
-
-void Logger::flush(raw_ostream &OS) {
-  std::string Buff;
-  LoggerData->flush(&Buff);
-  OS << Buff;
-}
-
-void Logger::flushLogs(raw_ostream &OS,
-                       const StringMap<std::unique_ptr<Logger>> &Loggers) {
-  google::protobuf::Struct Msg;
-  for (const auto &NamedLogger : Loggers) {
-    tensorflow::SequenceExample SE;
-    const auto &Logger = NamedLogger.second;
-    std::string Unencoded;
-    if (Logger->LoggerData->getNrRecords() > 0)
-      Logger->flush(&Unencoded);
-
-    (*Msg.mutable_fields())[NamedLogger.first().str()]
-        .mutable_string_value()
-        ->append(ProtobufTextMode ? Unencoded : encodeBase64(Unencoded));
-  }
-
-  std::string OutStr;
-  serialize(Msg, &OutStr);
-  OS << OutStr;
-}
 #endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/TrainingLogger.cpp

+//===- TrainingLogger.cpp - mlgo feature/reward logging -------------------===//
+//
+// 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 logging infrastructure for extracting features and
+// rewards for mlgo policy training.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/Config/config.h"
+#if defined(LLVM_HAVE_TF_API)
+
+#include "llvm/ADT/Twine.h"
+#include "llvm/Analysis/Utils/TrainingLogger.h"
+#include "llvm/Support/Base64.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "google/protobuf/struct.pb.h"
+#include "google/protobuf/text_format.h"
+#include "tensorflow/core/example/example.pb.h"
+#include <cassert>
+#include <numeric>
+
+using namespace llvm;
+
+using google::protobuf::Message;
+using google::protobuf::TextFormat;
+
+static cl::opt<bool>
+    ProtobufTextMode("tfutils-text-log", cl::init(false), cl::Hidden,
+                     cl::desc("Output textual (human-readable) protobuf."));
+
+namespace {
+
+void serialize(const Message &SE, std::string *OutStr) {
+  if (ProtobufTextMode) {
+    TextFormat::PrintToString(SE, OutStr);
+  } else {
+    *OutStr = SE.SerializeAsString();
+  }
+}
+} // namespace
+
+namespace llvm {
+
+class LoggerDataImpl {
+  const std::vector<LoggedFeatureSpec> LoggedFeatureSpecs;
+  const TensorSpec RewardSpec;
+  const bool IncludeReward;
+
+  std::vector<tensorflow::FeatureList> FeatureLists;
+  tensorflow::FeatureList Reward;
+
+  bool isSelfConsistent(const tensorflow::SequenceExample &SE,
+                        size_t NrRecords) const {
+    bool Ret = true;
+    for (const auto &TSpecs : LoggedFeatureSpecs) {
+      const auto &Name = TSpecs.getLoggingName();
+      const auto &FL = SE.feature_lists().feature_list().at(Name).feature();
+      if (NrRecords != static_cast<size_t>(FL.size())) {
+        dbgs() << "[TF-UTILS]: " << Name << " has missing records. Expected "
+               << NrRecords << " got " << FL.size() << "\n";
+        Ret = false;
+      }
+    }
+    if (IncludeReward && static_cast<size_t>(SE.feature_lists()
+                                                 .feature_list()
+                                                 .at(RewardSpec.name())
+                                                 .feature()
+                                                 .size()) != NrRecords) {
+      dbgs() << "[TF-UTILS]: reward is missing records.\n";
+      Ret = false;
+    }
+    return Ret;
+  }
+
+  void transferLog(tensorflow::SequenceExample &SE) {
+    auto *FL = SE.mutable_feature_lists()->mutable_feature_list();
+    if (IncludeReward)
+      (*FL)[RewardSpec.name()] = std::move(Reward);
+    assert(FeatureLists.size() == LoggedFeatureSpecs.size());
+    for (size_t I = 0; I < FeatureLists.size(); ++I) {
+      const auto &LFS = LoggedFeatureSpecs[I];
+      (*FL)[LFS.getLoggingName()] = std::move(FeatureLists[I]);
+    }
+  }
+
+public:
+  LoggerDataImpl(const std::vector<LoggedFeatureSpec> &LoggedSpecs,
+                 const TensorSpec &RewardSpec, bool IncludeReward)
+      : LoggedFeatureSpecs(LoggedSpecs), RewardSpec(RewardSpec),
+        IncludeReward(IncludeReward), FeatureLists(LoggedFeatureSpecs.size()) {}
+
+  // flush the logged info to a stream and clear the log contents.
+  void flush(std::string *Str) {
+    size_t NrRecords = getNrRecords();
+    (void)NrRecords;
+    tensorflow::SequenceExample SE;
+    transferLog(SE);
+    assert(isSelfConsistent(SE, NrRecords));
+    serialize(SE, Str);
+  }
+
+  char *addNewTensor(size_t FeatureID) {
+    const auto &Spec = LoggedFeatureSpecs[FeatureID].Spec;
+    if (Spec.isElementType<float>()) {
+      auto *RF = FeatureLists[FeatureID]
+                     .add_feature()
+                     ->mutable_float_list()
+                     ->mutable_value();
+      RF->Resize(Spec.getElementCount(), 0.0);
+      return reinterpret_cast<char *>(RF->mutable_data());
+    } else if (Spec.isElementType<int32_t>() || Spec.isElementType<int64_t>()) {
+      auto *RF = FeatureLists[FeatureID]
+                     .add_feature()
+                     ->mutable_int64_list()
+                     ->mutable_value();
+      RF->Resize(Spec.getElementCount(), 0);
+      return reinterpret_cast<char *>(RF->mutable_data());
+    }
+    llvm_unreachable("Unsupported tensor type.");
+  }
+
+  template <typename T> void logReward(T Value) {
+    assert(IncludeReward);
+    if (RewardSpec.isElementType<float>())
+      Reward.add_feature()->mutable_float_list()->add_value(Value);
+    else if (RewardSpec.isElementType<int32_t>() ||
+             RewardSpec.isElementType<int64_t>())
+      Reward.add_feature()->mutable_int64_list()->add_value(Value);
+    else
+      llvm_unreachable("Unsupported tensor type.");
+  }
+
+  size_t getNrRecords() const {
+    return FeatureLists.empty() ? 0 : FeatureLists[0].feature().size();
+  }
+};
+} // namespace llvm
+
+Logger::Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
+               const TensorSpec &RewardSpec, bool IncludeReward)
+    : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
+      IncludeReward(IncludeReward),
+      LoggerData(std::make_unique<LoggerDataImpl>(FeatureSpecs, RewardSpec,
+                                                  IncludeReward)) {}
+
+Logger::~Logger() {}
+
+#define LOG_REWARD(NAME, TYPE)                                                 \
+  void Logger::log##NAME##Reward(TYPE Value) {                                 \
+    assert(IncludeReward);                                                     \
+    LoggerData->logReward(Value);                                              \
+  }
+
+LOG_REWARD(Float, float)
+LOG_REWARD(Int32, int32_t)
+LOG_REWARD(Int64, int64_t)
+#undef LOG_REWARD
+
+#define LOG_FINAL_REWARD(NAME, TYPE)                                           \
+  void Logger::log##NAME##FinalReward(TYPE Value) {                            \
+    assert(RewardSpec.isElementType<TYPE>());                                  \
+    for (size_t I = 1; I < LoggerData->getNrRecords(); ++I)                    \
+      log##NAME##Reward(0);                                                    \
+    log##NAME##Reward(Value);                                                  \
+  }
+
+LOG_FINAL_REWARD(Float, float)
+LOG_FINAL_REWARD(Int32, int32_t)
+LOG_FINAL_REWARD(Int64, int64_t)
+#undef LOG_FINAL_REWARD
+
+void Logger::logFloatValue(size_t FeatureID, const float *Value) {
+  assert(FeatureSpecs[FeatureID].Spec.isElementType<float>());
+  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
+}
+
+void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
+  assert(FeatureSpecs[FeatureID].Spec.isElementType<int64_t>());
+  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
+}
+
+void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
+  assert(FeatureSpecs[FeatureID].Spec.isElementType<int32_t>());
+  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
+}
+
+void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
+  const auto &Spec = FeatureSpecs[FeatureID].Spec;
+  char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID);
+  if (Spec.isElementType<int32_t>())
+    for (size_t I = 0; I < Spec.getElementCount(); ++I)
+      (reinterpret_cast<int64_t *>(Buff))[I] =
+          static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]);
+  else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>())
+    std::memcpy(Buff, RawData,
+                Spec.getElementCount() * Spec.getElementByteSize());
+  else
+    llvm_unreachable("Unsupported tensor type");
+}
+
+char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) {
+  return reinterpret_cast<char *>(LoggerData->addNewTensor(FeatureID));
+}
+
+void Logger::flush(std::string *Str) { LoggerData->flush(Str); }
+
+void Logger::flush(raw_ostream &OS) {
+  std::string Buff;
+  LoggerData->flush(&Buff);
+  OS << Buff;
+}
+
+void Logger::flushLogs(raw_ostream &OS,
+                       const StringMap<std::unique_ptr<Logger>> &Loggers) {
+  google::protobuf::Struct Msg;
+  for (const auto &NamedLogger : Loggers) {
+    tensorflow::SequenceExample SE;
+    const auto &Logger = NamedLogger.second;
+    std::string Unencoded;
+    if (Logger->LoggerData->getNrRecords() > 0)
+      Logger->flush(&Unencoded);
+
+    (*Msg.mutable_fields())[NamedLogger.first().str()]
+        .mutable_string_value()
+        ->append(ProtobufTextMode ? Unencoded : encodeBase64(Unencoded));
+  }
+
+  std::string OutStr;
+  serialize(Msg, &OutStr);
+  OS << OutStr;
+}
+#endif // defined(LLVM_HAVE_TF_API)

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 #include "AllocationOrder.h"
 #include "RegAllocEvictionAdvisor.h"
 #include "RegAllocGreedy.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TF_API)
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
+#include "llvm/Analysis/Utils/TrainingLogger.h"
 #endif
 #include "llvm/Analysis/ReleaseModeModelRunner.h"
 #include "llvm/CodeGen/CalcSpillWeights.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"

llvm/unittests/Analysis/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   Support
   TransformUtils
   )
 
-set(MLGO_TESTS TFUtilsTest.cpp)
+set(MLGO_TESTS TFUtilsTest.cpp TrainingLoggerTest.cpp)
 if (DEFINED LLVM_HAVE_TF_API)
   LIST(APPEND EXTRA_TESTS ${MLGO_TESTS})
 else()
   LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS})
 endif()
 
 add_llvm_unittest_with_input_files(AnalysisTests
   AliasAnalysisTest.cpp

llvm/unittests/Analysis/TFUtilsTest.cpp

 //===- TFUtilsTest.cpp - test for TFUtils ---------------------------------===//
 //
 // 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/Utils/TFUtils.h"
-#include "google/protobuf/struct.pb.h"
-#include "tensorflow/core/example/example.pb.h"
-#include "tensorflow/core/example/feature.pb.h"
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Path.h"
TEST(TFUtilsTest, UnsupportedFeature) {
   float Ret = Evaluator->evaluate<float>();
   EXPECT_EQ(static_cast<int64_t>(Ret), 80);
   // The input vector should be unchanged
   for (auto I = 0; I < KnownSize; ++I)
     EXPECT_EQ(V[I], 1);
   for (auto I = 0; I < 2 * 5; ++I)
     EXPECT_FLOAT_EQ(F[I], 3.14 + I);
 }
-
-#define PROTO_CHECKER(FNAME, TYPE, INDEX, EXP)                                 \
-  do {                                                                         \
-    const auto &V = Expected.feature_lists()                                   \
-                        .feature_list()                                        \
-                        .at(FNAME)                                             \
-                        .feature(INDEX)                                        \
-                        .TYPE()                                                \
-                        .value();                                              \
-    for (auto I = 0; I < V.size(); ++I)                                        \
-      EXPECT_EQ(V.at(I), EXP[I]);                                              \
-  } while (false)
-
-TEST(TFUtilsTest, Logger) {
-  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);
-  const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
-  const int64_t F01[]{2, 3};
-
-  L.logFloatValue(0, F00);
-  L.logInt64Value(1, F01);
-  L.logFloatReward(3.4);
-  const float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
-  const int64_t F11[]{-2, -3};
-  L.logFloatValue(0, F10);
-  L.logInt64Value(1, F11);
-  L.logFloatReward(-3.0);
-  std::string Result;
-  raw_string_ostream OS(Result);
-  L.flush(OS);
-
-  tensorflow::SequenceExample Expected;
-  ASSERT_TRUE(Expected.ParseFromString(Result));
-  PROTO_CHECKER("the_float", float_list, 0, F00);
-  PROTO_CHECKER("the_float", float_list, 1, F10);
-  PROTO_CHECKER("alternate_name", int64_list, 0, F01);
-  PROTO_CHECKER("alternate_name", int64_list, 1, F11);
-  float R0[]{3.4};
-  float R1[]{-3.0};
-  PROTO_CHECKER("reward", float_list, 0, R0);
-  PROTO_CHECKER("reward", float_list, 1, R1);
-}
-
-TEST(TFUtilsTest, LoggerInt32FeaturesAndReward) {
-  std::vector<LoggedFeatureSpec> Features;
-  Features.push_back(
-      {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
-  Features.push_back({TensorSpec::createSpec<int32_t>("the_int", {2}),
-                      std::string("alternate_name")});
-
-  auto Rewards = TensorSpec::createSpec<int32_t>("reward", {1});
-  Logger L(Features, Rewards, true);
-  const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
-  const int32_t F01[]{2, 3};
-
-  L.logFloatValue(0, F00);
-  L.logInt32Value(1, F01);
-  L.logInt32Reward(3);
-  const float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
-  const int32_t F11[]{-2, -3};
-  L.logFloatValue(0, F10);
-  L.logInt32Value(1, F11);
-  L.logInt32Reward(-3);
-  std::string Result;
-  raw_string_ostream OS(Result);
-  L.flush(OS);
-
-  tensorflow::SequenceExample Expected;
-  ASSERT_TRUE(Expected.ParseFromString(Result));
-  PROTO_CHECKER("the_float", float_list, 0, F00);
-  PROTO_CHECKER("the_float", float_list, 1, F10);
-  PROTO_CHECKER("alternate_name", int64_list, 0, F01);
-  PROTO_CHECKER("alternate_name", int64_list, 1, F11);
-  int32_t R0[]{3};
-  int32_t R1[]{-3};
-  PROTO_CHECKER("reward", int64_list, 0, R0);
-  PROTO_CHECKER("reward", int64_list, 1, R1);
-}
-
-TEST(TFUtilsTest, LoggerNoReward) {
-  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);
-  const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
-  const int64_t F01[]{2, 3};
-
-  L.logFloatValue(0, F00);
-  L.logInt64Value(1, F01);
-  const float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
-  const int64_t F11[]{-2, -3};
-  L.logFloatValue(0, F10);
-  L.logInt64Value(1, F11);
-
-  std::string Result;
-  raw_string_ostream OS(Result);
-  L.flush(OS);
-  tensorflow::SequenceExample Expected;
-  ASSERT_TRUE(Expected.ParseFromString(Result));
-  PROTO_CHECKER("the_float", float_list, 0, F00);
-  PROTO_CHECKER("the_float", float_list, 1, F10);
-  PROTO_CHECKER("alternate_name", int64_list, 0, F01);
-  PROTO_CHECKER("alternate_name", int64_list, 1, F11);
-}
-
-TEST(TFUtilsTest, LoggerFinalReward) {
-  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 (int64_t I = 0; I < 3; ++I) {
-    float F = static_cast<float>(I);
-    L.logFloatValue(0, &F);
-    L.logInt64Value(1, &I);
-  }
-  L.logFloatFinalReward(3.14);
-  std::string Result;
-  raw_string_ostream OS(Result);
-  L.flush(OS);
-  const float Zero[]{0.0};
-  const float R[]{3.14};
-  tensorflow::SequenceExample Expected;
-  ASSERT_TRUE(Expected.ParseFromString(Result));
-  PROTO_CHECKER("reward", float_list, 0, Zero);
-  PROTO_CHECKER("reward", float_list, 1, Zero);
-  PROTO_CHECKER("reward", float_list, 2, R);
-}
-
-TEST(TFUtilsTest, LoggerGroup) {
-  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});
-  StringMap<std::unique_ptr<Logger>> Loggers;
-  std::vector<std::string> Names{"a", "b"};
-  size_t Bump = 0;
-  for (auto Name : Names) {
-    auto L = std::make_unique<Logger>(Features, Rewards, true);
-    for (int64_t I = 0; I < 3; ++I) {
-      float F = static_cast<float>(I) + Bump;
-      L->logFloatValue(0, &F);
-      L->logInt64Value(1, &I);
-    }
-    L->logFloatFinalReward(3.14 + Bump);
-    Loggers.insert(std::make_pair(Name, std::move(L)));
-  }
-  std::string Result;
-  raw_string_ostream OS(Result);
-  Logger::flushLogs(OS, Loggers);
-  google::protobuf::Struct Expected;
-  ASSERT_TRUE(Expected.ParseFromString(Result));
-  EXPECT_EQ(Expected.fields_size(), 2);
-  EXPECT_TRUE(Expected.fields().contains("a"));
-  EXPECT_TRUE(Expected.fields().contains("b"));
-}

llvm/unittests/Analysis/TrainingLoggerTest.cpp

-//===- TFUtilsTest.cpp - test for TFUtils ---------------------------------===//
+//===- TrainingLoggerTest.cpp - test for TrainingLogger -------------------===//
 //
 // 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/Utils/TFUtils.h"
+#include "llvm/Analysis/Utils/TrainingLogger.h"
 #include "google/protobuf/struct.pb.h"
 #include "tensorflow/core/example/example.pb.h"
 #include "tensorflow/core/example/feature.pb.h"
-#include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 extern const char *TestMainArgv0;
 
-// NOTE! This test model is currently also used by test/Transforms/Inline/ML tests
+// NOTE! This test model is currently also used by test/Transforms/Inline/ML
+// tests
 //- relevant if updating this model.
-static std::string getModelPath() {
-  SmallString<128> InputsDir = unittest::getInputFileDirectory(TestMainArgv0);
-  llvm::sys::path::append(InputsDir, "ir2native_x86_64_model");
-  return std::string(InputsDir);
-}
-
-// Test observable behavior when no model is provided.
-TEST(TFUtilsTest, NoModel) {
-  TFModelEvaluator Evaluator("", {}, {});
-  EXPECT_FALSE(Evaluator.isValid());
-}
-
-// Test we can correctly load a savedmodel and evaluate it.
-TEST(TFUtilsTest, LoadAndExecuteTest) {
-  // We use the ir2native model for test. We know it has one feature of
-  // dimension (1, 214)
-  const static int64_t KnownSize = 214;
-  std::vector<TensorSpec> InputSpecs{TensorSpec::createSpec<int32_t>(
-      "serving_default_input_1", {1, KnownSize})};
-  std::vector<TensorSpec> OutputSpecs{
-      TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
-
-  TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
-  EXPECT_TRUE(Evaluator.isValid());
-
-  int32_t *V = Evaluator.getInput<int32_t>(0);
-  // Fill it up with 1's, we know the output.
-  for (auto I = 0; I < KnownSize; ++I) {
-    V[I] = 1;
-  }
-  {
-    auto ER = Evaluator.evaluate();
-    EXPECT_TRUE(ER.hasValue());
-    float Ret = *ER->getTensorValue<float>(0);
-    EXPECT_EQ(static_cast<int64_t>(Ret), 80);
-    EXPECT_EQ(ER->getUntypedTensorValue(0),
-              reinterpret_cast<const void *>(ER->getTensorValue<float>(0)));
-  }
-  // The input vector should be unchanged
-  for (auto I = 0; I < KnownSize; ++I) {
-    EXPECT_EQ(V[I], 1);
-  }
-  // Zero-out the unused position '0' of the instruction histogram, which is
-  // after the first 9 calculated values. Should the the same result.
-  V[9] = 0;
-  {
-    auto ER = Evaluator.evaluate();
-    EXPECT_TRUE(ER.hasValue());
-    float Ret = *ER->getTensorValue<float>(0);
-    EXPECT_EQ(static_cast<int64_t>(Ret), 80);
-  }
-}
-
-// Test incorrect input setup
-TEST(TFUtilsTest, EvalError) {
-  // We use the ir2native model for test. We know it has one feature of
-  // dimension (1, 214)
-  const static int64_t KnownSize = 213;
-  std::vector<TensorSpec> InputSpecs{TensorSpec::createSpec<int32_t>(
-      "serving_default_input_1", {1, KnownSize})};
-  std::vector<TensorSpec> OutputSpecs{
-      TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
-
-  TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
-  EXPECT_TRUE(Evaluator.isValid());
-
-  int32_t *V = Evaluator.getInput<int32_t>(0);
-  // Fill it up with 1's, we know the output.
-  for (auto I = 0; I < KnownSize; ++I) {
-    V[I] = 1;
-  }
-  auto ER = Evaluator.evaluate();
-  EXPECT_FALSE(ER.hasValue());
-  EXPECT_FALSE(Evaluator.isValid());
-}
-
-TEST(TFUtilsTest, UnsupportedFeature) {
-  const static int64_t KnownSize = 214;
-  std::vector<TensorSpec> InputSpecs{
-      TensorSpec::createSpec<int32_t>("serving_default_input_1",
-                                      {1, KnownSize}),
-      TensorSpec::createSpec<float>("this_feature_does_not_exist", {2, 5})};
-
-  LLVMContext Ctx;
-  auto Evaluator = ModelUnderTrainingRunner::createAndEnsureValid(
-      Ctx, getModelPath(), "StatefulPartitionedCall", InputSpecs,
-      {LoggedFeatureSpec{
-          TensorSpec::createSpec<float>("StatefulPartitionedCall", {1}),
-          None}});
-  int32_t *V = Evaluator->getTensor<int32_t>(0);
-  // Fill it up with 1s, we know the output.
-  for (auto I = 0; I < KnownSize; ++I)
-    V[I] = 1;
-
-  float *F = Evaluator->getTensor<float>(1);
-  for (auto I = 0; I < 2 * 5; ++I)
-    F[I] = 3.14 + I;
-  float Ret = Evaluator->evaluate<float>();
-  EXPECT_EQ(static_cast<int64_t>(Ret), 80);
-  // The input vector should be unchanged
-  for (auto I = 0; I < KnownSize; ++I)
-    EXPECT_EQ(V[I], 1);
-  for (auto I = 0; I < 2 * 5; ++I)
-    EXPECT_FLOAT_EQ(F[I], 3.14 + I);
-}
 
 #define PROTO_CHECKER(FNAME, TYPE, INDEX, EXP)                                 \
   do {                                                                         \
     const auto &V = Expected.feature_lists()                                   \
                         .feature_list()                                        \
                         .at(FNAME)                                             \
                         .feature(INDEX)                                        \
                         .TYPE()                                                \
                         .value();                                              \
     for (auto I = 0; I < V.size(); ++I)                                        \
       EXPECT_EQ(V.at(I), EXP[I]);                                              \
   } while (false)
 
-TEST(TFUtilsTest, Logger) {
+TEST(TrainingLoggerTest, Logger) {
   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);
TEST(TrainingLoggerTest, Logger) {
   PROTO_CHECKER("alternate_name", int64_list, 0, F01);
   PROTO_CHECKER("alternate_name", int64_list, 1, F11);
   float R0[]{3.4};
   float R1[]{-3.0};
   PROTO_CHECKER("reward", float_list, 0, R0);
   PROTO_CHECKER("reward", float_list, 1, R1);
 }
 
-TEST(TFUtilsTest, LoggerInt32FeaturesAndReward) {
+TEST(TrainingLoggerTest, LoggerInt32FeaturesAndReward) {
   std::vector<LoggedFeatureSpec> Features;
   Features.push_back(
       {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
   Features.push_back({TensorSpec::createSpec<int32_t>("the_int", {2}),
                       std::string("alternate_name")});
 
   auto Rewards = TensorSpec::createSpec<int32_t>("reward", {1});
   Logger L(Features, Rewards, true);
TEST(TrainingLoggerTest, LoggerInt32FeaturesAndReward) {
   PROTO_CHECKER("alternate_name", int64_list, 0, F01);
   PROTO_CHECKER("alternate_name", int64_list, 1, F11);
   int32_t R0[]{3};
   int32_t R1[]{-3};
   PROTO_CHECKER("reward", int64_list, 0, R0);
   PROTO_CHECKER("reward", int64_list, 1, R1);
 }
 
-TEST(TFUtilsTest, LoggerNoReward) {
+TEST(TrainingLoggerTest, LoggerNoReward) {
   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);
TEST(TrainingLoggerTest, LoggerNoReward) {
   tensorflow::SequenceExample Expected;
   ASSERT_TRUE(Expected.ParseFromString(Result));
   PROTO_CHECKER("the_float", float_list, 0, F00);
   PROTO_CHECKER("the_float", float_list, 1, F10);
   PROTO_CHECKER("alternate_name", int64_list, 0, F01);
   PROTO_CHECKER("alternate_name", int64_list, 1, F11);
 }
 
-TEST(TFUtilsTest, LoggerFinalReward) {
+TEST(TrainingLoggerTest, LoggerFinalReward) {
   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 (int64_t I = 0; I < 3; ++I) {
     float F = static_cast<float>(I);
     L.logFloatValue(0, &F);
     L.logInt64Value(1, &I);
   }
   L.logFloatFinalReward(3.14);
   std::string Result;
   raw_string_ostream OS(Result);
   L.flush(OS);
   const float Zero[]{0.0};
   const float R[]{3.14};
   tensorflow::SequenceExample Expected;
   ASSERT_TRUE(Expected.ParseFromString(Result));
   PROTO_CHECKER("reward", float_list, 0, Zero);
   PROTO_CHECKER("reward", float_list, 1, Zero);
   PROTO_CHECKER("reward", float_list, 2, R);
 }
 
-TEST(TFUtilsTest, LoggerGroup) {
+TEST(TrainingLoggerTest, LoggerGroup) {
   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});
   StringMap<std::unique_ptr<Logger>> Loggers;
   std::vector<std::string> Names{"a", "b"};
   size_t Bump = 0;