in grappler/convert/convert_graph.cc [146:464]
tensorflow::Status ConvertSubGraphToNeuronNodeDef(SubGraphParams& sg_params) {
string neuron_op_name =
tensorflow::strings::StrCat("neuron_op_", sg_params.neuron_op_index);
VLOG(1) << "Start Node building ...." << neuron_op_name;
std::vector<string> input_names;
std::vector<tensorflow::DataType> input_dtypes;
std::vector<tensorflow::PartialTensorShape> input_shapes;
std::vector<NodeBuilder::NodeOut> input_nodes;
// map of nodename -> nodef of nodes to be added subgraph graphdef
std::unordered_map<std::string, NodeDef> subgraph_nodes;
// map of (name -> vector of datatype)
// This is to store all the nodes in subgraph that is named as output_nodes
// for the graph
std::map<string, std::vector<std::pair<DataType, PartialTensorShape>>>
main_graph_output_nodes;
// map of output_name -> max tensor index, transforming list of output_names
// of graph to map i.e list output_names : A:1, A: 3, B:0 => map: A->3, B->0
std::unordered_map<string, int> map_out_names_to_index;
for (auto out_tensor : *sg_params.output_names) {
auto tensor_vec = split_tensor(out_tensor);
map_out_names_to_index[tensor_vec[0]] = std::max(
map_out_names_to_index[tensor_vec[0]], std::stoi(tensor_vec[1]));
}
// Adding placeholder for all incoming edges, since
// placeholders do not need to specify inputs.
std::unordered_map<std::string, std::string>
outside_name_port_to_placeholder_name;
for (auto incoming_edge : sg_params.subgraph_incoming_edges) {
if (incoming_edge->IsControlEdge()) {
VLOG(2) << "Control edge: src-> " << incoming_edge->src()->name()
<< " dst-> " << incoming_edge->dst()->name();
// Need to remove this input from the node. Control edges incoming to
// the subgraph should point to Neuron op directly. The inputs in nodes
// taking control edges from outside should be removed.
NodeDef inside_node_def;
std::vector<string> input_list;
auto inside_node_name = incoming_edge->dst()->name();
if (subgraph_nodes.find(inside_node_name) == subgraph_nodes.end()) {
inside_node_def = incoming_edge->dst()->def();
} else {
inside_node_def = subgraph_nodes[inside_node_name];
}
for (int i = 0; i < inside_node_def.input_size(); ++i) {
auto input_name = inside_node_def.input(i);
if (input_name[0] != '^') {
input_list.push_back(input_name);
}
}
// Clearing the input and adding back all the inputs except
// control inputs.
inside_node_def.clear_input();
for (auto input : input_list) {
inside_node_def.add_input(input);
}
// adding the new nodedef to the list.
subgraph_nodes[incoming_edge->dst()->name()] = inside_node_def;
continue;
}
int outside_node_id = incoming_edge->src()->id();
int outside_node_index = incoming_edge->src_output();
int inside_node_id = incoming_edge->dst()->id();
int inside_node_index = incoming_edge->dst_input();
tensorflow::Node* inside_node = sg_params.graph->FindNodeId(inside_node_id);
tensorflow::Node* outside_node =
sg_params.graph->FindNodeId(outside_node_id);
VLOG(2) << "edge-> src:" << outside_node->name()
<< " idx: " << outside_node_index << " dst:" << inside_node->name()
<< " idx: " << inside_node_index;
NodeDef placeholder_def;
std::string outside_node_port =
outside_node->name() + std::to_string(outside_node_index);
std::string placeholder_name;
if (outside_name_port_to_placeholder_name.count(outside_node_port)) {
placeholder_name =
outside_name_port_to_placeholder_name[outside_node_port];
} else {
placeholder_name = sg_params.graph->NewName(outside_node_port);
outside_name_port_to_placeholder_name[outside_node_port] =
placeholder_name;
PartialTensorShape out_shape;
std::vector<PartialTensorShape> v_output_shape;
std::vector<PartialTensorShape> v_inferred_shape;
PartialTensorShape inferred_shape;
AttrSlice attrs = outside_node->attrs();
DataType out_dtype = outside_node->output_type(outside_node_index);
if (IsPlaceholder(outside_node->def()) &&
GetNodeAttr(attrs, "shape", &out_shape).ok()) {
} else if (GetNodeAttr(attrs, "_output_shapes", &v_output_shape).ok()) {
out_shape = v_output_shape[outside_node_index];
}
if (GetNodeAttr(attrs, kNeuronInferredShapes, &v_inferred_shape).ok()) {
inferred_shape = v_inferred_shape[outside_node_index];
} else {
inferred_shape = out_shape;
}
TF_RETURN_IF_ERROR(NodeDefBuilder(placeholder_name, "Placeholder")
.Attr("dtype", out_dtype)
.Attr("shape", out_shape)
.Attr(kNeuronInferredShapes, {inferred_shape})
.Finalize(&placeholder_def));
// for creating Neuron op node, storing:
// input_names -> Attr(input_names) -> name of input nodes inside subgraph
// input_dtypes -> Attr(input_dtypes) -> dtype of input nodes inside
// subgraph input_shapes -> Attr(input_shapes) -> shape of input nodes
// inside subgraph
input_names.push_back(placeholder_name + ":0");
input_dtypes.push_back(out_dtype);
input_shapes.push_back(inferred_shape);
// input_nodes -> Input(input_nodes) ->. list of nodes with index (outside
// the subgraph) feeding into the subgraph.
input_nodes.push_back(
NodeBuilder::NodeOut(outside_node, outside_node_index));
subgraph_nodes[placeholder_name] = placeholder_def;
}
// Changing the inside nodes taking inputs from outside to point to
// placeholder created above. If Nodedef of the node is already added to
// the map, update the same nodedef, else create a nodedef from the node.
NodeDef inside_node_def =
subgraph_nodes.find(inside_node->name()) != subgraph_nodes.end()
? subgraph_nodes[inside_node->name()]
: inside_node->def();
inside_node_def.set_input(inside_node_index, placeholder_name.c_str());
VLOG(2) << "Input node def:" << inside_node_def.DebugString();
VLOG(2) << " Place holder def :" << placeholder_def.DebugString();
// Storing/updating values in nodedef map
subgraph_nodes[inside_node->name()] = inside_node_def;
}
// Debug code
VLOG(2) << " Incoming nodes in graphdef";
for (auto it = subgraph_nodes.begin(); it != subgraph_nodes.end(); it++) {
VLOG(2) << "name:" << it->first << "\n Nodedef:\n"
<< it->second.DebugString();
}
// Adding all the interior nodes to the list.
for (auto node_id : *sg_params.subgraph_node_ids) {
tensorflow::Node* node = sg_params.graph->FindNodeId(node_id);
// this is to handle main graph output nodes in this segment.
// filling map of (tensor_name , vector(dtype)) i.e A ->DT_FLOAT,->INT8
// dtype order will match the the index.
if (map_out_names_to_index.count(node->name())) {
int num_outputs = map_out_names_to_index[node->name()];
std::vector<PartialTensorShape> v_partial_shape;
Status status =
GetNodeAttr(node->attrs(), "_output_shapes", &v_partial_shape);
if (!status.ok()) {
status =
GetNodeAttr(node->attrs(), kNeuronInferredShapes, &v_partial_shape);
}
if (!status.ok()) {
v_partial_shape.resize(num_outputs);
}
for (auto port = 0; port <= num_outputs; ++port) {
DataType tensor_dtype = node->output_type(port);
main_graph_output_nodes[node->name()].push_back(
std::make_pair(tensor_dtype, v_partial_shape[port]));
VLOG(1) << " ConvertSubGraphToNeuronNodeDef node: making pair: ("
<< node->name() << " , " << tensor_dtype << " )";
}
}
if (subgraph_nodes.find(node->name()) == subgraph_nodes.end()) {
subgraph_nodes[node->name()] = node->def();
}
}
// Inserting all the nodedef to graphdef for subgraph
GraphDef subgraph_graph_def;
for (auto it = subgraph_nodes.begin(); it != subgraph_nodes.end(); it++) {
(*subgraph_graph_def.add_node()) = it->second;
}
VLOG(2) << "Neuron subgraph graphdef: " << subgraph_graph_def.DebugString();
std::string in_graph_def_string = "";
if (!subgraph_graph_def.SerializeToString(&in_graph_def_string)) {
return tensorflow::errors::InvalidArgument("Failed to serialize subgraph");
}
// Gather output metadata
VLOG(2) << "Neuron op: " << neuron_op_name
<< " sg_params.output_inds: " << sg_params.output_inds->size();
std::vector<std::string> neuron_node_output_names;
std::vector<tensorflow::DataType> neuron_node_output_dtypes;
std::vector<PartialTensorShape> neuron_node_output_shapes;
for (const std::pair<int, int>& output : *sg_params.output_inds) {
int node_id = output.first;
int output_idx = output.second;
Node* node = sg_params.graph->FindNodeId(node_id);
std::string op_name = node->name();
std::string tensor_name = op_name + ":" + std::to_string(output_idx);
VLOG(2) << "op_name: " << op_name;
VLOG(2) << "Output tensor name: " << tensor_name;
if (!main_graph_output_nodes.count(op_name)) {
neuron_node_output_names.push_back(tensor_name);
DataType tf_dtype = node->output_type(output_idx);
neuron_node_output_dtypes.push_back(tf_dtype);
std::vector<PartialTensorShape> v_partial_shape;
if (GetNodeAttr(node->attrs(), "_output_shapes", &v_partial_shape).ok()) {
neuron_node_output_shapes.push_back(v_partial_shape[output_idx]);
} else if (GetNodeAttr(node->attrs(), kNeuronInferredShapes,
&v_partial_shape)
.ok()) {
neuron_node_output_shapes.push_back(v_partial_shape[output_idx]);
} else {
neuron_node_output_shapes.emplace_back();
}
}
}
auto start_index = neuron_node_output_names.size();
// Pushing the outputs of graph (if in the subgraph) to the output for the
// subgraph. As we will create IdentityN connecting to the these nodes.
for (auto name_type : main_graph_output_nodes) {
uint counter = 0;
auto name = name_type.first;
for (auto dtype_shape : name_type.second) {
auto dtype = dtype_shape.first;
auto shape = dtype_shape.second;
VLOG(1) << "tensorname: " << name << " dtype: " << dtype;
std::string tensor_name = name + ":" + std::to_string(counter++);
neuron_node_output_names.push_back(tensor_name);
neuron_node_output_dtypes.push_back(dtype);
neuron_node_output_shapes.push_back(shape);
}
}
VLOG(1) << "Finished op preparation";
StringPieceHasher hasher;
std::string hash_string = "";
for (auto const& s : input_names) {
hash_string += s;
}
for (auto const& s : neuron_node_output_names) {
hash_string += s;
}
for (auto const& s : neuron_node_output_dtypes) {
hash_string += s;
}
std::stringstream stream;
stream << std::hex << hasher(hash_string);
std::string hex_string(stream.str());
VLOG(2) << "String to be hashed " << hash_string << "Hashed String "
<< hasher(hash_string) << "Hex Rep" << hex_string;
neuron_op_name = "neuron_op_" + hex_string;
VLOG(1) << "Hashed neuron_op_name: " << neuron_op_name;
Node* neuron_node;
TF_CHECK_OK(NodeBuilder(neuron_op_name, "NeuronOp")
.Input(input_nodes)
.Attr("graph_def", in_graph_def_string)
.Attr("input_names", input_names)
.Attr("input_dtypes", input_dtypes)
.Attr("input_shapes", input_shapes)
.Attr("output_names", neuron_node_output_names)
.Attr("output_dtypes", neuron_node_output_dtypes)
.Attr("output_shapes", neuron_node_output_shapes)
.Attr(kNeuronInferredShapes, neuron_node_output_shapes)
.Finalize(sg_params.graph, &neuron_node));
sg_params.neuron_node = neuron_node;
// Creating identityN node corresponding to any output nodes(if any) in this
// subgraph
// start_index is the index of Neuron op output to which IdentityN node should
// be connected to.
VLOG(1) << "start_index: " << start_index;
std::vector<tensorflow::NodeBuilder::NodeOut> identity_inputs;
// Iterate through the output node list found.
for (auto name_index : main_graph_output_nodes) {
identity_inputs.clear();
auto name = name_index.first;
VLOG(1) << " indentity inputs: name:" << name;
VLOG(1) << " max index: " << map_out_names_to_index[name];
for (size_t i = 0; i < main_graph_output_nodes[name].size(); ++i) {
VLOG(1) << "start_index: " << start_index;
identity_inputs.push_back(
NodeBuilder::NodeOut(neuron_node, start_index++));
}
Node* node;
TF_CHECK_OK(NodeBuilder(name, "IdentityN")
.Input(identity_inputs)
.Finalize(sg_params.graph, &node));
VLOG(1) << " New output IdentityN node: " << node->def().DebugString();
}
VLOG(1) << "Created new node ..............";
VLOG(2) << " new node: " << neuron_node->def().DebugString();
return tensorflow::Status::OK();
}