[BOLT] Track fragment info for all split fragments

To generate all symbols correctly, it is necessary to record the address
of each fragment. This patch moves the address info for the main and
cold fragments from BinaryFunction to FunctionFragment, where this data
is recorded for all fragments.

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D132051

Author: Fabian Parzefall

bolt/include/bolt/Core/BinaryFunction.h

@@ -232,9 +232,6 @@ class BinaryFunction {
   /// Size of the function in the output file.
   uint64_t OutputSize{0};
 
-  /// Offset in the file.
-  uint64_t FileOffset{0};
-
   /// Maximum size this function is allowed to have.
   uint64_t MaxSize{std::numeric_limits<uint64_t>::max()};
 
@@ -345,14 +342,6 @@ class BinaryFunction {
   /// This attribute is only valid when hasCFG() == true.
   bool HasCanonicalCFG{true};
 
-  /// The address for the code for this function in codegen memory.
-  /// Used for functions that are emitted in a dedicated section with a fixed
-  /// address. E.g. for functions that are overwritten in-place.
-  uint64_t ImageAddress{0};
-
-  /// The size of the code in memory.
-  uint64_t ImageSize{0};
-
   /// Name for the section this function code should reside in.
   std::string CodeSectionName;
 
@@ -1060,7 +1049,9 @@ class BinaryFunction {
   }
 
   /// Return offset of the function body in the binary file.
-  uint64_t getFileOffset() const { return FileOffset; }
+  uint64_t getFileOffset() const {
+    return getLayout().getMainFragment().getFileOffset();
+  }
 
   /// Return (original) byte size of the function.
   uint64_t getSize() const { return Size; }
@@ -1698,7 +1689,7 @@ class BinaryFunction {
                                              int64_t NewOffset);
 
   BinaryFunction &setFileOffset(uint64_t Offset) {
-    FileOffset = Offset;
+    getLayout().getMainFragment().setFileOffset(Offset);
     return *this;
   }
 
@@ -1785,20 +1776,24 @@ class BinaryFunction {
   uint16_t getMaxColdAlignmentBytes() const { return MaxColdAlignmentBytes; }
 
   BinaryFunction &setImageAddress(uint64_t Address) {
-    ImageAddress = Address;
+    getLayout().getMainFragment().setImageAddress(Address);
     return *this;
   }
 
   /// Return the address of this function' image in memory.
-  uint64_t getImageAddress() const { return ImageAddress; }
+  uint64_t getImageAddress() const {
+    return getLayout().getMainFragment().getImageAddress();
+  }
 
   BinaryFunction &setImageSize(uint64_t Size) {
-    ImageSize = Size;
+    getLayout().getMainFragment().setImageSize(Size);
     return *this;
   }
 
   /// Return the size of this function' image in memory.
-  uint64_t getImageSize() const { return ImageSize; }
+  uint64_t getImageSize() const {
+    return getLayout().getMainFragment().getImageSize();
+  }
 
   /// Return true if the function is a secondary fragment of another function.
   bool isFragment() const { return IsFragment; }
@@ -2302,33 +2297,6 @@ class BinaryFunction {
   bool isAArch64Veneer() const;
 
   virtual ~BinaryFunction();
-
-  /// Info for fragmented functions.
-  class FragmentInfo {
-  private:
-    uint64_t Address{0};
-    uint64_t ImageAddress{0};
-    uint64_t ImageSize{0};
-    uint64_t FileOffset{0};
-
-  public:
-    uint64_t getAddress() const { return Address; }
-    uint64_t getImageAddress() const { return ImageAddress; }
-    uint64_t getImageSize() const { return ImageSize; }
-    uint64_t getFileOffset() const { return FileOffset; }
-
-    void setAddress(uint64_t VAddress) { Address = VAddress; }
-    void setImageAddress(uint64_t Address) { ImageAddress = Address; }
-    void setImageSize(uint64_t Size) { ImageSize = Size; }
-    void setFileOffset(uint64_t Offset) { FileOffset = Offset; }
-  };
-
-  /// Cold fragment of the function.
-  FragmentInfo ColdFragment;
-
-  FragmentInfo &cold() { return ColdFragment; }
-
-  const FragmentInfo &cold() const { return ColdFragment; }
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS,

bolt/include/bolt/Core/FunctionLayout.h

@@ -83,6 +83,20 @@ class FunctionFragment {
   unsigned StartIndex;
   unsigned Size = 0;
 
+  /// Output address for the fragment.
+  uint64_t Address = 0;
+
+  /// The address for the code for this fragment in codegen memory. Used for
+  /// functions that are emitted in a dedicated section with a fixed address,
+  /// e.g. for functions that are overwritten in-place.
+  uint64_t ImageAddress = 0;
+
+  /// The size of the code in memory.
+  uint64_t ImageSize = 0;
+
+  /// Offset in the file.
+  uint64_t FileOffset = 0;
+
   FunctionFragment(FunctionLayout &Layout, FragmentNum Num);
   FunctionFragment(const FunctionFragment &) = default;
   FunctionFragment(FunctionFragment &&) = default;
@@ -97,6 +111,15 @@ class FunctionFragment {
   }
   bool isSplitFragment() const { return !isMainFragment(); }
 
+  uint64_t getAddress() const { return Address; }
+  void setAddress(uint64_t Value) { Address = Value; }
+  uint64_t getImageAddress() const { return ImageAddress; }
+  void setImageAddress(uint64_t Address) { ImageAddress = Address; }
+  uint64_t getImageSize() const { return ImageSize; }
+  void setImageSize(uint64_t Size) { ImageSize = Size; }
+  uint64_t getFileOffset() const { return FileOffset; }
+  void setFileOffset(uint64_t Offset) { FileOffset = Offset; }
+
   unsigned size() const { return Size; };
   bool empty() const { return size() == 0; };
   iterator begin();

bolt/lib/Core/BinaryFunction.cpp

@@ -425,19 +425,19 @@ void BinaryFunction::print(raw_ostream &OS, std::string Annotation,
       Sep = "\n                ";
     }
   }
-  OS << "\n  Number      : "   << FunctionNumber
-     << "\n  State       : "   << CurrentState
-     << "\n  Address     : 0x" << Twine::utohexstr(Address)
-     << "\n  Size        : 0x" << Twine::utohexstr(Size)
-     << "\n  MaxSize     : 0x" << Twine::utohexstr(MaxSize)
-     << "\n  Offset      : 0x" << Twine::utohexstr(FileOffset)
-     << "\n  Section     : "   << SectionName
-     << "\n  Orc Section : "   << getCodeSectionName()
-     << "\n  LSDA        : 0x" << Twine::utohexstr(getLSDAAddress())
-     << "\n  IsSimple    : "   << IsSimple
-     << "\n  IsMultiEntry: "   << isMultiEntry()
-     << "\n  IsSplit     : "   << isSplit()
-     << "\n  BB Count    : "   << size();
+  OS << "\n  Number      : " << FunctionNumber;
+  OS << "\n  State       : " << CurrentState;
+  OS << "\n  Address     : 0x" << Twine::utohexstr(Address);
+  OS << "\n  Size        : 0x" << Twine::utohexstr(Size);
+  OS << "\n  MaxSize     : 0x" << Twine::utohexstr(MaxSize);
+  OS << "\n  Offset      : 0x" << Twine::utohexstr(getFileOffset());
+  OS << "\n  Section     : " << SectionName;
+  OS << "\n  Orc Section : " << getCodeSectionName();
+  OS << "\n  LSDA        : 0x" << Twine::utohexstr(getLSDAAddress());
+  OS << "\n  IsSimple    : " << IsSimple;
+  OS << "\n  IsMultiEntry: " << isMultiEntry();
+  OS << "\n  IsSplit     : " << isSplit();
+  OS << "\n  BB Count    : " << size();
 
   if (HasFixedIndirectBranch)
     OS << "\n  HasFixedIndirectBranch : true";
@@ -473,8 +473,8 @@ void BinaryFunction::print(raw_ostream &OS, std::string Annotation,
     for (const BinaryBasicBlock *BB : Layout.blocks())
       OS << LS << BB->getName();
   }
-  if (ImageAddress)
-    OS << "\n  Image       : 0x" << Twine::utohexstr(ImageAddress);
+  if (getImageAddress())
+    OS << "\n  Image       : 0x" << Twine::utohexstr(getImageAddress());
   if (ExecutionCount != COUNT_NO_PROFILE) {
     OS << "\n  Exec Count  : " << ExecutionCount;
     OS << "\n  Profile Acc : " << format("%.1f%%", ProfileMatchRatio * 100.0f);
@@ -4062,7 +4062,7 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
       setOutputDataAddress(BaseAddress + DataOffset);
     }
     if (isSplit()) {
-      for (const FunctionFragment &FF : getLayout().getSplitFragments()) {
+      for (FunctionFragment &FF : getLayout().getSplitFragments()) {
         ErrorOr<BinarySection &> ColdSection =
             getCodeSection(FF.getFragmentNum());
         // If fragment is empty, cold section might not exist
@@ -4084,8 +4084,8 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
         const uint64_t ColdStartOffset =
             Layout.getSymbolOffset(*ColdStartSymbol);
         const uint64_t ColdEndOffset = Layout.getSymbolOffset(*ColdEndSymbol);
-        cold().setAddress(ColdBaseAddress + ColdStartOffset);
-        cold().setImageSize(ColdEndOffset - ColdStartOffset);
+        FF.setAddress(ColdBaseAddress + ColdStartOffset);
+        FF.setImageSize(ColdEndOffset - ColdStartOffset);
         if (hasConstantIsland()) {
           const uint64_t DataOffset =
               Layout.getSymbolOffset(*getFunctionColdConstantIslandLabel());
@@ -4112,44 +4112,39 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
   if (getLayout().block_empty())
     return;
 
-  assert((getLayout().isHotColdSplit() ||
-          (cold().getAddress() == 0 && cold().getImageSize() == 0 &&
-           BC.HasRelocations)) &&
-         "Function must be split two ways or cold fragment must have no "
-         "address (only in relocation mode)");
-
-  BinaryBasicBlock *PrevBB = nullptr;
   for (FunctionFragment &FF : getLayout().fragments()) {
+    if (FF.empty())
+      continue;
+
     const uint64_t FragmentBaseAddress =
         getCodeSection(isSimple() ? FF.getFragmentNum() : FragmentNum::main())
             ->getOutputAddress();
+
+    BinaryBasicBlock *PrevBB = nullptr;
     for (BinaryBasicBlock *const BB : FF) {
       assert(BB->getLabel()->isDefined() && "symbol should be defined");
       if (!BC.HasRelocations) {
-        if (BB->isSplit()) {
-          assert(FragmentBaseAddress == cold().getAddress());
-        } else {
+        if (BB->isSplit())
+          assert(FragmentBaseAddress == FF.getAddress());
+        else
           assert(FragmentBaseAddress == getOutputAddress());
-        }
       }
+
       const uint64_t BBOffset = Layout.getSymbolOffset(*BB->getLabel());
       const uint64_t BBAddress = FragmentBaseAddress + BBOffset;
       BB->setOutputStartAddress(BBAddress);
 
-      if (PrevBB) {
-        uint64_t PrevBBEndAddress = BBAddress;
-        if (BB->isSplit() != PrevBB->isSplit())
-          PrevBBEndAddress = getOutputAddress() + getOutputSize();
-        PrevBB->setOutputEndAddress(PrevBBEndAddress);
-      }
+      if (PrevBB)
+        PrevBB->setOutputEndAddress(BBAddress);
       PrevBB = BB;
 
       BB->updateOutputValues(Layout);
     }
+
+    PrevBB->setOutputEndAddress(PrevBB->isSplit()
+                                    ? FF.getAddress() + FF.getImageSize()
+                                    : getOutputAddress() + getOutputSize());
   }
-  PrevBB->setOutputEndAddress(PrevBB->isSplit()
-                                  ? cold().getAddress() + cold().getImageSize()
-                                  : getOutputAddress() + getOutputSize());
 }
 
 DebugAddressRangesVector BinaryFunction::getOutputAddressRanges() const {
@@ -4165,8 +4160,9 @@ DebugAddressRangesVector BinaryFunction::getOutputAddressRanges() const {
                             getOutputAddress() + getOutputSize());
   if (isSplit()) {
     assert(isEmitted() && "split function should be emitted");
-    OutputRanges.emplace_back(cold().getAddress(),
-                              cold().getAddress() + cold().getImageSize());
+    for (const FunctionFragment &FF : getLayout().getSplitFragments())
+      OutputRanges.emplace_back(FF.getAddress(),
+                                FF.getAddress() + FF.getImageSize());
   }
 
   if (isSimple())

bolt/lib/Core/FunctionLayout.cpp

@@ -204,10 +204,14 @@ bool FunctionLayout::update(const ArrayRef<BinaryBasicBlock *> NewLayout) {
 
 void FunctionLayout::clear() {
   Blocks = BasicBlockListType();
-  for (FunctionFragment *const FF : Fragments)
-    delete FF;
-  Fragments = FragmentListType();
-  addFragment();
+  // If the binary does not have relocations and is not split, the function will
+  // be written to the output stream at its original file offset (see
+  // `RewriteInstance::rewriteFile`). Hence, when the layout is cleared, retain
+  // the main fragment, so that this information is not lost.
+  std::for_each(Fragments.begin() + 1, Fragments.end(),
+                [](FunctionFragment *const FF) { delete FF; });
+  Fragments = FragmentListType{Fragments.front()};
+  getMainFragment().Size = 0;
 }
 
 const BinaryBasicBlock *

bolt/lib/Profile/BoltAddressTranslation.cpp

@@ -73,29 +73,27 @@ void BoltAddressTranslation::write(const BinaryContext &BC, raw_ostream &OS) {
     LLVM_DEBUG(dbgs() << "Function name: " << Function.getPrintName() << "\n");
     LLVM_DEBUG(dbgs() << " Address reference: 0x"
                       << Twine::utohexstr(Function.getOutputAddress()) << "\n");
+
     MapTy Map;
-    const bool IsSplit = Function.isSplit();
-    for (const BinaryBasicBlock *const BB : Function.getLayout().blocks()) {
-      if (IsSplit && BB->isCold())
-        break;
+    for (const BinaryBasicBlock *const BB :
+         Function.getLayout().getMainFragment())
       writeEntriesForBB(Map, *BB, Function.getOutputAddress());
-    }
-    Maps.insert(std::pair<uint64_t, MapTy>(Function.getOutputAddress(), Map));
+    Maps.emplace(Function.getOutputAddress(), std::move(Map));
 
-    if (!IsSplit)
+    if (!Function.isSplit())
       continue;
 
-    // Cold map
-    Map.clear();
+    // Split maps
     LLVM_DEBUG(dbgs() << " Cold part\n");
-    for (const BinaryBasicBlock *const BB : Function.getLayout().blocks()) {
-      if (!BB->isCold())
-        continue;
-      writeEntriesForBB(Map, *BB, Function.cold().getAddress());
+    for (const FunctionFragment &FF :
+         Function.getLayout().getSplitFragments()) {
+      Map.clear();
+      for (const BinaryBasicBlock *const BB : FF)
+        writeEntriesForBB(Map, *BB, FF.getAddress());
+
+      Maps.emplace(FF.getAddress(), std::move(Map));
+      ColdPartSource.emplace(FF.getAddress(), Function.getOutputAddress());
     }
-    Maps.insert(std::pair<uint64_t, MapTy>(Function.cold().getAddress(), Map));
-    ColdPartSource.insert(std::pair<uint64_t, uint64_t>(
-        Function.cold().getAddress(), Function.getOutputAddress()));
   }
 
   const uint32_t NumFuncs = Maps.size();