[BOLT] Track fragment info for all split fragments
authorFabian Parzefall <parzefall@fb.com>
Thu, 25 Aug 2022 01:07:06 +0000 (18:07 -0700)
committerFabian Parzefall <parzefall@fb.com>
Thu, 25 Aug 2022 01:07:09 +0000 (18:07 -0700)
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

bolt/include/bolt/Core/BinaryFunction.h
bolt/include/bolt/Core/FunctionLayout.h
bolt/lib/Core/BinaryFunction.cpp
bolt/lib/Core/FunctionLayout.cpp
bolt/lib/Profile/BoltAddressTranslation.cpp
bolt/lib/Rewrite/RewriteInstance.cpp

index c17662b..6b1ded0 100644 (file)
@@ -232,9 +232,6 @@ private:
   /// 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 @@ private:
   /// 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 @@ public:
   }
 
   /// 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 @@ public:
                                              int64_t NewOffset);
 
   BinaryFunction &setFileOffset(uint64_t Offset) {
-    FileOffset = Offset;
+    getLayout().getMainFragment().setFileOffset(Offset);
     return *this;
   }
 
@@ -1785,20 +1776,24 @@ public:
   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 @@ public:
   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,
index 5f4cf8f..8eeabd9 100644 (file)
@@ -83,6 +83,20 @@ private:
   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 @@ public:
   }
   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();
index 6bd06f0..cde56ff 100644 (file)
@@ -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())
index d37c430..2ccf82c 100644 (file)
@@ -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 *
index 3059623..f475f17 100644 (file)
@@ -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();
index 6b050db..565b3e3 100644 (file)
@@ -3726,38 +3726,38 @@ void RewriteInstance::mapCodeSections(RuntimeDyld &RTDyld) {
     if (!Function.isSplit())
       continue;
 
-    for (const FunctionFragment &FF :
-         Function.getLayout().getSplitFragments()) {
-      ErrorOr<BinarySection &> ColdSection =
-          Function.getCodeSection(FF.getFragmentNum());
-      assert(ColdSection && "cannot find section for cold part");
-      // Cold fragments are aligned at 16 bytes.
-      NextAvailableAddress = alignTo(NextAvailableAddress, 16);
-      BinaryFunction::FragmentInfo &ColdPart = Function.cold();
-      if (TooLarge) {
-        // The corresponding FDE will refer to address 0.
-        ColdPart.setAddress(0);
-        ColdPart.setImageAddress(0);
-        ColdPart.setImageSize(0);
-        ColdPart.setFileOffset(0);
-      } else {
-        ColdPart.setAddress(NextAvailableAddress);
-        ColdPart.setImageAddress(ColdSection->getAllocAddress());
-        ColdPart.setImageSize(ColdSection->getOutputSize());
-        ColdPart.setFileOffset(getFileOffsetForAddress(NextAvailableAddress));
-        ColdSection->setOutputAddress(ColdPart.getAddress());
-      }
+    assert(Function.getLayout().isHotColdSplit() &&
+           "Cannot allocate more than two fragments per function in "
+           "non-relocation mode.");
+
+    FunctionFragment &FF =
+        Function.getLayout().getFragment(FragmentNum::cold());
+    ErrorOr<BinarySection &> ColdSection =
+        Function.getCodeSection(FF.getFragmentNum());
+    assert(ColdSection && "cannot find section for cold part");
+    // Cold fragments are aligned at 16 bytes.
+    NextAvailableAddress = alignTo(NextAvailableAddress, 16);
+    if (TooLarge) {
+      // The corresponding FDE will refer to address 0.
+      FF.setAddress(0);
+      FF.setImageAddress(0);
+      FF.setImageSize(0);
+      FF.setFileOffset(0);
+    } else {
+      FF.setAddress(NextAvailableAddress);
+      FF.setImageAddress(ColdSection->getAllocAddress());
+      FF.setImageSize(ColdSection->getOutputSize());
+      FF.setFileOffset(getFileOffsetForAddress(NextAvailableAddress));
+      ColdSection->setOutputAddress(FF.getAddress());
+    }
 
-      LLVM_DEBUG(dbgs() << "BOLT: mapping cold fragment 0x"
-                        << Twine::utohexstr(ColdPart.getImageAddress())
-                        << " to 0x" << Twine::utohexstr(ColdPart.getAddress())
-                        << " with size "
-                        << Twine::utohexstr(ColdPart.getImageSize()) << '\n');
-      RTDyld.reassignSectionAddress(ColdSection->getSectionID(),
-                                    ColdPart.getAddress());
+    LLVM_DEBUG(
+        dbgs() << formatv(
+            "BOLT: mapping cold fragment {0:x+} to {1:x+} with size {2:x+}\n",
+            FF.getImageAddress(), FF.getAddress(), FF.getImageSize()));
+    RTDyld.reassignSectionAddress(ColdSection->getSectionID(), FF.getAddress());
 
-      NextAvailableAddress += ColdPart.getImageSize();
-    }
+    NextAvailableAddress += FF.getImageSize();
   }
 
   // Add the new text section aggregating all existing code sections.
@@ -4518,20 +4518,22 @@ void RewriteInstance::updateELFSymbolTable(
       ICFSymbol.st_shndx = ICFParent->getCodeSection()->getIndex();
       Symbols.emplace_back(ICFSymbol);
     }
