[CallGraph] Refine call graph for indirect calls with !callees metadata

For indirect call sites having a small set of possible callees,
!callees metadata can be used to indicate what those callees are.
This patch updates the call graph and lazy call graph analyses so
that they consider this metadata when encountering call sites. For
the call graph, it adds a new external call graph node to the graph
for each unique !callees metadata node. A call graph edge connects
an indirect call site with the external node associated with the
!callees metadata that is attached to it. And there is an edge from
this external node to each of the callees indicated by the metadata.
Similarly, for the lazy call graph, the patch adds Ref edges from a
caller to the possible callees indicated by the metadata.

The primary purpose of the patch is to facilitate iterating over the
functions in a module such that all of the callees indicated by a
given !callees metadata node will be visited prior to the functions
containing call sites annotated by that node. This property is
required by optimizations performing a bottom-up traversal of the
SCC DAG. For example, the inliner can be made to inline through an
indirect call. If the call site is annotated with !callees metadata,
this patch ensures that the inliner will have visited all of the
callees prior to the caller, allowing it to reliably compute the
cost of inlining one or more of the potential callees.

Original patch by @mssimpso. I've made some small changes to get it
to apply, build, and pass tests on the top of tree, as well as
some minor tweaks to formatting and functionality.

Subscribers: mehdi_amini, hiraditya, llvm-commits, mssimpso

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D39339

llvm-svn: 369025

3 files changed, 44 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/CGSCCPassManager.cpp b/llvm/lib/Analysis/CGSCCPassManager.cpp
index a0b3f83cca6..54886d45a1d 100644
--- a/llvm/lib/Analysis/CGSCCPassManager.cpp
+++ b/llvm/lib/Analysis/CGSCCPassManager.cpp
@@ -449,7 +449,7 @@ LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
   // irrelevant.
   for (Instruction &I : instructions(F))
     if (auto CS = CallSite(&I))
-      if (Function *Callee = CS.getCalledFunction())
+      for (Function *Callee : CS.getKnownCallees()) {
         if (Visited.insert(Callee).second && !Callee->isDeclaration()) {
           Node &CalleeN = *G.lookup(*Callee);
           Edge *E = N->lookup(CalleeN);
@@ -467,6 +467,7 @@ LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
           if (!E->isCall())
             PromotedRefTargets.insert(&CalleeN);
         }
+      }
 
   // Now walk all references.
   for (Instruction &I : instructions(F))
diff --git a/llvm/lib/Analysis/CallGraph.cpp b/llvm/lib/Analysis/CallGraph.cpp
index 70aeb1a688e..a7bb16299a5 100644
--- a/llvm/lib/Analysis/CallGraph.cpp
+++ b/llvm/lib/Analysis/CallGraph.cpp
@@ -37,6 +37,7 @@ CallGraph::CallGraph(Module &M)
 
 CallGraph::CallGraph(CallGraph &&Arg)
     : M(Arg.M), FunctionMap(std::move(Arg.FunctionMap)),
