| 1369 | } // namespace |
| 1370 | |
| 1371 | Status MemoryOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 1372 | GraphDef* optimized_graph) { |
| 1373 | std::set<int> nodes_to_relax; |
| 1374 | TF_RETURN_IF_ERROR(FindAssignNodesToRelax(item.graph, &nodes_to_relax)); |
| 1375 | |
| 1376 | bool run_recomputation_pass = |
| 1377 | (optimization_level_ == RewriterConfig::RECOMPUTATION_HEURISTICS || |
| 1378 | optimization_level_ == RewriterConfig::HEURISTICS || |
| 1379 | optimization_level_ == RewriterConfig::MANUAL); |
| 1380 | if (!run_recomputation_pass && nodes_to_relax.empty() && item.fetch.empty()) { |
| 1381 | return errors::Aborted("Nothing to do."); |
| 1382 | } |
| 1383 | |
| 1384 | GrapplerItem optimized_item(item); |
| 1385 | RelaxAssignNodes(nodes_to_relax, &optimized_item.graph); |
| 1386 | |
| 1387 | if (run_recomputation_pass) { |
| 1388 | RecomputationRewritingPass(optimization_level_, |
| 1389 | recomputation_targets_name_scope_, |
| 1390 | &optimized_item.graph, item); |
| 1391 | } |
| 1392 | |
| 1393 | std::unordered_set<string> skip_list; |
| 1394 | // Bound the number of rewrite passes to avoid long processing times on graphs |
| 1395 | // that simply won't fit in memory. |
| 1396 | // SchedulingPass() and SwappingPass() rely on defined fetches in order to |
| 1397 | // infer the memory usage, so skip optimization if there are no fetches. |
| 1398 | if (!item.fetch.empty()) { |
| 1399 | bool updated_graph = true; |
| 1400 | for (int i = 0; i < 25 && updated_graph; ++i) { |
| 1401 | GRAPPLER_RETURN_IF_DEADLINE_EXCEEDED(); |
| 1402 | updated_graph = false; |
| 1403 | if ((optimization_level_ == RewriterConfig::DEFAULT_MEM_OPT || |
| 1404 | optimization_level_ == RewriterConfig::SCHEDULING_HEURISTICS || |
| 1405 | optimization_level_ == RewriterConfig::HEURISTICS) && |
| 1406 | cluster != nullptr) { |
| 1407 | updated_graph |= SchedulingPass(cluster, &optimized_item); |
| 1408 | } |
| 1409 | |
| 1410 | GRAPPLER_RETURN_IF_DEADLINE_EXCEEDED(); |
| 1411 | if ((optimization_level_ == RewriterConfig::DEFAULT_MEM_OPT || |
| 1412 | optimization_level_ == RewriterConfig::SWAPPING_HEURISTICS || |
| 1413 | optimization_level_ == RewriterConfig::HEURISTICS || |
| 1414 | optimization_level_ == RewriterConfig::MANUAL) && |
| 1415 | cluster != nullptr) { |
| 1416 | updated_graph |= SwappingPass(optimization_level_, cluster, |
| 1417 | &optimized_item, &skip_list); |
| 1418 | } |
| 1419 | } |
| 1420 | } |
| 1421 | |
| 1422 | optimized_graph->Swap(&optimized_item.graph); |
| 1423 | return Status::OK(); |
| 1424 | } |
| 1425 | |
| 1426 | void MemoryOptimizer::Feedback(Cluster* cluster, const GrapplerItem& item, |
| 1427 | const GraphDef& optimized_graph, double result) { |
nothing calls this directly
no test coverage detected