1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DEPENDENCY_OPTIMIZER_H_ 17 #define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DEPENDENCY_OPTIMIZER_H_ 18 19 #include <unordered_set> 20 #include "tensorflow/core/grappler/optimizers/graph_optimizer.h" 21 #include "tensorflow/core/grappler/utils.h" 22 #include "tensorflow/core/protobuf/rewriter_config.pb.h" 23 24 namespace tensorflow { 25 namespace grappler { 26 27 // Optimize TF computations by removing control dependencies or re-arranging 28 // them to shorten the critical path for a model step or enable other 29 // optimizations, such as removing nodes that are effectively noops. 30 class DependencyOptimizer : public GraphOptimizer { 31 public: 32 DependencyOptimizer() : opt_level_(RewriterConfig::ON) {} 33 explicit DependencyOptimizer(RewriterConfig::Toggle opt_level) 34 : opt_level_(opt_level) {} 35 ~DependencyOptimizer() override {} 36 37 string name() const override { return "dependency_optimizer"; }; 38 39 Status Optimize(Cluster* cluster, const GrapplerItem& item, 40 GraphDef* optimized_graph) override; 41 42 void Feedback(Cluster* cluster, const GrapplerItem& item, 43 const GraphDef& optimized_graph, double result) override; 44 45 private: 46 // Returns true if node is not an Identity node or if it is an Identity 47 // that is safe to remove. 48 bool SafeToRemoveIdentity(const NodeDef& node); 49 // Returns true if it is safe to convert node to NoOp. 50 bool SafeToConvertToNoOp(const NodeDef& node); 51 // Removes all duplicate control dependencies. 52 void CleanControlInputs(); 53 // Builds a map from the &optimized_graph_->node(i) to i. 54 void BuildNodeToIdx(); 55 // Tries to optimize the node with the given index, possibly additional 56 // optimizations by inserting nodes in nodes_to_simplify, and pruning nodes by 57 // inserting them in nodes_to_delete. 58 void OptimizeNode(int node_idx, SetVector<int>* nodes_to_simplify, 59 std::set<int>* nodes_to_delete); 60 // Eliminates redundant control dependencies by computing the transitive 61 // reduction of the graph. 62 Status TransitiveReduction(); 63 // Main driver of dependency optimizations. 64 Status OptimizeDependencies(); 65 66 RewriterConfig::Toggle opt_level_; 67 bool fetch_nodes_known_; 68 std::unordered_set<string> nodes_to_preserve_; 69 std::unique_ptr<NodeMap> node_map_; 70 std::unordered_map<const NodeDef*, int> node_to_idx_; 71 GraphDef* optimized_graph_; // Not owned. 72 }; 73 74 } // end namespace grappler 75 } // end namespace tensorflow 76 77 #endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DEPENDENCY_OPTIMIZER_H_ 78
