[CodeGen][StaticDataSplitter]Support constant pool partitioning

llvm/include/llvm/Analysis/StaticDataProfileInfo.h

@@ -0,0 +1,68 @@
+#ifndef LLVM_ANALYSIS_STATICDATAPROFILEINFO_H
+#define LLVM_ANALYSIS_STATICDATAPROFILEINFO_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/Pass.h"
+
+namespace llvm {
+
+/// A class that holds the constants that represent static data and their
+/// profile information and provides methods to operate on them.
+class StaticDataProfileInfo {
+public:
+  /// Accummulate the profile count of a constant that will be lowered to static
+  /// data sections.
+  DenseMap<const Constant *, uint64_t> ConstantProfileCounts;
+
+  /// Keeps track of the constants that are seen at least once without profile
+  /// counts.
+  DenseSet<const Constant *> ConstantWithoutCounts;
+
+public:
+  StaticDataProfileInfo() = default;
+
+  /// If \p Count is not nullopt, add it to the profile count of the constant \p
+  /// C in a saturating way, and clamp the count to \p getInstrMaxCountValue if
+  /// the result exceeds it. Otherwise, mark the constant as having no profile
+  /// count.
+  void addConstantProfileCount(const Constant *C,
+                               std::optional<uint64_t> Count);
+
+  /// If \p C has a count, return it. Otherwise, return std::nullopt.
+  std::optional<uint64_t> getConstantProfileCount(const Constant *C) const;
+
+  /// Return true if the constant \p C is seen at least once without profiles.
+  bool hasUnknownCount(const Constant *C) const {
+    return ConstantWithoutCounts.count(C);
+  }
+};
+
+/// This wraps the StaticDataProfileInfo object as an immutable pass, for a
+/// backend pass to operate on.
+class StaticDataProfileInfoWrapperPass : public ImmutablePass {
+public:
+  static char ID;
+  StaticDataProfileInfoWrapperPass();
+  bool doInitialization(Module &M) override;
+  bool doFinalization(Module &M) override;
+
+  StaticDataProfileInfo &getStaticDataProfileInfo() { return *Info; }
+  const StaticDataProfileInfo &getStaticDataProfileInfo() const {
+    return *Info;
+  }
+
+  /// This pass provides StaticDataProfileInfo for reads/writes but does not
+  /// modify \p M or other analysis. All analysis are preserved.
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+  }
+
+private:
+  std::unique_ptr<StaticDataProfileInfo> Info;
+};
+
+} // namespace llvm
+
+#endif // LLVM_ANALYSIS_STATICDATAPROFILEINFO_H

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -18,6 +18,8 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/Analysis/StaticDataProfileInfo.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/CodeGen/DwarfStringPoolEntry.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
@@ -132,6 +134,12 @@ class AsmPrinter : public MachineFunctionPass {
   /// default, this is equal to CurrentFnSym.
   MCSymbol *CurrentFnSymForSize = nullptr;
 
+  /// Provides the profile information for constants.
+  const StaticDataProfileInfo *SDPI = nullptr;
+
+  /// The profile summary information.
+  const ProfileSummaryInfo *PSI = nullptr;
+
   /// Map a basic block section ID to the begin and end symbols of that section
   ///  which determine the section's range.
   struct MBBSectionRange {
@@ -330,6 +338,10 @@ class AsmPrinter : public MachineFunctionPass {
     DwarfUsesRelocationsAcrossSections = Enable;
   }
 
+  // Returns a section suffix (hot or unlikely) for the constant if profiles
+  // are available. Returns empty string otherwise.
+  StringRef getConstantSectionSuffix(const Constant *C) const;
+
   //===------------------------------------------------------------------===//
   // XRay instrumentation implementation.
   //===------------------------------------------------------------------===//

llvm/include/llvm/CodeGen/Passes.h

@@ -71,10 +71,15 @@ namespace llvm {
   /// using profile information.
   MachineFunctionPass *createMachineFunctionSplitterPass();
 
-  /// createStaticDataSplitterPass - This pass partitions a static data section
-  /// into a hot and cold section using profile information.
+  /// createStaticDataSplitterPass - This is a machine-function pass that
+  /// categorizes static data hotness using profile information.
   MachineFunctionPass *createStaticDataSplitterPass();
 
+  /// createStaticDataAnnotatorPASS - This is a module pass that reads from
+  /// StaticDataProfileInfoWrapperPass and annotates the section prefix of
+  /// global variables.
+  ModulePass *createStaticDataAnnotatorPass();
+
   /// MachineFunctionPrinter pass - This pass prints out the machine function to
   /// the given stream as a debugging tool.
   MachineFunctionPass *

llvm/include/llvm/CodeGen/TargetLoweringObjectFileImpl.h

@@ -68,6 +68,12 @@ class TargetLoweringObjectFileELF : public TargetLoweringObjectFile {
                                    const Constant *C,
                                    Align &Alignment) const override;
 
+  /// Similar to the function above, but append \p SectionSuffix to the section
+  /// name.
+  MCSection *getSectionForConstant(const DataLayout &DL, SectionKind Kind,
+                                   const Constant *C, Align &Alignment,
+                                   StringRef SectionSuffix) const override;
+
   MCSection *getExplicitSectionGlobal(const GlobalObject *GO, SectionKind Kind,
                                       const TargetMachine &TM) const override;
 

llvm/include/llvm/InitializePasses.h

@@ -203,6 +203,8 @@ void initializeMachineLoopInfoWrapperPassPass(PassRegistry &);
 void initializeMachineModuleInfoWrapperPassPass(PassRegistry &);
 void initializeMachineOptimizationRemarkEmitterPassPass(PassRegistry &);
 void initializeMachineOutlinerPass(PassRegistry &);
+void initializeStaticDataProfileInfoWrapperPassPass(PassRegistry &);
+void initializeStaticDataAnnotatorPass(PassRegistry &);
 void initializeMachinePipelinerPass(PassRegistry &);
 void initializeMachinePostDominatorTreeWrapperPassPass(PassRegistry &);
 void initializeMachineRegionInfoPassPass(PassRegistry &);

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -220,6 +220,7 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
 #define DUMMY_MACHINE_MODULE_PASS(NAME, PASS_NAME)
 #endif
 DUMMY_MACHINE_MODULE_PASS("machine-outliner", MachineOutlinerPass)
+DUMMY_MACHINE_MODULE_PASS("static-data-annotator", StaticDataAnnotator)
 DUMMY_MACHINE_MODULE_PASS("pseudo-probe-inserter", PseudoProbeInserterPass)
 DUMMY_MACHINE_MODULE_PASS("mir-debugify", DebugifyMachineModule)
 DUMMY_MACHINE_MODULE_PASS("mir-check-debugify", CheckDebugMachineModulePass)

llvm/include/llvm/Target/TargetLoweringObjectFile.h

@@ -104,6 +104,13 @@ class TargetLoweringObjectFile : public MCObjectFileInfo {
                                            SectionKind Kind, const Constant *C,
                                            Align &Alignment) const;
 
+  /// Similar to the function above, but append \p SectionSuffix to the section
+  /// name.
+  virtual MCSection *getSectionForConstant(const DataLayout &DL,
+                                           SectionKind Kind, const Constant *C,
+                                           Align &Alignment,
+                                           StringRef SectionSuffix) const;
+
   virtual MCSection *
   getSectionForMachineBasicBlock(const Function &F,
                                  const MachineBasicBlock &MBB,

llvm/lib/Analysis/CMakeLists.txt

@@ -126,6 +126,7 @@ add_llvm_component_library(LLVMAnalysis
   ScalarEvolutionAliasAnalysis.cpp
   ScalarEvolutionDivision.cpp
   ScalarEvolutionNormalization.cpp
+  StaticDataProfileInfo.cpp
   StackLifetime.cpp
   StackSafetyAnalysis.cpp
   StructuralHash.cpp

llvm/lib/Analysis/StaticDataProfileInfo.cpp

@@ -0,0 +1,50 @@
+#include "llvm/Analysis/StaticDataProfileInfo.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/ProfileData/InstrProf.h"
+#include <sys/types.h>
+
+using namespace llvm;
+void StaticDataProfileInfo::addConstantProfileCount(
+    const Constant *C, std::optional<uint64_t> Count) {
+  if (!Count) {
+    ConstantWithoutCounts.insert(C);
+    return;
+  }
+  uint64_t &OriginalCount = ConstantProfileCounts[C];
+  OriginalCount += llvm::SaturatingAdd(*Count, OriginalCount);
+  // Clamp the count to getInstrMaxCountValue. InstrFDO reserves a few
+  // large values for special use.
+  if (OriginalCount > getInstrMaxCountValue())
+    OriginalCount = getInstrMaxCountValue();
+}
+
+std::optional<uint64_t>
+StaticDataProfileInfo::getConstantProfileCount(const Constant *C) const {
+  auto I = ConstantProfileCounts.find(C);
+  if (I == ConstantProfileCounts.end())
+    return std::nullopt;
+  return I->second;
+}
+
+bool StaticDataProfileInfoWrapperPass::doInitialization(Module &M) {
+  Info.reset(new StaticDataProfileInfo());
+  return false;
+}
+
+bool StaticDataProfileInfoWrapperPass::doFinalization(Module &M) {
+  Info.reset();
+  return false;
+}
+
+INITIALIZE_PASS(StaticDataProfileInfoWrapperPass, "static-data-profile-info",
+                "Static Data Profile Info", false, true)
+
+StaticDataProfileInfoWrapperPass::StaticDataProfileInfoWrapperPass()
+    : ImmutablePass(ID) {
+  initializeStaticDataProfileInfoWrapperPassPass(
+      *PassRegistry::getPassRegistry());
+}
+
+char StaticDataProfileInfoWrapperPass::ID = 0;

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -2769,6 +2769,19 @@ namespace {
 
 } // end anonymous namespace
 
+StringRef AsmPrinter::getConstantSectionSuffix(const Constant *C) const {
+  if (TM.Options.EnableStaticDataPartitioning && C && SDPI && PSI) {
+    if (auto Count = SDPI->getConstantProfileCount(C)) {
+      if (PSI->isHotCount(*Count))
+        return "hot";
+
+      if (PSI->isColdCount(*Count) && !SDPI->hasUnknownCount(C))
+        return "unlikely";
+    }
+  }
+  return "";
+}
+
 /// EmitConstantPool - Print to the current output stream assembly
 /// representations of the constants in the constant pool MCP. This is
 /// used to print out constants which have been "spilled to memory" by
@@ -2792,7 +2805,7 @@ void AsmPrinter::emitConstantPool() {
       C = CPE.Val.ConstVal;
 
     MCSection *S = getObjFileLowering().getSectionForConstant(
-        getDataLayout(), Kind, C, Alignment);
+        getDataLayout(), Kind, C, Alignment, getConstantSectionSuffix(C));
 
     // The number of sections are small, just do a linear search from the
     // last section to the first.

llvm/lib/CodeGen/CMakeLists.txt

@@ -228,6 +228,7 @@ add_llvm_component_library(LLVMCodeGen
   StackProtector.cpp
   StackSlotColoring.cpp
   StaticDataSplitter.cpp
+  StaticDataAnnotator.cpp
   SwiftErrorValueTracking.cpp
   SwitchLoweringUtils.cpp
   TailDuplication.cpp

llvm/lib/CodeGen/CodeGen.cpp

@@ -132,6 +132,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeStackProtectorPass(Registry);
   initializeStackSlotColoringLegacyPass(Registry);
   initializeStaticDataSplitterPass(Registry);
+  initializeStaticDataAnnotatorPass(Registry);
   initializeStripDebugMachineModulePass(Registry);
   initializeTailDuplicateLegacyPass(Registry);
   initializeTargetPassConfigPass(Registry);

llvm/lib/CodeGen/StaticDataAnnotator.cpp

@@ -0,0 +1,119 @@
+//===- StaticDataAnnotator - Annotate static data's section prefix --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// To reason about module-wide data hotness in a module granularity, this file
+// implements a module pass StaticDataAnnotator to work coordinately with the
+// StaticDataSplitter pass.
+//
+// The StaticDataSplitter pass is a machine function pass. It analyzes data
+// hotness based on code and adds counters in the StaticDataProfileInfo.
+// The StaticDataAnnotator pass is a module pass. It iterates global variables
+// in the module, looks up counters from StaticDataProfileInfo and sets the
+// section prefix based on profiles.
+//
+// The three-pass structure is implemented for practical reasons, to work around
+// the limitation that a module pass based on legacy pass manager cannot make
+// use of MachineBlockFrequencyInfo analysis. In the future, we can consider
+// porting the StaticDataSplitter pass to a module-pass using the new pass
+// manager framework. That way, analysis are lazily computed as opposed to
+// eagerly scheduled, and a module pass can use MachineBlockFrequencyInfo.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/Analysis/StaticDataProfileInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+
+#define DEBUG_TYPE "static-data-annotator"
+
+using namespace llvm;
+
+class StaticDataAnnotator : public ModulePass {
+public:
+  static char ID;
+
+  StaticDataProfileInfo *SDPI = nullptr;
+  const ProfileSummaryInfo *PSI = nullptr;
+
+  StaticDataAnnotator() : ModulePass(ID) {
+    initializeStaticDataAnnotatorPass(*PassRegistry::getPassRegistry());
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<StaticDataProfileInfoWrapperPass>();
+    AU.addRequired<ProfileSummaryInfoWrapperPass>();
+    AU.setPreservesAll();
+    ModulePass::getAnalysisUsage(AU);
+  }
+
+  StringRef getPassName() const override { return "Static Data Annotator"; }
+
+  bool runOnModule(Module &M) override;
+};
+
+// Returns true if the global variable already has a section prefix that is the
+// same as `Prefix`.
+static bool alreadyHasSectionPrefix(const GlobalVariable &GV,
+                                    StringRef Prefix) {
+  std::optional<StringRef> SectionPrefix = GV.getSectionPrefix();
+  return SectionPrefix && (*SectionPrefix == Prefix);
+}
+
+bool StaticDataAnnotator::runOnModule(Module &M) {
+  SDPI = &getAnalysis<StaticDataProfileInfoWrapperPass>()
+              .getStaticDataProfileInfo();
+  PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
+
+  if (!PSI->hasProfileSummary())
+    return false;
+
+  bool Changed = false;
+  for (auto &GV : M.globals()) {
+    if (GV.isDeclarationForLinker())
+      continue;
+
+    // Skip global variables without profile counts. The module may not be
+    // profiled or instrumented.
+    auto Count = SDPI->getConstantProfileCount(&GV);
+    if (!Count)
+      continue;
+
+    if (PSI->isHotCount(*Count) && !alreadyHasSectionPrefix(GV, "hot")) {
+      // The variable counter is hot, set 'hot' section prefix if the section
+      // prefix isn't hot already.
+      GV.setSectionPrefix("hot");
+      Changed = true;
+    } else if (PSI->isColdCount(*Count) && !SDPI->hasUnknownCount(&GV) &&
+               !alreadyHasSectionPrefix(GV, "unlikely")) {
+      // The variable counter is cold, set 'unlikely' section prefix when
+      // 1) the section prefix isn't unlikely already, and
+      // 2) the variable is not seen without profile counts. The reason is that
+      // a variable without profile counts doesn't have all its uses profiled,
+      // for example when a function is not instrumented, or not sampled (new
+      // code paths).
+      GV.setSectionPrefix("unlikely");
+      Changed = true;
+    }
+  }
+
+  return Changed;
+}
+
+char StaticDataAnnotator::ID = 0;
+
+INITIALIZE_PASS(StaticDataAnnotator, DEBUG_TYPE, "Static Data Annotator", false,
+                false)
+
+ModulePass *llvm::createStaticDataAnnotatorPass() {
+  return new StaticDataAnnotator();
+}