[SDP]Introduce StaticDataSplitter pass and implemenet jump table splitting

mingmingl-llvm · mingmingl-llvm · commit 5d207e9f4738 · 2025-01-08T14:28:43.000-08:00
diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
@@ -997,6 +997,12 @@ class MachineBasicBlock
   /// no changes occurred in the meantime.
   bool canSplitCriticalEdge(const MachineBasicBlock *Succ) const;
 
+  /// Return an index for MachineJumpTableInfo if \p this basic block ends with
+  /// an indirect jump using a jump table, otherwise -1.
+  /// This function is a thin wrapper and forward calls to the per-target method
+  /// `TargetInstrInfo::getjumpTableIndex`.
+  int getJumpTableIndex() const;
+
   void pop_front() { Insts.pop_front(); }
   void pop_back() { Insts.pop_back(); }
   void push_back(MachineInstr *MI) { Insts.push_back(MI); }
diff --git a/llvm/include/llvm/CodeGen/MachineFunction.h b/llvm/include/llvm/CodeGen/MachineFunction.h
@@ -88,6 +88,15 @@ template <> struct ilist_callback_traits<MachineBasicBlock> {
   }
 };
 
+// The hotness of static data tracked by a MachineFunction and not represented
+// as a global object in the module IR / MIR. Typical examples are
+// MachineJumpTableInfo and MachineConstantPool.
+enum class DataHotness {
+  Unknown,
+  Cold,
+  Hot,
+};
+
 /// MachineFunctionInfo - This class can be derived from and used by targets to
 /// hold private target-specific information for each MachineFunction.  Objects
 /// of type are accessed/created with MF::getInfo and destroyed when the
