CallGraph.h

Go to the documentation of this file.

//===- CallGraph.h - CallGraph analysis for MLIR ----------------*- C++ -*-===//
//
// Part of the MLIR 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
//
//===----------------------------------------------------------------------===//
//
// This file contains an analysis for computing the multi-level callgraph from a
// given top-level operation. This nodes within this callgraph are defined by
// the `CallOpInterface` and `CallableOpInterface` operation interfaces defined
// in CallInterface.td.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_ANALYSIS_CALLGRAPH_H
#define MLIR_ANALYSIS_CALLGRAPH_H

#include "mlir/Support/LLVM.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PointerIntPair.h"
#include "llvm/ADT/SetVector.h"

namespace mlir {
struct CallInterfaceCallable;
class Operation;
class Region;

//===----------------------------------------------------------------------===//
// CallGraphNode
//===----------------------------------------------------------------------===//

class CallGraphNode {
public:
  class Edge {
    enum class Kind {
      // An 'Abstract' edge represents an opaque, non-operation, reference
      // between this node and the target. Edges of this type are only valid
      // from the external node, as there is no valid connection to an operation
      // in the module.
      Abstract,

      // A 'Call' edge represents a direct reference to the target node via a
      // call-like operation within the callable region of this node.
      Call,

      // A 'Child' edge is used when the region of target node is defined inside
      // of the callable region of this node. This means that the region of this
      // node is an ancestor of the region for the target node. As such, this
      // edge cannot be used on the 'external' node.
      Child,
    };

  public:
    bool isAbstract() const { return targetAndKind.getInt() == Kind::Abstract; }

    bool isCall() const { return targetAndKind.getInt() == Kind::Call; }

    bool isChild() const { return targetAndKind.getInt() == Kind::Child; }

    CallGraphNode *getTarget() const { return targetAndKind.getPointer(); }

    bool operator==(const Edge &edge) const {
      return targetAndKind == edge.targetAndKind;
    }

  private:
    Edge(CallGraphNode *node, Kind kind) : targetAndKind(node, kind) {}
    explicit Edge(llvm::PointerIntPair<CallGraphNode *, 2, Kind> targetAndKind)
        : targetAndKind(targetAndKind) {}

    llvm::PointerIntPair<CallGraphNode *, 2, Kind> targetAndKind;

    // Provide access to the constructor and Kind.
    friend class CallGraphNode;
  };

  bool isExternal() const;

  Region *getCallableRegion() const;

  void addAbstractEdge(CallGraphNode *node);

  void addCallEdge(CallGraphNode *node);

  void addChildEdge(CallGraphNode *child);

  using iterator = SmallVectorImpl<Edge>::const_iterator;
  iterator begin() const { return edges.begin(); }
  iterator end() const { return edges.end(); }

  bool hasChildren() const;

private:
  struct EdgeKeyInfo {
    using BaseInfo =
        DenseMapInfo<llvm::PointerIntPair<CallGraphNode *, 2, Edge::Kind>>;

    static Edge getEmptyKey() { return Edge(BaseInfo::getEmptyKey()); }
    static Edge getTombstoneKey() { return Edge(BaseInfo::getTombstoneKey()); }
    static unsigned getHashValue(const Edge &edge) {
      return BaseInfo::getHashValue(edge.targetAndKind);
    }
    static bool isEqual(const Edge &lhs, const Edge &rhs) { return lhs == rhs; }
  };

  CallGraphNode(Region *callableRegion) : callableRegion(callableRegion) {}

  void addEdge(CallGraphNode *node, Edge::Kind kind);

  Region *callableRegion;

  llvm::SetVector<Edge, SmallVector<Edge, 4>,
                  llvm::SmallDenseSet<Edge, 4, EdgeKeyInfo>>
      edges;

  // Provide access to private methods.
  friend class CallGraph;
};

//===----------------------------------------------------------------------===//
// CallGraph
//===----------------------------------------------------------------------===//

class CallGraph {
  using NodeMapT = llvm::MapVector<Region *, std::unique_ptr<CallGraphNode>>;

  class NodeIterator final
      : public llvm::mapped_iterator<
            NodeMapT::const_iterator,
            CallGraphNode *(*)(const NodeMapT::value_type &)> {
    static CallGraphNode *unwrap(const NodeMapT::value_type &value) {
      return value.second.get();
    }

  public:
    NodeIterator(NodeMapT::const_iterator it)
        : llvm::mapped_iterator<
              NodeMapT::const_iterator,
              CallGraphNode *(*)(const NodeMapT::value_type &)>(it, &unwrap) {}
  };

public:
  CallGraph(Operation *op);

  CallGraphNode *getOrAddNode(Region *region, CallGraphNode *parentNode);

  CallGraphNode *lookupNode(Region *region) const;

  CallGraphNode *getExternalNode() const {
    return const_cast<CallGraphNode *>(&externalNode);
  }

  CallGraphNode *resolveCallable(CallInterfaceCallable callable,
                                 Operation *from = nullptr) const;

  using iterator = NodeIterator;
  iterator begin() const { return nodes.begin(); }
  iterator end() const { return nodes.end(); }

  void dump() const;
  void print(raw_ostream &os) const;

private:
  NodeMapT nodes;

  CallGraphNode externalNode;
};

} // end namespace mlir

namespace llvm {
// Provide graph traits for traversing call graphs using standard graph
// traversals.
template <> struct GraphTraits<const mlir::CallGraphNode *> {
  using NodeRef = mlir::CallGraphNode *;
  static NodeRef getEntryNode(NodeRef node) { return node; }

  static NodeRef unwrap(const mlir::CallGraphNode::Edge &edge) {
    return edge.getTarget();
  }

  // ChildIteratorType/begin/end - Allow iteration over all nodes in the graph.
  using ChildIteratorType =
      mapped_iterator<mlir::CallGraphNode::iterator, decltype(&unwrap)>;
  static ChildIteratorType child_begin(NodeRef node) {
    return {node->begin(), &unwrap};
  }
  static ChildIteratorType child_end(NodeRef node) {
    return {node->end(), &unwrap};
  }
};

template <>
struct GraphTraits<const mlir::CallGraph *>
    : public GraphTraits<const mlir::CallGraphNode *> {
  static NodeRef getEntryNode(const mlir::CallGraph *cg) {
    return cg->getExternalNode();
  }

  // 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(); }
};
} // end namespace llvm

#endif // MLIR_ANALYSIS_CALLGRAPH_H