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

[rudkx]

Summary

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.

[mssimpso]

• Incorporated the change into the lazy call graph as well. The metadata is represented as ref edges there.
• Added a helper in CallSite to get the functions listed in the metadata.

[mssimpso]

Ping.

[mssimpso]

Ping.

[mssimpso]

Updated to work with the latest revision of D39869.

[rudkx]

@mssimpso Is this something you're still interested in pushing forward?

I was just looking at how we could use callee metadata to improve code, and was going to implement a patch like this to update the call graph so that we had more accurate bottom-up code generation, and then started poking around and found this patch.

[mssimpso]

In D39339#1596064, @rudkx wrote:

@mssimpso Is this something you're still interested in pushing forward?

I was just looking at how we could use callee metadata to improve code, and was going to implement a patch like this to update the call graph so that we had more accurate bottom-up code generation, and then started poking around and found this patch.

I'm no longer working on this, so feel free to take it over if you like. The purpose of the patch was to refine the iteration order over functions, and there was a related patch (D39869) that used the !callees metadata for inlining.

[rudkx]

I'll take a look at getting this or something like it in.

[rudkx]

Apply the original changes to top of tree.

[bogner]

llvm/include/llvm/Analysis/CallGraph.h

88 ↗

(On Diff #215097)

Does this really need to be void *? AFAICT it's always CallBase * or null. If that's not correct, it'd probably still be better to use PointerUnion or something.

llvm/lib/Analysis/CallGraph.cpp

182–189 ↗

(On Diff #215097)

Maybe I'm misreading this, but isn't CGN == nullptr here, given the assignment just above?

[rudkx]

llvm/include/llvm/Analysis/CallGraph.h

88 ↗

(On Diff #215097)

These are actually keyed by MDNode *. I suspect the intent here was to make it more generally useful for creating dummy nodes mapped from various keys. I will change it to MDNode * and it can be generalized in the future if needed.

llvm/lib/Analysis/CallGraph.cpp

182–189 ↗

(On Diff #215097)

GCN should be a non-null CallGraphNode holding a Function * that has a nullptr value.

[bogner]

llvm/lib/Analysis/CallGraph.cpp

182–189 ↗

(On Diff #215097)

Ah right. In that case, doesn't the node get destroyed when the unique_ptr goes out of scope here? You may have meant to use .release(), but in that case who owns this object?

[rudkx]

Address review feedback regarding using void* as a key.

[rudkx]

llvm/lib/Analysis/CallGraph.cpp

182–189 ↗

(On Diff #215097)

GCN is a reference to the unique_ptr<T> in the map, and unique_ptr has a move assignment operator, so it transfers ownership to the map node right after it is allocated.

[bogner]

LGTM

llvm/lib/Analysis/CallGraph.cpp

182–189 ↗

(On Diff #215097)

Ah, I see. The use of auto & is subtle and I didn't notice that we were actually assigning a reference here. Nevermind.

include/llvm/Analysis/CallGraph.h

 /// transformations.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_CALLGRAPH_H
 #define LLVM_ANALYSIS_CALLGRAPH_H
 
 #include "llvm/ADT/GraphTraits.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include <cassert>

 ///
 /// The core call graph itself can also be updated to reflect changes to the IR.
 class CallGraph {
   Module &M;
 
   using FunctionMapTy =
       std::map<const Function *, std::unique_ptr<CallGraphNode>>;
 
+  /// \brief A type for maintaining dummy nodes.
+  ///
+  /// Dummy nodes (i.e., nodes having a null function) include, for example,
+  /// those created to represent !callees metadata. We use a void pointer as
+  /// the key to allow for various kinds of dummy nodes. We use a MapVector to
+  /// ensure a deterministic iteration order (there's no good way to sort dummy
+  /// nodes). A deterministic iteration order is primarily useful for printing.
+  using DummyNodeMapTy =
+      MapVector<const void *, std::unique_ptr<CallGraphNode>>;
+
   /// \brief A map from \c Function* to \c CallGraphNode*.
   FunctionMapTy FunctionMap;
 
+  /// \brief A map for maintaining dummy nodes.
+  DummyNodeMapTy DummyNodeMap;
+
   /// \brief This node has edges to all external functions and those internal
   /// functions that have their address taken.
   CallGraphNode *ExternalCallingNode;
 
   /// \brief This node has edges to it from all functions making indirect calls
   /// or calling an external function.
   std::unique_ptr<CallGraphNode> CallsExternalNode;
 
   /// \brief Replace the function represented by this node by another.
   ///
   /// This does not rescan the body of the function, so it is suitable when
   /// splicing the body of one function to another while also updating all
   /// callers from the old function to the new.
   void spliceFunction(const Function *From, const Function *To);
 
   /// \brief Add a function to the call graph, and link the node to all of the
   /// functions that it calls.
   void addToCallGraph(Function *F);
 
+  /// \brief Return the dummy node associated with the given pointer.
+  CallGraphNode *getOrInsertDummyNode(const void *P);
+
 public:
   explicit CallGraph(Module &M);
   CallGraph(CallGraph &&Arg);
   ~CallGraph();
 
   void print(raw_ostream &OS) const;
   void dump() const;
 

include/llvm/IR/CallSite.h

#undef CALLSITE_DELEGATE_SETTER
   bool hasArgument(const Value *Arg) const {
     for (arg_iterator AI = this->arg_begin(), E = this->arg_end(); AI != E;
          ++AI)
       if (AI->get() == Arg)
         return true;
     return false;
   }
 
+  /// Return the set of functions this call site is known to call. For direct
+  /// calls, the returned set contains the called function. For indirect calls,
+  /// this function collects known callees from !callees metadata, if present.
+  SmallVector<FunTy *, 1> getKnownCallees() {
+    SmallVector<FunTy *, 1> Callees;
+    InstrTy *Inst = getInstruction();
+    if (getCalledFunction()) {
+      // If the call site is direct, just add the called function to the set.
+      Callees.push_back(getCalledFunction());
+    } else if (auto *Node = Inst->getMetadata(LLVMContext::MD_callees)) {
+      // Otherwise, if the call site is indirect, collect the known callees from
+      // !callees metadata if present.
+      for (const MDOperand &Op : Node->operands())
+        if (auto *MDConstant = mdconst::extract_or_null<Constant>(Op))
+          Callees.push_back(cast<FunTy>(MDConstant->stripPointerCasts()));
+      if (!Callees.empty()) {
+        // Ensure a stable iteration order over the known callees. To do this,
+        // we sort the functions by name.
+        std::sort(Callees.begin(), Callees.end(),
+                  [](const Function *RHS, const Function *LHS) {
+                    return RHS->getName() < LHS->getName();
+                  });
+        Callees.erase(std::unique(Callees.begin(), Callees.end()),
+                      Callees.end());
+      }
+    }
+    return Callees;
+  }
+
 private:
   IterTy getCallee() const {
     if (isCall()) // Skip Callee
       return cast<CallInst>(getInstruction())->op_end() - 1;
     else // Skip BB, BB, Callee
       return cast<InvokeInst>(getInstruction())->op_end() - 3;
   }
 };

lib/Analysis/CGSCCPassManager.cpp

LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
   SmallSetVector<Node *, 4> PromotedRefTargets;
   SmallSetVector<Node *, 4> DemotedCallTargets;
 
   // First walk the function and handle all called functions. We do this first
   // because if there is a single call edge, whether there are ref edges is
   // 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);
           // FIXME: We should really handle adding new calls. While it will
           // make downstream usage more complex, there is no fundamental
           // limitation and it will allow passes within the CGSCC to be a bit
           // more flexible in what transforms they can do. Until then, we
           // verify that new calls haven't been introduced.

lib/Analysis/CallGraph.cpp

: M(M), ExternalCallingNode(getOrInsertFunction(nullptr)),
       CallsExternalNode(llvm::make_unique<CallGraphNode>(nullptr)) {
   // Add every function to the call graph.
   for (Function &F : M)
     addToCallGraph(&F);
 }
 
 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();
   Arg.ExternalCallingNode = nullptr;
 }
 
 CallGraph::~CallGraph() {
   // CallsExternalNode is not in the function map, delete it explicitly.
   if (CallsExternalNode)
     CallsExternalNode->allReferencesDropped();
 
 // Reset all node's use counts to zero before deleting them to prevent an
 // assertion from firing.
 #ifndef NDEBUG
   for (auto &I : FunctionMap)
     I.second->allReferencesDropped();
+  for (auto &I : DummyNodeMap)
+    I.second->allReferencesDropped();
 #endif
 }
 
 void CallGraph::addToCallGraph(Function *F) {
   CallGraphNode *Node = getOrInsertFunction(F);
 
   // If this function has external linkage or has its address taken, anything
   // could call it.
   if (!F->hasLocalLinkage() || F->hasAddressTaken())
     ExternalCallingNode->addCalledFunction(CallSite(), Node);
 
   // If this function is not defined in this translation unit, it could call
   // anything.
   if (F->isDeclaration() && !F->isIntrinsic())
     Node->addCalledFunction(CallSite(), CallsExternalNode.get());
 
   // Look for calls by this function.
   for (BasicBlock &BB : *F)
     for (Instruction &I : BB) {
       if (auto CS = CallSite(&I)) {
         const Function *Callee = CS.getCalledFunction();
-        if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
+        MDNode *CalleesMDNode = I.getMetadata(LLVMContext::MD_callees);
+        if (!Callee && CalleesMDNode) {
+          // 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.
+          bool InitCallees = !DummyNodeMap.count(CalleesMDNode);
+          CallGraphNode *Dummy = getOrInsertDummyNode(CalleesMDNode);
+          Node->addCalledFunction(CS, Dummy);
+          if (InitCallees)
+            for (Function *PossibleCallee : CS.getKnownCallees())
+              if (!PossibleCallee->isIntrinsic())
+                Dummy->addCalledFunction(CallSite(),
+                                         getOrInsertFunction(PossibleCallee));
+        } 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.
           Node->addCalledFunction(CS, CallsExternalNode.get());
-        else if (!Callee->isIntrinsic())
+        } else if (!Callee->isIntrinsic()) {
           Node->addCalledFunction(CS, getOrInsertFunction(Callee));
         }
       }
     }
+}
 
 void CallGraph::print(raw_ostream &OS) const {
   // Print in a deterministic order by sorting CallGraphNodes by name.  We do
   // this here to avoid slowing down the non-printing fast path.
 
   SmallVector<CallGraphNode *, 16> Nodes;
   Nodes.reserve(FunctionMap.size());
 
   for (const auto &I : *this)
     Nodes.push_back(I.second.get());
 
   std::sort(Nodes.begin(), Nodes.end(),
             [](CallGraphNode *LHS, CallGraphNode *RHS) {
     if (Function *LF = LHS->getFunction())
       if (Function *RF = RHS->getFunction())
         return LF->getName() < RF->getName();
 
     return RHS->getFunction() != nullptr;
   });
 
   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)
 LLVM_DUMP_METHOD void CallGraph::dump() const { print(dbgs()); }
 #endif
 
 // removeFunctionFromModule - Unlink the function from this module, returning
 // it.  Because this removes the function from the module, the call graph node
 // is destroyed.  This is only valid if the function does not call any other
 // functions (ie, there are no edges in it's CGN).  The easiest way to do this
 // is to dropAllReferences before calling this.
 //
 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
 
   M.getFunctionList().remove(F);
   return F;
 }
 
 /// spliceFunction - Replace the function represented by this node by another.
CallGraphNode *CallGraph::getOrInsertFunction(const Function *F) {
   if (CGN)
     return CGN.get();
 
   assert((!F || F->getParent() == &M) && "Function not in current module!");
   CGN = llvm::make_unique<CallGraphNode>(const_cast<Function *>(F));
   return CGN.get();
 }
 
+CallGraphNode *CallGraph::getOrInsertDummyNode(const void *P) {
+  auto &CGN = DummyNodeMap[P];
+  if (CGN)
+    return CGN.get();
+  CGN = llvm::make_unique<CallGraphNode>(nullptr);
+  return CGN.get();
+}
+
 //===----------------------------------------------------------------------===//
 // Implementations of the CallGraphNode class methods.
 //
 
 void CallGraphNode::print(raw_ostream &OS) const {
   if (Function *F = getFunction())
     OS << "Call graph node for function: '" << F->getName() << "'";
   else
     OS << "Call graph node <<null function>>";
   
   OS << "<<" << this << ">>  #uses=" << getNumReferences() << '\n';
 
   for (const auto &I : *this) {
     OS << "  CS<" << I.first << "> calls ";
     if (Function *FI = I.second->getFunction())
       OS << "function '" << FI->getName() <<"'\n";
     else
-      OS << "external node\n";
+      OS << "<<null function>><<" << I.second << ">>\n";
   }
   OS << '\n';
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void CallGraphNode::dump() const { print(dbgs()); }
 #endif
 

lib/Analysis/LazyCallGraph.cpp

LazyCallGraph::EdgeSequence &LazyCallGraph::Node::populateSlow() {
   // used. For example, this could be done by moving the weak definition to
   // a strong (internal) definition and making the weak definition be an
   // alias. Then a test of the address of the weak function against the new
   // strong definition's address would be an effective way to determine the
   // safety of optimizing a direct call edge.
   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())
         if (Constant *C = dyn_cast<Constant>(Op))
           if (Visited.insert(C).second)
             Worklist.push_back(C);
     }
 

test/Analysis/CallGraph/callees-metadata.ll

+; RUN: opt < %s -print-callgraph -disable-output 2>&1 | FileCheck %s
+
+; CHECK:      Call graph node <<null function>><<{{.+}}>> #uses=0
+; CHECK-DAG:    CS<0x0> calls function 'main'
+; CHECK-DAG:    CS<0x0> calls function 'add'
+; CHECK-DAG:    CS<0x0> calls function 'sub'
+;
+; CHECK:      Call graph node for function: 'add'<<{{.+}}>> #uses=2
+;
+; CHECK:      Call graph node for function: 'main'<<{{.+}}>> #uses=1
+; CHECK-NEXT:   CS<{{.+}}> calls <<null function>><<[[CALLEES:.+]]>>
+;
+; CHECK:      Call graph node for function: 'sub'<<{{.+}}>> #uses=2
+;
+; CHECK:      Call graph node <<null function>><<[[CALLEES]]>> #uses=1
+; CHECK-DAG:    CS<0x0> calls function 'add'
+; CHECK-DAG:    CS<0x0> calls function 'sub'
+
+define i64 @main(i64 %x, i64 %y, i64 (i64, i64)* %binop) {
+  %tmp0 = call i64 %binop(i64 %x, i64 %y),
+     !callees !{i64 (i64, i64)* @add, i64 (i64, i64)* @sub}
+  ret i64 %tmp0
+}
+
+define i64 @add(i64 %x, i64 %y) {
+  %tmp0 = add i64 %x, %y
+  ret i64 %tmp0
+}
+
+define i64 @sub(i64 %x, i64 %y) {
+  %tmp0 = sub i64 %x, %y
+  ret i64 %tmp0
+}

test/Analysis/CallGraph/non-leaf-intrinsics.ll

entry:
   ret void
 }
 
 
 ; CHECK: Call graph node <<null function>>
 ; CHECK:  CS<0x0> calls function 'f'
 
 ; CHECK: Call graph node for function: 'calls_patchpoint'
-; CHECK-NEXT:  CS<[[addr_1:[^>]+]]> calls external node
+; CHECK-NEXT:  CS<[[addr_1:[^>]+]]> calls <<null function>>
 
 ; CHECK: Call graph node for function: 'calls_statepoint'
-; CHECK-NEXT:  CS<[[addr_0:[^>]+]]> calls external node
+; CHECK-NEXT:  CS<[[addr_0:[^>]+]]> calls <<null function>>

test/Analysis/LazyCallGraph/callees-metadata.ll

+; RUN: opt < %s -passes=print-lcg -disable-output 2>&1 | FileCheck %s
+
+; CHECK:      Edges in function: main
+; CHECK-DAG:    ref  -> add
+; CHECK-DAG:    ref  -> sub
+;
+; CHECK:      Edges in function: add
+;
+; CHECK:      Edges in function: sub
+;
+; CHECK:      RefSCC with 1 call SCCs:
+; CHECK-NEXT:   SCC with 1 functions:
+; CHECK-NEXT:     sub
+;
+; CHECK:      RefSCC with 1 call SCCs:
+; CHECK-NEXT:   SCC with 1 functions:
+; CHECK-NEXT:     add
+;
+; CHECK:      RefSCC with 1 call SCCs:
+; CHECK-NEXT:   SCC with 1 functions:
+; CHECK-NEXT:     main
+
+define i64 @main(i64 %x, i64 %y, i64 (i64, i64)* %binop) {
+  %tmp0 = call i64 %binop(i64 %x, i64 %y),
+     !callees !{i64 (i64, i64)* @add, i64 (i64, i64)* @sub}
+  ret i64 %tmp0
+}
+
+define i64 @add(i64 %x, i64 %y) {
+  %tmp0 = add i64 %x, %y
+  ret i64 %tmp0
+}
+
+define i64 @sub(i64 %x, i64 %y) {
+  %tmp0 = sub i64 %x, %y
+  ret i64 %tmp0
+}