resolve comments

Author: mingmingl-llvm

llvm/include/llvm/Analysis/StaticDataProfileInfo.h

@@ -3,6 +3,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/Pass.h"
 
@@ -20,6 +21,9 @@ class StaticDataProfileInfo {
   /// 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;
 
@@ -30,13 +34,16 @@ class StaticDataProfileInfo {
   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);
-  }
+  /// 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

llvm/lib/Analysis/StaticDataProfileInfo.cpp

@@ -1,4 +1,5 @@
 #include "llvm/Analysis/StaticDataProfileInfo.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/InitializePasses.h"
@@ -13,7 +14,7 @@ void StaticDataProfileInfo::addConstantProfileCount(
     return;
   }
   uint64_t &OriginalCount = ConstantProfileCounts[C];
-  OriginalCount += llvm::SaturatingAdd(*Count, OriginalCount);
+  OriginalCount = llvm::SaturatingAdd(*Count, OriginalCount);
   // Clamp the count to getInstrMaxCountValue. InstrFDO reserves a few
   // large values for special use.
   if (OriginalCount > getInstrMaxCountValue())
@@ -28,6 +29,20 @@ StaticDataProfileInfo::getConstantProfileCount(const Constant *C) const {
   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;

llvm/lib/CodeGen/StaticDataAnnotator.cpp

@@ -82,28 +82,12 @@ bool StaticDataAnnotator::runOnModule(Module &M) {
     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)
+    StringRef SectionPrefix = SDPI->getConstantSectionPrefix(&GV, PSI);
+    if (SectionPrefix.empty() || alreadyHasSectionPrefix(GV, SectionPrefix))
       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;
-    }
+    GV.setSectionPrefix(SectionPrefix);
+    Changed = true;
   }
 
   return Changed;