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