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

[Haocong Lu]

================
@@ -73,14 +90,31 @@ void CallGraphNode::addEdge(CallGraphNode *node, Edge::Kind kind) {
 /// edges are placed into the given callgraph object.
 static void computeCallGraph(Operation *op, CallGraph &cg,
                              SymbolTableCollection &symbolTable,
-                             CallGraphNode *parentNode, bool resolveCalls) {
-  if (CallOpInterface call = dyn_cast<CallOpInterface>(op)) {
-    // If there is no parent node, we ignore this operation. Even if this
-    // operation was a call, there would be no callgraph node to attribute it
-    // to.
-    if (resolveCalls && parentNode)
-      parentNode->addCallEdge(cg.resolveCallable(call, symbolTable));
-    return;
+                             CallGraphNode *parentNode, bool resolveSymbolRef) {
+  if (resolveSymbolRef) {
+    bool skipCalleeSymbolRef = false;
+    if (auto call = dyn_cast<CallOpInterface>(op)) {
+      // If there is no parent node, there would be no callgraph node to call it
+      // directly, we use a reference egde to represent this call.
+      if (!parentNode->isExternal()) {
+        parentNode->addCallEdge(cg.resolveCallable(call, symbolTable));
+        if (isa<SymbolRefAttr>(call.getCallableForCallee()))
+          skipCalleeSymbolRef = true;
+      }
+    }
+    // If an operation reference a callable symbol, create a reference edge.
+    op->getAttrDictionary().walk<WalkOrder::PreOrder>(
+        [&](SymbolRefAttr symbolRef) {
+          // Skip the first symbol reference if it is a resolved callee.
+          if (skipCalleeSymbolRef) {
+            skipCalleeSymbolRef = false;
+            return WalkResult::skip();
+          }
----------------
Luhaocong wrote:

looks better to me, thanks.

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