[BOLT] Add method for better function size estimation

Summary:
Add BinaryContext::calculateEmittedSize() that ephemerally emits code
to allow precise estimation of the function size. Relaxation and
macro-op alignment adjustments are taken into account.

(cherry picked from FBD13092139)

bolt/src/BinaryContext.cpp

@@ -15,8 +15,13 @@
 #include "llvm/ADT/Twine.h"
 #include "llvm/DebugInfo/DWARF/DWARFFormValue.h"
 #include "llvm/DebugInfo/DWARF/DWARFUnit.h"
+#include "llvm/MC/MCAssembler.h"
+#include "llvm/MC/MCAsmLayout.h"
 #include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCELFStreamer.h"
+#include "llvm/MC/MCObjectStreamer.h"
 #include "llvm/MC/MCObjectWriter.h"
+#include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/CommandLine.h"
@@ -1149,3 +1154,78 @@ BinaryContext::createInjectedBinaryFunction(const std::string &Name,
   setSymbolToFunctionMap(BF->getSymbol(), BF);
   return BF;
 }
+
+std::pair<size_t, size_t>
+BinaryContext::calculateEmittedSize(BinaryFunction &BF) {
+  // Adjust branch instruction to match the current layout.
+  BF.fixBranches();
+
+  // Create local MC context to isolate the effect of ephemeral code emission.
+  std::unique_ptr<MCObjectFileInfo> LocalMOFI =
+    llvm::make_unique<MCObjectFileInfo>();
+  std::unique_ptr<MCContext> LocalCtx =
+    llvm::make_unique<MCContext>(AsmInfo.get(), MRI.get(), LocalMOFI.get());
+  LocalMOFI->InitMCObjectFileInfo(*TheTriple, /*PIC=*/false, *LocalCtx);
+  auto *MAB = TheTarget->createMCAsmBackend(*STI, *MRI, MCTargetOptions());
+  auto *MCE = TheTarget->createMCCodeEmitter(*MII, *MRI, *LocalCtx);
+  SmallString<256> Code;
+  raw_svector_ostream VecOS(Code);
+
+  std::unique_ptr<MCStreamer> Streamer(TheTarget->createMCObjectStreamer(
+      *TheTriple, *LocalCtx, std::unique_ptr<MCAsmBackend>(MAB), VecOS,
+      std::unique_ptr<MCCodeEmitter>(MCE), *STI,
+      /* RelaxAll */ false,
+      /* IncrementalLinkerCompatible */ false,
+      /* DWARFMustBeAtTheEnd */ false));
+
+  Streamer->InitSections(false);
+
+  auto *Section = LocalMOFI->getTextSection();
+  Section->setHasInstructions(true);
+
+  auto *StartLabel = LocalCtx->getOrCreateSymbol("__hstart");
+  auto *EndLabel = LocalCtx->getOrCreateSymbol("__hend");
+  auto *ColdStartLabel = LocalCtx->getOrCreateSymbol("__cstart");
+  auto *ColdEndLabel = LocalCtx->getOrCreateSymbol("__cend");
+
+  Streamer->SwitchSection(Section);
+  Streamer->EmitLabel(StartLabel);
+  BF.emitBody(*Streamer, /*EmitColdPart = */false, /*EmitCodeOnly = */true);
+  Streamer->EmitLabel(EndLabel);
+
+  if (BF.isSplit()) {
+    auto *ColdSection =
+      LocalCtx->getELFSection(BF.getColdCodeSectionName(),
+                              ELF::SHT_PROGBITS,
+                              ELF::SHF_EXECINSTR | ELF::SHF_ALLOC);
+    ColdSection->setHasInstructions(true);
+
+    Streamer->SwitchSection(ColdSection);
+    Streamer->EmitLabel(ColdStartLabel);
+    BF.emitBody(*Streamer, /*EmitColdPart = */true, /*EmitCodeOnly = */true);
+    Streamer->EmitLabel(ColdEndLabel);
+  }
+
+  // To avoid calling MCObjectStreamer::flushPendingLabels() which is private.
+  Streamer->EmitBytes(StringRef(""));
+
+  auto &Assembler =
+      static_cast<MCObjectStreamer *>(Streamer.get())->getAssembler();
+  MCAsmLayout Layout(Assembler);
+  Assembler.layout(Layout);
+
+  const auto HotSize = Layout.getSymbolOffset(*EndLabel) -
+                       Layout.getSymbolOffset(*StartLabel);
+  const auto ColdSize = BF.isSplit() ? Layout.getSymbolOffset(*ColdEndLabel) -
+                                       Layout.getSymbolOffset(*ColdStartLabel)
+                                     : 0ULL;
+
+  // Clean-up the effect of the code emission.
+  for (const auto &Symbol : Assembler.symbols()) {
+    auto *MutableSymbol = const_cast<MCSymbol *>(&Symbol);
+    MutableSymbol->setUndefined();
+    MutableSymbol->setIsRegistered(false);
+  }
+
+  return std::make_pair(HotSize, ColdSize);
+}

