| 198 | } |
| 199 | |
| 200 | Status CondBuilder::AddOutputs() { |
| 201 | // Construct the then and else nodes. |
| 202 | TF_RETURN_IF_ERROR(then_call_builder_.Finalize(graph_, &then_call_node_)); |
| 203 | graph_->AddControlEdge(pivot_t_, then_call_node_); |
| 204 | TF_RETURN_IF_ERROR(else_call_builder_.Finalize(graph_, &else_call_node_)); |
| 205 | graph_->AddControlEdge(pivot_f_, else_call_node_); |
| 206 | |
| 207 | // Add Merge node for each data output of the If node. |
| 208 | std::vector<Node*> merges(then_call_node_->num_outputs()); |
| 209 | outputs_.resize(merges.size()); |
| 210 | for (int i = 0; i < then_call_node_->num_outputs(); ++i) { |
| 211 | TF_RETURN_IF_ERROR( |
| 212 | NodeBuilder(NewName("output"), "Merge", graph_->op_registry(), |
| 213 | &debug_info_) |
| 214 | .Input({NodeOut(then_call_node_, i), NodeOut(else_call_node_, i)}) |
| 215 | .Device(if_op_->requested_device()) |
| 216 | .Finalize(graph_, &merges[i])); |
| 217 | outputs_[i] = NodeOut(merges[i], 0); |
| 218 | } |
| 219 | |
| 220 | // Add a Merge node that will be used as a control dependency source for the |
| 221 | // lowered output node. This Merge node will guarantee that lowered else/then |
| 222 | // function calls will be executed even if they do not have data outputs. |
| 223 | // |
| 224 | // Furthermore it will guarantee that all function side effects will be |
| 225 | // executed, if the function will be inlined into the graph. Having data |
| 226 | // outputs is not enough, because they might become unused after inlining. |
| 227 | // |
| 228 | // We will use this node to rewrite outgoing control edges from lowered 'If' |
| 229 | // node. All data edges will read tensors directly from Merge nodes. |
| 230 | TF_RETURN_IF_ERROR(NodeBuilder(NewName("branch_executed"), "Merge", |
| 231 | graph_->op_registry(), &debug_info_) |
| 232 | .Input({pivot_t_, pivot_f_}) |
| 233 | .ControlInputs({then_call_node_, else_call_node_}) |
| 234 | .Device(if_op_->requested_device()) |
| 235 | .Finalize(graph_, &branch_executed_node_)); |
| 236 | |
| 237 | TF_RETURN_IF_ERROR(BuildLoweredIfOutput()); |
| 238 | |
| 239 | // Add outputs. |
| 240 | for (const Edge* e : if_op_->out_edges()) { |
| 241 | if (e->IsControlEdge()) { |
| 242 | graph_->AddControlEdge(branch_executed_node_, e->dst()); |
| 243 | } else { |
| 244 | // Feed the outputs directly from the merge nodes so that downstream ops |
| 245 | // can start before all the outputs have been computed. |
| 246 | graph_->AddEdge(merges[e->src_output()], 0, e->dst(), e->dst_input()); |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | return Status::OK(); |
| 251 | } |
| 252 | |
| 253 | Status CondBuilder::BuildLoweredIfOutput() { |
| 254 | // If outputs are empty, it means that we might have only output control |
no test coverage detected