For static data layout, update codegen passes to reconcile data hotness from two sources of profile information

Author: mingmingl-llvm

llvm/include/llvm/Analysis/StaticDataProfileInfo.h

@@ -26,6 +26,10 @@ class StaticDataProfileInfo {
   LLVM_ABI std::optional<uint64_t>
   getConstantProfileCount(const Constant *C) const;
 
+  LLVM_ABI std::optional<StringRef>
+  getDataHotnessBasedOnProfileCount(const Constant *C,
+                                    const ProfileSummaryInfo *PSI) const;
+
 public:
   StaticDataProfileInfo() = default;
 

llvm/lib/Analysis/StaticDataProfileInfo.cpp

@@ -28,11 +28,15 @@ StaticDataProfileInfo::getConstantProfileCount(const Constant *C) const {
   return I->second;
 }
 
-StringRef StaticDataProfileInfo::getConstantSectionPrefix(
+std::optional<StringRef>
+StaticDataProfileInfo::getDataHotnessBasedOnProfileCount(
     const Constant *C, const ProfileSummaryInfo *PSI) const {
   auto Count = getConstantProfileCount(C);
+  // The constant `C` doesn't have a profile count. `C` might be a external
+  // linkage global variable, whose PGO-based counter is not tracked within one
+  // IR module.
   if (!Count)
-    return "";
+    return std::nullopt;
   // The accummulated counter shows the constant is hot. Return 'hot' whether
   // this variable is seen by unprofiled functions or not.
   if (PSI->isHotCount(*Count))
@@ -41,14 +45,46 @@ StringRef StaticDataProfileInfo::getConstantSectionPrefix(
   // assign it to unlikely sections, even if the counter says 'cold'. So return
   // an empty prefix before checking whether the counter is cold.
   if (ConstantWithoutCounts.count(C))
-    return "";
+    return std::nullopt;
   // The accummulated counter shows the constant is cold. Return 'unlikely'.
-  if (PSI->isColdCount(*Count))
+  if (PSI->isColdCount(*Count)) {
     return "unlikely";
-  // The counter says lukewarm. Return an empty prefix.
+  }
   return "";
 }
 
+static StringRef reconcileHotness(StringRef SectionPrefix, StringRef Hotness) {
+  assert((SectionPrefix == "hot" || SectionPrefix == "unlikely") &&
+         "Section prefix must be 'hot' or 'unlikely'");
+
+  if (SectionPrefix == "hot" || Hotness == "hot")
+    return "hot";
+  assert(SectionPrefix == "unlikely" && "Section prefix must be 'unlikely'.");
+  return Hotness;
+}
+
+static StringRef
+reconcileOptionalHotness(std::optional<StringRef> SectionPrefix,
+                         std::optional<StringRef> Hotness) {
+  if (!SectionPrefix)
+    return Hotness.value_or("");
+  if (!Hotness)
+    return SectionPrefix.value_or("");
+
+  return reconcileHotness(*SectionPrefix, *Hotness);
+}
+
+StringRef StaticDataProfileInfo::getConstantSectionPrefix(
+    const Constant *C, const ProfileSummaryInfo *PSI) const {
+  std::optional<StringRef> HotnessBasedOnCount =
+      getDataHotnessBasedOnProfileCount(C, PSI);
+  if (const GlobalVariable *GV = dyn_cast<GlobalVariable>(C))
+    return reconcileOptionalHotness(GV->getSectionPrefix(),
+                                    HotnessBasedOnCount);
+
+  return HotnessBasedOnCount.value_or("");
+}
+
 bool StaticDataProfileInfoWrapperPass::doInitialization(Module &M) {
   Info.reset(new StaticDataProfileInfo());
   return false;

llvm/lib/CodeGen/StaticDataAnnotator.cpp

@@ -78,15 +78,6 @@ bool StaticDataAnnotator::runOnModule(Module &M) {
     if (GV.isDeclarationForLinker())
       continue;
 
-    // The implementation below assumes prior passes don't set section prefixes,
-    // and specifically do 'assign' rather than 'update'. So report error if a
-    // section prefix is already set.
-    if (auto maybeSectionPrefix = GV.getSectionPrefix();
-        maybeSectionPrefix && !maybeSectionPrefix->empty())
-      llvm::report_fatal_error("Global variable " + GV.getName() +
-                               " already has a section prefix " +
-                               *maybeSectionPrefix);
-
     StringRef SectionPrefix = SDPI->getConstantSectionPrefix(&GV, PSI);
     if (SectionPrefix.empty())
       continue;

llvm/test/CodeGen/X86/global-variable-partition.ll

@@ -90,17 +90,24 @@ target triple = "x86_64-unknown-linux-gnu"
 ; UNIQ-NEXT:   .section	.data.rel.ro.unlikely.,"aw",@progbits,unique,6
 ; AGG-NEXT:    .section	.data.rel.ro.unlikely.,"aw",@progbits
 
-; Currently static-data-splitter only analyzes access from code.
-; @bss2 and @data3 are indirectly accessed by code through @hot_relro_array
-; and @cold_relro_array. A follow-up item is to analyze indirect access via data
-; and prune the unlikely list.
-; For @bss2
+
+; @bss2 and @data3 are indirectly accessed via @hot_relro_array and
+; @cold_relro_array, and actually hot due to accesses via @hot_relro_array.
+; Under the hood, the static data splitter pass analyzes accesses from code but
+; won't aggressively propgate the hotness of @hot_relro_array into the array
+; elements -- instead, this pass reconciles the hotness information from both
+; global variable section prefix and PGO counters.
+
+; @bss2 has a section prefix 'hot' in the IR. StaticDataProfileInfo reconciles
+; it into a hot prefix.
 ; COMMON:      .type bss2,@object
-; SYM-NEXT:    .section	.bss.unlikely.bss2,"aw",@nobits
-; UNIQ-NEXT:   .section	.bss.unlikely.,"aw",@nobits,unique,7
-; AGG-NEXT:    .section	.bss.unlikely.,"aw",@nobits
+; SYM-NEXT:    .section	.bss.hot.bss2,"aw",@nobits
+; UNIQ-NEXT:   .section	.bss.hot.,"aw",@nobits,unique,7
+; AGG-NEXT:    .section	.bss.hot.,"aw",@nobits
 
-; For @data3
+; @data3 doesn't have data access profile coverage and thereby doesn't have a
+; section prefix. PGO counter analysis categorizes it as cold, so it will have
+; section name `.data.unlikely`.
 ; COMMON:      .type data3,@object
 ; SYM-NEXT:    .section	.data.unlikely.data3,"aw",@progbits
 ; UNIQ-NEXT:   .section	.data.unlikely.,"aw",@progbits,unique,8
@@ -133,7 +140,7 @@ target triple = "x86_64-unknown-linux-gnu"
 @cold_data = internal global i32 4
 @cold_data_custom_foo_section = internal global i32 100, section "foo"
 @cold_relro_array = internal constant [2 x ptr] [ptr @data3, ptr @bss2]
-@bss2 = internal global i32 0
+@bss2 = internal global i32 0, !section_prefix !17
 @data3 = internal global i32 3
 @data_with_unknown_hotness = private global i32 5
 @hot_data_custom_bar_section = internal global i32 101 #0
@@ -227,3 +234,4 @@ attributes #0 = {"data-section"="bar"}
 !14 = !{!"function_entry_count", i64 100000}
 !15 = !{!"function_entry_count", i64 1}
 !16 = !{!"branch_weights", i32 1, i32 99999}
+!17 = !{!"section_prefix", !"hot"}