[CodeGen] Introduce Static Data Splitter pass (llvm#122183)

https://discourse.llvm.org/t/rfc-profile-guided-static-data-partitioning/83744
proposes to partition static data sections.

This patch introduces a codegen pass. This patch produces jump table
hotness in the in-memory states (machine jump table info and entries).
Target-lowering and asm-printer consume the states and produce `.hot`
section suffix. The follow up PR
llvm#122215 implements such
changes.

---------

Co-authored-by: Ellis Hoag <[email protected]>

Author: mingmingl-llvm

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 MachineFunctionDataHotness {
+  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

llvm/include/llvm/CodeGen/MachineJumpTableInfo.h

@@ -28,15 +28,19 @@ namespace llvm {
 class MachineBasicBlock;
 class DataLayout;
 class raw_ostream;
+enum class MachineFunctionDataHotness;
 
 /// 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) {}
+  /// The hotness of MJTE is inferred from the hotness of the source basic
+  /// block(s) that reference it.
+  MachineFunctionDataHotness Hotness;
+
+  explicit MachineJumpTableEntry(const std::vector<MachineBasicBlock *> &M);
 };
 
 class MachineJumpTableInfo {
@@ -107,6 +111,11 @@ class MachineJumpTableInfo {
     return JumpTables;
   }
 
+  // Update machine jump table entry's hotness. Return true if the hotness is
+  // updated.
+  bool updateJumpTableEntryHotness(size_t JTI,
+                                   MachineFunctionDataHotness Hotness);
+
   /// RemoveJumpTable - Mark the specific index as being dead.  This will
   /// prevent it from being emitted.
   void RemoveJumpTable(unsigned Idx) {

llvm/include/llvm/CodeGen/Passes.h

@@ -71,6 +71,10 @@ namespace llvm {
   /// using profile information.
   MachineFunctionPass *createMachineFunctionSplitterPass();
 
+  /// createStaticDataSplitterPass - This pass partitions a static data section
+  /// 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 *

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 &);

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -235,6 +235,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)

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

llvm/lib/CodeGen/CodeGen.cpp

@@ -131,6 +131,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeStackMapLivenessPass(Registry);
   initializeStackProtectorPass(Registry);
   initializeStackSlotColoringPass(Registry);
+  initializeStaticDataSplitterPass(Registry);
   initializeStripDebugMachineModulePass(Registry);
   initializeTailDuplicateLegacyPass(Registry);
   initializeTargetPassConfigPass(Registry);

llvm/lib/CodeGen/MachineFunction.cpp

@@ -1311,6 +1311,10 @@ const unsigned MachineFunction::DebugOperandMemNumber = 1000000;
 //  MachineJumpTableInfo implementation
 //===----------------------------------------------------------------------===//
 
+MachineJumpTableEntry::MachineJumpTableEntry(
+    const std::vector<MachineBasicBlock *> &MBBs)
+    : MBBs(MBBs), Hotness(MachineFunctionDataHotness::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
@@ -1360,6 +1364,17 @@ unsigned MachineJumpTableInfo::createJumpTableIndex(
   return JumpTables.size()-1;
 }
 
+bool MachineJumpTableInfo::updateJumpTableEntryHotness(
+    size_t JTI, MachineFunctionDataHotness Hotness) {
+  assert(JTI < JumpTables.size() && "Invalid JTI!");
+  // Record the largest hotness value.
+  if (Hotness <= JumpTables[JTI].Hotness)
+    return false;
+
+  JumpTables[JTI].Hotness = Hotness;
+  return true;
+}
+
 /// If Old is the target of any jump tables, update the jump tables to branch
 /// to New instead.
 bool MachineJumpTableInfo::ReplaceMBBInJumpTables(MachineBasicBlock *Old,

llvm/lib/CodeGen/StaticDataSplitter.cpp

@@ -0,0 +1,181 @@
+//===- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// The pass uses branch profile data to assign hotness based section qualifiers
+// for the following types of static data:
+// - Jump tables
+// - Constant pools (TODO)
+// - Other module-internal data (TODO)
+//
+// For the original RFC of this pass please see
+// https://discourse.llvm.org/t/rfc-profile-guided-static-data-partitioning/83744
+
+#include "llvm/ADT/ScopeExit.h"
+#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. Option "
+          "-static-data-default-hotness specifies the hotness.");
+
+static cl::opt<MachineFunctionDataHotness> StaticDataDefaultHotness(
+    "static-data-default-hotness", cl::Hidden,
+    cl::desc("This option specifies the hotness of static data when profile "
+             "information is unavailable"),
+    cl::init(MachineFunctionDataHotness::Hot),
+    cl::values(clEnumValN(MachineFunctionDataHotness::Hot, "hot", "Hot"),
+               clEnumValN(MachineFunctionDataHotness::Cold, "cold", "Cold")));
+
+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(const 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();
+
+  return splitJumpTables(MF);
+}
+
+bool StaticDataSplitter::splitJumpTablesWithProfiles(
+    const MachineFunction &MF, MachineJumpTableInfo &MJTI) {
+  int NumChangedJumpTables = 0;
+
+  // Jump table could be used by either terminating instructions or
+  // non-terminating ones, so we walk all instructions and use
+  // `MachineOperand::isJTI()` to identify jump table operands.
+  // Similarly, `MachineOperand::isCPI()` can identify constant pool usages
+  // in the same loop.
+  for (const auto &MBB : MF) {
+    for (const MachineInstr &I : MBB) {
+      for (const MachineOperand &Op : I.operands()) {
+        if (!Op.isJTI())
+          continue;
+        const int JTI = Op.getIndex();
+        // This is not a source block of jump table.
+        if (JTI == -1)
+          continue;
+
+        auto Hotness = MachineFunctionDataHotness::Hot;
+
+        // Hotness is based on source basic block hotness.
+        // TODO: PSI APIs are about instruction hotness. Introduce API for data
+        // access hotness.
+        if (PSI->isColdBlock(&MBB, MBFI))
+          Hotness = MachineFunctionDataHotness::Cold;
+
+        if (MJTI.updateJumpTableEntryHotness(JTI, Hotness))
+          ++NumChangedJumpTables;
+      }
+    }
+  }
+  return NumChangedJumpTables > 0;
+}
+
+bool StaticDataSplitter::splitJumpTables(MachineFunction &MF) {
+  MachineJumpTableInfo *MJTI = MF.getJumpTableInfo();
+  if (!MJTI || MJTI->getJumpTables().empty())
+    return false;
+
+  const bool ProfileAvailable = PSI && PSI->hasProfileSummary() && MBFI &&
+                                MF.getFunction().hasProfileData();
+  auto statOnExit = llvm::make_scope_exit([&] {
+    if (!AreStatisticsEnabled())
+      return;
+
+    if (!ProfileAvailable) {
+      NumUnknownJumpTables += MJTI->getJumpTables().size();
+      return;
+    }
+
+    for (size_t JTI = 0; JTI < MJTI->getJumpTables().size(); JTI++) {
+      auto Hotness = MJTI->getJumpTables()[JTI].Hotness;
+      if (Hotness == MachineFunctionDataHotness::Hot) {
+        ++NumHotJumpTables;
+      } else {
+        assert(Hotness == MachineFunctionDataHotness::Cold &&
+               "A jump table is either hot or cold when profile information is "
+               "available.");
+        ++NumColdJumpTables;
+      }
+    }
+  });
+
+  // Place jump tables according to block hotness if function has profile data.
+  if (ProfileAvailable)
+    return splitJumpTablesWithProfiles(MF, *MJTI);
+
+  // If function profile is unavailable (e.g., module not instrumented, or new
+  // code paths lacking samples), -static-data-default-hotness specifies the
+  // hotness.
+  for (size_t JTI = 0; JTI < MJTI->getJumpTables().size(); JTI++)
+    MF.getJumpTableInfo()->updateJumpTableEntryHotness(
+        JTI, StaticDataDefaultHotness);
+
+  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();
+}

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -261,6 +261,11 @@ static cl::opt<bool>
     GCEmptyBlocks("gc-empty-basic-blocks", cl::init(false), cl::Hidden,
                   cl::desc("Enable garbage-collecting empty basic blocks"));
 
+static cl::opt<bool>
+    SplitStaticData("split-static-data", cl::Hidden, cl::init(false),
+                    cl::desc("Split static data sections into hot and cold "
+                             "sections using profile information"));
+
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.
 /// i.e. -disable-mypass=false has no effect.
@@ -1251,6 +1256,8 @@ void TargetPassConfig::addMachinePasses() {
       }
     }
     addPass(createMachineFunctionSplitterPass());
+    if (SplitStaticData)
+      addPass(createStaticDataSplitterPass());
   }
   // We run the BasicBlockSections pass if either we need BB sections or BB
   // address map (or both).