[SyntheticCounts] Rewrite the code using only graph traits.

[eraman]

Summary

The intent of this is to allow the code to be used with ThinLTO. In
Thinlink phase, a traditional Callgraph can not be computed even though
all the necessary information (nodes and edges of a call graph) is
available. This is due to the fact that CallGraph class is closely tied
to the IR. This patch first extends GraphTraits to add a CallGraphTraits
graph. This is then used to implement a version of counts propagation
on a generic callgraph.

[• dberlin]

I like the idea here a lot, but i don't understand why you need to make a new type of graph traits with things called call_edge_begin/end.
Can you explain the necessity here?
(IE why can you just not implement a graphtraits, why do you need a callgraphtraits)

[davidxl]

include/llvm/Analysis/CallGraph.h

541

use 'child_begin()/child_end()'?

546

Is it better to push this method into the base class ( template<> GraphTraits<Callgraph*> )?

include/llvm/Analysis/SyntheticCountsUtils.h

31

nit: I suggest not using 'Utils' in the name. Also the profile propagation engine will be used for partial profile case as well, so it is not completely Synthetic. May be just call it 'ProfileCountProp' or ProfCountPropagator'

[eraman]

include/llvm/Analysis/CallGraph.h

541

child_begin/end is used in CallGraphTraits to iterate over children nodes.

546

This is not really a graph trait and obviously doesn't make sense for other graphs. So an algorithm that works on a general graph which is specialized for GraphTraits cannot do GraphTraits<G>::has_function_def and so this logically belongs here.

[davidxl]

include/llvm/Analysis/CallGraph.h

541

ok. Does it make sense to push edge_begin method to the subclass instead?

546

ok. Makes sense

[• dberlin]

I don't understand why you need has_function_def. It looks very much like a thing you should just be using a filter_iterator for or something. It does not look like it belongs in the graphtraits at all.

[• dberlin]

Is it a thing that is even different for each callgraph implementation? (If not, why is it a trait?)
Even if it is, if you look, you'll see we don't try to encode random computed Node property differences in GraphTraits. Only the things that *every* graph algorithm needs (IE it caters to anything, not everything). Random node property differences are handled through other trait classes or filtering. This does not, at a glance, seem like a thing that literally every algorithm that is going to use a Callgraph needs. It looks like one that some may need. You even prove this yourself.

If you want to do something different here for CallGraphTraits, i'd like to understand what advantage it has. At least in this patch as currently written, there is no advantage to has_function_def existing as a CallGraphTrait.
It would work just existing fine outside of CallGraphTraits, and even as a templated function local to a single file.

[eraman]

Remove has_function_def from CallGraphTraits and fixup counts propagation code.

[davidxl]

lgtm

include/llvm/Analysis/CallGraph.h

static nodes_iterator nodes_begin(const CallGraph *CG) {
     return nodes_iterator(CG->begin(), &CGGetValuePtr);
   }
 
   static nodes_iterator nodes_end(const CallGraph *CG) {
     return nodes_iterator(CG->end(), &CGGetValuePtr);
   }
 };
 
