[NFC][MLInliner] Set up the logger outside the development mode advisor

[mtrofin]

Summary

This allows us to subsequently configure the logger for the case when we
use a model evaluator and want to log additional outputs.

[yundiqian]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

400–401

swap IsDoingInference and Logger so that the sequence matches?

[MaskRay]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

371

A common variable name of type error_code is EC. You may consider it

[mtrofin]

feedback

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

371

done

400–401

done

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

 /// Tensorflow SequenceExample
 /// (https://www.tensorflow.org/api_docs/python/tf/train/SequenceExample)
 /// protobuf (https://developers.google.com/protocol-buffers).
 /// Because this is a protobuf, we cannot just stream the events as they come.
 /// Internally, TrainingLogger stores data in column-major format, because that
 /// lines up with how TF SequenceExample represents it.
 class TrainingLogger final {
 public:
-  TrainingLogger();
+  TrainingLogger(StringRef LogFileName);
 
   /// Log one inlining event.
   void logInlineEvent(const InlineEvent &Event,
                       const MLModelRunner &ModelRunner);
 
   /// Print the stored tensors.
-  void print(raw_fd_ostream &OutFile);
+  void print();
 
 private:
   /// Write the values of one tensor as a list.
   template <typename T>
   void writeTensorValues(raw_fd_ostream &OutFile, const char *TensorData,
                          size_t ElemCount) const {
     OutFile << "[";
     const T *TypedData = reinterpret_cast<const T *>(TensorData);
for (const char *P = TensorData,
       OutFile << "      feature: { " << FieldName << ": { value: ";
       ValueWriter(P);
       OutFile << " } }\n";
     }
     OutFile << "    }\n";
     OutFile << "  }\n";
   }
 
+  StringRef LogFileName;
   std::vector<InlineFeatures> Features;
   std::vector<int64_t> DefaultDecisions;
   std::vector<int64_t> Decisions;
   std::vector<bool> Effects;
   std::vector<int64_t> Rewards;
 };
 
 /// An extension of the MLInlineAdvisor for the 'development' mode, targeting

 /// release mode. The expectation is that a well-trained model provides a good
 /// policy over a sufficiently diverse codebase, over many changes (i.e.
 /// training happens seldom).
 class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor {
 public:
   DevelopmentModeMLInlineAdvisor(
       Module &M, ModuleAnalysisManager &MAM,
       std::unique_ptr<MLModelRunner> ModelRunner,
-      std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference);
+      std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference,
+      std::unique_ptr<TrainingLogger> Logger);
 
   size_t getTotalSizeEstimate();
 
   virtual ~DevelopmentModeMLInlineAdvisor();
   void updateNativeSizeEstimate(int64_t Change) { CurrentNativeSize += Change; }
   void resetNativeSize(Function *F) {
     FAM.invalidate<InlineSizeEstimatorAnalysis>(*F);
   }
 
   std::unique_ptr<MLInlineAdvice>
   getMandatoryAdvice(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
   std::unique_ptr<MLInlineAdvice>
   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
 
   size_t getNativeSizeEstimate(const Function &F) const;
 
 private:
-  bool isLogging() const { return !TrainingLog.empty(); }
+  bool isLogging() const { return !!Logger; }
 
   std::function<bool(CallBase &)> GetDefaultAdvice;
-  TrainingLogger Logger;
+  std::unique_ptr<TrainingLogger> Logger;
   const bool IsDoingInference;
 
   const int32_t InitialNativeSize;
   int32_t CurrentNativeSize = 0;
 };
 
 /// A variant of MLInlineAdvice that tracks all non-trivial inlining
 /// decisions, for training/logging.
