1 /* Copyright 2018 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_IMPLEMENTATION_SELECTOR_H_ 17 #define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_IMPLEMENTATION_SELECTOR_H_ 18 19 #include <string> 20 21 #include "tensorflow/core/framework/op.h" 22 #include "tensorflow/core/grappler/costs/graph_properties.h" 23 #include "tensorflow/core/grappler/grappler_item.h" 24 #include "tensorflow/core/grappler/op_types.h" 25 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h" 26 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h" 27 #include "tensorflow/core/grappler/optimizers/function_api_info.h" 28 #include "tensorflow/core/lib/core/errors.h" 29 #include "tensorflow/core/lib/core/stringpiece.h" 30 #include "tensorflow/core/lib/strings/strcat.h" 31 #include "tensorflow/core/util/device_name_utils.h" 32 33 namespace tensorflow { 34 namespace grappler { 35 36 // This transformation replaces function calls by the appropriate function 37 // definition based on properties of the runtime system. For instance, 38 // we may choose one implementation over another if we have a GPU with 39 // enough memory available. 40 // 41 // It is a way for the programmer to specify alternative implementations 42 // of the same functionality in the graph, and let TensorFlow pick the 43 // most appropriate one at runtime. 44 // 45 // For instance, the python code might specify: 46 // @Defun(tf.float32, 47 // api_implements='plus_one', 48 // api_preferred_device='GPU') 49 // def plus_one_gpu(x): return x + 1.0 50 // 51 // @Defun(tf.float32, 52 // api_implements='plus_one') 53 // def plus_one_reference_implementation(x): return x + 1.0 54 // input = tf.constant(2.0, dtype=tf.float32) 55 // 56 // z = plus_one_reference_implementation(input) 57 // z = plus_one_gpu(input) 58 // print(sess.run(z)) 59 // 60 // At runtime, we will trim either `plus_one_gpu` or 61 // `plus_one_reference_implementation` based on the availability of the GPU. 62 // 63 // Available annotations: 64 // - api_implements(string): all functions mapping to the same 65 // string can be interchanged. For now, all functions must have the same 66 // signature and overloads are not allowed. Defuns within defuns are 67 // allowed. 68 // - api_preferred_device(string): sets which device is preferred. 69 class ImplementationSelector : public CustomGraphOptimizer { 70 public: 71 ImplementationSelector() = default; 72 ~ImplementationSelector() override = default; Init(const tensorflow::RewriterConfig_CustomGraphOptimizer * config)73 Status Init( 74 const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override { 75 return Status::OK(); 76 } name()77 string name() const override { 78 return "implementation_selector"; 79 } 80 81 // This call is not thread-safe. 82 Status Optimize(Cluster* cluster, const GrapplerItem& item, 83 GraphDef* optimized_graph) override; 84 85 // Does not take any feedback. Feedback(Cluster * cluster,const GrapplerItem & item,const GraphDef & optimized_graph,double result)86 void Feedback(Cluster* cluster, const GrapplerItem& item, 87 const GraphDef& optimized_graph, double result) override {} 88 89 private: 90 Status LoadFunctions(const GraphDef& graph); 91 Status MaybeOptimizeFunctionCall(NodeDef* node_def) const; 92 93 // Finds all call sites for functions, then replace with the appropriate 94 // implementation. 95 // There are two ways of calling functions: 96 // 1. By specifying an op name as a function name, and 97 // 2. Via the functional interface, where the function name appears as an 98 // Attr. 99 // 100 // There may be multiple call sites for a given function. The function body 101 // may call into another function, so a function might have to be duplicated. 102 // For simplicity, we do not change function bodies. Also, we do not change 103 // gradients. 104 Status SelectImplementation(GraphDef* graph) const; 105 106 std::unique_ptr<FunctionLibraryApiInfo> lib_info_; 107 108 TF_DISALLOW_COPY_AND_ASSIGN(ImplementationSelector); 109 }; 110 111 } // namespace grappler 112 } // namespace tensorflow 113 114 #endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_IMPLEMENTATION_SELECTOR_H_ 115