| 85 | } |
| 86 | } |
| 87 | Status ConcatCastFusing::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 88 | GraphDef* optimized_graph) { |
| 89 | Status status; |
| 90 | *optimized_graph = item.graph; |
| 91 | TF_RETURN_IF_ERROR(status); |
| 92 | utils::MutableGraphView graph_view(optimized_graph, &status); |
| 93 | TF_RETURN_IF_ERROR(status); |
| 94 | TF_RETURN_IF_ERROR(graph_view.SortTopologically(/*ignore_cycles=*/false, {})); |
| 95 | const int num_nodes = item.graph.node_size(); |
| 96 | // invalidated_nodes - nodes that have been changed into a fused op |
| 97 | // nodes_to_delete - nodes that were fused into a fused op and are not needed anymore |
| 98 | std::vector<bool> invalidated_nodes(num_nodes); |
| 99 | std::vector<bool> nodes_to_delete(num_nodes); |
| 100 | const GraphDef* graph = graph_view.graph(); |
| 101 | |
| 102 | VLOG(3) << "Before concat cast graph rewrites: " << graph->DebugString(); |
| 103 | |
| 104 | for (int i = 0; i < num_nodes; ++i) { |
| 105 | if (invalidated_nodes[i] || nodes_to_delete[i]) { |
| 106 | continue; |
| 107 | } |
| 108 | |
| 109 | ConcatWithCast base; |
| 110 | if (FindConcatWithCast(graph_view, i, &base)) { |
| 111 | const auto* node_view = graph_view.GetNode(i); |
| 112 | const auto& fused_node = graph->node(i); |
| 113 | VLOG(2) << "Optimizing fused concat cast node " << SummarizeNodeDef(fused_node); |
| 114 | |
| 115 | // Adding fused concat+cast |
| 116 | const NodeDef& concat = graph->node(base.concat_id); |
| 117 | const NodeDef& cast = graph->node(base.cast_id); |
| 118 | const std::size_t concat_num_inputs = node_view->NumRegularFanins(); |
| 119 | VLOG(2) << "Fuse " << concat.op() << " with Cast: " |
| 120 | << " cast_name=" << cast.name(); |
| 121 | |
| 122 | NodeDef fused_op; |
| 123 | fused_op.set_name(cast.name()); |
| 124 | fused_op.set_op("FusedConcatCast"); |
| 125 | fused_op.set_device(concat.device()); |
| 126 | for (size_t j = 0; j < concat_num_inputs - 1; ++j) { |
| 127 | fused_op.add_input(concat.input(j)); // inputs |
| 128 | } |
| 129 | fused_op.add_input(concat.input(concat_num_inputs - 1)); // axis |
| 130 | |
| 131 | auto* attr = fused_op.mutable_attr(); |
| 132 | auto& concat_attr = concat.attr(); |
| 133 | auto& cast_attr = cast.attr(); |
| 134 | (*attr)["N"] = concat_attr.at("N"); |
| 135 | (*attr)["Tidx"] = concat_attr.at("Tidx"); |
| 136 | |
| 137 | (*attr)["SrcT"] = cast_attr.at("SrcT"); |
| 138 | (*attr)["DstT"] = cast_attr.at("DstT"); |
| 139 | (*attr)["Truncate"] = cast_attr.at("Truncate"); |
| 140 | |
| 141 | utils::Mutation* mutation = graph_view.GetMutationBuilder(); |
| 142 | Status status; |
| 143 | mutation->AddNode(std::move(fused_op), &status); |
| 144 | TF_RETURN_IF_ERROR(status); |
nothing calls this directly
no test coverage detected