[mlgo] Stream the training data

[mtrofin]

Summary

This leverages the new logging format in that we don't need to buffer
the training data, we can just write it out.

[kazu]

LGTM. Just a minor comment about currentObservationID.

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

122–124

If you only care about currentObservationID().has_value(), I wonder if you can rename this function to isObservationInProgress or something and do:

return ObservationIDs.count(CurrentContext);

[mtrofin]

api

[mtrofin]

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

122–124

good point, done!

[MatzeB]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

I think I am hitting this for cases where RegAllocGreedy bails out early in the if (!hasVirtRegAlloc()) return false; check and never initializes the advisor.

Replacing with Log->switchCurrentContext() for my experiments...

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

ack. I'll take a look. Probably in that case (!hasVirtRegAlloc()) there's also no output to log, right?

[MatzeB]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

Guess I can also try to first test for hasObservationInProgress() and bail out on that before checking the function name... Rebuilding and testing right now...

[MatzeB]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

checking for hasObservationInProgress() first was not enough to avoid hitting those errors. (the "InProgress" part of the function name seems somewhat confusing anyway as this seems to be more of an "everHadAnObservation()" behavior rather than being in the middle of a startObservation() / endObservation() thingy...)

[MatzeB]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

I'm not sure it's really that early-exit in RAGreedy that makes me hit that error, but I am hitting it for multiple functions within many compilation units here; will continue digging next week...

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

Ack, thanks - so in the case you're seeing, we did have an observation but we get to logRewardIfNeeded for a different function first? (lmk if you can share a repro, happy to help)

[jacobhegna]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

468–469

I am hitting this error when training the regalloc priority model as well.

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

 #include "llvm/Config/llvm-config.h"
 
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/JSON.h"
 
 #include <memory>
+#include <optional>
 #include <vector>
 
 namespace llvm {
 
 /// Logging utility - given an ordered specification of features, and assuming
 /// a scalar reward, allow logging feature values and rewards.
 /// 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

 ///   logInt64Value(1, ...)
 ///   ...
 ///   logFloatReward(...)
 ///
 /// At the end, call print to generate the log.
 /// Alternatively, don't call logReward at the end of each event, just
 /// log{Float|Int32|Int64}FinalReward at the end.
 class Logger final {
+  std::unique_ptr<raw_ostream> OS;
   const std::vector<TensorSpec> FeatureSpecs;
   const TensorSpec RewardSpec;
   const bool IncludeReward;
-  std::vector<std::unique_ptr<char[]>> FeatureStorage;
-  std::vector<std::unique_ptr<char[]>> RewardStorage;
-  raw_ostream &dumpHeader(raw_ostream &OS) const;
-  raw_ostream &startContext(raw_ostream &OS, StringRef Name) const;
-  raw_ostream &startObservation(raw_ostream &OS, size_t Nr) const;
-  raw_ostream &writeOutcome(raw_ostream &OS, size_t CurrentObservationID) const;
-  char *addNewTensor(size_t FeatureID);
-  size_t getNrRecords() const;
+  StringMap<size_t> ObservationIDs;
+  std::string CurrentContext;
 
-  void logRewardImpl(const char *Value, size_t Size);
+  void writeHeader();
+  void writeTensor(const TensorSpec &Spec, const char *RawData) {
+    OS->write(RawData, Spec.getTotalTensorBufferSize());
+  }
+  void logRewardImpl(const char *RawData);
 
 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<TensorSpec> &FeatureSpecs,
-         const TensorSpec &RewardSpec, bool IncludeReward)
-      : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
-        IncludeReward(IncludeReward) {}
+  Logger(std::unique_ptr<raw_ostream> OS,
+         const std::vector<TensorSpec> &FeatureSpecs,
+         const TensorSpec &RewardSpec, bool IncludeReward);
+
+  void switchContext(StringRef Name);
+  void startObservation();
+  void endObservation();
+
+  const std::string &currentContext() const { return CurrentContext; }
+
+  bool hasObservationInProgress() const {
+    return ObservationIDs.find(CurrentContext) != ObservationIDs.end();
+  }
 
   template <typename T> void logReward(T Value) {
-    logRewardImpl(reinterpret_cast<const char *>(&Value), sizeof(T));
+    logRewardImpl(reinterpret_cast<const char *>(&Value));
+  }
+
+  void logTensorValue(size_t FeatureID, const char *RawData) {
+    writeTensor(FeatureSpecs[FeatureID], RawData);
   }
-  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.
-  raw_ostream &flush(raw_ostream &OS, bool WithHeader = true,
-                     StringRef Context = "default") const;
-
-  // Flush a set of logs that are produced from the same module, e.g.
-  // per-function regalloc traces.
-  static void flushLogs(raw_ostream &OS,
-                        const StringMap<std::unique_ptr<Logger>> &Loggers);
 };
 
 } // namespace llvm
 #endif // LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

 class TrainingLogger final {
 public:
   TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR);
 
   /// Log one inlining event.
   void logInlineEvent(const InlineEvent &Event,
                       const MLModelRunner &ModelRunner);
 
