[CodeGen][StaticDataPartitioning]Place local-linkage global variables in hot or unlikely prefixed sections based on profile information

llvm/include/llvm/Analysis/StaticDataProfileInfo.h

@@ -0,0 +1,75 @@
+#ifndef LLVM_ANALYSIS_STATICDATAPROFILEINFO_H
+#define LLVM_ANALYSIS_STATICDATAPROFILEINFO_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/Analysis/ProfileSummaryInfo.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;
+
+  /// If \p C has a count, return it. Otherwise, return std::nullopt.
+  std::optional<uint64_t> getConstantProfileCount(const Constant *C) const;
+
+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);
+
+  /// Return a section prefix for the constant \p C based on its profile count.
+  /// - If a constant doesn't have a counter, return an empty string.
+  /// - Otherwise,
+  ///   - If it has a hot count, return "hot".
+  ///   - If it is seen by unprofiled function, return an empty string.
+  ///   - If it has a cold count, return "unlikely".
+  ///   - Otherwise (e.g. it's used by lukewarm functions), return an empty
+  ///     string.
+  StringRef getConstantSectionPrefix(const Constant *C,
+                                     const ProfileSummaryInfo *PSI) const;
+};
+
+/// 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/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/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/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,65 @@
+#include "llvm/Analysis/StaticDataProfileInfo.h"
+#include "llvm/Analysis/ProfileSummaryInfo.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;
+}
+
+StringRef StaticDataProfileInfo::getConstantSectionPrefix(
+    const Constant *C, const ProfileSummaryInfo *PSI) const {
+  auto Count = getConstantProfileCount(C);
+  if (!Count)
+    return "";
+  if (PSI->isHotCount(*Count))
+    return "hot";
+  if (ConstantWithoutCounts.count(C))
+    return "";
+  if (PSI->isColdCount(*Count))
+    return "unlikely";
+  return "";
+}
+
+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/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,103 @@
+//===- 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;
+
+    StringRef SectionPrefix = SDPI->getConstantSectionPrefix(&GV, PSI);
+    if (SectionPrefix.empty() || alreadyHasSectionPrefix(GV, SectionPrefix))
+      continue;
+
+    GV.setSectionPrefix(SectionPrefix);
+    Changed = true;
+  }
+
+  return Changed;
+}
+
+char StaticDataAnnotator::ID = 0;
+
+INITIALIZE_PASS(StaticDataAnnotator, DEBUG_TYPE, "Static Data Annotator", false,
+                false)
+
+ModulePass *llvm::createStaticDataAnnotatorPass() {
+  return new StaticDataAnnotator();
+}