[llvm] [AsmPrinter][ELF] Support profile-guided section prefix for jump tables' (read-only) data sections (PR #122215)

[Snehasish Kumar via llvm-commits]

================
@@ -2876,42 +2875,111 @@ void AsmPrinter::emitJumpTableInfo() {
       MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
           MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference64,
       F);
+
+  std::vector<unsigned> JumpTableIndices;
+  if (!TM.Options.EnableStaticDataPartitioning) {
+    for (unsigned JTI = 0, JTSize = JT.size(); JTI < JTSize; ++JTI)
+      JumpTableIndices.push_back(JTI);
+    emitJumpTableImpl(
+        *MJTI,
+        llvm::make_range(JumpTableIndices.begin(), JumpTableIndices.end()),
+        JTInDiffSection);
+    return;
+  }
+
+  // When static data partitioning is enabled, collect jump table entries that
+  // go into the same section together to reduce the amount of section switch
+  // statements.
+  //
+  // Iterate all jump tables, put hot jump table indices towards the beginning
+  // of the vector, and cold jump table indices towards the end. Meanwhile
+  // retain the relative orders of original jump tables.
+  int NumHotJumpTables = 0, NextColdJumpTableIndex = JT.size() - 1;
+  JumpTableIndices.resize(JT.size());
+  for (unsigned JTI = 0, JTSize = JT.size(); JTI < JTSize; ++JTI) {
+    if (JT[JTI].Hotness == MachineFunctionDataHotness::Cold) {
+      JumpTableIndices[NextColdJumpTableIndex--] = JTI;
+    } else {
+      JumpTableIndices[NumHotJumpTables++] = JTI;
+    }
+  }
+
+  emitJumpTableImpl(
+      *MJTI,
+      llvm::make_range(JumpTableIndices.begin(),
+                       JumpTableIndices.begin() + NumHotJumpTables),
+
+      JTInDiffSection);
+
+  const int NumColdJumpTables = JT.size() - NumHotJumpTables;
+  assert(NumColdJumpTables >= 0 && "Invalid number of cold jump tables.");
+
+  // Reverse iterating cold jump table indices to emit in the original order.
+  emitJumpTableImpl(
+      *MJTI,
+      llvm::make_range(JumpTableIndices.rbegin(),
+                       JumpTableIndices.rbegin() + NumColdJumpTables),
+      JTInDiffSection);
+
+  return;
+}
+
+template <typename Iterator>
+void AsmPrinter::emitJumpTableImpl(
+    const MachineJumpTableInfo &MJTI,
+    const llvm::iterator_range<Iterator> &JumpTableIndices,
+    bool JTInDiffSection) {
+  if (JumpTableIndices.empty())
+    return;
+
+  const TargetLoweringObjectFile &TLOF = getObjFileLowering();
+  const Function &F = MF->getFunction();
+  const std::vector<MachineJumpTableEntry> &JT = MJTI.getJumpTables();
+  MCSection *JumpTableSection = nullptr;
+  if (TM.Options.EnableStaticDataPartitioning) {
+    JumpTableSection =
+        TLOF.getSectionForJumpTable(F, TM, &JT[*JumpTableIndices.begin()]);
+  } else {
+    JumpTableSection = TLOF.getSectionForJumpTable(F, TM);
+  }
+
+  const DataLayout &DL = MF->getDataLayout();
   if (JTInDiffSection) {
-    // Drop it in the readonly section.
-    MCSection *ReadOnlySection = TLOF.getSectionForJumpTable(F, TM);
-    OutStreamer->switchSection(ReadOnlySection);
+    OutStreamer->switchSection(JumpTableSection);
   }
 
-  emitAlignment(Align(MJTI->getEntryAlignment(DL)));
+  emitAlignment(Align(MJTI.getEntryAlignment(MF->getDataLayout())));
 
   // Jump tables in code sections are marked with a data_region directive
   // where that's supported.
   if (!JTInDiffSection)
     OutStreamer->emitDataRegion(MCDR_DataRegionJT32);
 
-  for (unsigned JTI = 0, e = JT.size(); JTI != e; ++JTI) {
-    const std::vector<MachineBasicBlock*> &JTBBs = JT[JTI].MBBs;
+  for (const unsigned JumpTableIndex : JumpTableIndices) {
+    ArrayRef<MachineBasicBlock *> JTBBs = JT[JumpTableIndex].MBBs;
 
     // If this jump table was deleted, ignore it.
-    if (JTBBs.empty()) continue;
+    if (JTBBs.empty())
----------------
snehasish wrote:

There are some formatting changes here (and below) that make it hard to reason about what actually changed. Any way to separate them out into a separate PR easily?

https://github.com/llvm/llvm-project/pull/122215