-  /// Print the stored tensors.
-  void print();
-
 private:
   StringRef LogFileName;
   const ModelUnderTrainingRunner *const MUTR;
   std::unique_ptr<Logger> L;
   BitVector Effects;
   /// Set these 2 clearly OOB, to make sure we set them later.
   size_t DefaultDecisionPos = std::numeric_limits<size_t>::max();
   size_t DecisionPos = std::numeric_limits<size_t>::max();
public:
   DevelopmentModeMLInlineAdvisor(
       Module &M, ModuleAnalysisManager &MAM,
       std::unique_ptr<MLModelRunner> ModelRunner,
       std::function<bool(CallBase &)> GetDefaultAdvice,
       std::unique_ptr<TrainingLogger> Logger);
 
   size_t getTotalSizeEstimate();
 
-  virtual ~DevelopmentModeMLInlineAdvisor();
   void updateNativeSizeEstimate(int64_t Change) {
     *CurrentNativeSize += Change;
   }
   void resetNativeSize(Function *F) {
     PreservedAnalyses PA = PreservedAnalyses::all();
     PA.abandon<InlineSizeEstimatorAnalysis>();
     FAM.invalidate(*F, PA);
   }
TrainingLogger::TrainingLogger(StringRef LogFileName,
   if (MUTR)
     append_range(FT, MUTR->extraOutputsForLoggingSpecs());
 
   DefaultDecisionPos = FT.size();
   FT.push_back(TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}));
 
   DecisionPos = FT.size();
   FT.push_back(TensorSpec::createSpec<int64_t>(DecisionName, {1}));
+  std::error_code EC;
+  auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
+  if (EC)
+    dbgs() << (EC.message() + ":" + TrainingLog);
 
   L = std::make_unique<Logger>(
-      FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
+      std::move(OS), FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
       InlineSizeEstimatorAnalysis::isEvaluatorRequested());
+  L->switchContext("");
 }
 
 /// Log one inlining event.
 void TrainingLogger::logInlineEvent(const InlineEvent &Event,
                                     const MLModelRunner &ModelRunner) {
+  L->startObservation();
   size_t CurrentFeature = 0;
-  for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) {
-    int64_t F = *ModelRunner.getTensor<int64_t>(CurrentFeature);
-    L->logInt64Value(CurrentFeature, &F);
-  }
+  for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature)
+    L->logTensorValue(CurrentFeature,
+                      reinterpret_cast<const char *>(
+                          ModelRunner.getTensorUntyped(CurrentFeature)));
 
   if (MUTR)
     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) {
       const char *RawData =
           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I));