bolt/src/BinaryContext.h

@@ -669,6 +669,13 @@ public:
   static std::vector<BinaryFunction *>
   getSortedFunctions(std::map<uint64_t, BinaryFunction> &BinaryFunctions);
 
+  /// Do the best effort to calculate the size of the function by emitting
+  /// its code, and relaxing branch instructions.
+  ///
+  /// Return the pair where the first size is for the main part, and the second
+  /// size is for the cold one.
+  std::pair<size_t, size_t> calculateEmittedSize(BinaryFunction &BF);
+
   /// Compute the native code size for a range of instructions.
   /// Note: this can be imprecise wrt the final binary since happening prior to
   /// relaxation, as well as wrt the original binary because of opcode

bolt/src/BinaryFunction.cpp

@@ -2639,8 +2639,9 @@ uint64_t BinaryFunction::getEditDistance() const {
                                                  BasicBlocksLayout);
 }
 
-void BinaryFunction::emitBody(MCStreamer &Streamer, bool EmitColdPart) {
-  if (EmitColdPart && hasConstantIsland())
+void BinaryFunction::emitBody(MCStreamer &Streamer, bool EmitColdPart,
+                              bool EmitCodeOnly) {
+  if (!EmitCodeOnly && EmitColdPart && hasConstantIsland())
     duplicateConstantIslands();
 
   for (auto BB : layout()) {
@@ -2672,6 +2673,9 @@ void BinaryFunction::emitBody(MCStreamer &Streamer, bool EmitColdPart) {
     for (auto I = BB->begin(), E = BB->end(); I != E; ++I) {
       auto &Instr = *I;
 
+      if (EmitCodeOnly && BC.MII->get(Instr.getOpcode()).isPseudo())
+        continue;
+
       // Handle pseudo instructions.
       if (BC.MIB->isEHLabel(Instr)) {
         const auto *Label = BC.MIB->getTargetSymbol(Instr);
@@ -2695,7 +2699,7 @@ void BinaryFunction::emitBody(MCStreamer &Streamer, bool EmitColdPart) {
         Streamer.EmitNeverAlignCodeAtEnd(/*Alignment to avoid=*/64);
       }
 
-      if (opts::UpdateDebugSections && UnitLineTable.first) {
+      if (!EmitCodeOnly && opts::UpdateDebugSections && UnitLineTable.first) {
         LastLocSeen = emitLineInfo(Instr.getLoc(), LastLocSeen);
       }
 
@@ -2704,7 +2708,8 @@ void BinaryFunction::emitBody(MCStreamer &Streamer, bool EmitColdPart) {
     }
   }
 
-  emitConstantIslands(Streamer, EmitColdPart);
+  if (!EmitCodeOnly)
+    emitConstantIslands(Streamer, EmitColdPart);
 }
 
 void BinaryFunction::emitBodyRaw(MCStreamer *Streamer) {

bolt/src/BinaryFunction.h

@@ -2074,7 +2074,8 @@ public:
 
   /// Emit function code. The caller is responsible for emitting function
   /// symbol(s) and setting the section to emit the code to.
-  void emitBody(MCStreamer &Streamer, bool EmitColdPart);
+  void emitBody(MCStreamer &Streamer, bool EmitColdPart,
+                bool EmitCodeOnly = false);
 
   /// Emit function as a blob with relocations and labels for relocations.
   void emitBodyRaw(MCStreamer *Streamer);