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: DependencyOptimizer()32 DependencyOptimizer() {} DependencyOptimizer(RewriterConfig::Toggle opt_level)33 explicit DependencyOptimizer(RewriterConfig::Toggle opt_level) {} ~DependencyOptimizer()34 ~DependencyOptimizer() override {} 35 name()36 string name() const override { return "dependency_optimizer"; }; 37 UsesFunctionLibrary()38 bool UsesFunctionLibrary() const override { return false; } 39 40 Status Optimize(Cluster* cluster, const GrapplerItem& item, 41 GraphDef* optimized_graph) override; 42 43 private: 44 // Returns true if bypassing node does not increase the number of edges or 45 // number of edges crossing a device boundary. 46 bool BypassingNodeIsBeneficial( 47 const NodeDef& node, const std::vector<NodeDef*>& input_nodes, 48 const std::vector<NodeDef*>& output_nodes) const; 49 int NumEdgesIfBypassed(const NodeDef& node, 50 const std::vector<NodeDef*>& output_nodes) const; 51 // Returns true if node is not an Identity node or if it is an Identity 52 // that is safe to remove. 53 bool SafeToRemoveIdentity(const NodeDef& node) const; 54 // Returns true if it is safe to convert node to NoOp. 55 bool SafeToConvertToNoOp(const NodeDef& node) const; 56 // Removes all duplicate control dependencies. 57 void CleanControlInputs(); 58 // Builds a map from the &optimized_graph_->node(i) to i. 59 void BuildNodeToIdx(); 60 // Tries to optimize the node with the given index, possibly additional 61 // optimizations by inserting nodes in nodes_to_simplify, and pruning nodes by 62 // inserting them in nodes_to_delete. 63 void OptimizeNode(int node_idx, SetVector<int>* nodes_to_simplify, 64 std::set<int>* nodes_to_delete); 65 // Eliminates redundant control dependencies by computing the transitive 66 // reduction of the graph. 67 Status TransitiveReduction(); 68 // Main driver of dependency optimizations. 69 Status OptimizeDependencies(); 70 // Replaces multiple cross-device control edges from the same device with a 71 // single control edge. If `host_granularity` is true then group control 72 // edges from all devices on the same host. 73 void GroupCrossDeviceControlEdges(bool host_granularity); 74 75 bool fetch_nodes_known_; 76 std::unordered_set<string> nodes_to_preserve_; 77 std::unique_ptr<NodeMap> node_map_; 78 std::unordered_map<const NodeDef*, int> node_to_idx_; 79 GraphDef* optimized_graph_; // Not owned. 80 }; 81 82 } // end namespace grappler 83 } // end namespace tensorflow 84 85 #endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DEPENDENCY_OPTIMIZER_H_ 86