+      DummyNodeMap(std::move(Arg.DummyNodeMap)),
       ExternalCallingNode(Arg.ExternalCallingNode),
       CallsExternalNode(std::move(Arg.CallsExternalNode)) {
   Arg.FunctionMap.clear();
@@ -53,6 +54,8 @@ CallGraph::~CallGraph() {
 #ifndef NDEBUG
   for (auto &I : FunctionMap)
     I.second->allReferencesDropped();
+  for (auto &N : DummyNodeMap)
+    N.second->allReferencesDropped();
 #endif
 }
 
@@ -74,7 +77,14 @@ void CallGraph::addToCallGraph(Function *F) {
     for (Instruction &I : BB) {
       if (auto *Call = dyn_cast<CallBase>(&I)) {
         const Function *Callee = Call->getCalledFunction();
-        if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
+        MDNode *CalleesMD = I.getMetadata(LLVMContext::MD_callees);
+        if (!Callee && CalleesMD)
+          // If an indirect call site has !callees metadata indicating its
+          // possible callees, we add an edge from the call site to a dummy
+          // node. When we construct the dummy node, we add edges from it to
+          // the functions indicated in the !callees metadata.
+          Node->addCalledFunction(Call, getOrInsertNodeForCalleesMD(CalleesMD));
+        else if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
           // Indirect calls of intrinsics are not allowed so no need to check.
           // We can be more precise here by using TargetArg returned by
           // Intrinsic::isLeaf.
@@ -105,6 +115,11 @@ void CallGraph::print(raw_ostream &OS) const {
 
   for (CallGraphNode *CN : Nodes)
     CN->print(OS);
+
+  // The iteration order of the DummyNodeMap is deterministic, so we don't need
+  // to sort the nodes. Just print them.
+  for (auto &Entry : DummyNodeMap)
+    Entry.second->print(OS);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
@@ -120,6 +135,12 @@ LLVM_DUMP_METHOD void CallGraph::dump() const { print(dbgs()); }
 Function *CallGraph::removeFunctionFromModule(CallGraphNode *CGN) {
   assert(CGN->empty() && "Cannot remove function from call "
          "graph if it references other functions!");
+
+  // If any dummy node references the node for the given function, we first
+  // need to remove those edges.
+  for (auto &Entry : DummyNodeMap)
+    Entry.second->removeAnyCallEdgeTo(CGN);
+
   Function *F = CGN->getFunction(); // Get the function for the call graph node
   FunctionMap.erase(F);             // Remove the call graph node from the map
 
@@ -154,6 +175,21 @@ CallGraphNode *CallGraph::getOrInsertFunction(const Function *F) {
   return CGN.get();
 }
 
+CallGraphNode *CallGraph::getOrInsertNodeForCalleesMD(MDNode *Callees) {
+  auto &CGN = DummyNodeMap[Callees];
+  if (CGN)
+    return CGN.get();
+  CGN = llvm::make_unique<CallGraphNode>(nullptr);
+  for (const MDOperand &Op : Callees->operands())
+    if (auto *MDConstant = mdconst::extract_or_null<Constant>(Op)) {
+      auto *F = cast<Function>(MDConstant);
+
+      assert(!F->isIntrinsic());
+      CGN->addCalledFunction(nullptr, getOrInsertFunction(F));
+    }
+  return CGN.get();
+}
+
 //===----------------------------------------------------------------------===//
 // Implementations of the CallGraphNode class methods.
 //
@@ -171,7 +207,7 @@ void CallGraphNode::print(raw_ostream &OS) const {
     if (Function *FI = I.second->getFunction())
       OS << "function '" << FI->getName() <<"'\n";
     else
-      OS << "external node\n";
+      OS << "<<null function>><<" << I.second << ">>\n";
   }
   OS << '\n';
 }
diff --git a/llvm/lib/Analysis/LazyCallGraph.cpp b/llvm/lib/Analysis/LazyCallGraph.cpp
index 08d6e76ea03..c832dc94877 100644
--- a/llvm/lib/Analysis/LazyCallGraph.cpp
+++ b/llvm/lib/Analysis/LazyCallGraph.cpp
@@ -100,12 +100,14 @@ LazyCallGraph::EdgeSequence &LazyCallGraph::Node::populateSlow() {
   for (BasicBlock &BB : *F)
     for (Instruction &I : BB) {
       if (auto CS = CallSite(&I))
-        if (Function *Callee = CS.getCalledFunction())
+        for (Function *Callee : CS.getKnownCallees())
           if (!Callee->isDeclaration())
             if (Callees.insert(Callee).second) {
               Visited.insert(Callee);
+              auto EdgeK = CS.getCalledFunction() ? LazyCallGraph::Edge::Call
+                                                  : LazyCallGraph::Edge::Ref;
               addEdge(Edges->Edges, Edges->EdgeIndexMap, G->get(*Callee),
-                      LazyCallGraph::Edge::Call);
+                      EdgeK);
             }
 
       for (Value *Op : I.operand_values())