| 189 | } |
| 190 | |
| 191 | Status CaseBuilder::AddOutputs() { |
| 192 | // Construct the call nodes for each branch. |
| 193 | call_nodes_.resize(num_branches_, nullptr); |
| 194 | for (int b = 0; b < num_branches_; b++) { |
| 195 | TF_RETURN_IF_ERROR( |
| 196 | branch_call_builders_[b].Finalize(graph_, &call_nodes_[b])); |
| 197 | graph_->AddControlEdge(pivots_[b], call_nodes_[b]); |
| 198 | } |
| 199 | |
| 200 | // Merge the outputs from the N branches (all branches have matching outputs). |
| 201 | const int num_outputs = call_nodes_[0]->num_outputs(); |
| 202 | std::vector<Node*> merges(num_outputs); |
| 203 | outputs_.resize(merges.size()); |
| 204 | for (int i = 0; i < num_outputs; ++i) { |
| 205 | std::vector<NodeOut> merge_input; |
| 206 | merge_input.reserve(num_branches_); |
| 207 | for (int j = 0; j < num_branches_; j++) { |
| 208 | merge_input.emplace_back(call_nodes_[j], i); |
| 209 | } |
| 210 | TF_RETURN_IF_ERROR(NodeBuilder(NewName("merge"), "Merge", |
| 211 | graph_->op_registry(), &debug_info_) |
| 212 | .Input(merge_input) |
| 213 | .Device(case_op_->requested_device()) |
| 214 | .Finalize(graph_, &merges[i])); |
| 215 | outputs_[i] = NodeOut(merges[i], 0); |
| 216 | } |
| 217 | |
| 218 | // Add a Merge node that will be used as a control dependency source for the |
| 219 | // lowered output node. This Merge node will guarantee that lowered else/then |
| 220 | // function calls will be executed even if they do not have data outputs. |
| 221 | // |
| 222 | // Furthermore it will guarantee that all function side effects will be |
| 223 | // executed, if the function will be inlined into the graph. Having data |
| 224 | // outputs is not enough, because they might become unused after inlining. |
| 225 | // |
| 226 | // We will use this node to rewrite outgoing control edges from lowered 'Case' |
| 227 | // node. All data edges will read tensors directly from Merge nodes. |
| 228 | std::vector<NodeOut> pivots(num_branches_); |
| 229 | for (int j = 0; j < num_branches_; j++) { |
| 230 | pivots[j] = NodeOut(pivots_[j]); |
| 231 | } |
| 232 | TF_RETURN_IF_ERROR(NodeBuilder(NewName("branch_executed"), "Merge", |
| 233 | graph_->op_registry(), &debug_info_) |
| 234 | .Input(pivots) |
| 235 | .ControlInputs(call_nodes_) |
| 236 | .Device(case_op_->requested_device()) |
| 237 | .Finalize(graph_, &branch_executed_node_)); |
| 238 | |
| 239 | TF_RETURN_IF_ERROR(BuildLoweredCaseOutput()); |
| 240 | |
| 241 | // Add outputs. |
| 242 | for (const Edge* e : case_op_->out_edges()) { |
| 243 | if (e->IsControlEdge()) { |
| 244 | graph_->AddControlEdge(branch_executed_node_, e->dst()); |
| 245 | } else { |
| 246 | // Feed the outputs directly from the merge nodes so that downstream ops |
| 247 | // can start before all the outputs have been computed. |
| 248 | graph_->AddEdge(merges[e->src_output()], 0, e->dst(), e->dst_input()); |
no test coverage detected