path: root/tensorflow/core/graph/graph.h

// A Graph describes a set of computations that are to be
// performed, as well as the dependencies between those
// compuations. The basic model is a DAG (directed acyclic graph) with
// * internal nodes representing computational operations to be performed;
// * edges represent dependencies, indicating the target may only be
//   executed once the source has completed; and
// * predefined "source" (start) and "sink" (finish) nodes -- the source
//   should be the only node that doesn't depend on anything, and the sink
//   should be the only node that nothing depends on.
//
// Note: Node ids are intended to be relatively dense in the
// 0..max_id range, but there may be gaps since ids won't be reused.
//
// Note: Some dependencies between operations are due to one operation
// consuming the output of another. In fact operations can produce
// multiple outputs and consume multiple inputs, and some
// optimizations will care about which specific outputs are connected
// to which specific inputs.  We therefore represent data dependency
// between output O of layer A and input I of layer B using
// "input index" and "output index" labels per edge.

#ifndef TENSORFLOW_GRAPH_GRAPH_H_
#define TENSORFLOW_GRAPH_GRAPH_H_

#include <functional>
#include <string>
#include <vector>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/edgeset.h"
#include "tensorflow/core/lib/core/arena.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/gtl/iterator_range.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/public/status.h"

namespace tensorflow {

class Edge;
class EdgeSetTest;
class Graph;
class Node;

class NeighborIter;  // Declared below
class NodeIter;      // Declared below

class Node {
 public:
  string DebugString() const;
  int id() const { return id_; }
  int cost_id() const { return cost_id_; }
  const string& name() const { return props_->node_def_.name(); }
  const string& type_string() const { return props_->node_def_.op(); }
  const NodeDef& def() const { return props_->node_def_; }
  const OpDef& op_def() const { return *props_->op_def_; }

  // input and output types
  int num_inputs() const { return props_->input_types_.size(); }
  DataType input_type(int i) const { return props_->input_types_[i]; }
  const DataTypeVector& input_types() const { return props_->input_types_; }

  int num_outputs() const { return props_->output_types_.size(); }
  DataType output_type(int o) const { return props_->output_types_[o]; }
  const DataTypeVector& output_types() const { return props_->output_types_; }

  // This gives the device the runtime has assigned this node to.  If
  // you want the device the user requested, use def().device() instead.
  // TODO(josh11b): Validate that the assigned_device, if not empty:
  // fully specifies a device, and satisfies def().device().
  // TODO(josh11b): Move device_name outside of Node into a NodeId->DeviceName
  // map.
  string assigned_device_name() const { return assigned_device_name_; }
  void set_assigned_device_name(const string& device_name) {
    assigned_device_name_ = device_name;
  }

  // Get the neighboring nodes via edges either in or out of this node.
  gtl::iterator_range<NeighborIter> in_nodes() const;
  gtl::iterator_range<NeighborIter> out_nodes() const;
  const EdgeSet& in_edges() const { return in_edges_; }
  const EdgeSet& out_edges() const { return out_edges_; }

  // Node type helpers.
  bool IsSource() const { return id() == 0; }
  bool IsSink() const { return id() == 1; }
  // Anything other than the special Source & Sink nodes.
  bool IsOp() const { return id() > 1; }

 private:
  friend class Graph;
  Node();
  ~Node();

  class Properties : public core::RefCounted {
   public:
    Properties(const OpDef* op_def, const NodeDef& node_def,
               const DataTypeSlice inputs, const DataTypeSlice outputs);

    const OpDef* op_def_;  // not owned
    const NodeDef node_def_;
    const DataTypeVector input_types_;
    const DataTypeVector output_types_;

   private:
    // Destructor invoked when last reference goes away via Unref()
    virtual ~Properties();
    TF_DISALLOW_COPY_AND_ASSIGN(Properties);
  };

  Properties* properties() const { return props_; }

  // Initialize() adopts a reference to props, and so is suitable if props was
  // just allocated or you call props->Ref() to increment the reference
  // count for a props being held by another Node.
  void Initialize(int id, int cost_id, Properties* props);
  // Releases memory from props_, in addition to restoring *this to its
  // uninitialized state.
  void Clear();

  int id_;       // -1 until Initialize() is called
  int cost_id_;  // -1 if there is no corresponding cost accounting node

  EdgeSet in_edges_;
  EdgeSet out_edges_;

  Properties* props_;

  // Name of device assigned to perform this computation.
  string assigned_device_name_;

  TF_DISALLOW_COPY_AND_ASSIGN(Node);
};

class Edge {
 public:
  Node* src() const { return src_; }
  Node* dst() const { return dst_; }
  int id() const { return id_; }

  // Return the number of the source output that produces the data
  // carried by this edge.  The special value kControlSlot is used
  // for control dependencies.
  int src_output() const { return src_output_; }

