CallGraph.cpp

Go to the documentation of this file.

//===-- CallGraph.cpp - LLZK-specific call graph implementation -*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
// The contents of this file are adapted from llvm/lib/Analysis/CallGraph.cpp
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "llzk/Analysis/CallGraph.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/LLZK/IR/AttributeHelper.h"
#include "llzk/Util/SymbolHelper.h"
 
#include <mlir/Analysis/CallGraph.h>
#include <mlir/IR/Operation.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Interfaces/CallInterfaces.h>
 
#include <llvm/ADT/DepthFirstIterator.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/Support/ErrorHandling.h>
 
namespace llzk {
 
using namespace function;
 
//===----------------------------------------------------------------------===//
// CallGraphNode
//===----------------------------------------------------------------------===//

bool CallGraphNode::isExternal() const { return !callableRegion; }

mlir::Region *CallGraphNode::getCallableRegion() const {
  assert(!isExternal() && "the external node has no callable region");
  return callableRegion;
}
 
FuncDefOp CallGraphNode::getCalledFunction() const {
  return mlir::dyn_cast<FuncDefOp>(getCallableRegion()->getParentOp());
}

void CallGraphNode::addAbstractEdge(CallGraphNode *node) {
  assert(isExternal() && "abstract edges are only valid on external nodes");
  addEdge(node, Edge::Kind::Abstract);
}

void CallGraphNode::addCallEdge(CallGraphNode *node) { addEdge(node, Edge::Kind::Call); }

void CallGraphNode::addChildEdge(CallGraphNode *child) { addEdge(child, Edge::Kind::Child); }

bool CallGraphNode::hasChildren() const {
  return llvm::any_of(edges, [](const Edge &edge) { return edge.isChild(); });
}

void CallGraphNode::addEdge(CallGraphNode *node, Edge::Kind kind) {
  edges.insert({this, node, kind});
}
 
//===----------------------------------------------------------------------===//
// CallGraph
//===----------------------------------------------------------------------===//

static void computeCallGraph(
    mlir::Operation *op, CallGraph &cg, mlir::SymbolTableCollection &symbolTable,
    CallGraphNode *parentNode, bool resolveCalls
) {
  if (mlir::CallOpInterface call = mlir::dyn_cast<mlir::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;
  }
 
  // Compute the callgraph nodes and edges for each of the nested operations.
  if (mlir::CallableOpInterface callable = mlir::dyn_cast<mlir::CallableOpInterface>(op)) {
    if (auto *callableRegion = callable.getCallableRegion()) {
      parentNode = cg.getOrAddNode(callableRegion, parentNode);
    } else {
      return;
    }
  }
 
  for (mlir::Region &region : op->getRegions()) {
    for (mlir::Operation &nested : region.getOps()) {
      computeCallGraph(&nested, cg, symbolTable, parentNode, resolveCalls);
    }
  }
}
 
CallGraph::CallGraph(mlir::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.
  mlir::SymbolTableCollection symbolTable;
  computeCallGraph(
      op, *this, symbolTable, /*parentNode=*/nullptr,
      /*resolveCalls=*/false
  );
  computeCallGraph(
      op, *this, symbolTable, /*parentNode=*/nullptr,
      /*resolveCalls=*/true
  );
}

CallGraphNode *CallGraph::getOrAddNode(mlir::Region *region, CallGraphNode *parentNode) {
  assert(
      region && mlir::isa<mlir::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 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());
    }
  }
  return node.get();
}

CallGraphNode *CallGraph::lookupNode(mlir::Region *region) const {
  const auto *it = nodes.find(region);
  return it == nodes.end() ? nullptr : it->second.get();
}

CallGraphNode *CallGraph::resolveCallable(
    mlir::CallOpInterface call, mlir::SymbolTableCollection &symbolTable
) const {
  auto res = llzk::resolveCallable<FuncDefOp>(symbolTable, call);
  if (mlir::succeeded(res)) {
    if (auto *node = lookupNode(res->get().getCallableRegion())) {
      return node;
    }
  }
 
  return getUnknownCalleeNode();
}

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()) {
        eraseNode(edge.getTarget());
      }
    }
  }
  // Erase any edges to this node from any other nodes.
  for (auto &it : nodes) {
    it.second->edges.remove_if([node](const CallGraphNode::Edge &edge) {
      return edge.getTarget() == node;
    });
  }
  nodes.erase(node->getCallableRegion());
}
 
//===----------------------------------------------------------------------===//
// Printing

void CallGraph::dump() const { print(llvm::errs()); }
void CallGraph::print(llvm::raw_ostream &os) const {
  os << "// ---- CallGraph ----\n";
 
  // Functor used to output the name for the given node.
  auto emitNodeName = [&](const CallGraphNode *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();
    if (!isNullOrEmpty(attrs)) {
      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";
}
 
} // namespace llzk