Link the instrumentation runtime on OSX

Summary: Link the instrumentation runtime on OSX.

(cherry picked from FBD24390019)

bolt/runtime/CMakeLists.txt

@@ -30,7 +30,7 @@ target_include_directories(bolt_rt_hugify PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
 install(TARGETS bolt_rt_instr DESTINATION lib)
 install(TARGETS bolt_rt_hugify DESTINATION lib)
 
-if (CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+if (CMAKE_CXX_COMPILER_ID MATCHES ".*Clang.*")
   add_library(bolt_rt_instr_osx STATIC
     instr.cpp
     ${CMAKE_CURRENT_BINARY_DIR}/config.h

bolt/runtime/common.h

@@ -1,9 +1,26 @@
 #if !defined(__APPLE__)
+
 #include <cstddef>
 #include <cstdint>
+
+#else
+
+typedef __SIZE_TYPE__ size_t;
+#define __SSIZE_TYPE__                                                         \
+  __typeof__(_Generic((__SIZE_TYPE__)0, unsigned long long int                 \
+                      : (long long int)0, unsigned long int                    \
+                      : (long int)0, unsigned int                              \
+                      : (int)0, unsigned short                                 \
+                      : (short)0, unsigned char                                \
+                      : (signed char)0))
+typedef __SSIZE_TYPE__ ssize_t;
+
+typedef unsigned long long uint64_t;
+
 #endif
 
 #include "config.h"
+
 #ifdef HAVE_ELF_H
 #include <elf.h>
 #endif
@@ -46,11 +63,25 @@
   "pop %%rbx\n"                                                                \
   "pop %%rax\n"
 
-#if !defined(__APPLE__)
-
 // Anonymous namespace covering everything but our library entry point
 namespace {
 
+#if defined(__APPLE__)
+
+uint64_t __write(uint64_t fd, const void *buf, uint64_t count) {
+  uint64_t ret;
+  const long write = 0x2000004;
+  __asm__ __volatile__("syscall;\n"
+                       "movq %%rax, %0;\n"
+                       : "=g"(ret)
+                       : /* rax */ "a"(write), /* rdi */ "D"(fd),
+                         /* rsi */ "S"(buf), /* rdx */ "d"(count)
+                       : "memory");
+  return ret;
+}
+
+#else
+
 // We use a stack-allocated buffer for string manipulation in many pieces of
 // this code, including the code that prints each line of the fdata file. This
 // buffer needs to accomodate large function names, but shouldn't be arbitrarily
@@ -330,6 +361,7 @@ class Lock {
 inline uint64_t alignTo(uint64_t Value, uint64_t Align) {
   return (Value + Align - 1) / Align * Align;
 }
-} // anonymous namespace
 
 #endif
+
+} // anonymous namespace

bolt/runtime/instr.cpp

@@ -1467,6 +1467,9 @@ extern "C" void __bolt_instr_fini() {
 
 // On OSX/iOS the final symbol name of an extern "C" function/variable contains
 // one extra leading underscore: _bolt_instr_setup -> __bolt_instr_setup.
-extern "C" __attribute((section("__TEXT,__setup"))) void _bolt_instr_setup() {}
+extern "C" __attribute((section("__TEXT,__setup"))) void _bolt_instr_setup() {
+  const char* Message = "Hello!\n";
+  __write(2, Message, 7);
+}
 
 #endif

bolt/src/ExecutableFileMemoryManager.cpp

@@ -46,23 +46,26 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(intptr_t Size,
     Ret = SectionMemoryManager::allocateCodeSection(Size, Alignment,
                                                     SectionID, SectionName);
   } else {
-    Ret = SectionMemoryManager::allocateDataSection(Size, Alignment,
-                                                    SectionID, SectionName,
-                                                    IsReadOnly);
+    Ret = SectionMemoryManager::allocateDataSection(Size, Alignment, SectionID,
+                                                    SectionName, IsReadOnly);
   }
 
-  const auto Flags = BinarySection::getFlags(IsReadOnly, IsCode, true);
   SmallVector<char, 256> Buf;
-  if (ObjectsLoaded > 0)
-    SectionName = (Twine(SectionName) + ".bolt.extra." + Twine(ObjectsLoaded))
-                      .toStringRef(Buf);
+  if (ObjectsLoaded > 0) {
+    if (BC.isELF()) {
+      SectionName = (Twine(SectionName) + ".bolt.extra." + Twine(ObjectsLoaded))
+                        .toStringRef(Buf);
+    } else if (BC.isMachO()) {
+      assert((SectionName == "__text" || SectionName == "__data" ||
+              SectionName == "__setup" || SectionName == "__cstring") &&
+             "Unexpected section in the instrumentation library");
+      SectionName = ("I" + Twine(SectionName)).toStringRef(Buf);
+    }
+  }
 
-  auto &Section = BC.registerOrUpdateSection(SectionName,
-                                             ELF::SHT_PROGBITS,
-                                             Flags,
-                                             Ret,
-                                             Size,
-                                             Alignment);
+  auto &Section = BC.registerOrUpdateSection(
+      SectionName, ELF::SHT_PROGBITS,
+      BinarySection::getFlags(IsReadOnly, IsCode, true), Ret, Size, Alignment);
   Section.setSectionID(SectionID);
   assert(Section.isAllocatable() &&
          "verify that allocatable is marked as allocatable");
@@ -72,7 +75,6 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(intptr_t Size,
                << " section : " << SectionName
                << " with size " << Size << ", alignment " << Alignment
                << " at 0x" << Ret << ", ID = " << SectionID << "\n");
-
   return Ret;
 }
 

bolt/src/MachORewriteInstance.cpp

@@ -45,6 +45,7 @@ extern cl::opt<bool> PrintReordered;
 extern cl::opt<bool> PrintSections;
 extern cl::opt<bool> PrintDisasm;
 extern cl::opt<bool> PrintCFG;
+extern cl::opt<std::string> RuntimeInstrumentationLib;
 extern cl::opt<unsigned> Verbosity;
 } // namespace opts
 
@@ -286,15 +287,17 @@ void MachORewriteInstance::runOptimizationPasses() {
   Manager.runPasses();
 }
 
-void MachORewriteInstance::mapExtraSections(orc::VModuleKey Key) {
+void MachORewriteInstance::mapInstrumentationSection(orc::VModuleKey Key, StringRef SectionName) {
   if (!opts::Instrument)
     return;
-  ErrorOr<BinarySection &> Counters = BC->getUniqueSectionByName("__counters");
-  if (!Counters) {
-    llvm::errs() << "Cannot find __counters section\n";
+  ErrorOr<BinarySection &> Section = BC->getUniqueSectionByName(SectionName);
+  if (!Section) {
+    llvm::errs() << "Cannot find " + SectionName + " section\n";
     exit(1);
   }
-  OLT->mapSectionAddress(Key, Counters->getSectionID(), Counters->getAddress());
+  if (!Section->hasValidSectionID())
+    return;
+  OLT->mapSectionAddress(Key, Section->getSectionID(), Section->getAddress());
 }
 
 void MachORewriteInstance::mapCodeSections(orc::VModuleKey Key) {
@@ -419,17 +422,47 @@ void MachORewriteInstance::emitAndLink() {
       },
       [&](orc::VModuleKey Key, const object::ObjectFile &Obj,
           const RuntimeDyld::LoadedObjectInfo &) {
-        assert(Key == K && "Linking multiple objects is unsupported");
-        mapCodeSections(Key);
-        mapExtraSections(Key);
+        if (Key == K) {
+          mapCodeSections(Key);
+          mapInstrumentationSection(Key, "__counters");
+        } else {
+          // TODO: Refactor addRuntimeLibSections to work properly on Mach-O
+          // and use it here.
+          mapInstrumentationSection(Key, "I__setup");
+          mapInstrumentationSection(Key, "I__data");
+          mapInstrumentationSection(Key, "I__text");
+          mapInstrumentationSection(Key, "I__cstring");
+        }
       },
       [&](orc::VModuleKey Key) {
-        assert(Key == K && "Linking multiple objects is unsupported");
       }));
 
   OLT->setProcessAllSections(true);
   cantFail(OLT->addObject(K, std::move(ObjectMemBuffer)));
   cantFail(OLT->emitAndFinalize(K));
+
+  if (auto *RtLibrary = BC->getRuntimeLibrary()) {
+    RtLibrary->link(*BC, ToolPath, *ES, *OLT);
+  }
+}
+
+void MachORewriteInstance::writeInstrumentationSection(StringRef SectionName,
+                                                       raw_pwrite_stream &OS) {
+  if (!opts::Instrument)
+    return;
+  ErrorOr<BinarySection &> Section = BC->getUniqueSectionByName(SectionName);
+  if (!Section) {
+    llvm::errs() << "Cannot find " + SectionName + " section\n";
+    exit(1);
+  }
+  if (!Section->hasValidSectionID())
+    return;
+  assert(Section->getInputFileOffset() &&
+         "Section input offset cannot be zero");
+  assert(Section->getAllocAddress() && "Section alloc address cannot be zero");
+  assert(Section->getOutputSize() && "Section output size cannot be zero");
+  OS.pwrite(reinterpret_cast<char *>(Section->getAllocAddress()),
+            Section->getOutputSize(), Section->getInputFileOffset());
 }
 
 void MachORewriteInstance::rewriteFile() {
@@ -459,6 +492,13 @@ void MachORewriteInstance::rewriteFile() {
               Function->getImageSize(), Function->getFileOffset());
   }
 
+  // TODO: Refactor addRuntimeLibSections to work properly on Mach-O and
+  // use it here.
+  writeInstrumentationSection("I__setup", OS);
+  writeInstrumentationSection("I__data", OS);
+  writeInstrumentationSection("I__text", OS);
+  writeInstrumentationSection("I__cstring", OS);
+
   Out->keep();
 }
 
@@ -470,6 +510,7 @@ void MachORewriteInstance::adjustCommandLineOptions() {
     opts::ForcePatch = true;
   opts::JumpTables = JTS_MOVE;
   opts::InstrumentCalls = false;
+  opts::RuntimeInstrumentationLib = "libbolt_rt_instr_osx.a";
 }
 
 void MachORewriteInstance::run() {

bolt/src/MachORewriteInstance.h

@@ -50,7 +50,7 @@ class MachORewriteInstance {
 
   static StringRef getOrgSecPrefix() { return ".bolt.org"; }
 
-  void mapExtraSections(orc::VModuleKey Key);
+  void mapInstrumentationSection(orc::VModuleKey Key, StringRef SectionName);
   void mapCodeSections(orc::VModuleKey Key);
 
   void adjustCommandLineOptions();
@@ -60,6 +60,8 @@ class MachORewriteInstance {
   void postProcessFunctions();
   void runOptimizationPasses();
   void emitAndLink();
+
+  void writeInstrumentationSection(StringRef SectionName, raw_pwrite_stream &OS);
   void rewriteFile();
 
 public:

bolt/src/Passes/Instrumentation.cpp

@@ -580,24 +580,6 @@ void Instrumentation::createAuxiliaryFunctions(BinaryContext &BC) {
       createSimpleFunction("__bolt_instr_default_ind_tailcall_handler",
                            BC.MIB->createInstrumentedNoopIndTailCallHandler());
 
-  // TODO: Remove this code once we start loading the runtime library for OSX.
-  if (BC.isMachO()) {
-    std::vector<MCInst> Instrs(8);
-    for (MCInst &Instruction : Instrs)
-      BC.MIB->createNoop(Instruction);
-    BC.MIB->createReturn(Instrs.back());
-    BinaryFunction *Placeholder = createSimpleFunction(
-        "__bolt_instr_setup_placeholder", std::move(Instrs));
-    ErrorOr<BinarySection &> SetupSection =
-        BC.getUniqueSectionByName("I__setup");
-    if (!SetupSection) {
-      llvm::errs() << "Cannot find I__setup section\n";
-      exit(1);
-    }
-    Placeholder->setOutputAddress(SetupSection->getAddress());
-    Placeholder->setFileOffset(SetupSection->getInputFileOffset());
-    Placeholder->setOriginSection(&*SetupSection);
-  }
 }
 
 void Instrumentation::setupRuntimeLibrary(BinaryContext &BC) {

bolt/src/RuntimeLibs/InstrumentationRuntimeLibrary.cpp

@@ -29,7 +29,7 @@ cl::opt<bool>
                cl::desc("instrument code to generate accurate profile data"),
                cl::ZeroOrMore, cl::cat(BoltOptCategory));
 
-static cl::opt<std::string> RuntimeInstrumentationLib(
+cl::opt<std::string> RuntimeInstrumentationLib(
     "runtime-instrumentation-lib",
     cl::desc("specify file name of the runtime instrumentation library"),
     cl::ZeroOrMore, cl::init("libbolt_rt_instr.a"), cl::cat(BoltOptCategory));
@@ -177,6 +177,9 @@ void InstrumentationRuntimeLibrary::link(BinaryContext &BC, StringRef ToolPath,
   auto LibPath = getLibPath(ToolPath, opts::RuntimeInstrumentationLib);
   loadLibraryToOLT(LibPath, ES, OLT);
 
+  if (BC.isMachO())
+    return;
+
   RuntimeFiniAddress =
       cantFail(OLT.findSymbol("__bolt_instr_fini", false).getAddress());
   if (!RuntimeFiniAddress) {