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_COMPILER_JIT_FLAGS_H_ 17 #define TENSORFLOW_COMPILER_JIT_FLAGS_H_ 18 19 #include <vector> 20 21 #include "tensorflow/core/platform/types.h" 22 #include "tensorflow/core/protobuf/config.pb.h" 23 #include "tensorflow/core/util/command_line_flags.h" 24 25 namespace tensorflow { 26 27 struct XlaAutoJitFlag { 28 // Control compilation of operators into XLA computations on CPU and GPU 29 // devices. 0 = use ConfigProto setting; -1 = off; 1 = on for things very 30 // likely to be improved; 2 = on for everything. 31 // 32 // If all non-CPU ops in the graph being optimized are placed on a single GPU 33 // and there is at least one node placed on that GPU then 34 // `optimization_level_single_gpu` applies. Otherwise 35 // `optimization_level_general` applies. 36 // 37 // Experimental. 38 int32 optimization_level_single_gpu; 39 int32 optimization_level_general; 40 }; 41 42 // Sets the xla_auto_jit_flag based on the given flag string. Supported syntax 43 // is: 44 // <number>: sets general and single_gpu setting to the provided number. 45 // single-gpu(<number>): sets the single_gpu setting to the provided number. 46 bool SetXlaAutoJitFlagFromFlagString(const string& value); 47 48 // Flags associated with the XLA bridge's mark_for_compilation_pass module. 49 struct MarkForCompilationPassFlags { 50 XlaAutoJitFlag xla_auto_jit_flag; 51 52 // Minimum number of operators in an XLA compilation. Ignored for operators 53 // placed on an XLA device or operators explicitly marked for compilation. 54 int32 tf_xla_min_cluster_size; 55 56 // Maximum number of operators in an XLA compilation. 57 int32 tf_xla_max_cluster_size; 58 59 // If non-empty, limit XLA clustering to the following TF operations. 60 string tf_xla_ops_to_cluster; 61 62 // Dump graphs during XLA compilation. 63 bool tf_xla_clustering_debug; 64 65 // Enables global JIT compilation for CPU via SessionOptions. 66 bool tf_xla_cpu_global_jit; 67 68 // "Compiler fuel" for clustering. Only this many ops will be marked as 69 // eligible for clustering. 70 int64 tf_xla_clustering_fuel; 71 72 // If tf_xla_disable_deadness_safety_checks_for_debugging is set to true then 73 // we do not do deadness related safety checks. This is unsound in general, 74 // but can be used as a debugging aid. 75 bool tf_xla_disable_deadness_safety_checks_for_debugging; 76 77 // If tf_xla_disable_resource_variable_safety_checks_for_debugging is set to 78 // true then we do not do safety checks to preserve TensorFlow's resource 79 // variable concurrency semantics. This is unsound in general, but can be 80 // used as a debugging aid. 81 bool tf_xla_disable_resource_variable_safety_checks_for_debugging; 82 }; 83 84 // Flags associated with the XLA bridge's xla_device module. 85 struct XlaDeviceFlags { 86 // Switch the CPU device into "on-demand" mode, where instead of 87 // autoclustering ops are compiled one by one just-in-time. 88 // Enabling this mode by a legacy flag is a temporary mechanism. When this 89 // feature is battle-tested, we will switch this to be a session option. 90 bool tf_xla_compile_on_demand; 91 92 // Enables "XLA" devices if this flag is set. 93 bool tf_xla_enable_xla_devices; 94 }; 95 96 // Flags common to the _Xla* ops and their kernels. 97 struct XlaOpsCommonFlags { 98 // If true, _XlaCompile always refuses to compile the cluster, which means the 99 // XLA clusters always run in the TF executor. Defaults to false. 100 bool tf_xla_always_defer_compilation; 101 }; 102 103 // Flags for the build_xla_ops pass. 104 struct BuildXlaOpsPassFlags { 105 // Enables lazy compilation for TF/XLA (only when auto-clustering) if true. 106 // Defaults to true. 107 bool tf_xla_enable_lazy_compilation; 108 109 // If true then insert Print nodes to print out values produced by XLA 110 // clusters. Useful for debugging. 111 bool tf_xla_print_cluster_outputs; 112 113 // If true, insert CheckNumerics nodes for every floating point typed input to 114 // an XLA cluster. 115 bool tf_xla_check_cluster_input_numerics; 116 117 // If true, insert CheckNumerics nodes for every floating point typed output 118 // from an XLA cluster. 119 bool tf_xla_check_cluster_output_numerics; 120 121 // Disables all constant folding. The primary use for this is for testing to 122 // guarantee that tests are run on XLA and not on TF's CPU implementation. 123 bool tf_xla_disable_constant_folding; 124 }; 125 126 // Flags for the IntroduceFloatingPointJitter pass. 127 struct IntroduceFloatingPointJitterPassFlags { 128 // The amount of jitter to introduce. This amount is added to each element in 129 // the tensors named in `tensor_names. 130 float jitter_amount; 131 132 // The Tensors to add the jitter to. The tensors are named in the TensorId 133 // format of <node name>:<output idx>. 134 std::vector<string> tensor_names; 135 }; 136 137 // Flags for common MLIR configurations. 138 struct MlirCommonFlags { 139 ConfigProto::Experimental::MlirBridgeRollout tf_mlir_enable_mlir_bridge; 140 }; 141 142 // Return a pointer to the DumpGraphFlags struct; 143 // repeated calls return the same pointer. 144 // This should be called only after Flags::Parse() has returned. 145 146 // Getters for flags structs defined above. The first call to any of these 147 // parses TF_XLA_FLAGS for all of them. Those functions which return a pointer 148 // always return the same pointer. 149 MarkForCompilationPassFlags* GetMarkForCompilationPassFlags(); 150 BuildXlaOpsPassFlags* GetBuildXlaOpsPassFlags(); 151 XlaDeviceFlags* GetXlaDeviceFlags(); 152 const XlaOpsCommonFlags& GetXlaOpsCommonFlags(); 153 154 const IntroduceFloatingPointJitterPassFlags& 155 GetIntroduceFloatingPointJitterPassFlags(); 156 157 MlirCommonFlags* GetMlirCommonFlags(); 158 159 // Appends the flag definitions associated with 160 // MarkForCompilationPassFlags/DumpGraphFlags to `flag_list`. 161 // 162 // Has the side-effect of parsing TF_XLA_FLAGS if that hasn't happened yet. 163 void AppendMarkForCompilationPassFlags( 164 std::vector<tensorflow::Flag>* flag_list); 165 166 // Disables XLA compilation, forces it to return an error message instead. Can 167 // be used by a server to ensure that JIT compilation is opt-in. 168 void DisableXlaCompilation(); 169 170 // Returns `false` unless `DisableXlaCompilation` was called. 171 bool FailOnXlaCompilation(); 172 173 } // namespace tensorflow 174 175 #endif // TENSORFLOW_COMPILER_JIT_FLAGS_H_ 176