Replaces each outgoing edge from `old_node` with a merge node that merges in the corresponding output from `new_node`.
| 91 | // Replaces each outgoing edge from `old_node` with a merge node that merges in |
| 92 | // the corresponding output from `new_node`. |
| 93 | void MergeOutgoingDataEdges(const Scope& s, Node* old_node, Node* new_node, |
| 94 | absl::string_view cluster_name, |
| 95 | const DebuggingOpts& debugging_opts) { |
| 96 | if (!s.status().ok()) { |
| 97 | return; |
| 98 | } |
| 99 | |
| 100 | std::vector<Output> merged_outputs(old_node->num_outputs(), Output(nullptr)); |
| 101 | |
| 102 | std::vector<const Edge*> data_edges; |
| 103 | absl::c_copy_if(old_node->out_edges(), std::back_inserter(data_edges), |
| 104 | [](const Edge* e) { return !e->IsControlEdge(); }); |
| 105 | |
| 106 | for (const Edge* e : data_edges) { |
| 107 | int oidx = e->src_output(); |
| 108 | Output merged_output = merged_outputs[oidx]; |
| 109 | if (merged_output.node() == nullptr) { |
| 110 | Output new_output(new_node, oidx); |
| 111 | if (debugging_opts.print_outputs) { |
| 112 | string cpu_device = "/job:localhost/replica:0/task:0/device:CPU:0"; |
| 113 | ops::Print print_op(s.WithOpName("print_", oidx) |
| 114 | .WithDevice(cpu_device) |
| 115 | .WithAssignedDevice(cpu_device), |
| 116 | new_output, {new_output}, |
| 117 | ops::Print::Attrs{} |
| 118 | .Message(absl::StrCat("output ", oidx, " from ", |
| 119 | old_node->name(), " is ")) |
| 120 | .FirstN(1000) |
| 121 | .Summarize(-1)); |
| 122 | new_output = print_op; |
| 123 | } |
| 124 | |
| 125 | if (debugging_opts.check_output_numerics && |
| 126 | DataTypeIsFloating(new_output.type())) { |
| 127 | ops::CheckNumerics check_numerics_op( |
| 128 | s.WithOpName("check_output_", oidx) |
| 129 | .WithDevice(new_node->requested_device()) |
| 130 | .WithAssignedDevice(new_node->assigned_device_name()), |
| 131 | new_output, |
| 132 | absl::StrCat("CheckNumerics failed for output ", oidx, "(", |
| 133 | new_output.name(), ") from cluster ", cluster_name)); |
| 134 | new_output = check_numerics_op; |
| 135 | } |
| 136 | |
| 137 | ops::_XlaMerge xla_merge_op(s.WithOpName("merge_oidx_", oidx), |
| 138 | Output(old_node, oidx), new_output); |
| 139 | merged_output = merged_outputs[oidx] = xla_merge_op.output; |
| 140 | } |
| 141 | |
| 142 | Node* dst = e->dst(); |
| 143 | int dst_idx = e->dst_input(); |
| 144 | |
| 145 | s.graph()->RemoveEdge(e); |
| 146 | s.graph()->AddEdge(merged_output.node(), merged_output.index(), dst, |
| 147 | dst_idx); |
| 148 | } |
| 149 | } |
| 150 |
no test coverage detected