reconcile hotnessusers/mingmingl-llvm/annotator

8 files changed, 138 insertions, 72 deletions

diff --git a/llvm/include/llvm/Analysis/StaticDataProfileInfo.h b/llvm/include/llvm/Analysis/StaticDataProfileInfo.h
index d7bbf71..fbbc13a 100644
--- a/llvm/include/llvm/Analysis/StaticDataProfileInfo.h
+++ b/llvm/include/llvm/Analysis/StaticDataProfileInfo.h
@@ -14,8 +14,10 @@ namespace llvm {
 /// 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.
+  /// A constant and its profile count.
+  /// A constant is tracked if both conditions are met:
+  ///   1) It has local (i.e., private or internal) linkage.
+  //    2) Its data kind is one of {.rodata, .data, .bss, .data.rel.ro}.
   DenseMap<const Constant *, uint64_t> ConstantProfileCounts;
 
   /// Keeps track of the constants that are seen at least once without profile
@@ -26,11 +28,24 @@ public:
   LLVM_ABI std::optional<uint64_t>
   getConstantProfileCount(const Constant *C) const;
 
-  LLVM_ABI std::optional<StringRef>
-  getDataHotnessBasedOnProfileCount(const Constant *C,
-                                    const ProfileSummaryInfo *PSI) const;
+  enum class StaticDataHotness : uint8_t {
+    Cold = 0,
+    LukewarmOrUnknown = 1,
+    Hot = 2,
+  };
+
+  LLVM_ABI StaticDataHotness getSectionHotnessUsingProfileCount(
+      const Constant *C, const ProfileSummaryInfo *PSI, uint64_t Count) const;
+  LLVM_ABI StaticDataHotness
+  getSectionHotnessUsingDAP(std::optional<StringRef> SectionPrefix) const;
+
+  LLVM_ABI StringRef hotnessToStr(StaticDataHotness Hotness) const;
+
+  bool HasDataAccessProf = false;
 
 public:
+  StaticDataProfileInfo(bool HasDataAccessProf)
+      : HasDataAccessProf(HasDataAccessProf) {}
   StaticDataProfileInfo() = default;
 
   /// If \p Count is not nullopt, add it to the profile count of the constant \p
@@ -40,14 +55,10 @@ public:
   LLVM_ABI 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.
+  /// Given a constant \p C, returns a section prefix.
+  /// If \p C is a global variable, the section prefix is the bigger one
+  /// between its existing section prefix and its use profile count. Otherwise,
+  /// the section prefix is based on its use profile count.
   LLVM_ABI StringRef getConstantSectionPrefix(
       const Constant *C, const ProfileSummaryInfo *PSI) const;
 };
diff --git a/llvm/include/llvm/IR/GlobalObject.h b/llvm/include/llvm/IR/GlobalObject.h
index 08a02b4..58138c79 100644
--- a/llvm/include/llvm/IR/GlobalObject.h
+++ b/llvm/include/llvm/IR/GlobalObject.h
@@ -121,9 +121,13 @@ public:
   /// appropriate default object file section.
   LLVM_ABI void setSection(StringRef S);
 
-  /// Set the section prefix for this global object.
+  /// Set the section prefix for this global object. If \p Prefix is empty,
+  /// the section prefix metadata will be cleared if it exists.
   LLVM_ABI void setSectionPrefix(StringRef Prefix);
 
+  /// If \p Prefix is different from existing prefix, update section prefix. Returns true if an update happens and false otherwise.
+  LLVM_ABI bool updateSectionPrefix(StringRef Prefix);
+
   /// Get the section prefix for this global object.
   LLVM_ABI std::optional<StringRef> getSectionPrefix() const;
 
diff --git a/llvm/lib/Analysis/StaticDataProfileInfo.cpp b/llvm/lib/Analysis/StaticDataProfileInfo.cpp
index b7ca654..05b1adf 100644
--- a/llvm/lib/Analysis/StaticDataProfileInfo.cpp
+++ b/llvm/lib/Analysis/StaticDataProfileInfo.cpp
@@ -1,11 +1,17 @@
 #include "llvm/Analysis/StaticDataProfileInfo.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/ProfileData/InstrProf.h"
 
 using namespace llvm;
