Add initial callgraph support.

Using the two call interfaces, CallOpInterface and CallableOpInterface, this change adds support for an initial multi-level CallGraph. This call graph builds a set of nodes for each callable region, and connects them via edges. An edge may be any of the following types:
* Abstract
  - An edge not produced by a call operation, used for connecting to internal nodes from external nodes.
* Call
  - A call edge is an edge defined via a call-like operation.
* Child
  - This is an artificial edge connecting nested callgraph nodes.

This callgraph will be used, and improved upon, to begin supporting more interesting interprocedural analyses and transformation. In a followup, this callgraph will be used to support more complex inlining support.

PiperOrigin-RevId: 270724968

1 files changed, 212 insertions, 0 deletions

diff --git a/mlir/lib/Analysis/CallGraph.cpp b/mlir/lib/Analysis/CallGraph.cpp
index 6eea929d5d2..2936e8a3f0e 100644
--- a/mlir/lib/Analysis/CallGraph.cpp
+++ b/mlir/lib/Analysis/CallGraph.cpp
@@ -19,8 +19,13 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Analysis/CallGraph.h"
 #include "mlir/Analysis/CallInterfaces.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/SCCIterator.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
@@ -29,3 +34,210 @@ using namespace mlir;
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/CallInterfaces.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// CallGraphNode
+//===----------------------------------------------------------------------===//
+
+/// Returns if this node refers to the indirect/external node.
+bool CallGraphNode::isExternal() const { return !callableRegion; }
+
+/// Return the callable region this node represents. This can only be called
+/// on non-external nodes.
+Region *CallGraphNode::getCallableRegion() const {
+  assert(!isExternal() && "the external node has no callable region");
+  return callableRegion;
+}
+
+/// Adds an reference edge to the given node. This is only valid on the
+/// external node.
+void CallGraphNode::addAbstractEdge(CallGraphNode *node) {
+  assert(isExternal() && "abstract edges are only valid on external nodes");
+  addEdge(node, Edge::Kind::Abstract);
+}
+
+/// Add an outgoing call edge from this node.
+void CallGraphNode::addCallEdge(CallGraphNode *node) {
+  addEdge(node, Edge::Kind::Call);
+}
+
+/// Adds a reference edge to the given child node.
+void CallGraphNode::addChildEdge(CallGraphNode *child) {
+  addEdge(child, Edge::Kind::Child);
+}
+
+/// Add an edge to 'node' with the given kind.
+void CallGraphNode::addEdge(CallGraphNode *node, Edge::Kind kind) {
+  edges.insert({node, kind});
+}
+
+//===----------------------------------------------------------------------===//
+// CallGraph
+//===----------------------------------------------------------------------===//
+
+/// Compute the set of callgraph nodes that are created by regions nested within
+/// 'op'.
+static void computeCallables(Operation *op, CallGraph &cg,
+                             CallGraphNode *parentNode) {
+  if (op->getNumRegions() == 0)
+    return;
+  if (auto callableOp = dyn_cast<CallableOpInterface>(op)) {
+    SmallVector<Region *, 1> callables;
+    callableOp.getCallableRegions(callables);
+    for (auto *callableRegion : callables)
+      cg.getOrAddNode(callableRegion, parentNode);
+  }
+}
+
+/// Recursively compute the callgraph edges for the given operation. Computed
+/// edges are placed into the given callgraph object.
+static void computeCallGraph(Operation *op, CallGraph &cg,
+                             CallGraphNode *parentNode) {
+  // Compute the callgraph nodes and edges for each of the nested operations.
+  auto isCallable = isa<CallableOpInterface>(op);
+  for (auto &region : op->getRegions()) {
+    // Check to see if this region is a callable node, if so this is the parent
+    // node of the nested region.
+    CallGraphNode *nestedParentNode;
+    if (!isCallable || !(nestedParentNode = cg.lookupNode(&region)))
+      nestedParentNode = parentNode;
+
+    // Iterate over the nested operations twice:
+    /// One to fully create nodes in the for each callable region of a nested
+    /// operation;
+    for (auto &block : region)
+      for (auto &nested : block)
+        computeCallables(&nested, cg, nestedParentNode);
+    /// And another to recursively compute the callgraph.
+    for (auto &block : region)
+      for (auto &nested : block)
+        computeCallGraph(&nested, cg, nestedParentNode);
+  }
+
+  // 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 (!parentNode)
+    return;
+
+  // If this is a call operation, resolve the callee.
+  if (auto call = dyn_cast<CallOpInterface>(op))
+    parentNode->addCallEdge(
+        cg.resolveCallable(call.getCallableForCallee(), op));
+}
+
+CallGraph::CallGraph(Operation *op) : externalNode(/*callableRegion=*/nullptr) {
+  computeCallGraph(op, *this, /*parentNode=*/nullptr);
+}
+
+/// Get or add a call graph node for the given region.
+CallGraphNode *CallGraph::getOrAddNode(Region *region,
+                                       CallGraphNode *parentNode) {
+  assert(region && isa<CallableOpInterface>(region->getParentOp()) &&
+         "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)
+      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(riverriddle) 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());
+  }
+  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.
+CallGraphNode *CallGraph::resolveCallable(CallInterfaceCallable callable,
+                                          Operation *from) const {
+  // Get the callee operation from the callable.
+  Operation *callee;
+  if (auto symbolRef = callable.dyn_cast<SymbolRefAttr>())
+    callee = SymbolTable::lookupNearestSymbolFrom(from, symbolRef.getValue());
+  else
+    callee = callable.get<Value *>()->getDefiningOp();
+
+  // If the callee is non-null and is a valid callable object, try to get the
+  // called region from it.
+  if (callee && callee->getNumRegions()) {
+    if (auto callableOp = dyn_cast_or_null<CallableOpInterface>(callee)) {
+      if (auto *node = lookupNode(callableOp.getCallableRegion(callable)))
+        return node;
+    }
+  }
+
+  // If we don't have a valid direct region, this is an external call.
+  return getExternalNode();
+}
+
+/// Dump the 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>";
+      return;
+    }
+
+    auto *callableRegion = node->getCallableRegion();
+    auto *parentOp = callableRegion->getParentOp();
+    os << "'" << callableRegion->getParentOp()->getName() << "' - Region #"
+       << callableRegion->getRegionNumber();
+    if (auto attrs = parentOp->getAttrList().getDictionary())
+      os << " : " << attrs;
+  };
+
+  for (auto &nodeIt : nodes) {
+    const CallGraphNode *node = nodeIt.second.get();
+
+    // Dump the header for this node.
+    os << "// - Node : ";
+    emitNodeName(node);
+    os << "\n";
+
+    // Emit each of the edges.
+    for (auto &edge : *node) {
+      os << "// -- ";
+      if (edge.isCall())
+        os << "Call";
+      else if (edge.isChild())
+        os << "Child";
+
+      os << "-Edge : ";
+      emitNodeName(edge.getTarget());
+      os << "\n";
+    }
+    os << "//\n";
+  }
+
+  os << "// -- SCCs --\n";
+
+  for (auto &scc : make_range(llvm::scc_begin(this), llvm::scc_end(this))) {
+    os << "// - SCC : \n";
+    for (auto &node : scc) {
+      os << "// -- Node :";
+      emitNodeName(node);
+      os << "\n";
+    }
+    os << "\n";
+  }
+
+  os << "// -------------------\n";
+}