  // Return the number of the destination input that consumes the data
  // carried by this edge.  The special value kControlSlot is used
  // for control dependencies.
  int dst_input() const { return dst_input_; }

  // Return true iff this is an edge that indicates a control-flow
  // (as opposed to a data-flow) dependency.
  bool IsControlEdge() const;

 private:
  Edge() {}

  friend class EdgeSetTest;
  friend class Graph;
  Node* src_;
  Node* dst_;
  int id_;
  int src_output_;
  int dst_input_;
};

// Thread compatible but not thread safe.
class Graph {
 public:
  // Constructs a graph with a single SOURCE (always id kSourceId) and a
  // single SINK (always id kSinkId) node, and an edge from SOURCE->SINK.
  //
  // The graph can hold ops found in registry.
  explicit Graph(const OpRegistryInterface* registry);
  ~Graph();

  static const int kControlSlot = -1;

  // Adds a new node to this graph, and returns it. Infers the Op and
  // input/output types for the node. *this owns the returned instance.
  // Returns nullptr and sets *status on error.
  Node* AddNode(const NodeDef& node_def, Status* status);

  // Copies *node, which may belong to another graph, to a new node,
  // which is returned.  Does not copy any edges.  *this owns the
  // returned instance.
  Node* CopyNode(Node* node);

  // Remove a node from this graph, including all edges from or to it.
  // *node should not be accessed after calling this function.
  // REQUIRES: node->IsOp()
  void RemoveNode(Node* node);

  // Add an edge that connects the xth output of "source" to the yth input
  // of "dest".
  const Edge* AddEdge(Node* source, int x, Node* dest, int y);

  // Add a control-edge (no data flows along this edge) that
  // connects "source" to "dest".
  const Edge* AddControlEdge(Node* source, Node* dest) {
    return AddEdge(source, kControlSlot, dest, kControlSlot);
  }

  // Removes edge from the graph.
  // REQUIRES: The edge must exist.
  void RemoveEdge(const Edge* edge);

  // Returns one more than the maximum id assigned to any node.
  int num_node_ids() const { return nodes_.size(); }

  // Serialize to a GraphDef.
  void ToGraphDef(GraphDef* graph_def) const;

  // Generate new node name with the specified prefix that is unique
  // across this graph.
  string NewName(StringPiece prefix);

  // Access to the list of all nodes.  Example usage:
  //   for (Node* node : graph.nodes()) { ... }
  gtl::iterator_range<NodeIter> nodes() const;

  // Returns the node associated with an id, or nullptr if no node
  // with that id (the node with that id was removed and the id has
  // not yet been re-used). *this owns the returned instance.
  // REQUIRES: 0 <= id < num_node_ids().
  Node* FindNodeId(int id) const { return nodes_[id]; }

  // Returns one more than the maximum id assigned to any edge.
  int num_edge_ids() const { return edges_.size(); }

  // Returns the Edge associated with an id, or nullptr if no edge
  // with that id (the node with that id was removed and the id has
  // not yet been re-used). *this owns the returned instance.
  // REQUIRES: 0 <= id < num_node_ids().
  const Edge* FindEdgeId(int id) const { return edges_[id]; }

  // Access to the set of all edges.  Example usage:
  //   for (const Edge* e : graph.edges()) { ... }
  const EdgeSet& edges() const { return edge_set_; }

  // The pre-defined nodes.
  enum { kSourceId = 0, kSinkId = 1 };
  Node* source_node() const { return FindNodeId(kSourceId); }
  Node* sink_node() const { return FindNodeId(kSinkId); }

  const OpRegistryInterface* op_registry() const { return ops_; }

  // TODO(josh11b): uint64 hash() const;

 private:
  bool IsValidNode(Node* node) const;
  // If cost_node is non-null, then cost accounting (in CostModel)
  // will be associated with that node rather than the new one being
  // created.
  Node* AllocateNode(Node::Properties* props, const Node* cost_node);
  void ReleaseNode(Node* node);

  // Registry of all known ops.  Not owned.
  const OpRegistryInterface* const ops_;

  // Allocator which will give us good locality.
  core::Arena arena_;

  // Map from node ids to allocated nodes.  nodes_[id] may be nullptr if
  // the node with that id was removed from the graph.
  std::vector<Node*> nodes_;

  // Map from edge ids to allocated edges.  edges_[id] may be nullptr if
  // the edge with that id was removed from the graph.
  std::vector<Edge*> edges_;

  // For ease of iteration, we currently just keep a set of all live
  // edges.  May want to optimize by removing this copy.
  EdgeSet edge_set_;

  // Allocated but free nodes and edges.
  std::vector<Node*> free_nodes_;
  std::vector<Edge*> free_edges_;

  // For generating unique names.
  int name_counter_ = 0;

