[Mlir-commits] [mlir] [mlir][CallGraph] Add edges for callable symbol references in CallGraph (PR #116177)

[Haocong Lu]

================
@@ -93,42 +127,41 @@ static void computeCallGraph(Operation *op, CallGraph &cg,
 
   for (Region &region : op->getRegions())
     for (Operation &nested : region.getOps())
-      computeCallGraph(&nested, cg, symbolTable, parentNode, resolveCalls);
+      computeCallGraph(&nested, cg, symbolTable, parentNode, resolveSymbolRef);
 }
 
 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.
+  // resolve the calls and symbol references. 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);
+  computeCallGraph(op, *this, symbolTable, getEntryNode(),
+                   /*resolveSymbolRef=*/false);
+  computeCallGraph(op, *this, symbolTable, getEntryNode(),
+                   /*resolveSymbolRef=*/true);
 }
 
 /// Get or add a call graph node for the given region.
 CallGraphNode *CallGraph::getOrAddNode(Region *region,
                                        CallGraphNode *parentNode) {
-  assert(region && isa<CallableOpInterface>(region->getParentOp()) &&
+  Operation *parentOp = region->getParentOp();
+  assert(region && isa<CallableOpInterface>(parentOp) &&
          "expected parent operation to be callable");
   std::unique_ptr<CallGraphNode> &node = nodes[region];
   if (!node) {
     node.reset(new CallGraphNode(region));
 
     // Add this node to the given parent node if necessary.
-    if (parentNode) {
+    assert(parentNode && "expected non-empty parent node");
+    if (!parentNode->isExternal()) {
       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.
-      externalCallerNode.addAbstractEdge(node.get());
+    } else if (auto symbolOp = dyn_cast<SymbolOpInterface>(parentOp)) {
+      // Otherwise, connect all symbol nodes with public visibility
+      // to the entry node, which may be referenced externally.
+      if (!symbolOp.canDiscardOnUseEmpty())
+        parentNode->addRefEdge(node.get());
----------------
Luhaocong wrote:

Although the default implementation of `canDiscardOnUseEmpty` is `getVisibility() != Visibility::Public`, I think it's a more accurate condition to indicate whether this node may be referenced somewhere. e.g. A public symbol with  `canDiscardOnUseEmpty == true` must not be referenced by external node, but a private symbol with `canDiscardOnUseEmpty == false` should be considered as having potential references. 

https://github.com/llvm/llvm-project/pull/116177