[BOLT] Track fragment info for all split fragments

[FPar]

Summary

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.

[rafauler]

Thanks Fabian! Overall, very nice work, although I feel that the diff stack in general has been a bit light on tests that guarantee that the changes have the intended effect. I suggest one test in one of the comments below.

bolt/lib/Core/FunctionLayout.cpp

207–215

Can you explain why was that added a bit better?

bolt/lib/Rewrite/RewriteInstance.cpp

3728–3761

This will write fragments in this order, assuming we have two funcs with three fragments each:

Func1.main_frag Func2.main_frag Func1.frag2 Func1.frag3 Func2.frag2 Func2.frag3

I'm not sure this is what you want? Aren't you looking for grouping all fragments N together? (bear in mind that this code only runs for non-reloc mode, I think reloc mode is correct in grouping all fragments together)

4524–4538

Can we add a simple test to check that this is working as intended?

Mistakes in the symbol table update are nasty because the program will not crash and we may take a long time to detect (since it takes a human to read the symbol table and see that it doesn't make sense). For this case, a test is a good way of making sure this code won't break in the future.

4670–4675

This needs to be aligned with getBinaryFunctionContainingAddress. If this function returned that this address is part of this function, but we can't find it in any of its fragments, then it is an assertion. At this point in the stack, is getBinaryFunctionContainingAddress already considering fragments? (I would imagine so)

[FPar]

Address @rafauler's comments

• Clarify why main fragment cannot be cleared.
• Assert not using multiple fragments in non-relocation for mapping code sections.
• Conservatively checking address ranges when updating symbol section indices of secondoary entrypoints.

[FPar]

Thanks for the feedback @rafauler!

bolt/lib/Core/FunctionLayout.cpp

207–215

Better?

bolt/lib/Rewrite/RewriteInstance.cpp

3728–3761

You are right. Splitting into more than 2 fragments in non-relocation mode requires more work beyond this. I would rather have that in another diff so for now this checks that the function is split only 2-ways.

4670–4675

getBinaryFunctionContainingAddress considers the original address range of the function, so it is not aware of fragments at all. This is fine though, because at this point Symbol.st_value still has the original value.

I agree about the assert. This checks now explicitly that the secondary entrypoint can be found in any of the function's fragments.

[FPar]

bolt/lib/Rewrite/RewriteInstance.cpp

4524–4538

I have added a test at bolt/test/X86/fragmented-symbols.s in D132052 (easier to test in this diff). Does that cover what you were thinking of?

[rafauler]

LGTM

bolt/lib/Core/FunctionLayout.cpp

207–215

Yes, thanks

bolt/lib/Rewrite/RewriteInstance.cpp

4524–4538

Yes, that's great, thanks

bolt/include/bolt/Core/BinaryFunction.h

private:
   uint64_t Size;
 
   /// Address of the function in output.
   uint64_t OutputAddress{0};
 
   /// 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()};
 
   /// Alignment requirements for the function.
   uint16_t Alignment{2};
 
   /// Maximum number of bytes used for alignment of hot part of the function.
   uint16_t MaxAlignmentBytes{0};
private:
 
   /// True if there are no control-flow edges with successors in other functions
   /// (i.e. if tail calls have edges to function-local basic blocks).
   /// Set to false by SCTC. Dynostats can't be reliably computed for
   /// functions with non-canonical CFG.
   /// 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;
 
   /// Name for the corresponding cold code section.
   std::string ColdCodeSectionName;
 
   /// Parent function fragment for split function fragments.
   SmallPtrSet<BinaryFunction *, 1> ParentFragments;
public:
 
   /// Set the streaming order index for this function.
   void setIndex(uint32_t Idx) {
     assert(!hasValidIndex());
     Index = Idx;
   }
 
   /// 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; }
 
   /// Return the maximum size the body of the function could have.
   uint64_t getMaxSize() const { return MaxSize; }
 
   /// Return the number of emitted instructions for this function.
public:
   }
 
   void mutateCFIRegisterFor(const MCInst &Instr, MCPhysReg NewReg);
 
   const MCCFIInstruction *mutateCFIOffsetFor(const MCInst &Instr,
                                              int64_t NewOffset);
 
   BinaryFunction &setFileOffset(uint64_t Offset) {
-    FileOffset = Offset;
+    getLayout().getMainFragment().setFileOffset(Offset);
     return *this;
   }
 
   BinaryFunction &setSize(uint64_t S) {
     Size = S;
     return *this;
   }
 

   BinaryFunction &setMaxColdAlignmentBytes(uint16_t MaxAlignBytes) {
     MaxColdAlignmentBytes = MaxAlignBytes;
     return *this;
   }
 
   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; }
 
   /// Returns if the given function is a parent fragment of this function.
   bool isParentFragment(BinaryFunction *Parent) const {
     return ParentFragments.count(Parent);
   }

   /// and moves associated data to output location lists.
   DebugLocationsVector
   translateInputToOutputLocationList(const DebugLocationsVector &InputLL) const;
 
   /// Return true if the function is an AArch64 linker inserted veneer
   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,
                                const BinaryFunction &Function) {
   OS << Function.getPrintName();
   return OS;
 }
 

