[mlir][CallGraph] Add special call graph node for representing unknown callees

[zero9178]

Summary

The callgraph currently contains a special external node that is used both as the quasi caller for any externally callable as well as callees that could not be resolved.
This has one negative side effect however, which is the motivation for this patch: It leads to every externally callable which contains a call that could not be resolved (eg. an indirect call), to be put into one giant SCC when iterating over the SCCs of the call graph.

This patch fixes that issue by creating a second special callgraph node that acts as the callee for any unresolved callable. This breaks the cycles produced in the callgraph, yielding proper SCCs for all direct calls.

mlir/include/mlir/Analysis/CallGraph.h

public:
   /// corresponds to the direct node in the callgraph that contains the parent
   /// operation of `region`, or nullptr if there is no parent node.
   CallGraphNode *getOrAddNode(Region *region, CallGraphNode *parentNode);
 
   /// Lookup a call graph node for the given region, or nullptr if none is
   /// registered.
   CallGraphNode *lookupNode(Region *region) const;
 
-  /// Return the callgraph node representing the indirect-external callee.
-  CallGraphNode *getExternalNode() const {
-    return const_cast<CallGraphNode *>(&externalNode);
+  /// Return the callgraph node representing an external caller.
+  CallGraphNode *getExternalCallerNode() const {
+    return const_cast<CallGraphNode *>(&externalCallerNode);
+  }
+
+  /// Return the callgraph node representing an indirect callee.
+  CallGraphNode *getUnknownCalleeNode() const {
+    return const_cast<CallGraphNode *>(&unknownCalleeNode);
   }
 
   /// Resolve the callable for given callee to a node in the callgraph, or the
   /// external node if a valid node was not resolved. The provided symbol table
   /// is used when resolving calls that reference callables via a symbol
   /// reference.
   CallGraphNode *resolveCallable(CallOpInterface call,
                                  SymbolTableCollection &symbolTable) const;
public:
   /// Dump the graph in a human readable format.
   void dump() const;
   void print(raw_ostream &os) const;
 
 private:
   /// The set of nodes within the callgraph.
   NodeMapT nodes;
 
-  /// A special node used to indicate an external edges.
-  CallGraphNode externalNode;
+  /// A special node used to indicate an external caller.
+  CallGraphNode externalCallerNode;
+
+  /// A special node used to indicate an unknown callee.
+  CallGraphNode unknownCalleeNode;
 };
 
 } // namespace mlir
 
 namespace llvm {
 // Provide graph traits for traversing call graphs using standard graph
 // traversals.
 template <>
struct GraphTraits<const mlir::CallGraphNode *> {
   }
 };
 
 template <>
 struct GraphTraits<const mlir::CallGraph *>
     : public GraphTraits<const mlir::CallGraphNode *> {
   /// The entry node into the graph is the external node.
   static NodeRef getEntryNode(const mlir::CallGraph *cg) {
-    return cg->getExternalNode();
+    return cg->getExternalCallerNode();
   }
 
   // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
   using nodes_iterator = mlir::CallGraph::iterator;
   static nodes_iterator nodes_begin(mlir::CallGraph *cg) { return cg->begin(); }
   static nodes_iterator nodes_end(mlir::CallGraph *cg) { return cg->end(); }
 };
 } // namespace llvm
 
 #endif // MLIR_ANALYSIS_CALLGRAPH_H

mlir/lib/Analysis/CallGraph.cpp

else
       return;
   }
 
   for (Region &region : op->getRegions())
     for (Operation &nested : region.getOps())
       computeCallGraph(&nested, cg, symbolTable, parentNode, resolveCalls);
 }
 
-CallGraph::CallGraph(Operation *op) : externalNode(/*callableRegion=*/nullptr) {
+CallGraph::CallGraph(Operation *op)
+    : externalCallerNode(/*callableRegion=*/nullptr),
+      unknownCalleeNode(/*callableRegion=*/nullptr) {
   // Make two passes over the graph, one to compute the callables and one to
   // resolve the calls. We split these up as we may have nested callable objects
   // that need to be reserved before the calls.
   SymbolTableCollection symbolTable;
   computeCallGraph(op, *this, symbolTable, /*parentNode=*/nullptr,
                    /*resolveCalls=*/false);
   computeCallGraph(op, *this, symbolTable, /*parentNode=*/nullptr,
                    /*resolveCalls=*/true);
if (parentNode) {
       parentNode->addChildEdge(node.get());
     } else {
       // Otherwise, connect all callable nodes to the external node, this allows
       // for conservatively including all callable nodes within the graph.
       // FIXME This isn't correct, this is only necessary for callable nodes
       // that *could* be called from external sources. This requires extending
       // the interface for callables to check if they may be referenced
       // externally.
-      externalNode.addAbstractEdge(node.get());
+      externalCallerNode.addAbstractEdge(node.get());
     }
   }
   return node.get();
 }
 
 /// Lookup a call graph node for the given region, or nullptr if none is
 /// registered.
 CallGraphNode *CallGraph::lookupNode(Region *region) const {
   auto it = nodes.find(region);
   return it == nodes.end() ? nullptr : it->second.get();
 }
 
 /// Resolve the callable for given callee to a node in the callgraph, or the