+// FIXME: The traits here are not limited to callgraphs and can be moved
+// elsewhere including GraphTraits. They are left here because only algorithms
+// that operate on Callgraphs currently use them. If other algorithms operating
+// on a general graph need edge traversals, these can be moved.
+template <class CallGraphType>
+struct CallGraphTraits : public GraphTraits<CallGraphType> {
+  // Elements to provide:
+
+  // typedef EdgeRef           - Type of Edge token in the graph, which should
+  //                             be cheap to copy.
+  // typedef CallEdgeIteratorType - Type used to iterate over children edges in
+  //                             graph, dereference to a EdgeRef.
+
+  // static CallEdgeIteratorType call_edge_begin(NodeRef)
+  // static CallEdgeIteratorType call_edge_end  (NodeRef)
+  //     Return iterators that point to the beginning and ending of the call
+  //     edges list for the given callgraph node.
+  //
+  // static NodeRef edge_dest(EdgeRef)
+  //     Return the destination node of an edge.
+
+  // If anyone tries to use this class without having an appropriate
+  // specialization, make an error.  If you get this error, it's because you
+  // need to include the appropriate specialization of GraphTraits<> for your
+  // graph, or you need to define it for a new graph type. Either that or
+  // your argument to XXX_begin(...) is unknown or needs to have the proper .h
+  // file #include'd.
+  using CallEdgeIteratorType =
+      typename CallGraphType::UnknownCallGraphTypeError;
+};
+
+template <class GraphType>
+iterator_range<typename CallGraphTraits<GraphType>::CallEdgeIteratorType>
+call_edges(const typename CallGraphTraits<GraphType>::NodeRef &G) {
+  return make_range(CallGraphTraits<GraphType>::call_edge_begin(G),
+                    CallGraphTraits<GraphType>::call_edge_end(G));
+}
+
+template <>
+struct CallGraphTraits<const CallGraph *>
+    : public GraphTraits<const CallGraph *> {
+  using EdgeRef = const CallGraphNode::CallRecord &;
+  using CallEdgeIteratorType = CallGraphNode::const_iterator;
+
+  static CallEdgeIteratorType call_edge_begin(NodeRef N) { return N->begin(); }
+  static CallEdgeIteratorType call_edge_end(NodeRef N) { return N->end(); }
+
+  static NodeRef edge_dest(EdgeRef E) { return E.second; }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_CALLGRAPH_H

include/llvm/Analysis/SyntheticCountsUtils.h

 // This file defines utilities for synthetic counts propagation.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_SYNTHETIC_COUNTS_UTILS_H
 #define LLVM_ANALYSIS_SYNTHETIC_COUNTS_UTILS_H
 
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/IR/CallSite.h"
+#include "llvm/Analysis/CallGraph.h"
 #include "llvm/Support/ScaledNumber.h"
 
 namespace llvm {
 
 class CallGraph;
 class Function;
 
+/// Class with methods to propagate synthetic entry counts.
+///
+/// This class is templated on the type of the call graph and designed to work
+/// with the traditional per-module callgraph and the summary callgraphs used in
+/// ThinLTO. This contains only static methods and alias templates.
+template <typename CallGraphType> class SyntheticCountsUtils {
+public:
   using Scaled64 = ScaledNumber<uint64_t>;
-void propagateSyntheticCounts(
-    const CallGraph &CG, function_ref<Scaled64(CallSite CS)> GetCallSiteRelFreq,
-    function_ref<uint64_t(Function *F)> GetCount,
-    function_ref<void(Function *F, uint64_t)> AddToCount);
+  using CGT = CallGraphTraits<CallGraphType>;
+  using NodeRef = typename CGT::NodeRef;
+  using EdgeRef = typename CGT::EdgeRef;
+  using SccTy = std::vector<NodeRef>;
+
+  using GetRelBBFreqTy = function_ref<Optional<Scaled64>(EdgeRef)>;
+  using GetCountTy = function_ref<uint64_t(NodeRef)>;
+  using AddCountTy = function_ref<void(NodeRef, uint64_t)>;
+
+  static void propagate(const CallGraphType &CG, GetRelBBFreqTy GetRelBBFreq,
+                        GetCountTy GetCount, AddCountTy AddCount);
+
+private:
+  static void propagateFromSCC(const SccTy &SCC, GetRelBBFreqTy GetRelBBFreq,
+                               GetCountTy GetCount, AddCountTy AddCount);
+};
 } // namespace llvm
 
 #endif

lib/Analysis/SyntheticCountsUtils.cpp

 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 
 using namespace llvm;
 
-// Given a set of functions in an SCC, propagate entry counts to functions
-// called by the SCC.
-static void
-propagateFromSCC(const SmallPtrSetImpl<Function *> &SCCFunctions,
-                 function_ref<Scaled64(CallSite CS)> GetCallSiteRelFreq,
-                 function_ref<uint64_t(Function *F)> GetCount,
-                 function_ref<void(Function *F, uint64_t)> AddToCount) {
-
-  SmallVector<CallSite, 16> CallSites;
-
-  // Gather all callsites in the SCC.
-  auto GatherCallSites = [&]() {
-    for (auto *F : SCCFunctions) {
-      assert(F && !F->isDeclaration());
-      for (auto &I : instructions(F)) {
-        if (auto CS = CallSite(&I)) {
-          CallSites.push_back(CS);
-        }
-      }
-    }
-  };
-
-  GatherCallSites();
-
-  // Partition callsites so that the callsites that call functions in the same
-  // SCC come first.
-  auto Mid = partition(CallSites, [&](CallSite &CS) {
-    auto *Callee = CS.getCalledFunction();
-    if (Callee)
-      return SCCFunctions.count(Callee);
-    // FIXME: Use the !callees metadata to propagate counts through indirect
-    // calls.
-    return 0U;
-  });
-
-  // For functions in the same SCC, update the counts in two steps:
-  // 1. Compute the additional count for each function by propagating the counts
-  // along all incoming edges to the function that originate from the same SCC
-  // and summing them up.
-  // 2. Add the additional counts to the functions in the SCC.
+// Given an SCC, propagate entry counts along the edge of the SCC nodes.
+template <typename CallGraphType>
+void SyntheticCountsUtils<CallGraphType>::propagateFromSCC(
+    const SccTy &SCC, GetRelBBFreqTy GetRelBBFreq, GetCountTy GetCount,
+    AddCountTy AddCount) {
+
+  SmallPtrSet<NodeRef, 8> SCCNodes;
+  SmallVector<std::pair<NodeRef, EdgeRef>, 8> SCCEdges, NonSCCEdges;
+
+  for (auto &Node : SCC)
+    SCCNodes.insert(Node);
+
+  // Partition the edges coming out of the SCC into those whose destination is
+  // in the SCC and the rest.
+  for (const auto &Node : SCCNodes) {
+    for (auto &E : call_edges<CallGraphType>(Node)) {
+      if (SCCNodes.count(CGT::edge_dest(E)))
+        SCCEdges.emplace_back(Node, E);
+      else
+        NonSCCEdges.emplace_back(Node, E);
+
+  // For nodes in the same SCC, update the counts in two steps:
+  // 1. Compute the additional count for each node by propagating the counts
+  // along all incoming edges to the node that originate from within the same
+  // SCC and summing them up.
+  // 2. Add the additional counts to the nodes in the SCC.
   // This ensures that the order of
-  // traversal of functions within the SCC doesn't change the final result.
+  // traversal of nodes within the SCC doesn't affect the final result.
 
-  DenseMap<Function *, uint64_t> AdditionalCounts;
-  for (auto It = CallSites.begin(); It != Mid; It++) {
-    auto &CS = *It;
-    auto RelFreq = GetCallSiteRelFreq(CS);
-    Function *Callee = CS.getCalledFunction();
-    Function *Caller = CS.getCaller();
+  DenseMap<NodeRef, uint64_t> AdditionalCounts;
+  for (auto &E : SCCEdges) {
+    auto OptRelFreq = GetRelBBFreq(E.second);
+    if (!OptRelFreq)
+      continue;
+    Scaled64 RelFreq = OptRelFreq.getValue();
+    auto Caller = E.first;
+    auto Callee = CGT::edge_dest(E.second);
     RelFreq *= Scaled64(GetCount(Caller), 0);
     uint64_t AdditionalCount = RelFreq.toInt<uint64_t>();
     AdditionalCounts[Callee] += AdditionalCount;
   }
 
-  // Update the counts for the functions in the SCC.
+  // Update the counts for the nodes in the SCC.
   for (auto &Entry : AdditionalCounts)
-    AddToCount(Entry.first, Entry.second);
+    AddCount(Entry.first, Entry.second);
 
-  // Now update the counts for functions not in SCC.
-  for (auto It = Mid; It != CallSites.end(); It++) {
-    auto &CS = *It;
-    auto Weight = GetCallSiteRelFreq(CS);
-    Function *Callee = CS.getCalledFunction();
-    Function *Caller = CS.getCaller();
-    Weight *= Scaled64(GetCount(Caller), 0);
-    AddToCount(Callee, Weight.toInt<uint64_t>());
+  // Now update the counts for nodes outside the SCC.
+  for (auto &E : NonSCCEdges) {
+    auto OptRelFreq = GetRelBBFreq(E.second);
+    if (!OptRelFreq)
+      continue;
+    Scaled64 RelFreq = OptRelFreq.getValue();
+    auto Caller = E.first;
+    auto Callee = CGT::edge_dest(E.second);
+    RelFreq *= Scaled64(GetCount(Caller), 0);
+    AddCount(Callee, RelFreq.toInt<uint64_t>());
   }
 }
 
-/// Propgate synthetic entry counts on a callgraph.
+/// Propgate synthetic entry counts on a callgraph \p CG.
 ///
 /// This performs a reverse post-order traversal of the callgraph SCC. For each
-/// SCC, it first propagates the entry counts to the functions within the SCC
+/// SCC, it first propagates the entry counts to the nodes within the SCC
 /// through call edges and updates them in one shot. Then the entry counts are
-/// propagated to functions outside the SCC.
-void llvm::propagateSyntheticCounts(
-    const CallGraph &CG, function_ref<Scaled64(CallSite CS)> GetCallSiteRelFreq,
-    function_ref<uint64_t(Function *F)> GetCount,
-    function_ref<void(Function *F, uint64_t)> AddToCount) {
-
-  SmallVector<SmallPtrSet<Function *, 8>, 16> SCCs;
-  for (auto I = scc_begin(&CG); !I.isAtEnd(); ++I) {
-    auto SCC = *I;
+/// propagated to nodes outside the SCC. This requires \p CallGraphTraits
+/// to have a specialization for \p CallGraphType.
 
-    SmallPtrSet<Function *, 8> SCCFunctions;
-    for (auto *Node : SCC) {
-      Function *F = Node->getFunction();
-      if (F && !F->isDeclaration()) {
-        SCCFunctions.insert(F);
-      }
-    }
-    SCCs.push_back(SCCFunctions);
-  }
+template <typename CallGraphType>
+void SyntheticCountsUtils<CallGraphType>::propagate(const CallGraphType &CG,
+                                                    GetRelBBFreqTy GetRelBBFreq,
+                                                    GetCountTy GetCount,
+                                                    AddCountTy AddCount) {
+  std::vector<SccTy> SCCs;
+
+  // Collect all the SCCs.
+  for (auto I = scc_begin(CG); !I.isAtEnd(); ++I)
+    SCCs.push_back(*I);
+
+  // The callgraph-scc needs to be visited in top-down order for propagation.
+  // The scc iterator returns the scc in bottom-up order, so reverse the SCCs
+  // and call propagateFromSCC.
+  for (auto &SCC : reverse(SCCs))
+    propagateFromSCC(SCC, GetRelBBFreq, GetCount, AddCount);
 
-  for (auto &SCCFunctions : reverse(SCCs))
-    propagateFromSCC(SCCFunctions, GetCallSiteRelFreq, GetCount, AddToCount);
-}
+template class llvm::SyntheticCountsUtils<const CallGraph *>;

lib/Transforms/IPO/SyntheticCountsPropagation.cpp

 PreservedAnalyses SyntheticCountsPropagation::run(Module &M,
                                                   ModuleAnalysisManager &MAM) {
   FunctionAnalysisManager &FAM =
       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   DenseMap<Function *, uint64_t> Counts;
   // Set initial entry counts.
   initializeCounts(M, [&](Function *F, uint64_t Count) { Counts[F] = Count; });
 
-  // Compute the relative block frequency for a callsite. Use scaled numbers
+  // Compute the relative block frequency for a call edge. Use scaled numbers
   // and not integers since the relative block frequency could be less than 1.
-  auto GetCallSiteRelFreq = [&](CallSite CS) {
+  auto GetCallSiteRelFreq = [&](const CallGraphNode::CallRecord &Edge) {
+    Optional<Scaled64> Res = None;
+    if (!Edge.first)
+      return Res;
+    assert(isa<Instruction>(Edge.first));
+    CallSite CS(cast<Instruction>(Edge.first));
     Function *Caller = CS.getCaller();
     auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(*Caller);
     BasicBlock *CSBB = CS.getInstruction()->getParent();
     Scaled64 EntryFreq(BFI.getEntryFreq(), 0);
     Scaled64 BBFreq(BFI.getBlockFreq(CSBB).getFrequency(), 0);
     BBFreq /= EntryFreq;
-    return BBFreq;
+    return Optional<Scaled64>(BBFreq);
   };
 
   CallGraph CG(M);
   // Propgate the entry counts on the callgraph.
-  propagateSyntheticCounts(
-      CG, GetCallSiteRelFreq, [&](Function *F) { return Counts[F]; },
-      [&](Function *F, uint64_t New) { Counts[F] += New; });
+  SyntheticCountsUtils<const CallGraph *>::propagate(
+      &CG, GetCallSiteRelFreq,
+      [&](const CallGraphNode *N) { return Counts[N->getFunction()]; },
+      [&](const CallGraphNode *N, uint64_t New) {
+        auto F = N->getFunction();
+        if (!F || F->isDeclaration())
+          return;
+        Counts[F] += New;
+      });
 
   // Set the counts as metadata.
   for (auto Entry : Counts)
     Entry.first->setEntryCount(
         ProfileCount(Entry.second, Function::PCT_Synthetic));
 
   return PreservedAnalyses::all();
 }