[BOLT] Output basic YAML profile in BAT mode

Relax assumptions that YAML output is not supported in BAT mode.
Set up basic infrastructure for emitting YAML for functions not covered
by BAT, such as from `.bolt.org.text` section (code identical to input binary
sans external refs), or non-rewritten functions in non-relocation mode (where
the function stays in the same section but BAT mapping is not emitted).

This diff only produces YAML profile for non-BAT functions (skipped,
non-simple). YAML profile for BAT functions is added in follow-up diffs:
- #76911 emits YAML profile with
  internal control flow information only (branch profile),
- #76896 adds cross-function profile
  (calls profile).

Test Plan: Added bolt/test/X86/bolt-address-translation-yaml.test

Reviewers: ayermolo, dcci, maksfb, rafaelauler

Reviewed By: rafaelauler

Pull Request: #76910

Author: aaupov

bolt/docs/BAT.md

@@ -14,9 +14,8 @@ binary onto the original binary.
 # Usage
 `--enable-bat` flag controls the generation of BAT section. Sampled profile
 needs to be passed along with the optimized binary containing BAT section to
-`perf2bolt` which reads BAT section and produces fdata profile for the original
-binary. Note that YAML profile generation is not supported since BAT doesn't
-contain the metadata for input functions.
+`perf2bolt` which reads BAT section and produces profile for the original
+binary.
 
 # Internals
 ## Section contents

bolt/include/bolt/Profile/BoltAddressTranslation.h

