[mlgo][regalloc] Handle training case when no regalloc happens.

[mtrofin]

Summary

There's an early-exit case for regalloc when we don't even get a chance
to ask for an advisor (priority or eviction), and switch the context.
Then, when we want to log the reward for that function (==the one with
the early exit case), we hit the error case where the function's name
doesn't match the last-seen context.

There are a few possible fixes, one would be to just switch context when
output-ing the reward, which would be correct. This patch opts for the
alternative where we check any loging happened in the first place - just
to re-validate that no function would have been regaloc-ed without first
log-ing its reward.

[MatzeB]

LGTM

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

477

Ah, I see now, why things didn't work for me when I just moved the existing check :)

[jacobhegna]

This resolved my issues in training.

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

public:
 
   void switchContext(StringRef Name);
   void startObservation();
   void endObservation();
   void flush() { OS->flush(); }
 
   const std::string &currentContext() const { return CurrentContext; }
 
+  /// Check if there is at least an observation for `currentContext()`.
   bool hasObservationInProgress() const {
-    return ObservationIDs.find(CurrentContext) != ObservationIDs.end();
+    return hasAnyObservationForContext(CurrentContext);
+  }
+
+  /// Check if there is at least an observation for the context `Ctx`.
+  bool hasAnyObservationForContext(StringRef Ctx) const {
+    return ObservationIDs.find(Ctx) != ObservationIDs.end();
   }
 
   template <typename T> void logReward(T Value) {
     logRewardImpl(reinterpret_cast<const char *>(&Value));
   }
 
   void logTensorValue(size_t FeatureID, const char *RawData) {
     writeTensor(FeatureSpecs[FeatureID], RawData);
   }
 };
 
 } // namespace llvm
 #endif // LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

public:
   }
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
-    if (!Log)
+    if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
       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.");

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

DevelopmentModePriorityAdvisorAnalysis()
       : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Development) {}
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
-    if (!Log)
+    if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
       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.");

llvm/test/CodeGen/MLRegalloc/empty-log.ll

+; REQUIRES: have_tflite
+; REQUIRES: x86_64-linux
+;
+; Check that we can log more than 1 function.
+;
+; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
+; RUN:   -regalloc-training-log=%t1 < %s
+; RUN: FileCheck --input-file %t1 %s
+
+; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-priority-advisor=development \
+; RUN:   -regalloc-priority-training-log=%t2 < %s
+; RUN: FileCheck --input-file %t2 %s
+
+declare void @f();
+
+define void @f1(i64 %lhs, i64 %rhs, i64* %addr) {
+  ret void
+}
+
+define void @f2(i64 %lhs, i64 %rhs, i64* %addr) {
+  %sum = add i64 %lhs, %rhs
+  call void @f();
+  store i64 %sum, i64* %addr
+  ret void
+}
+
+define void @f3(i64 %lhs, i64 %rhs, i64* %addr) {
+  ret void
+}
+
+; CHECK-NOT:  {"context":"f1"}
+; CHECK:      {"context":"f2"}
+; CHECK-NOT:  {"context":"f3"}