diff --git a/llvm/include/llvm/CodeGen/MachineJumpTableInfo.h b/llvm/include/llvm/CodeGen/MachineJumpTableInfo.h
@@ -28,15 +28,17 @@ namespace llvm {
 class MachineBasicBlock;
 class DataLayout;
 class raw_ostream;
+enum class DataHotness;
 
 /// MachineJumpTableEntry - One jump table in the jump table info.
 ///
 struct MachineJumpTableEntry {
   /// MBBs - The vector of basic blocks from which to create the jump table.
   std::vector<MachineBasicBlock*> MBBs;
 
-  explicit MachineJumpTableEntry(const std::vector<MachineBasicBlock*> &M)
-  : MBBs(M) {}
+  DataHotness Hotness;
+
+  explicit MachineJumpTableEntry(const std::vector<MachineBasicBlock *> &M);
 };
 
 class MachineJumpTableInfo {
@@ -107,6 +109,8 @@ class MachineJumpTableInfo {
     return JumpTables;
   }
 
+  void updateJumpTableHotness(size_t JTI, DataHotness Hotness);
+
   /// RemoveJumpTable - Mark the specific index as being dead.  This will
   /// prevent it from being emitted.
   void RemoveJumpTable(unsigned Idx) {
diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
@@ -71,6 +71,10 @@ namespace llvm {
   /// using profile information.
   MachineFunctionPass *createMachineFunctionSplitterPass();
 
+  /// createStaticDataSplitterPass - This pass partions static data sections
+  /// into a hot and cold section using profile information.
+  MachineFunctionPass *createStaticDataSplitterPass();
+
   /// MachineFunctionPrinter pass - This pass prints out the machine function to
   /// the given stream as a debugging tool.
   MachineFunctionPass *
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
@@ -293,6 +293,7 @@ void initializeSpeculativeExecutionLegacyPassPass(PassRegistry &);
 void initializeSpillPlacementWrapperLegacyPass(PassRegistry &);
 void initializeStackColoringLegacyPass(PassRegistry &);
 void initializeStackFrameLayoutAnalysisPassPass(PassRegistry &);
+void initializeStaticDataSplitterPass(PassRegistry &);
 void initializeStackMapLivenessPass(PassRegistry &);
 void initializeStackProtectorPass(PassRegistry &);
 void initializeStackSafetyGlobalInfoWrapperPassPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -236,6 +236,7 @@ DUMMY_MACHINE_FUNCTION_PASS("livedebugvalues", LiveDebugValuesPass)
 DUMMY_MACHINE_FUNCTION_PASS("lrshrink", LiveRangeShrinkPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-combiner", MachineCombinerPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-cp", MachineCopyPropagationPass)
+DUMMY_MACHINE_FUNCTION_PASS("static-data-splitter", StaticDataSplitter)
 DUMMY_MACHINE_FUNCTION_PASS("machine-function-splitter", MachineFunctionSplitterPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-latecleanup", MachineLateInstrsCleanupPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-sanmd", MachineSanitizerBinaryMetadata)
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
@@ -226,6 +226,7 @@ add_llvm_component_library(LLVMCodeGen
   StackMaps.cpp
   StackProtector.cpp
   StackSlotColoring.cpp
+  StaticDataSplitter.cpp
   SwiftErrorValueTracking.cpp
   SwitchLoweringUtils.cpp
   TailDuplication.cpp
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
@@ -130,6 +130,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeStackMapLivenessPass(Registry);
   initializeStackProtectorPass(Registry);
   initializeStackSlotColoringPass(Registry);
+  initializeStaticDataSplitterPass(Registry);
   initializeStripDebugMachineModulePass(Registry);
   initializeTailDuplicateLegacyPass(Registry);
   initializeTargetPassConfigPass(Registry);
diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -1426,6 +1426,10 @@ bool MachineBasicBlock::canSplitCriticalEdge(
   return true;
 }
 
+int MachineBasicBlock::getJumpTableIndex() const {
+  return findJumpTableIndex(*this);
+}
+
 /// Prepare MI to be removed from its bundle. This fixes bundle flags on MI's
 /// neighboring instructions so the bundle won't be broken by removing MI.
 static void unbundleSingleMI(MachineInstr *MI) {
diff --git a/llvm/lib/CodeGen/MachineFunction.cpp b/llvm/lib/CodeGen/MachineFunction.cpp
@@ -1291,6 +1291,10 @@ const unsigned MachineFunction::DebugOperandMemNumber = 1000000;
 //  MachineJumpTableInfo implementation
 //===----------------------------------------------------------------------===//
 
+MachineJumpTableEntry::MachineJumpTableEntry(
+    const std::vector<MachineBasicBlock *> &MBBs)
+    : MBBs(MBBs), Hotness(DataHotness::Unknown) {}
+
 /// Return the size of each entry in the jump table.
 unsigned MachineJumpTableInfo::getEntrySize(const DataLayout &TD) const {
   // The size of a jump table entry is 4 bytes unless the entry is just the
@@ -1340,6 +1344,15 @@ unsigned MachineJumpTableInfo::createJumpTableIndex(
   return JumpTables.size()-1;
 }
 
+void MachineJumpTableInfo::updateJumpTableHotness(size_t JTI,
+                                                  DataHotness Hotness) {
+  assert(JTI < JumpTables.size() && "Invalid JTI!");
+  // Note record the largest hotness is important for mergable data (constant
+  // pools). Even if jump table instances are not merged, record the largest
+  // value seen fwiw.
+  JumpTables[JTI].Hotness = std::max(JumpTables[JTI].Hotness, Hotness);
+}
+
 /// If Old is the target of any jump tables, update the jump tables to branch
 /// to New instead.
 bool MachineJumpTableInfo::ReplaceMBBInJumpTables(MachineBasicBlock *Old,
diff --git a/llvm/lib/CodeGen/StaticDataSplitter.cpp b/llvm/lib/CodeGen/StaticDataSplitter.cpp
@@ -0,0 +1,153 @@
+//===- StaticDataSplitter.cpp ---------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass uses profile information to partition static data sections into
+// hot and cold ones. It begins to split jump tables based on profile, and
+// subsequent patches will handle constant pools and other module internal data.
+//
+// For the original RFC of this pass please see
+// https://discourse.llvm.org/t/rfc-profile-guided-static-data-partitioning/83744.
+
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/CodeGen/MBFIWrapper.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
+#include "llvm/CodeGen/MachineConstantPool.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "static-data-splitter"
+
+STATISTIC(NumHotJumpTables, "Number of hot jump tables seen");
+STATISTIC(NumColdJumpTables, "Number of cold jump tables seen");
+STATISTIC(NumUnknownJumpTables,
+          "Number of jump tables with unknown hotness. Such jump tables will "
+          "be placed in the hot-suffixed section by default.");
+
+class StaticDataSplitter : public MachineFunctionPass {
+  const MachineBranchProbabilityInfo *MBPI = nullptr;
+  const MachineBlockFrequencyInfo *MBFI = nullptr;
+  const ProfileSummaryInfo *PSI = nullptr;
+
+  // Returns true iff any jump table is hot-cold categorized.
+  bool splitJumpTables(MachineFunction &MF);
+
+  // Same as above but works on functions with profile information.
+  bool splitJumpTablesWithProfiles(MachineFunction &MF,
+                                   MachineJumpTableInfo &MJTI);
+
+public:
+  static char ID;
+
+  StaticDataSplitter() : MachineFunctionPass(ID) {
+    initializeStaticDataSplitterPass(*PassRegistry::getPassRegistry());
+  }
+
+  StringRef getPassName() const override { return "Static Data Splitter"; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    MachineFunctionPass::getAnalysisUsage(AU);
+    AU.addRequired<MachineBranchProbabilityInfoWrapperPass>();
+    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
+    AU.addRequired<ProfileSummaryInfoWrapperPass>();
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+
+bool StaticDataSplitter::runOnMachineFunction(MachineFunction &MF) {
+  MBPI = &getAnalysis<MachineBranchProbabilityInfoWrapperPass>().getMBPI();
+  MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+  PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
+
+  // Split jump tables based on profile information. Subsequent patches will
+  // handle other data types like constant pools, module-internal data, etc.
+  return splitJumpTables(MF);
+}
+
+bool StaticDataSplitter::splitJumpTablesWithProfiles(
+    MachineFunction &MF, MachineJumpTableInfo &MJTI) {
+  int NumChangedJumpTables = 0;
+  // Regard a jump table as hot by default. If the source and all of destination
+  // blocks are cold, regard the jump table as cold.
+  DataHotness Hotness = DataHotness::Hot;
+  for (const auto &MBB : MF) {
+    // IMPORTANT, `getJumpTableIndex` is a thin wrapper around per-target
+    // interface `TargetInstrInfo::getjumpTableIndex`, and only X86 implements
+    // it so far.
+    const int JTI = MBB.getJumpTableIndex();
+    // This is not a source block of jump table.
+    if (JTI == -1)
+      continue;
+
+    bool AllBlocksCold = true;
+
+    if (!PSI->isColdBlock(&MBB, MBFI))
+      AllBlocksCold = false;
+
+    for (const MachineBasicBlock *MBB : MJTI.getJumpTables()[JTI].MBBs)
+      if (!PSI->isColdBlock(MBB, MBFI))
+        AllBlocksCold = false;
+
+    if (AllBlocksCold) {
+      Hotness = DataHotness::Cold;
+      ++NumColdJumpTables;
+    } else {
+      ++NumHotJumpTables;
+    }
+
+    MF.getJumpTableInfo()->updateJumpTableHotness(JTI, Hotness);
+    ++NumChangedJumpTables;
+  }
+  return NumChangedJumpTables > 0;
+}
+
+bool StaticDataSplitter::splitJumpTables(MachineFunction &MF) {
+  MachineJumpTableInfo *MJTI = MF.getJumpTableInfo();
+  if (!MJTI || MJTI->getJumpTables().empty())
+    return false;
+
+  // Place jump tables according to block hotness if block counters are
+  // available. Check function entry count because BFI depends on it to derive
+  // block counters.
+  if (PSI && PSI->hasProfileSummary() && MBFI &&
+      MF.getFunction().getEntryCount())
+    return splitJumpTablesWithProfiles(MF, *MJTI);
+
+  // Conservatively place all jump tables in the hot-suffixed section if profile
+  // information for the function is not available, or the target doesn't
+  // implement `TargetInstrInfo::getJumpTableIndex` yet.
+  for (size_t JTI = 0; JTI < MJTI->getJumpTables().size(); JTI++)
+    MF.getJumpTableInfo()->updateJumpTableHotness(JTI, DataHotness::Hot);
+
+  NumUnknownJumpTables += MJTI->getJumpTables().size();
+  return true;
+}
+
+char StaticDataSplitter::ID = 0;
+
+INITIALIZE_PASS_BEGIN(StaticDataSplitter, DEBUG_TYPE, "Split static data",
+                      false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineBranchProbabilityInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MachineBlockFrequencyInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
+INITIALIZE_PASS_END(StaticDataSplitter, DEBUG_TYPE, "Split static data", false,
+                    false)
+
+MachineFunctionPass *llvm::createStaticDataSplitterPass() {
+  return new StaticDataSplitter();
+}
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -1256,6 +1256,7 @@ void TargetPassConfig::addMachinePasses() {
                "performance.\n";
       }
     }
+    addPass(createStaticDataSplitterPass());
     addPass(createMachineFunctionSplitterPass());
   }
   // We run the BasicBlockSections pass if either we need BB sections or BB
diff --git a/llvm/test/CodeGen/X86/jump-table-partition.ll b/llvm/test/CodeGen/X86/jump-table-partition.ll
