Revert "[BOLT] Track fragment info for all split fragments"

This reverts commit 7e25481.

Author: Fabian Parzefall

bolt/include/bolt/Core/BinaryFunction.h

@@ -232,6 +232,9 @@ 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()};
 
@@ -342,6 +345,14 @@ 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;
 
@@ -1049,9 +1060,7 @@ class BinaryFunction {
   }
 
   /// Return offset of the function body in the binary file.
-  uint64_t getFileOffset() const {
-    return getLayout().getMainFragment().getFileOffset();
-  }
+  uint64_t getFileOffset() const { return FileOffset; }
 
   /// Return (original) byte size of the function.
   uint64_t getSize() const { return Size; }
@@ -1689,7 +1698,7 @@ class BinaryFunction {
                                              int64_t NewOffset);
 
   BinaryFunction &setFileOffset(uint64_t Offset) {
-    getLayout().getMainFragment().setFileOffset(Offset);
+    FileOffset = Offset;
     return *this;
   }
 
@@ -1776,24 +1785,20 @@ class BinaryFunction {
   uint16_t getMaxColdAlignmentBytes() const { return MaxColdAlignmentBytes; }
 
   BinaryFunction &setImageAddress(uint64_t Address) {
-    getLayout().getMainFragment().setImageAddress(Address);
+    ImageAddress = Address;
     return *this;
   }
 
   /// Return the address of this function' image in memory.
-  uint64_t getImageAddress() const {
-    return getLayout().getMainFragment().getImageAddress();
-  }
+  uint64_t getImageAddress() const { return ImageAddress; }
 
   BinaryFunction &setImageSize(uint64_t Size) {
-    getLayout().getMainFragment().setImageSize(Size);
+    ImageSize = Size;
     return *this;
   }
 
   /// Return the size of this function' image in memory.
-  uint64_t getImageSize() const {
-    return getLayout().getMainFragment().getImageSize();
-  }
+  uint64_t getImageSize() const { return ImageSize; }
 
   /// Return true if the function is a secondary fragment of another function.
   bool isFragment() const { return IsFragment; }
@@ -2297,6 +2302,33 @@ 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,20 +83,6 @@ 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;
@@ -111,15 +97,6 @@ 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;
-  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();
+  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();
 
   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 (getImageAddress())
-    OS << "\n  Image       : 0x" << Twine::utohexstr(getImageAddress());
+  if (ImageAddress)
+    OS << "\n  Image       : 0x" << Twine::utohexstr(ImageAddress);
   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 (FunctionFragment &FF : getLayout().getSplitFragments()) {
+      for (const 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);
-        FF.setAddress(ColdBaseAddress + ColdStartOffset);
-        FF.setImageSize(ColdEndOffset - ColdStartOffset);
+        cold().setAddress(ColdBaseAddress + ColdStartOffset);
+        cold().setImageSize(ColdEndOffset - ColdStartOffset);
         if (hasConstantIsland()) {
           const uint64_t DataOffset =
               Layout.getSymbolOffset(*getFunctionColdConstantIslandLabel());
@@ -4112,39 +4112,44 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
   if (getLayout().block_empty())
     return;
 
-  for (FunctionFragment &FF : getLayout().fragments()) {
-    if (FF.empty())
-      continue;
+  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()) {
     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 == FF.getAddress());
-        else
+        if (BB->isSplit()) {
+          assert(FragmentBaseAddress == cold().getAddress());
+        } else {
           assert(FragmentBaseAddress == getOutputAddress());
+        }
       }
-
       const uint64_t BBOffset = Layout.getSymbolOffset(*BB->getLabel());
       const uint64_t BBAddress = FragmentBaseAddress + BBOffset;
       BB->setOutputStartAddress(BBAddress);
 
-      if (PrevBB)
-        PrevBB->setOutputEndAddress(BBAddress);
+      if (PrevBB) {
+        uint64_t PrevBBEndAddress = BBAddress;
+        if (BB->isSplit() != PrevBB->isSplit())
+          PrevBBEndAddress = getOutputAddress() + getOutputSize();
+        PrevBB->setOutputEndAddress(PrevBBEndAddress);
+      }
       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 {
@@ -4160,9 +4165,8 @@ DebugAddressRangesVector BinaryFunction::getOutputAddressRanges() const {
                             getOutputAddress() + getOutputSize());
   if (isSplit()) {
     assert(isEmitted() && "split function should be emitted");
-    for (const FunctionFragment &FF : getLayout().getSplitFragments())
-      OutputRanges.emplace_back(FF.getAddress(),
-                                FF.getAddress() + FF.getImageSize());
+    OutputRanges.emplace_back(cold().getAddress(),
+                              cold().getAddress() + cold().getImageSize());
   }
 
   if (isSimple())

bolt/lib/Core/FunctionLayout.cpp

@@ -204,14 +204,10 @@ bool FunctionLayout::update(const ArrayRef<BinaryBasicBlock *> NewLayout) {
 
 void FunctionLayout::clear() {
   Blocks = BasicBlockListType();
-  // 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;
+  for (FunctionFragment *const FF : Fragments)
+    delete FF;
+  Fragments = FragmentListType();
+  addFragment();
 }
 
 const BinaryBasicBlock *

bolt/lib/Profile/BoltAddressTranslation.cpp

@@ -73,27 +73,29 @@ 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;
-    for (const BinaryBasicBlock *const BB :
-         Function.getLayout().getMainFragment())
+    const bool IsSplit = Function.isSplit();
+    for (const BinaryBasicBlock *const BB : Function.getLayout().blocks()) {
+      if (IsSplit && BB->isCold())
+        break;
       writeEntriesForBB(Map, *BB, Function.getOutputAddress());
-    Maps.emplace(Function.getOutputAddress(), std::move(Map));
+    }
+    Maps.insert(std::pair<uint64_t, MapTy>(Function.getOutputAddress(), Map));
 
-    if (!Function.isSplit())
+    if (!IsSplit)
       continue;
 
-    // Split maps
+    // Cold map
+    Map.clear();
     LLVM_DEBUG(dbgs() << " Cold part\n");
-    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());
+    for (const BinaryBasicBlock *const BB : Function.getLayout().blocks()) {
+      if (!BB->isCold())
+        continue;
+      writeEntriesForBB(Map, *BB, Function.cold().getAddress());
     }
+    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();