bolt/include/bolt/Core/FunctionLayout.h

using const_iterator =
                            const BinaryBasicBlock>;
 
 private:
   FunctionLayout *Layout;
   FragmentNum Num;
   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;
   FunctionFragment &operator=(const FunctionFragment &) = default;
   FunctionFragment &operator=(FunctionFragment &&) = default;
   ~FunctionFragment() = default;
 
 public:
   FragmentNum getFragmentNum() const { return Num; }
   bool isMainFragment() const {
     return getFragmentNum() == FragmentNum::main();
   }
   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();
   const_iterator begin() const;
   iterator end();
   const_iterator end() const;
   const BinaryBasicBlock *front() const;
 

bolt/lib/Core/BinaryFunction.cpp

void BinaryFunction::print(raw_ostream &OS, std::string Annotation,
   if (AllNames.size() > 1) {
     OS << "\n  All names   : ";
     const char *Sep = "";
     for (const StringRef &Name : AllNames) {
       OS << Sep << Name;
       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";
   if (HasUnknownControlFlow)
     OS << "\n  Unknown CF  : true";
   if (getPersonalityFunction())
     OS << "\n  Personality : " << getPersonalityFunction()->getName();
   if (IsFragment)
void BinaryFunction::print(raw_ostream &OS, std::string Annotation,
   if (FrameInstructions.size())
     OS << "\n  CFI Instrs  : " << FrameInstructions.size();
   if (!Layout.block_empty()) {
     OS << "\n  BB Layout   : ";
     ListSeparator LS;
     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);
   }
 
   if (opts::PrintDynoStats && !getLayout().block_empty()) {
     OS << '\n';
     DynoStats dynoStats = getDynoStats(*this);
if (BC.HasRelocations || isInjected()) {
     setOutputAddress(BaseAddress + StartOffset);
     setOutputSize(EndOffset - StartOffset);
     if (hasConstantIsland()) {
       const uint64_t DataOffset =
           Layout.getSymbolOffset(*getFunctionConstantIslandLabel());
       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
         if (FF.empty() && ColdSection.getError())
           continue;
         const uint64_t ColdBaseAddress = ColdSection->getOutputAddress();
 
         const MCSymbol *ColdStartSymbol = getSymbol(FF.getFragmentNum());
         // If fragment is empty, symbol might have not been emitted
         if (FF.empty() && (!ColdStartSymbol || !ColdStartSymbol->isDefined()) &&
             !hasConstantIsland())
           continue;
         assert(ColdStartSymbol && ColdStartSymbol->isDefined() &&
                "split function should have defined cold symbol");
         const MCSymbol *ColdEndSymbol =
             getFunctionEndLabel(FF.getFragmentNum());
         assert(ColdEndSymbol && ColdEndSymbol->isDefined() &&
                "split function should have defined cold end symbol");
         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());
           setOutputColdDataAddress(ColdBaseAddress + DataOffset);
         }
       }
     }
   } else {
void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
   // Output ranges should match the input if the body hasn't changed.
   if (!isSimple() && !BC.HasRelocations)
     return;
 
   // AArch64 may have functions that only contains a constant island (no code).
   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()
-                                  ? cold().getAddress() + cold().getImageSize()
+                                    ? FF.getAddress() + FF.getImageSize()
                                     : getOutputAddress() + getOutputSize());
   }
+}
 
 DebugAddressRangesVector BinaryFunction::getOutputAddressRanges() const {
   DebugAddressRangesVector OutputRanges;
 
   if (isFolded())
     return OutputRanges;
 
   if (IsFragment)
     return OutputRanges;
 
   OutputRanges.emplace_back(getOutputAddress(),
                             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())
     return OutputRanges;
 
   for (BinaryFunction *Frag : Fragments) {
     assert(!Frag->isSimple() &&
            "fragment of non-simple function should also be non-simple");

bolt/lib/Core/FunctionLayout.cpp

for (BinaryBasicBlock *const BB : NewLayout) {
     addBasicBlock(BB);
   }
 
   return true;
 }
 
 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 *
 FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
                                    bool IgnoreSplits) const {
   const block_const_iterator BBPos = find(blocks(), BB);
   if (BBPos == block_end())
     return nullptr;
 

bolt/lib/Profile/BoltAddressTranslation.cpp

for (auto &BFI : BC.getBinaryFunctions()) {
     // We don't need a translation table if the body of the function hasn't
     // changed
     if (Function.isIgnored() || (!BC.HasRelocations && !Function.isSimple()))
       continue;
 
     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();
   OS.write(reinterpret_cast<const char *>(&NumFuncs), 4);
   LLVM_DEBUG(dbgs() << "Writing " << NumFuncs << " functions for BAT.\n");
   for (auto &MapEntry : Maps) {
     const uint64_t Address = MapEntry.first;
     MapTy &Map = MapEntry.second;

bolt/lib/Rewrite/RewriteInstance.cpp

if (opts::JumpTables == JTS_BASIC) {
                           << " to 0x" << Twine::utohexstr(JT->getAddress())
                           << '\n');
         RTDyld.reassignSectionAddress(Section.getSectionID(), JT->getAddress());
       }
     }
 
     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());
