| 378 | } |
| 379 | |
| 380 | void SingleMachine::MergeCosts(CostGraphDef* graph_costs, |
| 381 | const CostGraphDef& init_costs, |
| 382 | const CostGraphDef& queue_costs) { |
| 383 | graph_costs->mutable_node()->Reserve(graph_costs->node_size() + |
| 384 | init_costs.node_size() + |
| 385 | queue_costs.node_size()); |
| 386 | std::unordered_set<string> nodes_seen; |
| 387 | int queue_costs_id_offset = graph_costs->node_size(); |
| 388 | for (const auto& node : graph_costs->node()) { |
| 389 | nodes_seen.insert(node.name()); |
| 390 | if (node.id() >= queue_costs_id_offset) { |
| 391 | queue_costs_id_offset = node.id() + 1; |
| 392 | } |
| 393 | } |
| 394 | |
| 395 | int init_costs_id_offset = queue_costs_id_offset + queue_costs.node_size(); |
| 396 | // The costs obtained by running the main graph could be more stable than |
| 397 | // the one we get from the queue runners since the queue runners run |
| 398 | // asynchronously. |
| 399 | for (const auto& node : queue_costs.node()) { |
| 400 | if (nodes_seen.find(node.name()) != nodes_seen.end()) { |
| 401 | continue; |
| 402 | } |
| 403 | |
| 404 | auto* new_node = graph_costs->add_node(); |
| 405 | new_node->MergeFrom(node); |
| 406 | |
| 407 | new_node->set_id(node.id() + queue_costs_id_offset); |
| 408 | if (new_node->id() >= init_costs_id_offset) { |
| 409 | init_costs_id_offset = new_node->id() + 1; |
| 410 | } |
| 411 | |
| 412 | for (auto& input_info : *new_node->mutable_input_info()) { |
| 413 | input_info.set_preceding_node(input_info.preceding_node() + |
| 414 | queue_costs_id_offset); |
| 415 | } |
| 416 | for (auto& control_input : *new_node->mutable_control_input()) { |
| 417 | control_input += queue_costs_id_offset; |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | // Don't overwrite the costs with that generated during initialization since |
| 422 | // these are possibly outdated. |
| 423 | for (const auto& node : init_costs.node()) { |
| 424 | if (nodes_seen.find(node.name()) != nodes_seen.end()) { |
| 425 | continue; |
| 426 | } |
| 427 | |
| 428 | auto* new_node = graph_costs->add_node(); |
| 429 | new_node->MergeFrom(node); |
| 430 | |
| 431 | new_node->set_id(node.id() + init_costs_id_offset); |
| 432 | for (auto& input_info : *new_node->mutable_input_info()) { |
| 433 | input_info.set_preceding_node(input_info.preceding_node() + |
| 434 | init_costs_id_offset); |
| 435 | } |
| 436 | for (auto& control_input : *new_node->mutable_control_input()) { |
| 437 | control_input += init_costs_id_offset; |