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

[David Li via llvm-commits]

================
@@ -0,0 +1,179 @@
+//===- 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(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(
+    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.
+        if (PSI->isColdBlock(&MBB, MBFI))
----------------
david-xl wrote:

PSI->isColdBlock(..) is about instruction coldness, it might be worth decouple this and introduce a new API to query data access hotness with option control. Can be done a a follow up patch.

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