| 295 | } // end namespace internal |
| 296 | |
| 297 | Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 298 | GraphDef* optimized_graph) { |
| 299 | *optimized_graph = item.graph; |
| 300 | |
| 301 | // Skip all TPU graphs. |
| 302 | if (IsTPUGraphDef(*optimized_graph)) { |
| 303 | return Status::OK(); |
| 304 | } |
| 305 | |
| 306 | GraphProperties properties(item); |
| 307 | GraphView graph(optimized_graph); |
| 308 | |
| 309 | gtl::FlatSet<string> devices; |
| 310 | if (cluster) { |
| 311 | const std::vector<string> device_names = cluster->GetDeviceNames(); |
| 312 | devices.insert(device_names.begin(), device_names.end()); |
| 313 | } else { |
| 314 | devices = {"/device:CPU:0"}; |
| 315 | } |
| 316 | |
| 317 | const bool has_device_cpu = devices.find("/device:CPU:0") != devices.end(); |
| 318 | |
| 319 | // Topologically sort the graph, so that we traverse the nodes in order. This |
| 320 | // will help us discover producer->consumer chains of Host ops. |
| 321 | TF_RETURN_IF_ERROR(TopologicalSort(optimized_graph)); |
| 322 | |
| 323 | // All the Const nodes, and their original devices in topological order. |
| 324 | std::vector<std::pair<NodeDef*, string>> const_nodes; |
| 325 | |
| 326 | for (auto& node : *optimized_graph->mutable_node()) { |
| 327 | GRAPPLER_RETURN_IF_DEADLINE_EXCEEDED(); |
| 328 | bool is_candidate = false; |
| 329 | TF_RETURN_IF_ERROR( |
| 330 | internal::IsNodeHostCandidate(graph, &properties, node, &is_candidate)); |
| 331 | if (!is_candidate) { |
| 332 | continue; |
| 333 | } |
| 334 | |
| 335 | string device = |
| 336 | internal::TryFindHostDevice(devices, has_device_cpu, node.device()); |
| 337 | if (!device.empty()) { |
| 338 | // Keep track of all Const nodes that we swapped. |
| 339 | if (IsConstant(node)) { |
| 340 | const_nodes.emplace_back(&node, node.device()); |
| 341 | } |
| 342 | *node.mutable_device() = std::move(device); |
| 343 | } |
| 344 | } |
| 345 | |
| 346 | // Traverse all `const_nodes`, and map them back to GPU greedily. |
| 347 | for (auto& it : const_nodes) { |
| 348 | GRAPPLER_RETURN_IF_DEADLINE_EXCEEDED(); |
| 349 | NodeDef* node = it.first; |
| 350 | const string& device = it.second; |
| 351 | |
| 352 | // Check all the consumers of this node, if any of them are not on CPU, swap |
| 353 | // this node back onto the original device. |
| 354 | for (const GraphView::InputPort& fanout : graph.GetFanouts(*node, false)) { |
nothing calls this directly
no test coverage detected