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(); }

// def() provides the NodeDef the user supplied, but the specifics

// of this Node may have changed due to placement, optimization, etc.

// In particular:

// * def().name() will match name();

// * def().op() will match type_string() and op_def().name();

// * def().input() is not reliable, use "in_edges()" below instead;

// * def().device() is the "user's requested device" and may not match

// the actual assigned device, see assigned_device_name() below;

// * def().attr() is authoritative.

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 assigned_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; }

// Node class helpers

bool IsSwitch() const { return class_ == NC_SWITCH; }

bool IsMerge() const { return class_ == NC_MERGE; }

bool IsEnter() const { return class_ == NC_ENTER; }

bool IsExit() const { return class_ == NC_EXIT; }

bool IsNextIteration() const { return class_ == NC_NEXT_ITERATION; }

bool IsLoopCond() const { return class_ == NC_LOOP_COND; }

bool IsControlTrigger() const { return class_ == NC_CONTROL_TRIGGER; }

bool IsSend() const { return class_ == NC_SEND || class_ == NC_HOST_SEND; }

bool IsRecv() const { return class_ == NC_RECV || class_ == NC_HOST_RECV; }

bool IsConstant() const { return class_ == NC_CONSTANT; }

bool IsVariable() const { return class_ == NC_VARIABLE; }

bool IsIdentity() const { return class_ == NC_IDENTITY; }

bool IsGetSessionHandle() const { return class_ == NC_GET_SESSION_HANDLE; }

bool IsGetSessionTensor() const { return class_ == NC_GET_SESSION_TENSOR; }

bool IsDeleteSessionTensor() const {

return class_ == NC_DELETE_SESSION_TENSOR;

}

bool IsControlFlow() const {

return (class_ != NC_OTHER) && // Fast path

(IsSwitch() || IsMerge() || IsEnter() || IsExit() ||

IsNextIteration());

}

bool IsHostSend() const { return class_ == NC_HOST_SEND; }

bool IsHostRecv() const { return class_ == NC_HOST_RECV; }

template <typename T>

void AddAttr(const string& name, const T& val) {

MaybeCopyOnWrite();

SetAttrValue(val, &((*props_->node_def_.mutable_attr())[name]));

}

void ClearAttr(const string& name);

// Returns into '*e' the edge connecting to the 'idx' input of this Node.

Status input_edge(int idx, const Edge** e) const;

// Returns into '*n' the node that has an output connected to the

// 'idx' input of this Node.

Status input_node(int idx, const Node** n) const;

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

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();

// Make a copy of the Node's props_ if props_ is shared with

// other nodes. This must be called before mutating properties,

// e.g. in AddAttr.

void MaybeCopyOnWrite();

// A set of mutually exclusive classes for different kinds of nodes,

// class_ is initialized in the Node::Initialize routine based on the

// node's type_string().

enum NodeClass {

NC_UNINITIALIZED,

NC_SWITCH,

NC_MERGE,

NC_ENTER,

NC_EXIT,

NC_NEXT_ITERATION,

NC_LOOP_COND,

NC_CONTROL_TRIGGER,

NC_SEND,

NC_HOST_SEND,

NC_RECV,

NC_HOST_RECV,

NC_CONSTANT,

NC_VARIABLE,

NC_IDENTITY,

NC_GET_SESSION_HANDLE,

NC_GET_SESSION_TENSOR,

NC_DELETE_SESSION_TENSOR,

NC_OTHER // Not a special kind of node

};

int id_; // -1 until Initialize() is called

int cost_id_; // -1 if there is no corresponding cost accounting node

NodeClass class_;

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);

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_;

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;

// The GraphDef version range of this graph (see graph.proto).

const VersionDef& versions() const { return versions_; }

void set_versions(const VersionDef& versions) { versions_ = versions; }

// 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);

// The number of live nodes in the graph.

//

// Because nodes can be removed from the graph, num_nodes() is often

// smaller than num_node_ids(). If one needs to create an array of

// nodes indexed by node ids, num_node_ids() should be used as the

// array's size.

int num_nodes() const { return num_nodes_; }

// The number of live edges in the graph.

//

// Because edges can be removed from the graph, num_edges() is often

// smaller than num_edge_ids(). If one needs to create an array of

// edges indexed by edge ids, num_edge_ids() should be used as the

// array's size.

int num_edges() const { return edges().size(); }

// Serialize the nodes starting at `from_node_id` to a GraphDef.

void ToGraphDefSubRange(GraphDef* graph_def, int from_node_id) const;

// 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 one more than the maximum id assigned to any node.

int num_node_ids() const { return nodes_.size(); }

// 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_;

// GraphDef versions

VersionDef versions_;

// 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_;

// Number of nodes alive.

int64 num_nodes_ = 0;

// 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);

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_;

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_;

// Note that if either src_output_ or dst_input_ is kControlSlot,

// so is the other one (AddEdge checks this).

return src_output_ == Graph::kControlSlot;