private:
   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() {
+TrainingLogger::TrainingLogger(StringRef LogFileName)
+    : LogFileName(LogFileName) {
   for (size_t I = 0; I < NumberOfFeatures; ++I) {
     Features.push_back(InlineFeatures());
   }
 }
 
 /// Log one inlining event.
 void TrainingLogger::logInlineEvent(const InlineEvent &Event,
                                     const MLModelRunner &ModelRunner) {
   for (size_t I = 0; I < NumberOfFeatures; ++I) {
     Features[I].push_back(ModelRunner.getFeature(I));
   }
   Decisions.push_back(Event.AdvisedDecision);
   Effects.push_back(Event.Effect);
   Rewards.push_back(Event.Reward);
   DefaultDecisions.push_back(Event.DefaultDecision);
 }
 
-void TrainingLogger::print(raw_fd_ostream &OutFile) {
+void TrainingLogger::print() {
+  std::error_code ErrorCode;
+  raw_fd_ostream OutFile(LogFileName, ErrorCode);
   size_t NumberOfRecords = Decisions.size();
   if (NumberOfRecords == 0)
     return;
 
   OutFile << "feature_lists: {\n";
   for (size_t I = 0; I < Features.size(); ++I)
     writeTensorsAsFeatureLists(
         OutFile, TensorSpec::createSpec<int64_t>(FeatureNameMap.at(I), {1}),
writeTensorsAsFeatureLists(OutFile,
                              Rewards.data(), NumberOfRecords);
 
   OutFile << "}\n";
 }
 
 DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
     Module &M, ModuleAnalysisManager &MAM,
     std::unique_ptr<MLModelRunner> ModelRunner,
-    std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference)
+    std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference,
+    std::unique_ptr<TrainingLogger> Logger)
     : MLInlineAdvisor(M, MAM, std::move(ModelRunner)),
-      GetDefaultAdvice(GetDefaultAdvice), IsDoingInference(IsDoingInference),
+      GetDefaultAdvice(GetDefaultAdvice), Logger(std::move(Logger)),
+      IsDoingInference(IsDoingInference),
       InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0),
       CurrentNativeSize(InitialNativeSize) {
   // We cannot have the case of neither inference nor logging.
   assert(IsDoingInference || isLogging());
 }
 
 DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
-  if (TrainingLog.empty())
-    return;
-  std::error_code ErrorCode;
-  raw_fd_ostream OutFile(TrainingLog, ErrorCode);
-  Logger.print(OutFile);
+  if (isLogging())
+    Logger->print();
 }
 
 size_t
 DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
   auto &R =
       FAM.getResult<InlineSizeEstimatorAnalysis>(const_cast<Function &>(F));
   if (!R) {
     F.getParent()->getContext().emitError(
         "Native size estimator is not present.");
     return 0;
   }
   return *R;
 }
 
 std::unique_ptr<MLInlineAdvice>
 DevelopmentModeMLInlineAdvisor::getMandatoryAdvice(
     CallBase &CB, OptimizationRemarkEmitter &ORE) {
   if (!isLogging())
     return MLInlineAdvisor::getMandatoryAdvice(CB, ORE);
   return std::make_unique<LoggingMLInlineAdvice>(
       /*Advisor=*/this,
-      /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/true, /*Logger=*/Logger,
+      /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/true, /*Logger=*/*Logger,
       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
       /*CalleeSizeEstimateBefore=*/
       getNativeSizeEstimate(*CB.getCalledFunction()),
       /*DefaultDecision=*/true, /*Mandatory*/ true);
 }
 
 std::unique_ptr<MLInlineAdvice>
 DevelopmentModeMLInlineAdvisor::getAdviceFromModel(
     CallBase &CB, OptimizationRemarkEmitter &ORE) {
   if (IsDoingInference && !isLogging())
     return MLInlineAdvisor::getAdviceFromModel(CB, ORE);
 
   bool DefaultAdvice = GetDefaultAdvice(CB);
   auto Recommendation = IsDoingInference ? ModelRunner->run() : DefaultAdvice;
   return std::make_unique<LoggingMLInlineAdvice>(
       /*Advisor=*/this,
       /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation,
-      /*Logger=*/Logger,
+      /*Logger=*/*Logger,
       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
       /*CalleeSizeEstimateBefore=*/
       getNativeSizeEstimate(*CB.getCalledFunction()),
       /*DefaultDecision=*/DefaultAdvice);
 }
 
 size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() {
   size_t Ret = 0;
else {
     Runner = std::make_unique<ModelUnderTrainingRunner>(
         Ctx, TFModelUnderTrainingPath);
     if (!Runner) {
       Ctx.emitError("Could not load the policy model from the provided path");
       return nullptr;
     }
     IsDoingInference = true;
   }
+  std::unique_ptr<TrainingLogger> Logger;
+  if (!TrainingLog.empty())
+    Logger = std::make_unique<TrainingLogger>(TrainingLog);
+
   return std::make_unique<DevelopmentModeMLInlineAdvisor>(
-      M, MAM, std::move(Runner), GetDefaultAdvice, IsDoingInference);
+      M, MAM, std::move(Runner), GetDefaultAdvice, IsDoingInference,
+      std::move(Logger));
 }
 #endif // defined(LLVM_HAVE_TF_API)