-      L->logSpecifiedTensorValue(CurrentFeature, RawData);
+      L->logTensorValue(CurrentFeature, RawData);
       ++CurrentFeature;
     }
 
   assert(CurrentFeature == DefaultDecisionPos);
-  L->logInt64Value(DefaultDecisionPos, &Event.DefaultDecision);
-  L->logInt64Value(DecisionPos, &Event.AdvisedDecision);
+  L->logTensorValue(DefaultDecisionPos,
+                    reinterpret_cast<const char *>(&Event.DefaultDecision));
+  L->logTensorValue(DecisionPos,
+                    reinterpret_cast<const char *>(&Event.AdvisedDecision));
+  L->endObservation();
   if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
-    L->logInt64Reward(Event.Reward);
+    L->logReward(Event.Reward);
 
   // For debugging / later use
   Effects.push_back(Event.Effect);
 }
 
-void TrainingLogger::print() {
-  std::error_code EC;
-  raw_fd_ostream OutFile(LogFileName, EC);
-  L->flush(OutFile);
-}
-
 DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
     Module &M, ModuleAnalysisManager &MAM,
     std::unique_ptr<MLModelRunner> ModelRunner,
     std::function<bool(CallBase &)> GetDefaultAdvice,
     std::unique_ptr<TrainingLogger> Logger)
     : MLInlineAdvisor(M, MAM, std::move(ModelRunner)),
       GetDefaultAdvice(GetDefaultAdvice),
       IsDoingInference(isa<ModelUnderTrainingRunner>(getModelRunner())),
       Logger(std::move(Logger)),
       InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0),
       CurrentNativeSize(InitialNativeSize) {
   // We cannot have the case of neither inference nor logging.
   assert(IsDoingInference || isLogging());
 }
 
-DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
-  if (isLogging())
-    Logger->print();
-}
-
 std::optional<size_t>
 DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
   if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
     return std::nullopt;
   auto &R =
       FAM.getResult<InlineSizeEstimatorAnalysis>(const_cast<Function &>(F));
   if (!R) {
     F.getParent()->getContext().emitError(

llvm/lib/Analysis/TrainingLogger.cpp

 
 using namespace llvm;
 
 // FIXME(mtrofin): remove the flag altogether
 static cl::opt<bool>
     UseSimpleLogger("tfutils-use-simplelogger", cl::init(true), cl::Hidden,
                     cl::desc("Output simple (non-protobuf) log."));
 
-raw_ostream &Logger::dumpHeader(raw_ostream &OS) const {
-  json::OStream JOS(OS);
+void Logger::writeHeader() {
+  json::OStream JOS(*OS);
   JOS.object([&]() {
     JOS.attributeArray("features", [&]() {
       for (const auto &TS : FeatureSpecs)
         TS.toJSON(JOS);
     });
     if (IncludeReward) {
       JOS.attributeBegin("score");
       RewardSpec.toJSON(JOS);
       JOS.attributeEnd();
     }
   });
-  OS << "\n";
-  return OS;
+  *OS << "\n";
 }
 
-raw_ostream &Logger::startContext(raw_ostream &OS, StringRef Name) const {
-  json::OStream JOS(OS);
+void Logger::switchContext(StringRef Name) {
+  CurrentContext = Name.str();
+  json::OStream JOS(*OS);
   JOS.object([&]() { JOS.attribute("context", Name); });
-  OS << "\n";
-  return OS;
+  *OS << "\n";
 }
 
-raw_ostream &Logger::startObservation(raw_ostream &OS, size_t Nr) const {
-  json::OStream JOS(OS);
-  JOS.object([&]() { JOS.attribute("observation", static_cast<int64_t>(Nr)); });
-  OS << "\n";
-  return OS;
-}
-
-raw_ostream &Logger::writeOutcome(raw_ostream &OS,
-                                  size_t CurrentObservationID) const {
-  if (IncludeReward) {
-    OS << "\n";
-    json::OStream JOS(OS);
-    JOS.object([&]() {
-      JOS.attribute("outcome", static_cast<int64_t>(CurrentObservationID));
-    });
-    OS << "\n";
-    OS.write(RewardStorage[CurrentObservationID].get(),
-             RewardSpec.getTotalTensorBufferSize());
-  }
-  OS << "\n";
-  return OS;
-}
-
-char *Logger::addNewTensor(size_t FeatureID) {
-  return FeatureStorage
-      .emplace_back(
-          new char[FeatureSpecs[FeatureID].getTotalTensorBufferSize()])
-      .get();
-}
-
-size_t Logger::getNrRecords() const {
-  assert(FeatureStorage.size() % FeatureSpecs.size() == 0);
-  return FeatureStorage.size() / FeatureSpecs.size();
+void Logger::startObservation() {
+  auto I = ObservationIDs.insert({CurrentContext, 0});
+  size_t NewObservationID = I.second ? 0 : ++I.first->second;
+  json::OStream JOS(*OS);
+    JOS.attribute("observation", static_cast<int64_t>(NewObservationID));
+  *OS << "\n";
 }
 
-void Logger::logRewardImpl(const char *Value, size_t Size) {
-  std::memcpy(RewardStorage.emplace_back(new char[Size]).get(), Value, Size);
-}
-
-raw_ostream &Logger::flush(raw_ostream &OS, bool WithHeader,
-                           StringRef Context) const {
-  if (WithHeader)
-    dumpHeader(OS);
-  startContext(OS, Context);
-  size_t CurrentObservationID = 0;
-  for (size_t I = 0; I < FeatureStorage.size(); ++I) {
-    size_t TensorID = I % FeatureSpecs.size();
-    if (TensorID == 0) {
-      CurrentObservationID = I / FeatureSpecs.size();
-      startObservation(OS, CurrentObservationID);
-    }
-    OS.write(FeatureStorage[I].get(),
-             FeatureSpecs[TensorID].getTotalTensorBufferSize());
-    if (TensorID == FeatureSpecs.size() - 1) {
-      writeOutcome(OS, CurrentObservationID);
-    }
-  }
-  return OS;
-}
-
-#define LOG_REWARD(NAME, TYPE)                                                 \
-  void Logger::log##NAME##Reward(TYPE Value) {                                 \
-    assert(IncludeReward);                                                     \
-    (void)IncludeReward;                                                       \
-    logReward(Value);                                                          \
-  }
+void Logger::endObservation() { *OS << "\n"; }
 
-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 < 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].isElementType<float>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
-  assert(FeatureSpecs[FeatureID].isElementType<int64_t>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
-  assert(FeatureSpecs[FeatureID].isElementType<int32_t>());
-  logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
-  const auto &Spec = FeatureSpecs[FeatureID];
-  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 *>(addNewTensor(FeatureID));
+void Logger::logRewardImpl(const char *RawData) {
+  assert(IncludeReward);
+  json::OStream JOS(*OS);
+  JOS.object([&]() {
+    JOS.attribute("outcome", static_cast<int64_t>(
+                                 ObservationIDs.find(CurrentContext)->second));
+  });
+  *OS << "\n";
+  writeTensor(RewardSpec, RawData);
+  *OS << "\n";
 }
 
-void Logger::flushLogs(raw_ostream &OS,
-                       const StringMap<std::unique_ptr<Logger>> &Loggers) {
-  bool IsFirst = true;
-  for (const auto &NamedLogger : Loggers) {
-    NamedLogger.second->flush(OS, IsFirst, NamedLogger.first());
-    IsFirst = false;
-  }
+Logger::Logger(std::unique_ptr<raw_ostream> OS,
+               const std::vector<TensorSpec> &FeatureSpecs,
+               const TensorSpec &RewardSpec, bool IncludeReward)
+    : OS(std::move(OS)), FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
+      IncludeReward(IncludeReward) {
+  writeHeader();
 }

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

if (EnableDevelopmentFeatures) {
           TensorSpec::createSpec<float>("action_reward", {1})};
     }
   }
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
-  /// get the logger for the given function, or nullptr if we didn't collect
-  /// one. This is used to inject the score by the RegAllocScoring pass.
-  Logger *getLogger(const MachineFunction &MF) const {
-    auto I = LogMap.find(MF.getName());
-    if (I == LogMap.end())
-      return nullptr;
-    return I->second.get();
-  }
-
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
-    if (auto *Log = this->getLogger(MF))
-      Log->logFloatFinalReward(GetReward());
+    if (!Log)
+      return;
+    // The function pass manager would run all the function passes for a
+    // function, so we assume the last context belongs to this function. If
+    // this invariant ever changes, we can implement at that time switching
+    // contexts. At this point, it'd be an error
+    if (Log->currentContext() != MF.getName()) {
+      MF.getFunction().getContext().emitError(
+          "The training log context shouldn't have had changed.");
+    }
+    if (Log->hasObservationInProgress())
+      Log->logReward<float>(GetReward());
   }
 
 private:
   std::vector<TensorSpec> InputFeatures;
   std::vector<TensorSpec> TrainingInputFeatures;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineLoopInfo>();
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
-  // Save all the logs (when requested).
-  bool doFinalization(Module &M) override {
-    if (TrainingLog.empty())
-      return false;
-    std::error_code EC;
-    auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
-    if (EC) {
-      M.getContext().emitError(EC.message() + ":" + TrainingLog);
-      return false;
-    }
-    Logger::flushLogs(*OS, LogMap);
-    return false;
-  }
-
-  std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    LLVMContext &Ctx = MF.getFunction().getContext();
+  bool doInitialization(Module &M) override {
+    LLVMContext &Ctx = M.getContext();
     if (ModelUnderTraining.empty() && TrainingLog.empty()) {
       Ctx.emitError("Regalloc development mode should be requested with at "
                     "least logging enabled and/or a training model");
-      return nullptr;
+      return false;
     }
-    if (!Runner) {
-      if (ModelUnderTraining.empty())
-        Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
-      else
-        Runner = ModelUnderTrainingRunner::createAndEnsureValid(
-            Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
-      if (!Runner) {
-        Ctx.emitError("Regalloc: could not set up the model runner");
-        return nullptr;
-      }
+      return false;
+    if (TrainingLog.empty())
+      return false;
+    std::error_code EC;
+    auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
+    if (EC) {
+      M.getContext().emitError(EC.message() + ":" + TrainingLog);
+      return false;
     }
-
-    Logger *Log = nullptr;
-    if (!TrainingLog.empty()) {
-      std::vector<TensorSpec> LFS = InputFeatures;
-      if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
-        append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
-      // We always log the output; in particular, if we're not evaluating, we
-      // don't have an output spec json file. That's why we handle the
-      // 'normal' output separately.
-      LFS.push_back(Output);
-      auto I = LogMap.insert(std::make_pair(
-          MF.getFunction().getName(),
-          std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
-      assert(I.second);
-      Log = I.first->second.get();
-    }
+
+    Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
+                                   /*IncludeReward*/ true);
+    return false;
+
+  std::unique_ptr<RegAllocEvictionAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    if (!Runner)
+      return nullptr;
+    if (Log)
+      Log->switchContext(MF.getName());
     return std::make_unique<DevelopmentModeEvictAdvisor>(
         MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
-        getAnalysis<MachineLoopInfo>(), Log);
+        getAnalysis<MachineLoopInfo>(), Log.get());
   }
 
   std::unique_ptr<MLModelRunner> Runner;
-  StringMap<std::unique_ptr<Logger>> LogMap;
+  std::unique_ptr<Logger> Log;
 };
 
 #endif //#ifdef LLVM_HAVE_TFLITE
 } // namespace
 
 float MLEvictAdvisor::getInitialQueueSize(const MachineFunction &MF) {
   auto &MRI = MF.getRegInfo();
   float Ret = 0.0;
if (isa<ModelUnderTrainingRunner>(getRunner())) {
     else
       for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit);
            I != E; ++I, ++Ret)
         if (*I == PhysReg)
           break;
   }
   if (TrainingLog.empty())
     return Ret;
+  // TODO(mtrofin): when we support optional rewards, this can go away. In the
+  // meantime, we log the "pretend" reward (0) for the previous observation
+  // before starting a new one.
+  if (Log->hasObservationInProgress())
+    Log->logReward<float>(0.0);
+
+  Log->startObservation();
   size_t CurrentFeature = 0;
   size_t FeatureCount = EnableDevelopmentFeatures
                             ? FeatureIDs::FeaturesWithDevelopmentCount
                             : FeatureIDs::FeatureCount;
   for (; CurrentFeature < FeatureCount; ++CurrentFeature) {
-    Log->logSpecifiedTensorValue(
-        CurrentFeature, reinterpret_cast<const char *>(
+    Log->logTensorValue(CurrentFeature,
+                        reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
          ++I, ++CurrentFeature)
-      Log->logSpecifiedTensorValue(
+      Log->logTensorValue(
           CurrentFeature,
           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
   // The output is right after the features and the extra outputs
-  Log->logInt64Value(CurrentFeature, &Ret);
+  Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
+  Log->endObservation();
   return Ret;
 }
 
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
   std::optional<float> CachedReward;
   auto GetReward = [&]() {
     if (!CachedReward)
       CachedReward = static_cast<float>(

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

 public:
   DevelopmentModePriorityAdvisorAnalysis()
       : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Development) {}
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
-  /// get the logger for the given function, or nullptr if we didn't collect
-  /// one. This is used to inject the score by the RegAllocScoring pass.
-  Logger *getLogger(const MachineFunction &MF) const {
-    auto I = LogMap.find(MF.getName());
-    if (I == LogMap.end())
-      return nullptr;
-    return I->second.get();
-  }
-
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
-    if (auto *Log = this->getLogger(MF))
-      Log->logFloatFinalReward(GetReward());
+    if (!Log)
+      return;
+    // The function pass manager would run all the function passes for a
+    // function, so we assume the last context belongs to this function. If
+    // this invariant ever changes, we can implement at that time switching
+    // contexts. At this point, it'd be an error
+    if (Log->currentContext() != MF.getName()) {
+      MF.getFunction().getContext().emitError(
+          "The training log context shouldn't have had changed.");
+    }
+    if (Log->hasObservationInProgress())
+      Log->logReward<float>(GetReward());
   }
 
 private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     AU.addRequired<SlotIndexes>();
     RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
   // Save all the logs (when requested).
-  bool doFinalization(Module &M) override {
-    if (TrainingLog.empty())
-      return false;
-    std::error_code EC;
-    auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
-    if (EC) {
-      M.getContext().emitError(EC.message() + ":" + TrainingLog);
-      return false;
-    }
-    Logger::flushLogs(*OS, LogMap);
-    return false;
-  }
-
-  std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-
-    LLVMContext &Ctx = MF.getFunction().getContext();
+  bool doInitialization(Module &M) override {
+    LLVMContext &Ctx = M.getContext();
     if (ModelUnderTraining.empty() && TrainingLog.empty()) {
       Ctx.emitError("Regalloc development mode should be requested with at "
                     "least logging enabled and/or a training model");
-      return nullptr;
+      return false;
     }
-    if (!Runner) {
-      if (ModelUnderTraining.empty())
-        Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
-      else
-        Runner = ModelUnderTrainingRunner::createAndEnsureValid(
-            Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
-      if (!Runner) {
-        Ctx.emitError("Regalloc: could not set up the model runner");
-        return nullptr;
-      }
+      return false;
+    if (TrainingLog.empty())
+      return false;
+    std::error_code EC;
+    auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
+    if (EC) {
+      M.getContext().emitError(EC.message() + ":" + TrainingLog);
+      return false;
     }
-
-    Logger *Log = nullptr;
-    if (!TrainingLog.empty()) {
-      std::vector<TensorSpec> LFS = InputFeatures;
-      if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
-        append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
-      // We always log the output; in particular, if we're not evaluating, we
-      // don't have an output spec json file. That's why we handle the
-      // 'normal' output separately.
-      LFS.push_back(Output);
-      auto I = LogMap.insert(std::make_pair(
-          MF.getFunction().getName(),
-          std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
-      assert(I.second);
-      Log = I.first->second.get();
+
+    Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
+                                   /*IncludeReward*/ true);
+    return false;
+  }
+
+  std::unique_ptr<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    if (!Runner)
+      return nullptr;
+    if (Log) {
+      Log->switchContext(MF.getName());
     }
 
     return std::make_unique<DevelopmentModePriorityAdvisor>(
-        MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log);
+        MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log.get());
   }
 
   std::unique_ptr<MLModelRunner> Runner;
-  StringMap<std::unique_ptr<Logger>> LogMap;
+  std::unique_ptr<Logger> Log;
 };
 #endif //#ifdef LLVM_HAVE_TFLITE
 
 } // namespace llvm
 
 RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() {
   return new ReleaseModePriorityAdvisorAnalysis();
 }
if (isa<ModelUnderTrainingRunner>(getRunner())) {
     Prio = MLPriorityAdvisor::getPriorityImpl(LI);
   } else {
     Prio = getDefaultAdvisor().getPriority(LI);
   }
 
   if (TrainingLog.empty())
     return Prio;
 
+  // TODO(mtrofin): when we support optional rewards, this can go away. In the
+  // meantime, we log the "pretend" reward (0) for the previous observation
+  // before starting a new one.
+  if (Log->hasObservationInProgress())
+    Log->logReward<float>(0.0);
+
+  Log->startObservation();
   size_t CurrentFeature = 0;
   for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
-    Log->logSpecifiedTensorValue(
-        CurrentFeature, reinterpret_cast<const char *>(
+    Log->logTensorValue(CurrentFeature,
+                        reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
 
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
          ++I, ++CurrentFeature)
-      Log->logSpecifiedTensorValue(
+      Log->logTensorValue(
           CurrentFeature,
           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
   }
 
   float Ret = static_cast<float>(Prio);
-  Log->logFloatValue(CurrentFeature, &Ret);
+  Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
+  Log->endObservation();
 
   return static_cast<unsigned>(Prio);
 }
 
 #endif // #ifdef LLVM_HAVE_TFLITE

llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll

 ; CHECK-NEXT: observation: 1
 ; CHECK-NEXT: mask:
 ; NOML:      observation: 16
 ; ML:      observation: 83
 ; ML: reward: 37.06
 ; NOML: reward: 36.64
 
 
-; CHECK-TWO-FCTS: context: SyFgetsCopy
+; CHECK-TWO-FCTS: context: SyFgets
 ; CHECK-TWO-FCTS-NEXT: observation: 0
 ; CHECK-TWO-FCTS-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
 ; CHECK-TWO-FCTS: index_to_evict: 12
 ; CHECK-TWO-FCTS: observation: 16
 ; CHECK-TWO-FCTS: reward: 36.64
-; CHECK-TWO-FCTS: context: SyFgets
+; CHECK-TWO-FCTS: context: SyFgetsCopy
 ; CHECK-TWO-FCTS-NEXT: observation: 0
 ; CHECK-TWO-FCTS-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
 ; CHECK-TWO-FCTS: index_to_evict: 12
 ; CHECK-TWO-FCTS: observation: 16
 ; CHECK-TWO-FCTS: reward: 36.64