  TF_DISALLOW_COPY_AND_ASSIGN(Graph);
};

// TODO(josh11b): We may want to support keeping an index on various
// node/edge attributes in a graph, particularly node names.

// Helper routines

inline bool IsSwitch(const Node* node) {
  return node->type_string() == "Switch" || node->type_string() == "RefSwitch";
}

inline bool IsMerge(const Node* node) { return node->type_string() == "Merge"; }

inline bool IsEnter(const Node* node) {
  return node->type_string() == "Enter" || node->type_string() == "RefEnter";
}

inline bool IsExit(const Node* node) { return node->type_string() == "Exit"; }

inline bool IsNextIteration(const Node* node) {
  return node->type_string() == "NextIteration";
}

inline bool IsLoopCond(const Node* node) {
  return node->type_string() == "LoopCond";
}

inline bool IsControlTrigger(const Node* node) {
  return node->type_string() == "ControlTrigger";
}

inline bool IsSend(const Node* node) {
  return node->type_string() == "_Send" || node->type_string() == "_HostSend";
}

inline bool IsRecv(const Node* node) {
  return node->type_string() == "_Recv" || node->type_string() == "_HostRecv";
}

// True for Nodes that mediate the transfer of values between processes.
inline bool IsTransferNode(const Node* n) { return IsSend(n) || IsRecv(n); }

inline bool IsConstant(const Node* node) {
  return node->type_string() == "Const" || node->type_string() == "HostConst";
}

inline bool IsVariable(const Node* node) {
  return node->type_string() == "Variable";
}

inline bool IsIdentity(const Node* node) {
  return (node->type_string() == "Identity" ||
          node->type_string() == "RefIdentity");
}

// Returns true iff 'n' is a control flow node.
inline bool IsControlFlow(const Node* n) {
  return IsSwitch(n) || IsMerge(n) || IsEnter(n) || IsExit(n) ||
         IsNextIteration(n);
}

inline bool IsHostMemoryPreserving(const Node* node) {
  return IsIdentity(node) || IsControlFlow(node);
}

// Iterator for stepping through the nodes of a graph.
class NodeIter {
 public:
  NodeIter(const Graph* graph, int id);
  bool operator==(const NodeIter& rhs);
  bool operator!=(const NodeIter& rhs);
  void operator++();
  Node* operator*();
  Node* operator->();

 private:
  // Invariant: id_ == graph_->num_node_ids() || graph_->FindId(id_) != nullptr
  const Graph* graph_;
  int id_;
};

// Iterator for stepping through the neighbors of a node.
class NeighborIter {
 public:
  NeighborIter(EdgeSet::const_iterator iter, bool incoming);
  bool operator==(const NeighborIter& rhs);
  bool operator!=(const NeighborIter& rhs);
  void operator++();
  Node* operator*();
  Node* operator->();

 private:
  EdgeSet::const_iterator iter_;
  bool incoming_;
};

// IMPLEMENTATION DETAILS, PLEASE IGNORE

inline NodeIter::NodeIter(const Graph* graph, int id)
    : graph_(graph), id_(id) {}

inline bool NodeIter::operator==(const NodeIter& rhs) {
  DCHECK(graph_ == rhs.graph_);
  return id_ == rhs.id_;
}

inline bool NodeIter::operator!=(const NodeIter& rhs) {
  return !(*this == rhs);
}

inline void NodeIter::operator++() {
  while (1) {
    DCHECK_LE(id_, graph_->num_node_ids());
    ++id_;
    if (id_ >= graph_->num_node_ids() || graph_->FindNodeId(id_) != nullptr) {
      return;
    }
  }
}

inline Node* NodeIter::operator*() { return graph_->FindNodeId(id_); }

inline Node* NodeIter::operator->() { return graph_->FindNodeId(id_); }

inline NeighborIter::NeighborIter(EdgeSet::const_iterator iter, bool incoming)
    : iter_(iter), incoming_(incoming) {}

inline bool NeighborIter::operator==(const NeighborIter& rhs) {
  return iter_ == rhs.iter_ && incoming_ == rhs.incoming_;
}

inline bool NeighborIter::operator!=(const NeighborIter& rhs) {
  return !(*this == rhs);
}

inline void NeighborIter::operator++() { ++iter_; }

inline Node* NeighborIter::operator*() {
  const Edge* e = *iter_;
  return incoming_ ? e->src() : e->dst();
}

inline Node* NeighborIter::operator->() {
  const Edge* e = *iter_;
  return incoming_ ? e->src() : e->dst();
}

inline bool Edge::IsControlEdge() const {
  // Note that if either src_output_ or dst_input_ is kControlSlot,
  // so is the other one (AddEdge checks this).
  return src_output_ == Graph::kControlSlot;
}

}  // namespace tensorflow

#endif  // TENSORFLOW_GRAPH_GRAPH_H_