allow to disable one broadcast optimization in algebraic_simplifier

PiperOrigin-RevId: 685914179

Author: tensorflower-gardener

third_party/xla/xla/service/algebraic_simplifier.cc

@@ -5122,7 +5122,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleBroadcast(
   if (options_.is_layout_sensitive()) {
     return absl::OkStatus();
   }
-  if (ShapeUtil::HasDegenerateDimensions(operand->shape())) {
+  if (options_.enable_broadcast_degenerate_dimension() &&
+      ShapeUtil::HasDegenerateDimensions(operand->shape())) {
     auto new_operand = operand->AddInstruction(HloInstruction::CreateReshape(
         ShapeUtil::DropDegenerateDimensions(operand->shape()), operand));
     std::vector<int64_t> new_dims;

third_party/xla/xla/service/algebraic_simplifier.h

@@ -313,6 +313,15 @@ class AlgebraicSimplifierOptions {
   }
   bool enable_fast_math() const { return enable_fast_math_; }
 
+  void set_enable_broadcast_degenerate_dimension(
+      bool enable_broadcast_degenerate_dimension) {
+    enable_broadcast_degenerate_dimension_ =
+        enable_broadcast_degenerate_dimension;
+  }
+  bool enable_broadcast_degenerate_dimension() const {
+    return enable_broadcast_degenerate_dimension_;
+  }
+
  private:
   // Metadata struct can be used to store any metadata information encapsulated
   // with the AlgebraicSimplifierOptions that can be later used in an
@@ -354,6 +363,7 @@ class AlgebraicSimplifierOptions {
   bool use_convert_constant_folding_{false};
   bool disable_dynamic_slice_to_slice_conversion_{false};
   bool enable_fast_math_{false};
+  bool enable_broadcast_degenerate_dimension_{true};
   Metadata metadata_;
 };