Have GpuScatterExpander always rewrite Scatter when TF op determinism…

… is enabled.

Scatter is nondeterministic, so this causes TF ops that use Scatter to be made deterministic when TF op determinism is used.

Also raise an error when TF op determinism is enabled and SelectAndScatter is used, as SelectAndScatter is nondeterministic.

PiperOrigin-RevId: 416901938
Change-Id: I40d4a31508f5f81f006c3591c7930928159ac481

tensorflow/compiler/jit/xla_compilation_cache.cc

@@ -50,6 +50,7 @@ limitations under the License.
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/public/version.h"
 #include "tensorflow/core/tpu/tpu_defs.h"
+#include "tensorflow/core/util/determinism.h"
 #include "tensorflow/core/util/dump_graph.h"
 
 namespace tensorflow {
@@ -232,6 +233,9 @@ Status XlaCompilationCache::BuildExecutable(
   build_options.set_alias_passthrough_params(options.alias_passthrough_params);
   build_options.mutable_debug_options()->set_xla_detailed_logging_and_dumping(
       options.detailed_logging);
+  if (tensorflow::OpDeterminismRequired()) {
+    build_options.mutable_debug_options()->set_xla_gpu_deterministic_ops(true);
+  }
   TF_ASSIGN_OR_RETURN(
       auto executables,
       client_->Compile(*result.computation, argument_layouts, build_options));

tensorflow/compiler/tf2xla/kernels/pooling_ops.cc

@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/determinism.h"
 #include "tensorflow/core/util/tensor_format.h"
 
 namespace tensorflow {
@@ -319,6 +320,13 @@ class MaxPoolGradOp : public XlaOpKernel {
     OP_REQUIRES(ctx, padding_ != EXPLICIT,
                 errors::Unimplemented(
                     "XLA does not support maxpoolgrad with explicit padding."));
+    // When determinism is enabled, the use of SelectAndScatter causes a generic
+    // error to be raised. We raise a more informative error here before
+    // SelectAndScatter is used.
+    OP_REQUIRES(
+        ctx, !tensorflow::OpDeterminismRequired(),
+        errors::Unimplemented("GPU MaxPool gradient ops do not yet have a "
+                              "deterministic XLA implementation."));
   }
 
   int num_dims() const { return num_spatial_dims_ + 2; }

tensorflow/compiler/xla/service/gpu/gpu_compiler.cc

@@ -349,7 +349,13 @@ Status GpuCompiler::OptimizeHloModule(
     // handle it.
     pipeline.AddPass<ZeroSizedHloElimination>();
 
-    pipeline.AddPass<GpuScatterExpander>();
+    if (debug_options.xla_gpu_deterministic_ops()) {
+      // Scatter is nondeterministic, so eliminate all Scatters.
+      pipeline.AddPass<ScatterExpander>(ScatterExpander::kEliminateAllScatters);
+    } else {
+      // Only Scatters unsupported on XLA:GPU are eliminated.
+      pipeline.AddPass<GpuScatterExpander>();
+    }
     // TODO(phawkins): replace QR and Eigh decompositions with calls to
     // cuSOLVER.
     pipeline.AddPass<QrExpander>();

tensorflow/python/kernel_tests/array_ops/scatter_nd_ops_test.py

@@ -471,7 +471,6 @@ def tearDown(self):
     super().tearDown()
     config.disable_op_determinism()
 
-  @test_util.disable_xla("Scatter ND is not deterministic with XLA")
   def testDeterminism(self):
     ref = variables.Variable(array_ops.zeros([1]))
     indices = array_ops.zeros([100000, 1], dtypes.int32)
@@ -815,7 +814,6 @@ def tearDown(self):
     super().tearDown()
     config.disable_op_determinism()
 
-  @test_util.disable_xla("Scatter ND is not deterministic with XLA")
   def testDeterminism(self):
     indices = array_ops.zeros([100000, 1], dtypes.int32)
     values = np.random.randn(100000)