From 125a2261e7f9f917289f872acff6e91333114361 Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <alexshap@fb.com>
Date: Fri, 29 Jan 2021 16:37:07 -0800
Subject: [PATCH] Add support for reading profile on Mach-O

Summary: Add support for reading profile on Mach-O.

Reviewed By: rafaelauler

Differential Revision: D25777049

fbshipit-source-id: ce9bf5616ca
---
 src/MachORewriteInstance.cpp | 60 ++++++++++++++++++++++++++++++++++++
 src/MachORewriteInstance.h   |  9 ++++++
 src/llvm-bolt.cpp            |  5 +++
 3 files changed, 74 insertions(+)

diff --git a/src/MachORewriteInstance.cpp b/src/MachORewriteInstance.cpp
index 9d62374d6f69..f67ce091bf0c 100644
--- a/src/MachORewriteInstance.cpp
+++ b/src/MachORewriteInstance.cpp
@@ -14,6 +14,7 @@
 #include "BinaryEmitter.h"
 #include "BinaryFunction.h"
 #include "BinaryPassManager.h"
+#include "DataReader.h"
 #include "ExecutableFileMemoryManager.h"
 #include "JumpTable.h"
 #include "Passes/Instrumentation.h"
@@ -64,6 +65,42 @@ MachORewriteInstance::MachORewriteInstance(object::MachOObjectFile *InputFile,
                                DWARFContext::defaultErrorHandler, "",
                                false))) {}
 
+Error MachORewriteInstance::setProfile(StringRef Filename) {
+  if (!sys::fs::exists(Filename))
+    return errorCodeToError(make_error_code(errc::no_such_file_or_directory));
+
+  if (ProfileReader) {
+    // Already exists
+    return make_error<StringError>(
+        Twine("multiple profiles specified: ") + ProfileReader->getFilename() +
+        " and " + Filename, inconvertibleErrorCode());
+  }
+
+  ProfileReader = llvm::make_unique<DataReader>(Filename);
+  return Error::success();
+}
+
+void MachORewriteInstance::preprocessProfileData() {
+  if (!ProfileReader)
+    return;
+  if (auto E = ProfileReader->preprocessProfile(*BC.get()))
+    report_error("cannot pre-process profile", std::move(E));
+}
+
+void MachORewriteInstance::processProfileDataPreCFG() {
+  if (!ProfileReader)
+    return;
+  if (auto E = ProfileReader->readProfilePreCFG(*BC.get()))
+    report_error("cannot read profile pre-CFG", std::move(E));
+}
+
+void MachORewriteInstance::processProfileData() {
+  if (!ProfileReader)
+    return;
+  if (auto E = ProfileReader->readProfile(*BC.get()))
+    report_error("cannot read profile", std::move(E));
+}
+
 void MachORewriteInstance::readSpecialSections() {
   for (const auto &Section : InputFile->sections()) {
     StringRef SectionName;
@@ -252,6 +289,14 @@ void MachORewriteInstance::disassembleFunctions() {
     Function.disassemble();
     if (opts::PrintDisasm)
       Function.print(outs(), "after disassembly", true);
+  }
+}
+
+void MachORewriteInstance::buildFunctionsCFG() {
+  for (auto &BFI : BC->getBinaryFunctions()) {
+    BinaryFunction &Function = BFI.second;
+    if (!Function.isSimple())
+      continue;
     if (!Function.buildCFG(/*AllocId*/ 0)) {
       errs() << "BOLT-WARNING: failed to build CFG for the function "
              << Function << "\n";
@@ -522,12 +567,27 @@ void MachORewriteInstance::adjustCommandLineOptions() {
 
 void MachORewriteInstance::run() {
   adjustCommandLineOptions();
+
   readSpecialSections();
+
   discoverFileObjects();
+
+  preprocessProfileData();
+
   disassembleFunctions();
+
+  processProfileDataPreCFG();
+
+  buildFunctionsCFG();
+
+  processProfileData();
+
   postProcessFunctions();
+
   runOptimizationPasses();
+
   emitAndLink();
+
   rewriteFile();
 }
 
diff --git a/src/MachORewriteInstance.h b/src/MachORewriteInstance.h
index 60a4e557ff79..99e7a8908832 100644
--- a/src/MachORewriteInstance.h
+++ b/src/MachORewriteInstance.h
@@ -15,6 +15,7 @@
 #define LLVM_TOOLS_LLVM_BOLT_MACHO_REWRITE_INSTANCE_H
 
 #include "NameResolver.h"
+#include "ProfileReaderBase.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/Object/MachO.h"
 #include <memory>
@@ -48,6 +49,11 @@ class MachORewriteInstance {
 
   std::unique_ptr<ToolOutputFile> Out;
 
+  std::unique_ptr<ProfileReaderBase> ProfileReader;
+  void preprocessProfileData();
+  void processProfileDataPreCFG();
+  void processProfileData();
+
   static StringRef getOrgSecPrefix() { return ".bolt.org"; }
 
   void mapInstrumentationSection(orc::VModuleKey Key, StringRef SectionName);
@@ -57,6 +63,7 @@ class MachORewriteInstance {
   void readSpecialSections();
   void discoverFileObjects();
   void disassembleFunctions();
+  void buildFunctionsCFG();
   void postProcessFunctions();
   void runOptimizationPasses();
   void emitAndLink();
@@ -68,6 +75,8 @@ class MachORewriteInstance {
   MachORewriteInstance(object::MachOObjectFile *InputFile, StringRef ToolPath);
   ~MachORewriteInstance();
 
+  Error setProfile(StringRef FileName);
+
   /// Run all the necessary steps to read, optimize and rewrite the binary.
   void run();
 };
diff --git a/src/llvm-bolt.cpp b/src/llvm-bolt.cpp
index 0ea26b1fd592..03f40a5a064a 100644
--- a/src/llvm-bolt.cpp
+++ b/src/llvm-bolt.cpp
@@ -323,6 +323,11 @@ int main(int argc, char **argv) {
       RI.run();
     } else if (auto *O = dyn_cast<MachOObjectFile>(&Binary)) {
       MachORewriteInstance MachORI(O, ToolPath);
+
+      if (!opts::InputDataFilename.empty())
+        if (auto E = MachORI.setProfile(opts::InputDataFilename))
+          report_error(opts::InputDataFilename, std::move(E));
+
       MachORI.run();
     } else {
       report_error(opts::InputFilename, object_error::invalid_file_type);