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_META_OPTIMIZER_H_ 17 #define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_META_OPTIMIZER_H_ 18 19 #include "tensorflow/core/common_runtime/device_set.h" 20 #include "tensorflow/core/framework/device_base.h" 21 #include "tensorflow/core/framework/function.h" 22 #include "tensorflow/core/graph/graph.h" 23 #include "tensorflow/core/grappler/grappler_item.h" 24 #include "tensorflow/core/grappler/optimizers/graph_optimizer.h" 25 #include "tensorflow/core/grappler/verifiers/graph_verifier.h" 26 #include "tensorflow/core/lib/core/status.h" 27 #include "tensorflow/core/protobuf/config.pb.h" 28 #include "tensorflow/core/protobuf/rewriter_config.pb.h" 29 #include "tensorflow/core/protobuf/verifier_config.pb.h" 30 31 namespace tensorflow { 32 namespace grappler { 33 34 // Run the other grappler optimizers based on the specified rewriter config. 35 class MetaOptimizer : public GraphOptimizer { 36 public: 37 MetaOptimizer(DeviceBase* cpu_device, const ConfigProto& cfg); 38 ~MetaOptimizer() override = default; 39 name()40 string name() const override { return "meta_optimizer"; }; 41 UsesFunctionLibrary()42 bool UsesFunctionLibrary() const override { return true; } 43 Optimize(Cluster * cluster,const GrapplerItem & item,GraphDef * optimized_graph)44 Status Optimize(Cluster* cluster, const GrapplerItem& item, 45 GraphDef* optimized_graph) override { 46 GrapplerItem copy(item); 47 return OptimizeConsumeItem(cluster, std::move(copy), optimized_graph); 48 } 49 50 Status OptimizeConsumeItem(Cluster* cluster, GrapplerItem&& item, 51 GraphDef* optimized_graph); 52 53 string GetResultString() const; 54 55 void PrintResult(); 56 57 private: 58 std::unique_ptr<GraphOptimizer> MakeNewOptimizer( 59 const string& optimizer, const std::set<string>& device_types) const; 60 61 // When grappler should lower control flow to V1 switch/merge style nodes. 62 bool LowerControlFlow() const; 63 64 // Initialize active optimizers from RewriterConfig toggles. 65 Status InitializeOptimizers( 66 const std::set<string>& device_types, 67 std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const; 68 // Initialize active optimizers from RewriterConfig optimizer names. 69 Status InitializeOptimizersByName( 70 const std::set<string>& device_types, 71 std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const; 72 // Initialize active optimizers from RewriterConfig.custom_optimizers. 73 Status InitializeCustomGraphOptimizers( 74 const std::set<string>& device_types, 75 const std::set<string>& pre_initialized_optimizers, 76 std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const; 77 Status InitializePluginGraphOptimizers( 78 const std::set<string>& device_types, 79 std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const; 80 // Returns the config for a custom graph optimizer. Null if none was found. 81 const RewriterConfig::CustomGraphOptimizer* GetCustomGraphOptimizerConfig( 82 const string& name) const; 83 84 // Initialize active verifiers from the RewriterConfig toggles. 85 void InitializeVerifiers( 86 std::vector<std::unique_ptr<GraphVerifier>>* inter_optimizer_verifiers, 87 std::vector<std::unique_ptr<GraphVerifier>>* post_optimization_verifiers) 88 const; 89 90 void PrintUserAndPluginConfigs(const std::set<string>& device_types) const; 91 92 // Run optimization pass over a single GrapplerItem. Meta optimizer might run 93 // multiple such passes: 1) for the main graph 2) for the function library 94 Status OptimizeGraph(Cluster* cluster, GrapplerItem&& item, 95 GraphDef* optimized_graph); 96 97 DeviceBase* const cpu_device_; // may be NULL 98 ConfigProto config_proto_; 99 RewriterConfig& cfg_; 100 bool xla_auto_clustering_on_; 101 102 struct OptimizerResult { 103 string optimizer_name; 104 string message; 105 Status status; 106 }; 107 108 struct GraphOptimizationResult { GraphOptimizationResultGraphOptimizationResult109 explicit GraphOptimizationResult(const string& id) : id(id) {} 110 string id; 111 std::vector<OptimizerResult> results; 112 }; 113 114 Status RunOptimizer(GraphOptimizer* optimizer, Cluster* cluster, 115 GrapplerItem* optimized_item, GraphDef* optimized_graph, 116 GraphOptimizationResult* optimization_result); 117 118 std::vector<GraphOptimizationResult> optimization_results_; 119 }; 120 121 bool MetaOptimizerEnabled(const ConfigProto& cfg); 122 123 // Run the meta optimizer. 124 // 125 // If <cpu_device> is non-null, it is the device to be used for executing ops 126 // during constant folding; if NULL, a new device is created for doing constant 127 // folding. For performance, it is recommended to pass in an existing cpu_device 128 // when possible. 129 Status RunMetaOptimizer(GrapplerItem&& item, const ConfigProto& cfg, 130 DeviceBase* cpu_device, Cluster* cluster, 131 GraphDef* optimized_graph); 132 133 // Wrapper around RunMetaOptimizer convenient for optimizing 134 // function graphs. 135 // 136 // Runs grappler optimizations on `g` based on `config_proto`. 137 // `ret_node_names`: a vector of node names whose outputs are returned, 138 // aka fetches. when `g` represent a function, these are _Retval nodes. 139 // `lib`: function library to use with `g`. 140 // `device_set`: the set of devices that graph can refer to. 141 // `cpu_device`: the CPU device. 142 // `config_proto`: Grapper configuration. 143 // `grappler_item_id': Grappler item id (e.g. optimized function name). 144 // `optimization_options`: Grappler optimization constraints that are known only 145 // at runtime. 146 // 147 // **g is a graph constructed based on the runtime library 'lib'. 148 // OptimizeGraph mutates **g extensively and replaces '*g' with a 149 // complete copy. Therefore, the caller should not keep any references 150 // to nodes *g. 151 Status OptimizeGraph( 152 std::vector<string> ret_node_names, std::vector<string> keep_node_names, 153 FunctionLibraryDefinition* lib, const DeviceSet& device_set, 154 Device* cpu_device, const ConfigProto& config_proto, 155 const string& grappler_item_id, 156 const GrapplerItem::OptimizationOptions& optimization_options, 157 std::unique_ptr<tensorflow::Graph>* g); 158 159 } // namespace grappler 160 } // namespace tensorflow 161 162 #endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_META_OPTIMIZER_H_ 163