+
+extern cl::opt<bool> AnnotateStaticDataSectionPrefix;
+
 void StaticDataProfileInfo::addConstantProfileCount(
     const Constant *C, std::optional<uint64_t> Count) {
   if (!Count) {
@@ -28,65 +34,79 @@ StaticDataProfileInfo::getConstantProfileCount(const Constant *C) const {
   return I->second;
 }
 
-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 std::nullopt;
+StaticDataProfileInfo::StaticDataHotness
+StaticDataProfileInfo::getSectionHotnessUsingProfileCount(
+    const Constant *C, const ProfileSummaryInfo *PSI, uint64_t Count) const {
   // The accummulated counter shows the constant is hot. Return 'hot' whether
   // this variable is seen by unprofiled functions or not.
-  if (PSI->isHotCount(*Count))
-    return "hot";
+  if (PSI->isHotCount(Count))
+    return StaticDataHotness::Hot;
   // The constant is not hot, and seen by unprofiled functions. We don't want to
   // 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 std::nullopt;
+    return StaticDataHotness::LukewarmOrUnknown;
   // The accummulated counter shows the constant is cold. Return 'unlikely'.
-  if (PSI->isColdCount(*Count)) {
-    return "unlikely";
-  }
-  return "";
-}
+  if (PSI->isColdCount(Count))
+    return StaticDataHotness::Cold;
 
-static StringRef reconcileHotness(StringRef SectionPrefix, StringRef Hotness) {
-  assert((SectionPrefix == "hot" || SectionPrefix == "unlikely") &&
-         "Section prefix must be 'hot' or 'unlikely'");
+  return StaticDataHotness::LukewarmOrUnknown;
+}
 
-  if (SectionPrefix == "hot" || Hotness == "hot")
+StringRef StaticDataProfileInfo::hotnessToStr(
+    StaticDataProfileInfo::StaticDataHotness Hotness) const {
+  switch (Hotness) {
+  case StaticDataProfileInfo::StaticDataHotness::Cold:
+    return "unlikely";
+  case StaticDataProfileInfo::StaticDataHotness::Hot:
     return "hot";
-  assert(SectionPrefix == "unlikely" && "Section prefix must be 'unlikely'.");
-  return Hotness;
+  default:
+    return "";
+  }
 }
 
-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);
+StaticDataProfileInfo::StaticDataHotness
+StaticDataProfileInfo::getSectionHotnessUsingDAP(
+    std::optional<StringRef> MaybeSectionPrefix) const {
+  if (!MaybeSectionPrefix)
+    return StaticDataProfileInfo::StaticDataHotness::LukewarmOrUnknown;
+  StringRef Prefix = *MaybeSectionPrefix;
+  assert((Prefix == "hot" || Prefix == "unlikely") &&
+         "Expect section_prefix to be one of hot or unlikely");
+  return Prefix == "hot" ? StaticDataProfileInfo::StaticDataHotness::Hot
+                         : StaticDataProfileInfo::StaticDataHotness::Cold;
 }
 
 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);
+  auto Count = getConstantProfileCount(C);
 
-  return HotnessBasedOnCount.value_or("");
+  if (HasDataAccessProf) {
+    // Module flag `HasDataAccessProf` is 1 -> empty section prefix means uknown
+    // hotness except for string literals.
+    if (const GlobalVariable *GV = dyn_cast<GlobalVariable>(C);
+        GV && !GV->getName().starts_with(".str")) {
+      auto HotnessFromDAP = getSectionHotnessUsingDAP(GV->getSectionPrefix());
+
+      if (!Count)
+        return hotnessToStr(HotnessFromDAP);
+
+      auto HotnessFromPGO = getSectionHotnessUsingProfileCount(C, PSI, *Count);
+      return hotnessToStr(std::max(HotnessFromDAP, HotnessFromPGO));
+    }
+  }
+
+  if (!Count)
+    return "";
+  return hotnessToStr(getSectionHotnessUsingProfileCount(C, PSI, *Count));
 }
 
 bool StaticDataProfileInfoWrapperPass::doInitialization(Module &M) {
-  Info.reset(new StaticDataProfileInfo());
+  bool HasDataAccessProf = false;
+  if (auto *MD = mdconst::extract_or_null<ConstantInt>(
+          M.getModuleFlag("HasDataAccessProf")))
+    HasDataAccessProf = MD->getZExtValue();
+  Info.reset(new StaticDataProfileInfo(HasDataAccessProf));
   return false;
 }
 
diff --git a/llvm/lib/CodeGen/StaticDataAnnotator.cpp b/llvm/lib/CodeGen/StaticDataAnnotator.cpp
index fdf48e1..72735d2 100644
--- a/llvm/lib/CodeGen/StaticDataAnnotator.cpp
+++ b/llvm/lib/CodeGen/StaticDataAnnotator.cpp
@@ -31,6 +31,8 @@
 #include "llvm/Analysis/StaticDataProfileInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Analysis.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
@@ -79,11 +81,7 @@ bool StaticDataAnnotator::runOnModule(Module &M) {
       continue;
 
     StringRef SectionPrefix = SDPI->getConstantSectionPrefix(&GV, PSI);
-    if (SectionPrefix.empty())
-      continue;
-
-    GV.setSectionPrefix(SectionPrefix);
-    Changed = true;
+    Changed |= GV.updateSectionPrefix(SectionPrefix);
   }
 
   return Changed;