-/// external node if a valid node was not resolved.
+/// unknown callee node if a valid node was not resolved.
 CallGraphNode *
 CallGraph::resolveCallable(CallOpInterface call,
                            SymbolTableCollection &symbolTable) const {
   Operation *callable = call.resolveCallable(&symbolTable);
   if (auto callableOp = dyn_cast_or_null<CallableOpInterface>(callable))
     if (auto *node = lookupNode(callableOp.getCallableRegion()))
       return node;
 
-  // If we don't have a valid direct region, this is an external call.
-  return getExternalNode();
+  return getUnknownCalleeNode();
 }
 
 /// Erase the given node from the callgraph.
 void CallGraph::eraseNode(CallGraphNode *node) {
   // Erase any children of this node first.
   if (node->hasChildren()) {
     for (const CallGraphNode::Edge &edge : llvm::make_early_inc_range(*node))
       if (edge.isChild())

 
 /// Dump the graph in a human readable format.
 void CallGraph::dump() const { print(llvm::errs()); }
 void CallGraph::print(raw_ostream &os) const {
   os << "// ---- CallGraph ----\n";
 
   // Functor used to output the name for the given node.
   auto emitNodeName = [&](const CallGraphNode *node) {
-    if (node->isExternal()) {
-      os << "<External-Node>";
+    if (node == getExternalCallerNode()) {
+      os << "<External-Caller-Node>";
+      return;
+    }
+    if (node == getUnknownCalleeNode()) {
+      os << "<Unknown-Callee-Node>";
       return;
     }
 
     auto *callableRegion = node->getCallableRegion();
     auto *parentOp = callableRegion->getParentOp();
     os << "'" << callableRegion->getParentOp()->getName() << "' - Region #"
        << callableRegion->getRegionNumber();
     auto attrs = parentOp->getAttrDictionary();

mlir/test/Analysis/test-callgraph.mlir

 // RUN: mlir-opt %s -test-print-callgraph -split-input-file -allow-unregistered-dialect 2>&1 | FileCheck %s
 
 // CHECK-LABEL: Testing : "simple"
 module attributes {test.name = "simple"} {
 
   // CHECK: Node{{.*}}func_a
   func.func @func_a() {
     return
   }
 
   func.func private @func_b()
 
   // CHECK: Node{{.*}}func_c
-  // CHECK-NEXT: Call-Edge{{.*}}External-Node
+  // CHECK-NEXT: Call-Edge{{.*}}Unknown-Callee-Node
   func.func @func_c() {
     call @func_b() : () -> ()
     return
   }
 
   // CHECK: Node{{.*}}func_d
   // CHECK-NEXT: Call-Edge{{.*}}func_c
   func.func @func_d() {
module attributes {test.name = "nested"} {
 
   // CHECK: Node{{.*}}func_b
   // CHECK: Call-Edge{{.*}}func_a
   func.func @func_b() {
     "test.conversion_call_op"() { callee = @nested_module::@func_a } : () -> ()
     return
   }
 }
+
+// -----
+
+// CHECK-LABEL: Testing : "SCC"
+// CHECK: SCCs
+module attributes {test.name = "SCC"} {
+  // CHECK: SCC :
+  // CHECK-NEXT: Node{{.*}}Unknown-Callee-Node
+
+  // CHECK: SCC :
+  // CHECK-NEXT: Node{{.*}}foo
+  func.func @foo(%arg0 : () -> ()) {
+    call_indirect %arg0() : () -> ()
+    return
+  }
+
+  // CHECK: SCC :
+  // CHECK-NEXT: Node{{.*}}bar
+  func.func @bar(%arg1 : () -> ()) {
+    call_indirect %arg1() : () -> ()
+    return
+  }
+
+  // CHECK: SCC :
+  // CHECK-NEXT: Node{{.*}}External-Caller-Node
+}
+