-      }
-
-      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());
-
-      NextAvailableAddress += ColdPart.getImageSize();
+    assert(Function.getLayout().isHotColdSplit() &&
+           "Cannot allocate more than two fragments per function in "
+           "non-relocation mode.");
+
+    FunctionFragment &FF =
+        Function.getLayout().getFragment(FragmentNum::cold());
+      FF.setAddress(0);
+      FF.setImageAddress(0);
+      FF.setImageSize(0);
+      FF.setFileOffset(0);
+      FF.setAddress(NextAvailableAddress);
+      FF.setImageAddress(ColdSection->getAllocAddress());
+      FF.setImageSize(ColdSection->getOutputSize());
+      FF.setFileOffset(getFileOffsetForAddress(NextAvailableAddress));
+      ColdSection->setOutputAddress(FF.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 += FF.getImageSize();
   }
 
   // Add the new text section aggregating all existing code sections.
   // This is pseudo-section that serves a purpose of creating a corresponding
   // entry in section header table.
   int64_t NewTextSectionSize =
       NextAvailableAddress - NewTextSectionStartAddress;
   if (NewTextSectionSize) {
     const unsigned Flags = BinarySection::getFlags(/*IsReadOnly=*/true,
if (Function.isFolded()) {
           AddToStrTab(Twine(cantFail(FunctionSymbol.getName(StringSection)))
                           .concat(".icf.0")
                           .toStringRef(Buf));
       ICFSymbol.st_value = ICFParent->getOutputAddress();
       ICFSymbol.st_size = ICFParent->getOutputSize();
       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()) {
+        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 = Function.cold().getAddress();
-        NewColdSym.st_size = Function.cold().getImageSize();
+          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()) {
       uint64_t DataMark = Function.getOutputDataAddress();
       uint64_t CISize = getConstantIslandSize(Function);
       uint64_t CodeMark = DataMark + CISize;
       ELFSymTy DataMarkSym = FunctionSymbol;
       DataMarkSym.st_name = AddToStrTab("$d");
       DataMarkSym.st_value = DataMark;
       DataMarkSym.st_size = 0;
if (Function) {
                "Adding symbols based on cold fragment when there are more than "
                "2 fragments");
         const uint64_t OutputAddress =
             Function->translateInputToOutputAddress(Symbol.st_value);
 
         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()
                              ? nullptr
                              : BC->getBinaryDataAtAddress(Symbol.st_value);
         if (BD && BD->isMoved() && !BD->isJumpTable()) {
           assert((!BD->getSize() || !Symbol.st_size ||
                   Symbol.st_size == BD->getSize()) &&
if (Function->isSplit()) {
       assert(Function->getLayout().isHotColdSplit() &&
              "Adding symbols based on cold fragment when there are more than "
              "2 fragments");
       ELFSymTy NewColdSym = NewSymbol;
       NewColdSym.setType(ELF::STT_NOTYPE);
       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);
     }
   }
 
   assert((!NumHotTextSymsUpdated || NumHotTextSymsUpdated == 2) &&
          "either none or both __hot_start/__hot_end symbols were expected");
   assert((!NumHotDataSymsUpdated || NumHotDataSymsUpdated == 2) &&
          "either none or both __hot_data_start/__hot_data_end symbols were "
if (Function->getImageSize() > Function->getMaxSize()) {
           JumpTable *JT = JTI.second;
           BinarySection &Section = JT->getOutputSection();
           BC->deregisterSection(Section);
         }
       }
       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.
     if (opts::Verbosity >= 2)
       outs() << "BOLT: rewriting function \"" << *Function << "\"\n";
 
     OS.pwrite(reinterpret_cast<char *>(Function->getImageAddress()),
               Function->getImageSize(), Function->getFileOffset());
if (!Function->isSplit()) {
         outs() << "BOLT: maximum number of functions reached\n";
         break;
       }
       continue;
     }
 
     // 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) {
       outs() << "BOLT: maximum number of functions reached\n";
       break;
     }
   }