Towards building a SCC using SCCIterator (Tarjan) for global variablesusers/mingmingl-llvm/spr/pre-scc

5 files changed, 386 insertions, 22 deletions

diff --git a/llvm/include/llvm/Analysis/StaticDataProfileInfo.h b/llvm/include/llvm/Analysis/StaticDataProfileInfo.h
index 9e2e5fbfc676..ebaa0aa9509c 100644
--- a/llvm/include/llvm/Analysis/StaticDataProfileInfo.h
+++ b/llvm/include/llvm/Analysis/StaticDataProfileInfo.h
@@ -9,6 +9,12 @@
 
 namespace llvm {
 
+enum class DataHotness {
+  kCold,
+  kUnknown,
+  kHot,
+};
+
 /// A class that holds the constants that represent static data and their
 /// profile information and provides methods to operate on them.
 class StaticDataProfileInfo {
@@ -44,6 +50,11 @@ public:
   ///     string.
   StringRef getConstantSectionPrefix(const Constant *C,
                                      const ProfileSummaryInfo *PSI) const;
+
+  DataHotness getConstantHotness(const Constant *C,
+                                 const ProfileSummaryInfo *PSI) const;
+
+  static StringRef hotnessToStr(DataHotness hotness);
 };
 
 /// This wraps the StaticDataProfileInfo object as an immutable pass, for a
diff --git a/llvm/lib/Analysis/StaticDataProfileInfo.cpp b/llvm/lib/Analysis/StaticDataProfileInfo.cpp
index a435aa00c2e9..8384978b9a47 100644
--- a/llvm/lib/Analysis/StaticDataProfileInfo.cpp
+++ b/llvm/lib/Analysis/StaticDataProfileInfo.cpp
@@ -29,25 +29,42 @@ StaticDataProfileInfo::getConstantProfileCount(const Constant *C) const {
   return I->second;
 }
 
+StringRef StaticDataProfileInfo::hotnessToStr(DataHotness hotness) {
+  switch (hotness) {
+  case DataHotness::kCold:
+    return "unlikely";
+  case DataHotness::kHot:
+    return "hot";
+  case DataHotness::kUnknown:
+    return "";
+  }
+}
+
 StringRef StaticDataProfileInfo::getConstantSectionPrefix(
     const Constant *C, const ProfileSummaryInfo *PSI) const {
+  return hotnessToStr(getConstantHotness(C, PSI));
+}
+
+DataHotness
+StaticDataProfileInfo::getConstantHotness(const Constant *C,
+                                          const ProfileSummaryInfo *PSI) const {
   auto Count = getConstantProfileCount(C);
   if (!Count)
-    return "";
+    return DataHotness::kUnknown;
   // 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";
+    return DataHotness::kHot;
   // 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 "";
+    return DataHotness::kUnknown;
   // The accummulated counter shows the constant is cold. Return 'unlikely'.
   if (PSI->isColdCount(*Count))
-    return "unlikely";
+    return DataHotness::kCold;
   // The counter says lukewarm. Return an empty prefix.
-  return "";
+  return DataHotness::kUnknown;
 }
 
 bool StaticDataProfileInfoWrapperPass::doInitialization(Module &M) {
diff --git a/llvm/lib/CodeGen/StaticDataAnnotator.cpp b/llvm/lib/CodeGen/StaticDataAnnotator.cpp
index edf85aef41c8..62321b5d74fc 100644
--- a/llvm/lib/CodeGen/StaticDataAnnotator.cpp
+++ b/llvm/lib/CodeGen/StaticDataAnnotator.cpp
@@ -27,6 +27,8 @@
 // eagerly scheduled, and a module pass can use MachineBlockFrequencyInfo.
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/GraphTraits.h"
+#include "llvm/ADT/SCCIterator.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/StaticDataProfileInfo.h"
 #include "llvm/CodeGen/Passes.h"
@@ -35,15 +37,224 @@
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <set>
+
 #define DEBUG_TYPE "static-data-annotator"
 
 using namespace llvm;
 
-/// A module pass which iterates global variables in the module and annotates
-/// their section prefixes based on profile-driven analysis.
+static cl::opt<std::string> GlobalVariableRefGraphDotFile(
+    "global-var-ref-graph-dot-file", cl::init(""),
+    cl::desc("Global variable reference graph to dot file"));
+
+static cl::opt<bool>
+    EnableSCCHotnessPropagation("enable-scc-hotness-propagation",
+                                cl::init(true),
+                                cl::desc("Enable SCC hotness propagation"));
+
+namespace {
+
+struct RefGraphNode;
+
+struct RefGraphEdge {
+  RefGraphEdge(RefGraphNode *Src, RefGraphNode *Dst) : Src(Src), Dst(Dst) {}
+
+  RefGraphNode *Src = nullptr;
+  RefGraphNode *Dst = nullptr;
+
+  operator RefGraphNode *() const { return Dst; }
+};
+
+struct RefGraphNode {
+  RefGraphNode(GlobalVariable *Var, int ID) : Var(Var), VarID(ID) {}
+
+  struct EdgeComparator {
+    bool operator()(const RefGraphEdge &L, const RefGraphEdge &R) const {
+      return L.Dst->VarID < R.Dst->VarID;
+    }
+  };
+
+  using Edge = RefGraphEdge;
+  using RefEdges = std::set<RefGraphEdge, EdgeComparator>;
+  using iterator = RefEdges::iterator;
+  using const_iterator = RefEdges::const_iterator;
+
+  GlobalVariable *Var;
+
+  int VarID;
+
+  RefEdges Edges;
+};
+
+class RefGraph {
+public:
+  using iterator = RefGraphNode::iterator;
+  RefGraph(Module &M) : M(M) {}
+
+  void buildGraph();
+
+  iterator begin() { return Root->Edges.begin(); }
+  iterator end() { return Root->Edges.end(); }
+
+  RefGraphNode *getEntryNode() { return Root.get(); }
+
+  void exportToDot(raw_ostream &OS);
+
+private:
+  void addEdge(GlobalVariable *Src, GlobalVariable *Dst);
+
+  RefGraphNode *getOrCreateNode(GlobalVariable *GV);
+
+  void computeGVDependencies(Value *V, SmallPtrSetImpl<GlobalVariable *> &Deps);
+
+  void updateGVDependencies(GlobalVariable *GV);
+
+  DenseMap<GlobalVariable *, std::unique_ptr<RefGraphNode>> GVToNode;
+
+  Module &M;
+  std::unique_ptr<RefGraphNode> Root; // A dummy root
+
+  std::unordered_map<Constant *, SmallPtrSet<GlobalVariable *, 4>>
+      ConstantDepsCache;
+};
+
+void RefGraph::addEdge(GlobalVariable *Src, GlobalVariable *Dst) {
+  auto SrcNode = getOrCreateNode(Src);
+  auto DstNode = getOrCreateNode(Dst);
+  SrcNode->Edges.insert(RefGraphEdge(SrcNode, DstNode));
+}
+
+RefGraphNode *RefGraph::getOrCreateNode(GlobalVariable *GV) {
+  if (GV == nullptr)
+    return Root.get();
+  auto &Node = GVToNode[GV];
+  if (Node == nullptr)
+    Node = std::make_unique<RefGraphNode>(GV, GVToNode.size());
+  return Node.get();
+}
+
+void RefGraph::computeGVDependencies(Value *V,
+                                     SmallPtrSetImpl<GlobalVariable *> &Deps) {
+  if (auto *GV = dyn_cast<GlobalVariable>(V)) {
+    Deps.insert(GV);
+  } else if (auto *CE = dyn_cast<Constant>(V)) {
+    auto [Where, Inserted] = ConstantDepsCache.try_emplace(CE);
+    auto &LocalDeps = Where->second;
+    if (Inserted)
+      for (User *CEUser : CE->users())
+        computeGVDependencies(CEUser, LocalDeps);
+    Deps.insert_range(LocalDeps);
+  }
+}
+
+void RefGraph::updateGVDependencies(GlobalVariable *GV) {
+  SmallPtrSet<GlobalVariable *, 4> Deps;
+  for (auto *User : GV->users())
+    computeGVDependencies(User, Deps);
+  Deps.erase(GV);
+  for (GlobalVariable *Dep : Deps)
+    addEdge(GV, Dep);
+}
+
+void RefGraph::buildGraph() {
+  Root = std::make_unique<RefGraphNode>(nullptr, -1);
+  for (GlobalVariable &GV : M.globals()) {
+    if (GV.isDeclarationForLinker())
+      continue;
+
+    updateGVDependencies(&GV);
+    addEdge(Root->Var, &GV);
+  }
+}
+
+void RefGraph::exportToDot(raw_ostream &OS) {
+  OS << "digraph {\n";
+  SmallVector<RefGraphNode *, 4> Nodes; // Non-owned pointers.
+  for (auto Iter = GVToNode.begin(); Iter != GVToNode.end(); Iter++)
+    Nodes.push_back(Iter->second.get());
+  // Sort nodes for deterministic output in tests.
+  std::sort(Nodes.begin(), Nodes.end(),
+            [](const RefGraphNode *L, const RefGraphNode *R) {
+              return L->VarID < R->VarID;
+            });
+  SmallVector<std::pair<int, int>> Edges;
+  for (auto Iter = GVToNode.begin(); Iter != GVToNode.end(); Iter++) {
+    RefGraphNode *Node = Iter->second.get();
+    for (const RefGraphEdge &Edge : Node->Edges) {
+      // OS << "\t" << Node->VarID << " -> " << Edge.Dst->VarID << "\n";
+      Edges.emplace_back(Node->VarID, Edge.Dst->VarID);
+    }
+  }
+  // Sort edges for deterministic output in tests.
+  std::sort(Edges.begin(), Edges.end(),
+            [](const std::pair<int, int> &L, const std::pair<int, int> &R) {
+              return L.first < R.first ||
+                     (L.first == R.first && L.second < R.second);
+            });
+  for (RefGraphNode *Node : Nodes) {
+    const GlobalVariable *Var = Node->Var;
+    OS << "\t" << Node->VarID << " [label=\"" << (Var ? Var->getName() : "Root")
+       << "\", style=filled, fillcolor=\"";
+    std::optional<StringRef> SectionPrefix = Var->getSectionPrefix();
+    if (SectionPrefix) {
+      if (*SectionPrefix == "hot")
+        OS << "darkgoldenrod1";
+      else {
+        assert(*SectionPrefix == "unlikely" && "Unknown section prefix");
+        OS << "cadetblue1";
+      }
+    } else {
+      OS << "lightgrey";
+    }
+    OS << "\", shape=\"";
+    if (Var && Var->hasLocalLinkage())
+      OS << "ellipse";
+    else
+      OS << "box";
+    OS << "\"]\n";
+  }
+
+  for (const auto &Edge : Edges)
+    OS << "\t" << Edge.first << " -> " << Edge.second << "\n";
+  OS << "}\n";
+}
+
+} // namespace
+
+namespace llvm {
+template <> struct GraphTraits<RefGraphNode *> {
+  using NodeType = RefGraphNode;
+  using NodeRef = RefGraphNode *;
+  using EdgeType = NodeType::Edge;
+  using ChildIteratorType = NodeType::const_iterator;
+
+  static NodeRef getEntryNode(NodeRef N) { return N; }
+  static ChildIteratorType child_begin(NodeRef N) { return N->Edges.begin(); }
+  static ChildIteratorType child_end(NodeRef N) { return N->Edges.end(); }
+};
+
+template <>
+struct GraphTraits<RefGraph *> : public GraphTraits<RefGraphNode *> {
+  static NodeRef getEntryNode(RefGraph *G) { return G->getEntryNode(); }
+
+  static ChildIteratorType nodes_begin(RefGraph *G) { return G->begin(); }
+
+  static ChildIteratorType nodes_end(RefGraph *G) { return G->end(); }
+};
+
+} // end namespace llvm
+
+/// A module pass which iterates global variables in the module and
+/// annotates their section prefixes based on profile-driven analysis.
 class StaticDataAnnotator : public ModulePass {
+  void propagateHotness(Module &M);
+
+  DenseMap<GlobalVariable *, DataHotness> GVHotness;
+
 public:
   static char ID;
 
@@ -66,10 +277,57 @@ public:
   bool runOnModule(Module &M) override;
 };
 
+void StaticDataAnnotator::propagateHotness(Module &M) {
+  RefGraph G(M);
+  G.buildGraph();
+
+  scc_iterator<RefGraph *> RefGraphI = scc_begin(&G);
+  // Iterate SCCs and propagate hotness.
+  int SccNum = 0;
+  while (!RefGraphI.isAtEnd()) {
+    ++SccNum;
+    auto Range = *RefGraphI;
+
+    DataHotness LocalHotness = DataHotness::kCold;
+    for (auto *Node : Range) {
+      if (Node != nullptr && Node->Var != nullptr) {
+        LocalHotness = std::max(LocalHotness, GVHotness[Node->Var]);
+      }
+    }
+
+    for (auto *Node : Range) {
+      DataHotness NodeVarHotness = LocalHotness;
+      if (Node != nullptr && Node->Var != nullptr) {
+        for (auto &Edge : Node->Edges) {
+          NodeVarHotness = std::max(NodeVarHotness, GVHotness[Edge.Dst->Var]);
+        }
+        if (NodeVarHotness != DataHotness::kUnknown) {
+
+          Node->Var->setSectionPrefix(
+              StaticDataProfileInfo::hotnessToStr(NodeVarHotness));
+        }
+      }
+    }
+    ++RefGraphI;
+  }
+
+  if (!GlobalVariableRefGraphDotFile.empty()) {
+    std::error_code EC;
+    raw_fd_ostream OSDot(GlobalVariableRefGraphDotFile, EC,
+                         sys::fs::OpenFlags::OF_Text);
+    if (EC)
+      report_fatal_error(Twine("Failed to open dot file ") +
+                         GlobalVariableRefGraphDotFile + ": " + EC.message() +
+                         "\n");
+    G.exportToDot(OSDot);
+  }
+}
+
 bool StaticDataAnnotator::runOnModule(Module &M) {
   SDPI = &getAnalysis<StaticDataProfileInfoWrapperPass>()
               .getStaticDataProfileInfo();
   PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
+  GVHotness.clear();
 
   if (!PSI->hasProfileSummary())
     return false;
@@ -87,13 +345,18 @@ bool StaticDataAnnotator::runOnModule(Module &M) {
       llvm::report_fatal_error("Global variable " + GV.getName() +
                                " already has a section prefix " +
                                *maybeSectionPrefix);
+    auto Hotness = SDPI->getConstantHotness(&GV, PSI);
+    GVHotness[&GV] = Hotness;
+    if (Hotness != DataHotness::kUnknown)
+      Changed = true;
+  }
 
-    StringRef SectionPrefix = SDPI->getConstantSectionPrefix(&GV, PSI);
-    if (SectionPrefix.empty())
-      continue;
-
-    GV.setSectionPrefix(SectionPrefix);
-    Changed = true;
+  if (EnableSCCHotnessPropagation) {
+    propagateHotness(M);
+  } else {
+    for (auto &[GV, Hotness] : GVHotness)
+      if (Hotness != DataHotness::kUnknown)
+        GV->setSectionPrefix(StaticDataProfileInfo::hotnessToStr(Hotness));
   }
 
   return Changed;
diff --git a/llvm/test/CodeGen/X86/global-variable-partition.ll b/llvm/test/CodeGen/X86/global-variable-partition.ll
index 91084d038cfe..be4cc2249fbe 100644
--- a/llvm/test/CodeGen/X86/global-variable-partition.ll
+++ b/llvm/test/CodeGen/X86/global-variable-partition.ll
@@ -12,14 +12,14 @@ target triple = "x86_64-unknown-linux-gnu"
 ; This RUN command sets `-data-sections=true -unique-section-names=true` so data
 ; sections are uniqufied by numbers.
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-split-machine-functions \
-; RUN:     -partition-static-data-sections=true -data-sections=true \
-; RUN:     -unique-section-names=true -relocation-model=pic \
+; RUN:     -partition-static-data-sections -data-sections \
+; RUN:     -unique-section-names -relocation-model=pic \
 ; RUN:     %s -o - 2>&1 | FileCheck %s --check-prefixes=SYM,COMMON --dump-input=always
 
 ; This RUN command sets `-data-sections=true -unique-section-names=false` so
 ; data sections are uniqufied by variable names.
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-split-machine-functions \
-; RUN:     -partition-static-data-sections=true -data-sections=true \
+; RUN:     -partition-static-data-sections -data-sections \
 ; RUN:     -unique-section-names=false -relocation-model=pic \
 ; RUN:     %s -o - 2>&1 | FileCheck %s --check-prefixes=UNIQ,COMMON --dump-input=always
 
@@ -96,15 +96,15 @@ target triple = "x86_64-unknown-linux-gnu"
 ; and prune the unlikely list.
 ; For @bss2
 ; COMMON:      .type bss2,@object
-; SYM-NEXT:    .section	.bss.unlikely.bss2,"aw",@nobits
-; UNIQ-NEXT:   .section	.bss.unlikely.,"aw",@nobits,unique,9
-; AGG-NEXT:    .section	.bss.unlikely.,"aw",@nobits
+; SYM-NEXT:    .section	.bss.hot.bss2,"aw",@nobits
+; UNIQ-NEXT:   .section	.bss.hot.,"aw",@nobits,unique,9
+; AGG-NEXT:    .section	.bss.hot.,"aw",@nobits
 
 ; For @data3
 ; COMMON:      .type data3,@object
-; SYM-NEXT:    .section	.data.unlikely.data3,"aw",@progbits
-; UNIQ-NEXT:   .section	.data.unlikely.,"aw",@progbits,unique,10
-; AGG-NEXT:    .section	.data.unlikely.,"aw",@progbits
+; SYM-NEXT:    .section	.data.hot.data3,"aw",@progbits
+; UNIQ-NEXT:   .section	.data.hot.,"aw",@progbits,unique,10
+; AGG-NEXT:    .section	.data.hot.,"aw",@progbits
 
 ; For @data_with_unknown_hotness
 ; SYM: 	       .type	.Ldata_with_unknown_hotness,@object          # @data_with_unknown_hotness
diff --git a/llvm/test/CodeGen/X86/global-variable-scc.ll b/llvm/test/CodeGen/X86/global-variable-scc.ll
new file mode 100644
index 000000000000..eddb01713a79
--- /dev/null
+++ b/llvm/test/CodeGen/X86/global-variable-scc.ll
@@ -0,0 +1,73 @@
+
+
+; This RUN command sets `-data-sections=true -unique-section-names=true` so data
+; sections are uniqufied by numbers.
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-split-machine-functions \
+; RUN:     -partition-static-data-sections=true -data-sections=true \
+; RUN:     -unique-section-names=true -relocation-model=pic \
+; RUN:     -global-var-ref-graph-dot-file=%t1.dot \
+; RUN:     %s -o - 2>&1 
+
+;; COM: Test section prefixes after giving the data a hotness.
+
+; RUN: cat %t1.dot | FileCheck --check-prefix=DOT %s
+
+; DOT:      digraph {
+; DOT-NEXT:     1 [label="scc1_var1", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     2 [label="scc1_var2", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     3 [label="scc2_var1", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     4 [label="scc2_var3", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     5 [label="scc2_var2", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     6 [label="scc3_var", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     7 [label="scc4_var", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     8 [label="scc5_var1", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     9 [label="scc5_var2", style=filled, fillcolor="lightgrey", shape="ellipse"]
+; DOT-NEXT:     1 -> 2
+; DOT-NEXT:     2 -> 1
+; DOT-NEXT:     3 -> 2
+; DOT-NEXT:     3 -> 4
+; DOT-NEXT:     4 -> 5
+; DOT-NEXT:     5 -> 3
+; DOT-NEXT:     6 -> 5
+; DOT-NEXT:     7 -> 5
+; DOT-NEXT:     8 -> 9
+; DOT-NEXT:     9 -> 8
+; DOT-NEXT: }
+
+
+; TODO: Make this cold scc
+@scc1_var1 = internal constant [1 x ptr][ptr @scc1_var2]
+@scc1_var2 = internal constant [2 x ptr][ptr @scc1_var1, ptr @scc2_var1]
+
+; TODO: In this scc, one hot, one cold, one unknown
+@scc2_var1 = internal constant [1 x ptr][ptr @scc2_var2]
+@scc2_var2 = internal constant [3 x ptr][ptr @scc2_var3, ptr @scc3_var, ptr @scc4_var]
+@scc2_var3 = internal constant [1 x ptr][ptr @scc2_var1]
+
+@scc3_var = internal constant i32 12345
+
+@scc4_var = private constant [5 x i8] c "abcde"
+
+; Have a un-named constant in the middle in this scc graph 
+@scc5_var1 = internal constant [1 x ptr][ptr @scc5_var2]
+@scc5_var2 = internal constant [1 x ptr][ptr @scc5_var1]
+
+
+!llvm.module.flags = !{!1}
+
+!1 = !{i32 1, !"ProfileSummary", !2}
+!2 = !{!3, !4, !5, !6, !7, !8, !9, !10}
+!3 = !{!"ProfileFormat", !"InstrProf"}
+!4 = !{!"TotalCount", i64 1460183}
+!5 = !{!"MaxCount", i64 849024}
+!6 = !{!"MaxInternalCount", i64 32769}
+!7 = !{!"MaxFunctionCount", i64 849024}
+!8 = !{!"NumCounts", i64 23627}
+!9 = !{!"NumFunctions", i64 3271}
+!10 = !{!"DetailedSummary", !11}
+!11 = !{!12, !13}
+!12 = !{i32 990000, i64 166, i32 73}
+!13 = !{i32 999999, i64 3, i32 1443}
+!14 = !{!"function_entry_count", i64 100000}
+!15 = !{!"function_entry_count", i64 1}
+!16 = !{!"branch_weights", i32 1, i32 99999}