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

llvm/include/llvm/IR/Function.h

@@ -346,12 +346,6 @@ class LLVM_ABI Function : public GlobalObject, public ilist_node<Function> {
   /// sample PGO, to enable the same inlines as the profiled optimized binary.
   DenseSet<GlobalValue::GUID> getImportGUIDs() const;
 
-  /// Set the section prefix for this function.
-  void setSectionPrefix(StringRef Prefix);
-
-  /// Get the section prefix for this function.
-  std::optional<StringRef> getSectionPrefix() const;
-
   /// hasGC/getGC/setGC/clearGC - The name of the garbage collection algorithm
   ///                             to use during code generation.
   bool hasGC() const {

llvm/include/llvm/IR/GlobalObject.h

@@ -124,6 +124,17 @@ class GlobalObject : public GlobalValue {
   /// appropriate default object file section.
   void setSection(StringRef S);
 
+  /// Set the section prefix for this global object.
+  void setSectionPrefix(StringRef Prefix);
+
+  /// Update the section prefix, unless the existing prefix is the same as
+  /// `KeepPrefix`.
+  bool updateSectionPrefix(StringRef Prefix,
+                           std::optional<StringRef> KeepPrefix = std::nullopt);
+
+  /// Get the section prefix for this global object.
+  std::optional<StringRef> getSectionPrefix() const;
+
   bool hasComdat() const { return getComdat() != nullptr; }
   const Comdat *getComdat() const { return ObjComdat; }
   Comdat *getComdat() { return ObjComdat; }

llvm/include/llvm/IR/MDBuilder.h

@@ -89,8 +89,8 @@ class MDBuilder {
   MDNode *createFunctionEntryCount(uint64_t Count, bool Synthetic,
                                    const DenseSet<GlobalValue::GUID> *Imports);
 
-  /// Return metadata containing the section prefix for a function.
-  MDNode *createFunctionSectionPrefix(StringRef Prefix);
+  /// Return metadata containing the section prefix for a global object.
+  MDNode *createGlobalObjectSectionPrefix(StringRef Prefix);
 
   /// Return metadata containing the pseudo probe descriptor for a function.
   MDNode *createPseudoProbeDesc(uint64_t GUID, uint64_t Hash, StringRef FName);

llvm/lib/CodeGen/StaticDataSplitter.cpp

@@ -9,13 +9,13 @@
 // The pass uses branch profile data to assign hotness based section qualifiers
 // for the following types of static data:
 // - Jump tables
+// - Module-internal global variables
 // - 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/APInt.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/MBFIWrapper.h"
@@ -27,9 +27,12 @@
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Target/TargetLoweringObjectFile.h"
 
 using namespace llvm;
 
@@ -46,6 +49,20 @@ class StaticDataSplitter : public MachineFunctionPass {
   const MachineBlockFrequencyInfo *MBFI = nullptr;
   const ProfileSummaryInfo *PSI = nullptr;
 
+  // If the global value is a local linkage global variable, return it.
+  // Otherwise, return nullptr.
+  const GlobalVariable *getLocalLinkageGlobalVariable(const GlobalValue *GV);
+
+  // Returns true if the global variable is in one of {.rodata, .bss, .data,
+  // .data.rel.ro} sections
+  bool inStaticDataSection(const GlobalVariable *GV, const TargetMachine &TM);
+
+  // Iterate all global variables in the module and update the section prefix
+  // of the module-internal data.
+  bool updateGlobalVariableSectionPrefix(MachineFunction &MF);
+
+  // Accummulated data profile count across machine functions in the module.
+  DenseMap<const GlobalVariable *, APInt> DataProfileCounts;
   // Update LLVM statistics for a machine function without profiles.
   void updateStatsWithoutProfiles(const MachineFunction &MF);
   // Update LLVM statistics for a machine function with profiles.
@@ -88,13 +105,16 @@ bool StaticDataSplitter::runOnMachineFunction(MachineFunction &MF) {
 
   bool Changed = partitionStaticDataWithProfiles(MF);
 
+  Changed |= updateGlobalVariableSectionPrefix(MF);
+
   updateStatsWithProfiles(MF);
   return Changed;
 }
 
 bool StaticDataSplitter::partitionStaticDataWithProfiles(MachineFunction &MF) {
   int NumChangedJumpTables = 0;
 
+  const TargetMachine &TM = MF.getTarget();
   MachineJumpTableInfo *MJTI = MF.getJumpTableInfo();
 
   // Jump table could be used by either terminating instructions or
@@ -105,6 +125,11 @@ bool StaticDataSplitter::partitionStaticDataWithProfiles(MachineFunction &MF) {
   for (const auto &MBB : MF) {
     for (const MachineInstr &I : MBB) {
       for (const MachineOperand &Op : I.operands()) {
+        if (!Op.isJTI() && !Op.isGlobal())
+          continue;
+
+        std::optional<uint64_t> Count = MBFI->getBlockProfileCount(&MBB);
+
         if (Op.isJTI()) {
           assert(MJTI != nullptr && "Jump table info is not available.");
           const int JTI = Op.getIndex();
@@ -117,18 +142,89 @@ bool StaticDataSplitter::partitionStaticDataWithProfiles(MachineFunction &MF) {
           // 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))
+          if (Count && PSI->isColdCount(*Count))
             Hotness = MachineFunctionDataHotness::Cold;
 
           if (MJTI->updateJumpTableEntryHotness(JTI, Hotness))
             ++NumChangedJumpTables;
+        } else if (Op.isGlobal()) {
+          // Find global variables with local linkage
+          const GlobalVariable *GV =
+              getLocalLinkageGlobalVariable(Op.getGlobal());
+          if (!GV || !inStaticDataSection(GV, TM))
+            continue;
+
+          // Acccumulate data profile count across machine function
+          // instructions.
+          // TODO: Analyze global variable's initializers.
+          if (Count) {
+            auto [It, Inserted] =
+                DataProfileCounts.try_emplace(GV, APInt(128, 0));
+            It->second += *Count;
+          }
         }
       }
     }
   }
   return NumChangedJumpTables > 0;
 }
 
+const GlobalVariable *
+StaticDataSplitter::getLocalLinkageGlobalVariable(const GlobalValue *GV) {
+  if (!GV || GV->isDeclarationForLinker())
+    return nullptr;
+
+  return GV->hasLocalLinkage() ? dyn_cast<GlobalVariable>(GV) : nullptr;
+}
+
+bool StaticDataSplitter::inStaticDataSection(const GlobalVariable *GV,
+                                             const TargetMachine &TM) {
+  assert(GV && "Caller guaranteed");
+
+  // Skip LLVM reserved symbols.
+  if (GV->getName().starts_with("llvm."))
+    return false;
+
+  SectionKind Kind = TargetLoweringObjectFile::getKindForGlobal(GV, TM);
+  return Kind.isData() || Kind.isReadOnly() || Kind.isReadOnlyWithRel() ||
+         Kind.isBSS();
+}
+
+bool StaticDataSplitter::updateGlobalVariableSectionPrefix(
+    MachineFunction &MF) {
+  bool Changed = false;
+  for (GlobalVariable &GV : MF.getFunction().getParent()->globals()) {
+    if (GV.isDeclarationForLinker())
+      continue;
+    // DataProfileCounts accumulates data profile count across all machine
+    // function instructions, and it can't model the indirect accesses through
+    // other global variables' initializers.
+    // TODO: Analyze the users of module-internal global variables and see
+    // through the users' initializers. Do not place a global variable into
+    // unlikely section if any of its users are potentially hot.
+    auto Iter = DataProfileCounts.find(&GV);
+    if (Iter == DataProfileCounts.end())
+      continue;
+
+    // StaticDataSplitter is made a machine function pass rather than a module
+    // pass because (Lazy)MachineBlockFrequencyInfo is a machine-function
+    // analysis pass and cannot be used for a legacy module pass.
+    // As a result, we use `DataProfileCounts` to accumulate data
+    // profile count across machine functions and update global variable section
+    // prefix once per machine function.
+    // FIXME: Make StaticDataSplitter a module pass under new pass manager
+    // framework, and set global variable section prefix once per module after
+    // analyzing all machine functions.
+    if (PSI->isColdCount(Iter->second.getZExtValue())) {
+      Changed |= GV.updateSectionPrefix("unlikely",
+                                        std::make_optional(StringRef("hot")));
+    } else if (PSI->isHotCount(Iter->second.getZExtValue())) {
+      Changed |= GV.updateSectionPrefix("hot");
+    }
+  }
+  return Changed;
+}
+
 void StaticDataSplitter::updateStatsWithProfiles(const MachineFunction &MF) {
   if (!AreStatisticsEnabled())
     return;

llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp

@@ -670,6 +670,7 @@ getELFSectionNameForGlobal(const GlobalObject *GO, SectionKind Kind,
   }
 
   bool HasPrefix = false;
+
   if (const auto *F = dyn_cast<Function>(GO)) {
     // Jump table hotness takes precedence over its enclosing function's hotness
     // if it's known. The function's section prefix is used if jump table entry
@@ -687,6 +688,11 @@ getELFSectionNameForGlobal(const GlobalObject *GO, SectionKind Kind,
       raw_svector_ostream(Name) << '.' << *Prefix;
       HasPrefix = true;
     }
+  } else if (const auto *GV = dyn_cast<GlobalVariable>(GO)) {
+    if (std::optional<StringRef> Prefix = GV->getSectionPrefix()) {
+      raw_svector_ostream(Name) << '.' << *Prefix;
+      HasPrefix = true;
+    }
   }
 
   if (UniqueSectionName) {

llvm/lib/IR/Function.cpp

@@ -1164,22 +1164,6 @@ DenseSet<GlobalValue::GUID> Function::getImportGUIDs() const {
   return R;
 }
 
-void Function::setSectionPrefix(StringRef Prefix) {
-  MDBuilder MDB(getContext());
-  setMetadata(LLVMContext::MD_section_prefix,
-              MDB.createFunctionSectionPrefix(Prefix));
-}
-
-std::optional<StringRef> Function::getSectionPrefix() const {
-  if (MDNode *MD = getMetadata(LLVMContext::MD_section_prefix)) {
-    assert(cast<MDString>(MD->getOperand(0))->getString() ==
-               "function_section_prefix" &&
-           "Metadata not match");
-    return cast<MDString>(MD->getOperand(1))->getString();
-  }
-  return std::nullopt;
-}
-
 bool Function::nullPointerIsDefined() const {
   return hasFnAttribute(Attribute::NullPointerIsValid);
 }

llvm/lib/IR/Globals.cpp

@@ -18,6 +18,7 @@
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -286,6 +287,35 @@ void GlobalObject::setSection(StringRef S) {
   setGlobalObjectFlag(HasSectionHashEntryBit, !S.empty());
 }
 
+void GlobalObject::setSectionPrefix(StringRef Prefix) {
+  MDBuilder MDB(getContext());
+  setMetadata(LLVMContext::MD_section_prefix,
+              MDB.createGlobalObjectSectionPrefix(Prefix));
+}
+
+bool GlobalObject::updateSectionPrefix(StringRef Prefix,
+                                       std::optional<StringRef> KeepPrefix) {
+  auto SectionPrefix = getSectionPrefix();
+  if (SectionPrefix && (*SectionPrefix == Prefix ||
+                        (KeepPrefix && *SectionPrefix == *KeepPrefix)))
+    return false;
+
+  setSectionPrefix(Prefix);
+  return true;
+}
+
+std::optional<StringRef> GlobalObject::getSectionPrefix() const {
+  if (MDNode *MD = getMetadata(LLVMContext::MD_section_prefix)) {
+    [[maybe_unused]] StringRef MDName =
+        cast<MDString>(MD->getOperand(0))->getString();
+    assert((MDName == "section_prefix" ||
+            (isa<Function>(this) && MDName == "function_section_prefix")) &&
+           "Metadata not match");
+    return cast<MDString>(MD->getOperand(1))->getString();
+  }
+  return std::nullopt;
+}
+
 bool GlobalValue::isNobuiltinFnDef() const {
   const Function *F = dyn_cast<Function>(this);
   if (!F || F->empty())

llvm/lib/IR/MDBuilder.cpp

@@ -87,9 +87,9 @@ MDNode *MDBuilder::createFunctionEntryCount(
   return MDNode::get(Context, Ops);
 }
 
-MDNode *MDBuilder::createFunctionSectionPrefix(StringRef Prefix) {
-  return MDNode::get(
-      Context, {createString("function_section_prefix"), createString(Prefix)});
+MDNode *MDBuilder::createGlobalObjectSectionPrefix(StringRef Prefix) {
+  return MDNode::get(Context,
+                     {createString("section_prefix"), createString(Prefix)});
 }
 
 MDNode *MDBuilder::createRange(const APInt &Lo, const APInt &Hi) {

llvm/test/CodeGen/X86/data-section-prefix.ll

@@ -0,0 +1,27 @@
+; RUN: llc -mtriple x86_64-linux-gnu -data-sections %s -o - | FileCheck %s --check-prefix=ELF
+; RUN: llc -mtriple x86_64-linux-gnu -unique-section-names=0 -data-sections %s -o - | FileCheck %s --check-prefix=ELF-NOUNIQ
+
+; RUN: llc -mtriple x86_64-windows-msvc -data-sections %s -o - | FileCheck %s --check-prefix=COFF-MSVC
+
+; ELF: .section .data.hot.foo,
+; ELF: .section .data.bar,
+; ELF: .section .bss.unlikely.baz,
+; ELF: .section .bss.quz,
+
+; ELF-NOUNIQ: .section    .data.hot.,"aw",@progbits,unique,1
+; ELF-NOUNIQ: .section    .data,"aw",@progbits,unique,2
+; ELF-NOUNIQ: .section    .bss.unlikely.,"aw",@nobits,unique,3
+; ELF-NOUNIQ: .section    .bss,"aw",@nobits,unique,4
+
+; COFF-MSVC: .section .data,"dw",one_only,foo
+; COFF-MSVC: .section .data,"dw",one_only,bar
+; COFF-MSVC: .section .bss,"bw",one_only,baz
+; COFF-MSVC: .section .bss,"bw",one_only,quz
+
+@foo = global i32 1, !section_prefix !0
+@bar = global i32 2
+@baz = global i32 0, !section_prefix !1
+@quz = global i32 0
+
+!0 = !{!"section_prefix", !"hot"}
+!1 = !{!"section_prefix", !"unlikely"}