@@ -122,6 +122,9 @@ class BoltAddressTranslation {
   /// Returns BF hash by function output address (after BOLT).
   size_t getBFHash(uint64_t OutputAddress) const;
 
+  /// True if a given \p Address is a function with translation table entry.
+  bool isBATFunction(uint64_t Address) const { return Maps.count(Address); }
+
 private:
   /// Helper to update \p Map by inserting one or more BAT entries reflecting
   /// \p BB for function located at \p FuncAddress. At least one entry will be

bolt/include/bolt/Profile/DataAggregator.h

@@ -463,6 +463,10 @@ class DataAggregator : public DataReader {
   /// Dump data structures into a file readable by llvm-bolt
   std::error_code writeAggregatedFile(StringRef OutputFilename) const;
 
+  /// Dump translated data structures into YAML
+  std::error_code writeBATYAML(BinaryContext &BC,
+                               StringRef OutputFilename) const;
+
   /// Filter out binaries based on PID
   void filterBinaryMMapInfo();
 

bolt/lib/Profile/DataAggregator.cpp

@@ -16,6 +16,7 @@
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Profile/BoltAddressTranslation.h"
 #include "bolt/Profile/Heatmap.h"
+#include "bolt/Profile/YAMLProfileWriter.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "bolt/Utils/Utils.h"
 #include "llvm/ADT/STLExtras.h"
@@ -85,6 +86,7 @@ MaxSamples("max-samples",
   cl::cat(AggregatorCategory));
 
 extern cl::opt<opts::ProfileFormatKind> ProfileFormat;
+extern cl::opt<std::string> SaveProfile;
 
 cl::opt<bool> ReadPreAggregated(
     "pa", cl::desc("skip perf and read data from a pre-aggregated file format"),
@@ -594,10 +596,21 @@ Error DataAggregator::readProfile(BinaryContext &BC) {
     convertBranchData(Function);
   }
 
-  if (opts::AggregateOnly &&
-      opts::ProfileFormat == opts::ProfileFormatKind::PF_Fdata) {
-    if (std::error_code EC = writeAggregatedFile(opts::OutputFilename))
-      report_error("cannot create output data file", EC);
+  if (opts::AggregateOnly) {
+    if (opts::ProfileFormat == opts::ProfileFormatKind::PF_Fdata)
+      if (std::error_code EC = writeAggregatedFile(opts::OutputFilename))
+        report_error("cannot create output data file", EC);
+
+    // BAT YAML is handled by DataAggregator since normal YAML output requires
+    // CFG which is not available in BAT mode.
+    if (usesBAT()) {
+      if (opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML)
+        if (std::error_code EC = writeBATYAML(BC, opts::OutputFilename))
+          report_error("cannot create output data file", EC);
+      if (!opts::SaveProfile.empty())
+        if (std::error_code EC = writeBATYAML(BC, opts::SaveProfile))
+          report_error("cannot create output data file", EC);
+    }
   }
 
   return Error::success();
@@ -2258,6 +2271,53 @@ DataAggregator::writeAggregatedFile(StringRef OutputFilename) const {
   return std::error_code();
 }
 
+std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
+                                             StringRef OutputFilename) const {
+  std::error_code EC;
+  raw_fd_ostream OutFile(OutputFilename, EC, sys::fs::OpenFlags::OF_None);
+  if (EC)
+    return EC;
+
+  yaml::bolt::BinaryProfile BP;
+
+  // Fill out the header info.
+  BP.Header.Version = 1;
+  BP.Header.FileName = std::string(BC.getFilename());
+  std::optional<StringRef> BuildID = BC.getFileBuildID();
+  BP.Header.Id = BuildID ? std::string(*BuildID) : "<unknown>";
+  BP.Header.Origin = std::string(getReaderName());
+  // Only the input binary layout order is supported.
+  BP.Header.IsDFSOrder = false;
+  // FIXME: Need to match hash function used to produce BAT hashes.
+  BP.Header.HashFunction = HashFunction::Default;
+
+  ListSeparator LS(",");
+  raw_string_ostream EventNamesOS(BP.Header.EventNames);
+  for (const StringMapEntry<std::nullopt_t> &EventEntry : EventNames)
+    EventNamesOS << LS << EventEntry.first().str();
+
+  BP.Header.Flags = opts::BasicAggregation ? BinaryFunction::PF_SAMPLE
+                                           : BinaryFunction::PF_LBR;
+
+  if (!opts::BasicAggregation) {
+    // Convert profile for functions not covered by BAT
+    for (auto &BFI : BC.getBinaryFunctions()) {
+      BinaryFunction &Function = BFI.second;
+      if (!Function.hasProfile())
+        continue;
+      if (BAT->isBATFunction(Function.getAddress()))
+        continue;
+      BP.Functions.emplace_back(
+          YAMLProfileWriter::convert(Function, /*UseDFS=*/false));
+    }
+  }
+
+  // Write the profile.
+  yaml::Output Out(OutFile, nullptr, 0);
+  Out << BP;
+  return std::error_code();
+}
+
 void DataAggregator::dump() const { DataReader::dump(); }
 
 void DataAggregator::dump(const LBREntry &LBR) const {

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -199,10 +199,7 @@ static cl::opt<cl::boolOrDefault> RelocationMode(
     "relocs", cl::desc("use relocations in the binary (default=autodetect)"),
     cl::cat(BoltCategory));
 
-static cl::opt<std::string>
-SaveProfile("w",
-  cl::desc("save recorded profile to a file"),
-  cl::cat(BoltOutputCategory));
+extern cl::opt<std::string> SaveProfile;
 
 static cl::list<std::string>
 SkipFunctionNames("skip-funcs",
@@ -732,6 +729,13 @@ Error RewriteInstance::run() {
   // Skip disassembling if we have a translation table and we are running an
   // aggregation job.
   if (opts::AggregateOnly && BAT->enabledFor(InputFile)) {
+    // YAML profile in BAT mode requires CFG for .bolt.org.text functions
+    if (!opts::SaveProfile.empty() ||
+        opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML) {
+      selectFunctionsToProcess();
+      disassembleFunctions();
+      buildFunctionsCFG();
+    }
     processProfileData();
     return Error::success();
   }
@@ -2027,14 +2031,6 @@ void RewriteInstance::adjustCommandLineOptions() {
 
   if (opts::Lite)
     BC->outs() << "BOLT-INFO: enabling lite mode\n";
-
-  if (!opts::SaveProfile.empty() && BAT->enabledFor(InputFile)) {
-    BC->errs()
-        << "BOLT-ERROR: unable to save profile in YAML format for input "
-           "file processed by BOLT. Please remove -w option and use branch "
-           "profile.\n";
-    exit(1);
-  }
 }
 
 namespace {
@@ -3126,12 +3122,13 @@ void RewriteInstance::processProfileData() {
     }
   }
 
-  if (!opts::SaveProfile.empty()) {
+  if (!opts::SaveProfile.empty() && !BAT->enabledFor(InputFile)) {
     YAMLProfileWriter PW(opts::SaveProfile);
     PW.writeProfile(*this);
   }
   if (opts::AggregateOnly &&
-      opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML) {
+      opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML &&
+      !BAT->enabledFor(InputFile)) {
     YAMLProfileWriter PW(opts::OutputFilename);
     PW.writeProfile(*this);
   }

bolt/lib/Utils/CommandLineOpts.cpp

@@ -162,6 +162,10 @@ cl::opt<ProfileFormatKind> ProfileFormat(
                clEnumValN(PF_YAML, "yaml", "dense YAML representation")),
     cl::ZeroOrMore, cl::Hidden, cl::cat(BoltCategory));
 
+cl::opt<std::string> SaveProfile("w",
+                                 cl::desc("save recorded profile to a file"),
+                                 cl::cat(BoltOutputCategory));
+
 cl::opt<bool> SplitEH("split-eh", cl::desc("split C++ exception handling code"),
                       cl::Hidden, cl::cat(BoltOptCategory));
 

bolt/test/X86/Inputs/blarge_new.preagg.txt

@@ -0,0 +1,81 @@
+B 40164b 401608 109 0
+B 401611 4017e0 115 0
+B 4017f0 401616 117 0
+B 401ba2 4015da 6 0
+B 4015d5 401b60 1 0
+B 40159a 401b60 5 0
+B 401b9d 401b70 615 2
+B 401b90 401b99 344 37
+B 401ba2 40159f 8 0
+B 4015b0 401070 9 0
+B 401544 4014a0 6 0
+B 40188a 401928 5 0
+B 40152a 4014b0 21 0
+B 40169e 40165b 2 0
+B 4014dd 401070 12 1
+B 401509 4014ec 2 2
+B 401510 401030 673 0
+B 4019de 401080 1 0
+B 401500 401070 22 0
+B 401921 4014d6 9 0
+B 4019b3 401080 3 0
+B 40162d 401070 113 0
+B 4014d1 401800 27 0
+B 401a3f 401080 1 0
+B 4018d2 401050 17 0
+B 401664 4017c0 2 0
+B 401680 401070 2 0
+B 4017d0 401669 2 0
+B 4018f7 40190d 9 0
+B 4015bc 401592 6 0
+B 401964 401090 5 0
+B 4015f8 4015cd 1 0
+B 4015ec 401070 6 0
+F 40165b 401664 2
+F 4017c0 4017d0 2
+F 401669 401680 2
+F 40190d 401921 9
+F 4014d6 4014dd 9
+F 401800 4018d2 17
+F 4018d7 4018f7 9
+F 40159f 4015b0 8
+F 401515 401544 6
+F 401070 401500 1
+F 401070 401070 157
+F 4014a0 4014d1 6
+F 401616 40162d 112
+F 4019e3 401a3f 1
+F 4014e2 401500 19
+F 401090 401090 5
+F 401030 401030 673
+F 401505 401510 668
+F 401616 4017f0 2
+F 401070 4015b0 1
+F 4015da 4015ec 6
+F 401b60 401b90 6
+F 4019b8 4019de 1
+F 401969 4019b3 3
+F 401505 401509 2
+F 401515 40152a 21
+F 401592 40159a 4
+F 401050 401050 17
+F 4015cd 4015d5 1
+F 401070 4014dd 1
+F 401b99 401ba2 8
+F 401b70 401b90 326
+F 401b99 401b9d 324
+F 401592 4015bc 1
+F 401608 401611 109
+F 401b70 401b9d 268
+F 4015b5 4015bc 5
+F 401b99 401b90 1
+F 401b70 401ba2 5
+F 401632 40164b 108
+F 401080 401080 5
+F 4014b0 4014d1 21
+F 4017e0 4017f0 115
+F 4015f1 4015f8 1
+F 401685 40169e 2
+F 401928 401964 5
+F 401800 40188a 5
+F 4014ec 401500 2