diff --git a/llvm/lib/IR/Globals.cpp b/llvm/lib/IR/Globals.cpp
index 11d33e2..0731fcb 100644
--- a/llvm/lib/IR/Globals.cpp
+++ b/llvm/lib/IR/Globals.cpp
@@ -289,11 +289,28 @@ void GlobalObject::setSection(StringRef S) {
 }
 
 void GlobalObject::setSectionPrefix(StringRef Prefix) {
+  if (Prefix.empty()) {
+    setMetadata(LLVMContext::MD_section_prefix, nullptr);
+    return;
+  }
   MDBuilder MDB(getContext());
   setMetadata(LLVMContext::MD_section_prefix,
               MDB.createGlobalObjectSectionPrefix(Prefix));
 }
 
+bool GlobalObject::updateSectionPrefix(StringRef Prefix) {
+  auto MD = getMetadata(LLVMContext::MD_section_prefix);
+  StringRef ExistingPrefix; // Empty by default.
+  if (MD != nullptr)
+    ExistingPrefix = cast<MDString>(MD->getOperand(1))->getString();
+
+  if (ExistingPrefix != Prefix) {
+    setSectionPrefix(Prefix);
+    return true;
+  }
+  return false;
+}
+
 std::optional<StringRef> GlobalObject::getSectionPrefix() const {
   if (MDNode *MD = getMetadata(LLVMContext::MD_section_prefix)) {
     [[maybe_unused]] StringRef MDName =
diff --git a/llvm/lib/ProfileData/DataAccessProf.cpp b/llvm/lib/ProfileData/DataAccessProf.cpp
index a1e686b..d16d0b6f 100644
--- a/llvm/lib/ProfileData/DataAccessProf.cpp
+++ b/llvm/lib/ProfileData/DataAccessProf.cpp
@@ -1,6 +1,7 @@
 #include "llvm/ProfileData/DataAccessProf.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ProfileData/InstrProf.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Errc.h"
@@ -9,6 +10,9 @@
 #include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
+cl::opt<bool> AnnotateStaticDataSectionPrefix(
+    "memprof-annotate-static-data-prefix", cl::init(false), cl::Hidden,
+    cl::desc("If true, annotate the static data section prefix"));
 namespace memprof {
 
 // If `Map` has an entry keyed by `Str`, returns the entry iterator. Otherwise,
diff --git a/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp b/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp
index ecb2f2d..789396c 100644
--- a/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp
@@ -41,6 +41,7 @@ using namespace llvm::memprof;
 #define DEBUG_TYPE "memprof"
 
 namespace llvm {
+extern cl::opt<bool> AnnotateStaticDataSectionPrefix;
 extern cl::opt<bool> PGOWarnMissing;
 extern cl::opt<bool> NoPGOWarnMismatch;
 extern cl::opt<bool> NoPGOWarnMismatchComdatWeak;
@@ -76,10 +77,6 @@ static cl::opt<unsigned> MinMatchedColdBytePercent(
     "memprof-matching-cold-threshold", cl::init(100), cl::Hidden,
     cl::desc("Min percent of cold bytes matched to hint allocation cold"));
 
-static cl::opt<bool> AnnotateStaticDataSectionPrefix(
-    "memprof-annotate-static-data-prefix", cl::init(false), cl::Hidden,
-    cl::desc("If true, annotate the static data section prefix"));
-
 // Matching statistics
 STATISTIC(NumOfMemProfMissing, "Number of functions without memory profile.");
 STATISTIC(NumOfMemProfMismatch,
@@ -797,7 +794,11 @@ bool MemProfUsePass::annotateGlobalVariables(
   if (!AnnotateStaticDataSectionPrefix || M.globals().empty())
     return false;
 
+  // The module flag helps codegen passes interpret empty section prefix:
+  // - 0 : empty section prefix is expected for each GV.
+  // - 1 : empty section prefix means the GV has unknown hotness.
   if (!DataAccessProf) {
+    M.addModuleFlag(Module::Warning, "HasDataAccessProf", 0U);
     M.getContext().diagnose(DiagnosticInfoPGOProfile(
         MemoryProfileFileName.data(),
         StringRef("Data access profiles not found in memprof. Ignore "
@@ -806,6 +807,8 @@ bool MemProfUsePass::annotateGlobalVariables(
     return false;
   }
 
+  M.addModuleFlag(Module::Warning, "HasDataAccessProf", 1);
+
   bool Changed = false;
   // Iterate all global variables in the module and annotate them based on
   // data access profiles. Note it's up to the linker to decide how to map input
diff --git a/llvm/test/CodeGen/X86/global-variable-partition.ll b/llvm/test/CodeGen/X86/global-variable-partition.ll
index 63d37f2..8a81481 100644
--- a/llvm/test/CodeGen/X86/global-variable-partition.ll
+++ b/llvm/test/CodeGen/X86/global-variable-partition.ll
@@ -109,9 +109,9 @@ target triple = "x86_64-unknown-linux-gnu"
 ; 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
-; AGG-NEXT:    .section	.data.unlikely.,"aw",@progbits
+; SYM-NEXT:    .section	.data.data3,"aw",@progbits
+; UNIQ-NEXT:   .section	.data,"aw",@progbits,unique,8
+; AGG-NEXT:    .data
 
 ; For @data_with_unknown_hotness
 ; SYM: 	       .type	.Ldata_with_unknown_hotness,@object          # @data_with_unknown_hotness
@@ -119,8 +119,6 @@ target triple = "x86_64-unknown-linux-gnu"
 ; UNIQ:        .section  .data,"aw",@progbits,unique,9
 ; The `.section` directive is omitted for .data with -unique-section-names=false.
 ; See MCSectionELF::shouldOmitSectionDirective for the implementation details.
-; AGG:         .data
-; COMMON:      .Ldata_with_unknown_hotness:
 
 ; For @hot_data_custom_bar_section
 ; It has an explicit section attribute 'var' and shouldn't have hot or unlikely suffix.
@@ -130,20 +128,27 @@ target triple = "x86_64-unknown-linux-gnu"
 ; UNIQ:        .section bar,"aw",@progbits
 ; AGG:         .section bar,"aw",@progbits
 
+; For @external_cold_bss
+; COMMON:      .type external_cold_bss,@object
+; SYM-NEXT:    .section	.bss.external_cold_bss,"aw",@nobits
+; UNIQ-NEXT:   .section	.bss,"aw",@nobits,unique,
+; AVG-NEXT:    .bss
+
 @.str = private unnamed_addr constant [5 x i8] c"hot\09\00", align 1
 @.str.1 = private unnamed_addr constant [10 x i8] c"%d\09%d\09%d\0A\00", align 1
 @hot_relro_array = internal constant [2 x ptr] [ptr @bss2, ptr @data3]
 @hot_data = internal global i32 5
 @hot_bss = internal global i32 0
 @.str.2 = private unnamed_addr constant [14 x i8] c"cold%d\09%d\09%d\0A\00", align 1
-@cold_bss = internal global i32 0
-@cold_data = internal global i32 4
+@cold_bss = internal global i32 0, !section_prefix !18
+@cold_data = internal global i32 4, !section_prefix !18
 @cold_data_custom_foo_section = internal global i32 100, section "foo"
-@cold_relro_array = internal constant [2 x ptr] [ptr @data3, ptr @bss2]
+@cold_relro_array = internal constant [2 x ptr] [ptr @data3, ptr @bss2], !section_prefix !18
 @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
+@external_cold_bss = global i32 0
 
 define void @cold_func(i32 %0) !prof !15 {
   %2 = load i32, ptr @cold_bss
@@ -156,7 +161,8 @@ define void @cold_func(i32 %0) !prof !15 {
   %9 = load i32, ptr @data_with_unknown_hotness
   %11 = load i32, ptr @hot_data
   %12 = load i32, ptr @cold_data_custom_foo_section
-  %13 = call i32 (...) @func_taking_arbitrary_param(ptr @.str.2, i32 %2, i32 %3, i32 %8, i32 %9, i32 %11, i32 %12)
+  %val = load i32, ptr @external_cold_bss
+  %13 = call i32 (...) @func_taking_arbitrary_param(ptr @.str.2, i32 %2, i32 %3, i32 %8, i32 %9, i32 %11, i32 %12, i32 %val)
   ret void
 }
 
@@ -216,8 +222,9 @@ declare i32 @func_taking_arbitrary_param(...)
 
 attributes #0 = {"data-section"="bar"}
 
-!llvm.module.flags = !{!1}
+!llvm.module.flags = !{!0, !1}
 
+!0 = !{i32 2, !"HasDataAccessProf", i32 1}
 !1 = !{i32 1, !"ProfileSummary", !2}
 !2 = !{!3, !4, !5, !6, !7, !8, !9, !10}
 !3 = !{!"ProfileFormat", !"InstrProf"}
@@ -235,3 +242,5 @@ attributes #0 = {"data-section"="bar"}
 !15 = !{!"function_entry_count", i64 1}
 !16 = !{!"branch_weights", i32 1, i32 99999}
 !17 = !{!"section_prefix", !"hot"}
+!18 = !{!"section_prefix", !"unlikely"}
+