Do not normalize broadcast operands on GPU.

There is no need to normalize broadcast operands on GPU, and it makes it harder to write matchers if they have to optionally match a bitcast. PiperOrigin-RevId: 514373072
vwbaker · Mar 6, 2023 · 5673903 · 5673903
1 parent 18534d1
commit 5673903
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 1 deletion.
diff --git a/xla/service/algebraic_simplifier.cc b/xla/service/algebraic_simplifier.cc
@@ -3665,7 +3665,8 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) {
   if (options_.is_layout_sensitive()) {
     return OkStatus();
   }
-  if (ShapeUtil::HasDegenerateDimensions(operand->shape())) {
+  if (options_.enable_normalize_broadcast_operand() &&
+      ShapeUtil::HasDegenerateDimensions(operand->shape())) {
     auto new_operand =
         operand->parent()->AddInstruction(HloInstruction::CreateReshape(
             ShapeUtil::DropDegenerateDimensions(operand->shape()), operand));

diff --git a/xla/service/algebraic_simplifier.h b/xla/service/algebraic_simplifier.h
@@ -186,6 +186,15 @@ class AlgebraicSimplifierOptions {
 
   bool enable_sink_broadcast() const { return enable_sink_broadcast_; }
 
+  void set_enable_normalize_broadcast_operand(
+      bool enable_normalize_broadcast_operand) {
+    enable_normalize_broadcast_operand_ = enable_normalize_broadcast_operand;
+  }
+
+  bool enable_normalize_broadcast_operand() const {
+    return enable_normalize_broadcast_operand_;
+  }
+
   // If true, min(x, NaN) = NaN.  If false, min(x, NaN) = x.
   //
   // TODO(b/209827141): Remove this and make minmax_propagate_nan uncondtionally
@@ -220,6 +229,7 @@ class AlgebraicSimplifierOptions {
   bool enable_reduce_of_reshape_{true};
   bool enable_negative_padding_replacement_{true};
   bool enable_sink_broadcast_{true};
+  bool enable_normalize_broadcast_operand_{true};
   int64_t very_small_gather_size_{4};
   bool minmax_propagate_nan_{true};
   Metadata metadata_;

diff --git a/xla/service/algebraic_simplifier_test.cc b/xla/service/algebraic_simplifier_test.cc
@@ -1045,6 +1045,23 @@ TEST_F(AlgebraicSimplifierTest, DegenerateDimsInOperandRemovedFromBroadcast) {
               GmockMatch(m::Broadcast(m::Reshape(m::Parameter(0)))));
 }
 
+// Test that Broadcast(x) where x has degenerate dimensions first does not
+// remove the degenerate dimensions if the corresponding option is disabled.
+TEST_F(AlgebraicSimplifierTest,
+       DegenerateDimsInOperandNotRemovedFromBroadcast) {
+  const char* kModuleStr = R"(
+    HloModule m
+    test {
+      c = f32[1,4] parameter(0)
+      ROOT b = f32[5,1,4,3] broadcast(c), dimensions={1,2}
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  AlgebraicSimplifierOptions options;
+  options.set_enable_normalize_broadcast_operand(false);
+  ASSERT_FALSE(AlgebraicSimplifier(options).Run(m.get()).value());
+}
+
 // Test to catch a crash where we were overshooting the reshaped_dimensions
 // vector.
 TEST_F(AlgebraicSimplifierTest, ArrayOvershootTest) {

diff --git a/xla/service/gpu/amdgpu_compiler.cc b/xla/service/gpu/amdgpu_compiler.cc
@@ -99,6 +99,7 @@ Status AMDGPUCompiler::OptimizeHloConvolutionCanonicalization(
   // AlgebraicSimplifier  We run algsimp to a fixed point.
   AlgebraicSimplifierOptions options;
   options.set_enable_conv_operand_swap(false);
+  options.set_enable_normalize_broadcast_operand(false);
   pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(options);
 
   pipeline.AddPass<HloConstantFolding>();

diff --git a/xla/service/gpu/gpu_compiler.cc b/xla/service/gpu/gpu_compiler.cc
@@ -381,6 +381,8 @@ Status GpuCompiler::OptimizeHloModule(
   layout_insensitive_algsimp_opts.set_minmax_propagate_nan(
       !debug_options.xla_gpu_enable_fast_min_max());
 
+  layout_insensitive_algsimp_opts.set_enable_normalize_broadcast_operand(false);
+
   if (gpu_target_config.platform_name == "ROCM") {
     layout_insensitive_algsimp_opts.set_enable_conv_operand_swap(false);
   }
@@ -806,6 +808,7 @@ Status GpuCompiler::OptimizeHloPostLayoutAssignment(
     options.set_supports_non_canonical_dots(false);
     options.set_is_layout_sensitive(true);
     options.set_enable_conv_operand_swap(false);
+    options.set_enable_normalize_broadcast_operand(false);
     // "slow" minmax means we propagate nan.
     options.set_minmax_propagate_nan(
         !debug_options.xla_gpu_enable_fast_min_max());
@@ -907,6 +910,7 @@ Status GpuCompiler::OptimizeHloPostLayoutAssignment(
     options.set_supports_non_canonical_dots(false);
     options.set_is_layout_sensitive(true);
     options.set_enable_conv_operand_swap(false);
+    options.set_enable_normalize_broadcast_operand(false);
     // "slow" minmax means we propagate nan.
     options.set_minmax_propagate_nan(
         !hlo_module->config().debug_options().xla_gpu_enable_fast_min_max());

diff --git a/xla/service/gpu/nvptx_compiler.cc b/xla/service/gpu/nvptx_compiler.cc
@@ -131,6 +131,7 @@ Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
 
   AlgebraicSimplifierOptions algsimp_options;
   algsimp_options.set_enable_conv_operand_swap(false);
+  algsimp_options.set_enable_normalize_broadcast_operand(false);
   pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(algsimp_options);
 
   // CudnnSimplifyPadding gets rid of some padding introduced by