From bd23bbf19c69467f41f0b1adbd7f39b8e1edfc6b Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Tue, 4 Feb 2020 12:43:13 -0800
Subject: [PATCH] Disable horizontal fusion by default

We have an internal customer who has asked us to do a critical systems freeze
until Feb 7.  We will re-enable the pass on Feb 7.

PiperOrigin-RevId: 293206387
Change-Id: I6b401fd68b2219c27136937f860dd21690d3e273
---
 tensorflow/compiler/xla/debug_options_flags.cc     |  6 ++++++
 .../compiler/xla/service/gpu/gpu_compiler.cc       | 14 ++++++++------
 .../xla/service/gpu/horizontal_fusion_test.cc      | 12 +++++++++++-
 tensorflow/compiler/xla/xla.proto                  |  6 +++++-
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/xla/debug_options_flags.cc b/tensorflow/compiler/xla/debug_options_flags.cc
index 9808dd3d092..699baa5731a 100644
--- a/tensorflow/compiler/xla/debug_options_flags.cc
+++ b/tensorflow/compiler/xla/debug_options_flags.cc
@@ -535,6 +535,12 @@ static void AllocateFlags() {
           bool_setter_for(&DebugOptions::set_xla_gpu_deterministic_reductions),
           flag_values->xla_gpu_deterministic_reductions(),
           "Always run deterministic reductions on GPU"),
+      // TODO(b/148871440): Enable by default on Feb 7.
+      tensorflow::Flag(
+          "xla_gpu_use_horizontal_fusion",
+          bool_setter_for(&DebugOptions::set_xla_gpu_use_horizontal_fusion),
+          flag_values->xla_gpu_use_horizontal_fusion(),
+          "Allow horizontal fusion for XLA GPU."),
   });
   ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", *flag_objects);
 }
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
index 6fe440dfdb0..271fbd85c62 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
@@ -280,12 +280,14 @@ Status GpuCompiler::OptimizeHloModule(
     fusion.AddPass<HloDCE>();
     TF_RETURN_IF_ERROR(fusion.Run(hlo_module).status());
 
-    HloPassPipeline horizontal_fusion("horizontal_fusion");
-    horizontal_fusion.AddPass<GpuHorizontalFusion>();
-    horizontal_fusion.AddPass<HloCSE>(/*is_layout_sensitive=*/true,
-                                      /*only_fusion_computations=*/true);
-    horizontal_fusion.AddPass<HloDCE>();
-    TF_RETURN_IF_ERROR(horizontal_fusion.Run(hlo_module).status());
+    if (hlo_module->config().debug_options().xla_gpu_use_horizontal_fusion()) {
+      HloPassPipeline horizontal_fusion("horizontal_fusion");
+      horizontal_fusion.AddPass<GpuHorizontalFusion>();
+      horizontal_fusion.AddPass<HloCSE>(/*is_layout_sensitive=*/true,
+                                        /*only_fusion_computations=*/true);
+      horizontal_fusion.AddPass<HloDCE>();
+      TF_RETURN_IF_ERROR(horizontal_fusion.Run(hlo_module).status());
+    }
   }
 
   return Status::OK();
diff --git a/tensorflow/compiler/xla/service/gpu/horizontal_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/horizontal_fusion_test.cc
index e1024f6017c..3642b76ffe9 100644
--- a/tensorflow/compiler/xla/service/gpu/horizontal_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/horizontal_fusion_test.cc
@@ -37,7 +37,17 @@ namespace {
 
 namespace op = xla::testing::opcode_matchers;
 
-class HorizontalFusionTest : public HloTestBase {};
+class HorizontalFusionTest : public HloTestBase {
+ protected:
+  HloModuleConfig GetModuleConfigWithHorizontalFusionEnabled() {
+    DebugOptions debug_options = GetDebugOptionsForTest();
+    debug_options.set_xla_gpu_use_horizontal_fusion(true);
+
+    HloModuleConfig config;
+    config.set_debug_options(debug_options);
+    return config;
+  }
+};
 
 TEST_F(HorizontalFusionTest, BasicTest) {
   auto module = ParseAndReturnVerifiedModule(R"(
diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto
index 10b77f947e5..da889cf87bf 100644
--- a/tensorflow/compiler/xla/xla.proto
+++ b/tensorflow/compiler/xla/xla.proto
@@ -260,7 +260,11 @@ message DebugOptions {
 
   // Guarantee run-to-run determinism from reductions on XLA:GPU.
   bool xla_gpu_deterministic_reductions = 130;
-  // Next id: 133
+
+  // Enable horizontal fusion on XLA GPU.
+  bool xla_gpu_use_horizontal_fusion = 133;
+
+  // Next id: 134
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.