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

[XLA] Respect TF_DETERMINISTIC_OPS environment variable for reductions

Browse Source 
PiperOrigin-RevId: 296094275
Change-Id: Iadcbf33d5d6432413c86d4d176865980de252eeb
This commit is contained in:
George Karpenkov 2020-02-19 17:20:19 -08:00 committed by TensorFlower Gardener
parent b56c66d883
commit 8b7a3db0b6
3 changed files with 15 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/xla/service/gpu/BUILD
View File

@ -1285,6 +1285,7 @@ cc_library(
":reduction_dimension_grouper",
":reduction_layout_normalizer",
":target_constants",
":tree_reduction_rewriter",
"//tensorflow/compiler/xla:statusor",
"//tensorflow/compiler/xla/service:algebraic_simplifier",
"//tensorflow/compiler/xla/service:hlo",

1
tensorflow/compiler/xla/service/gpu/amdgpu_compiler.cc
View File

@ -26,6 +26,7 @@ limitations under the License.
#include "tensorflow/compiler/xla/service/gpu/reduction_dimension_grouper.h"
#include "tensorflow/compiler/xla/service/gpu/reduction_layout_normalizer.h"
#include "tensorflow/compiler/xla/service/gpu/target_constants.h"
#include "tensorflow/compiler/xla/service/gpu/tree_reduction_rewriter.h"
#include "tensorflow/compiler/xla/service/hlo_constant_folding.h"
#include "tensorflow/compiler/xla/service/hlo_cse.h"
#include "tensorflow/compiler/xla/service/hlo_pass_fix.h"

14
tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
View File

@ -55,6 +55,7 @@ limitations under the License.
#include "tensorflow/core/platform/cuda_libdevice_path.h"
#include "tensorflow/core/platform/tracing.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/util/env_var.h"
#include "tensorflow/stream_executor/cuda/cuda_diagnostics.h"
#include "tensorflow/stream_executor/gpu/asm_compiler.h"
@ -151,6 +152,16 @@ Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
return Status::OK();
}
// TODO(cheshire): Duplication with gpu_conv_algorithm picker, figure out a
// right way to share this.
static bool RequireDeterminism() {
bool deterministic_ops = false;
TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar("TF_DETERMINISTIC_OPS",
/*default_val=*/false,
&deterministic_ops));
return deterministic_ops;
}
Status NVPTXCompiler::OptimizeHloPostLayoutAssignment(
HloModule* hlo_module, se::StreamExecutor* stream_exec,
se::DeviceMemoryAllocator* device_allocator) {
@ -172,7 +183,8 @@ Status NVPTXCompiler::OptimizeHloPostLayoutAssignment(
options.set_is_layout_sensitive(true);
pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(options);
if (hlo_module->config().debug_options().xla_gpu_deterministic_reductions()) {
if (RequireDeterminism() ||
hlo_module->config().debug_options().xla_gpu_deterministic_reductions()) {
pipeline.AddPass<HloPassFix<GpuTreeReductionRewriter>>();
}