-    if (Function.isSplit() && Function.cold().getAddress()) {
+    if (Function.isSplit()) {
       for (const FunctionFragment &FF :
            Function.getLayout().getSplitFragments()) {
-        ELFSymTy NewColdSym = FunctionSymbol;
-        const SmallString<256> SymbolName = formatv(
-            "{0}.cold.{1}", cantFail(FunctionSymbol.getName(StringSection)),
-            FF.getFragmentNum().get() - 1);
-        NewColdSym.st_name = AddToStrTab(SymbolName);
-        NewColdSym.st_shndx =
-            Function.getCodeSection(FF.getFragmentNum())->getIndex();
-        NewColdSym.st_value = Function.cold().getAddress();
-        NewColdSym.st_size = Function.cold().getImageSize();
-        NewColdSym.setBindingAndType(ELF::STB_LOCAL, ELF::STT_FUNC);
-        Symbols.emplace_back(NewColdSym);
+        if (FF.getAddress()) {
+          ELFSymTy NewColdSym = FunctionSymbol;
+          const SmallString<256> SymbolName = formatv(
+              "{0}.cold.{1}", cantFail(FunctionSymbol.getName(StringSection)),
+              FF.getFragmentNum().get() - 1);
+          NewColdSym.st_name = AddToStrTab(SymbolName);
+          NewColdSym.st_shndx =
+              Function.getCodeSection(FF.getFragmentNum())->getIndex();
+          NewColdSym.st_value = FF.getAddress();
+          NewColdSym.st_size = FF.getImageSize();
+          NewColdSym.setBindingAndType(ELF::STB_LOCAL, ELF::STT_FUNC);
+          Symbols.emplace_back(NewColdSym);
+        }
       }
     }
     if (Function.hasConstantIsland()) {
@@ -4656,11 +4658,26 @@ void RewriteInstance::updateELFSymbolTable(
         NewSymbol.st_value = OutputAddress;
         // Force secondary entry points to have zero size.
         NewSymbol.st_size = 0;
+
+        // Find fragment containing entrypoint
+        FunctionLayout::fragment_const_iterator FF = llvm::find_if(
+            Function->getLayout().fragments(), [&](const FunctionFragment &FF) {
+              uint64_t Lo = FF.getAddress();
+              uint64_t Hi = Lo + FF.getImageSize();
+              return Lo <= OutputAddress && OutputAddress < Hi;
+            });
+
+        if (FF == Function->getLayout().fragment_end()) {
+          assert(
+              OutputAddress >= Function->getCodeSection()->getOutputAddress() &&
+              OutputAddress < (Function->getCodeSection()->getOutputAddress() +
+                               Function->getCodeSection()->getOutputSize()) &&
+              "Cannot locate fragment containg secondary entrypoint");
+          FF = Function->getLayout().fragment_begin();
+        }
+
         NewSymbol.st_shndx =
-            OutputAddress >= Function->cold().getAddress() &&
-                    OutputAddress < Function->cold().getImageSize()
-                ? Function->getCodeSection(FragmentNum::cold())->getIndex()
-                : Function->getCodeSection()->getIndex();
+            Function->getCodeSection(FF->getFragmentNum())->getIndex();
       } else {
         // Check if the symbol belongs to moved data object and update it.
         BinaryData *BD = opts::ReorderData.empty()
@@ -4765,8 +4782,10 @@ void RewriteInstance::updateELFSymbolTable(
       SmallVector<char, 256> Buf;
       NewColdSym.st_name = AddToStrTab(
           Twine(Function->getPrintName()).concat(".cold.0").toStringRef(Buf));
-      NewColdSym.st_value = Function->cold().getAddress();
-      NewColdSym.st_size = Function->cold().getImageSize();
+      const FunctionFragment &ColdFF =
+          Function->getLayout().getFragment(FragmentNum::cold());
+      NewColdSym.st_value = ColdFF.getAddress();
+      NewColdSym.st_size = ColdFF.getImageSize();
       Symbols.emplace_back(NewColdSym);
     }
   }
@@ -5292,9 +5311,21 @@ void RewriteInstance::rewriteFile() {
       continue;
     }
 
-    if (Function->isSplit() && (Function->cold().getImageAddress() == 0 ||
-                                Function->cold().getImageSize() == 0))
+    const auto HasAddress = [](const FunctionFragment &FF) {
+      return FF.empty() ||
+             (FF.getImageAddress() != 0 && FF.getImageSize() != 0);
+    };
+    const bool SplitFragmentsHaveAddress =
+        llvm::all_of(Function->getLayout().getSplitFragments(), HasAddress);
+    if (Function->isSplit() && !SplitFragmentsHaveAddress) {
+      const auto HasNoAddress = [](const FunctionFragment &FF) {
+        return FF.getImageAddress() == 0 && FF.getImageSize() == 0;
+      };
+      assert(llvm::all_of(Function->getLayout().getSplitFragments(),
+                          HasNoAddress) &&
+             "Some split fragments have an address while others do not");
       continue;
+    }
 
     OverwrittenScore += Function->getFunctionScore();
     // Overwrite function in the output file.
@@ -5326,12 +5357,14 @@ void RewriteInstance::rewriteFile() {
 
     // Write cold part
     if (opts::Verbosity >= 2)
-      outs() << "BOLT: rewriting function \"" << *Function
-             << "\" (cold part)\n";
+      outs() << formatv("BOLT: rewriting function \"{0}\" (split parts)\n",
+                        *Function);
 
-    OS.pwrite(reinterpret_cast<char *>(Function->cold().getImageAddress()),
-              Function->cold().getImageSize(),
-              Function->cold().getFileOffset());
+    for (const FunctionFragment &FF :
+         Function->getLayout().getSplitFragments()) {
+      OS.pwrite(reinterpret_cast<char *>(FF.getImageAddress()),
+                FF.getImageSize(), FF.getFileOffset());
+    }
 
     ++CountOverwrittenFunctions;
     if (opts::MaxFunctions && CountOverwrittenFunctions == opts::MaxFunctions) {