experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA/Bridge] Provide a flag not to resolve compile time constants.

Browse Source 
For some models, resolving constant expressions in the bridge leads to
non-linear slowdown.

PiperOrigin-RevId: 280205828
Change-Id: Icaa0394d8045c31544b66123857da7cb1b1701b9
This commit is contained in:
George Karpenkov 2019-11-13 08:43:33 -08:00 committed by TensorFlower Gardener
parent ffe860458e
commit 2fb65c9308
3 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/compiler/jit/flags.cc
View File

@ -161,6 +161,9 @@ void AllocateAndParseFlags() {
Flag("tf_xla_always_defer_compilation",
&ops_flags->tf_xla_always_defer_compilation, ""),
Flag("tf_xla_noresolve_compile_time_constants",
&ops_flags->tf_xla_noresolve_compile_time_constants,
"Do not perform constant folding in XlaCompiler::CompileGraph"),
Flag("tf_introduce_floating_point_jitter_to_tensors",
setter_for_jitter_tensor_names, "",

8
tensorflow/compiler/jit/flags.h
View File

@ -91,6 +91,14 @@ struct XlaOpsCommonFlags {
// If true, _XlaCompile always refuses to compile the cluster, which means the
// XLA clusters always run in the TF executor. Defaults to false.
bool tf_xla_always_defer_compilation;
// If true, sets compile_options.resolve_compile_time_constants to false,
// which stops the bridge from using the HloEvaluator for constant resolution
// in XlaCompiler::CompileGraph.
//
// For some models, constant folding during compile graph experiences a
// non-linear blow up, which overshadows both compilation and execution.
bool tf_xla_noresolve_compile_time_constants;
};
// Flags for the build_xla_ops pass.

7
tensorflow/compiler/jit/kernels/xla_ops.cc
View File

@ -326,11 +326,8 @@ static Status CompileToLocalExecutable(
}
XlaCompiler::CompileOptions compile_options;
compile_options.is_entry_computation = true;
// If we resolve constants we never emit them on the device, meaning that if
// they are needed by a following computation the host has to transfer
// them. Not resolving constants is expected to be faster than resolving
// constants.
compile_options.resolve_compile_time_constants = true;
compile_options.resolve_compile_time_constants =
!GetXlaOpsCommonFlags().tf_xla_noresolve_compile_time_constants;
// Optimization: where possible, have the computation return a naked array
// rather than a one-element tuple.
compile_options.always_return_tuple = false;