From c1a6ae89945036e51fda2785c051f96ac1ddc5b9 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Tue, 30 Mar 2021 03:53:04 -0700
Subject: [PATCH] Generalize the HloBinaryElementwiseAdaptor

We can use it also for ternary ops like Select if we change the signature so
that a ValueRange is passed in.
Also remove special casing for HloComplexAdaptor. It can be handled with the
generic adaptor as well.

PiperOrigin-RevId: 365777493
---
 .../mhlo/transforms/map_chlo_to_hlo_op.h      | 32 +++++++------------
 .../mhlo/transforms/chlo_legalize_to_hlo.cc   |  9 +++---
 .../mhlo/transforms/transform_unranked_hlo.cc |  2 +-
 3 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/transforms/map_chlo_to_hlo_op.h b/include/mlir-hlo/Dialect/mhlo/transforms/map_chlo_to_hlo_op.h
index d9e637d..2b061c0 100644
--- a/include/mlir-hlo/Dialect/mhlo/transforms/map_chlo_to_hlo_op.h
+++ b/include/mlir-hlo/Dialect/mhlo/transforms/map_chlo_to_hlo_op.h
@@ -25,30 +25,22 @@ limitations under the License.
 namespace mlir {
 namespace chlo {
 
-struct HloComplexAdaptor {
-  static mhlo::ComplexOp CreateOp(BroadcastComplexOp from_op, Type result_type,
-                                  Value broadcasted_lhs, Value broadcasted_rhs,
-                                  OpBuilder &builder) {
-    return builder.create<mhlo::ComplexOp>(from_op.getLoc(), result_type,
-                                           broadcasted_lhs, broadcasted_rhs);
-  }
-};
 template <typename FromOpTy, typename ToOpTy>
-struct HloBinaryElementwiseAdaptor {
+struct HloNaryElementwiseAdaptor {
   static ToOpTy CreateOp(FromOpTy from_op, Type result_type,
-                         Value broadcasted_lhs, Value broadcasted_rhs,
-                         OpBuilder &builder) {
+                         ValueRange broadcasted_operands, OpBuilder &builder) {
     return builder.create<ToOpTy>(from_op.getLoc(), result_type,
-                                  broadcasted_lhs, broadcasted_rhs);
+                                  broadcasted_operands);
   }
 };
 struct HloCompareAdaptor {
   static mhlo::CompareOp CreateOp(BroadcastCompareOp from_op, Type result_type,
-                                  Value broadcasted_lhs, Value broadcasted_rhs,
+                                  ValueRange broadcasted_operands,
                                   OpBuilder &builder) {
     return builder.create<mhlo::CompareOp>(
-        from_op.getLoc(), result_type, broadcasted_lhs, broadcasted_rhs,
-        from_op.comparison_direction(), from_op.compare_typeAttr());
+        from_op.getLoc(), result_type, broadcasted_operands[0],
+        broadcasted_operands[1], from_op.comparison_direction(),
+        from_op.compare_typeAttr());
   }
 };
 
@@ -59,14 +51,15 @@ template <template <typename, typename, typename> class Pattern,
 void PopulateForBroadcastingBinaryOp(MLIRContext *context,
                                      OwningRewritePatternList *patterns,
                                      ConstructorArgs &&...args) {
-#define POPULATE_BCAST(ChloOp, HloOp)                                      \
-  patterns->insert<                                                        \
-      Pattern<ChloOp, HloOp, HloBinaryElementwiseAdaptor<ChloOp, HloOp>>>( \
+#define POPULATE_BCAST(ChloOp, HloOp)                                    \
+  patterns->insert<                                                      \
+      Pattern<ChloOp, HloOp, HloNaryElementwiseAdaptor<ChloOp, HloOp>>>( \
       context, args...);
 
   POPULATE_BCAST(BroadcastAddOp, mhlo::AddOp);
   POPULATE_BCAST(BroadcastAndOp, mhlo::AndOp);
   POPULATE_BCAST(BroadcastAtan2Op, mhlo::Atan2Op);
+  POPULATE_BCAST(BroadcastComplexOp, mhlo::ComplexOp);
   POPULATE_BCAST(BroadcastDivOp, mhlo::DivOp);
   POPULATE_BCAST(BroadcastMaxOp, mhlo::MaxOp);
   POPULATE_BCAST(BroadcastMinOp, mhlo::MinOp);
@@ -83,9 +76,6 @@ void PopulateForBroadcastingBinaryOp(MLIRContext *context,
   POPULATE_BCAST(BroadcastZetaOp, ZetaOp);
 
   // Broadcasting ops requiring special construction.
-  patterns
-      ->insert<Pattern<BroadcastComplexOp, mhlo::ComplexOp, HloComplexAdaptor>>(
-          context, args...);
   patterns
       ->insert<Pattern<BroadcastCompareOp, mhlo::CompareOp, HloCompareAdaptor>>(
           context, args...);
diff --git a/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo.cc b/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo.cc
index da883a6..916dc01 100644
--- a/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo.cc
+++ b/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo.cc
@@ -1119,9 +1119,8 @@ struct ConvertTrivialNonBroadcastBinaryOp
       }
     }
 
-    rewriter.replaceOp(
-        op, {Adaptor::CreateOp(op, op.getResult().getType(), operands[0],
-                               operands[1], rewriter)});
+    rewriter.replaceOp(op, {Adaptor::CreateOp(op, op.getResult().getType(),
+                                              operands, rewriter)});
     return success();
   }
 };
@@ -1214,8 +1213,8 @@ struct ConvertRankedDynamicBroadcastBinaryOp
         rewriter.getI64TensorAttr(rhs_broadcast_dimensions));
 
     // And generate the final non-broadcasted binary op.
-    Value final_result = Adaptor::CreateOp(op, result_type, broadcasted_lhs,
-                                           broadcasted_rhs, rewriter);
+    Value final_result = Adaptor::CreateOp(
+        op, result_type, {broadcasted_lhs, broadcasted_rhs}, rewriter);
     rewriter.create<shape::AssumingYieldOp>(loc, final_result);
     rewriter.replaceOp(op, {assuming_op.getResult(0)});
     return success();
diff --git a/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc b/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc
index e3ac0b6..6e2c9db 100644
--- a/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc
+++ b/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc
@@ -463,7 +463,7 @@ struct ConvertUnrankedDynamicBroadcastBinaryOp
     OpBuilder if_eq_shapes_builder =
         if_eq_shapes_op.getThenBodyBuilder(rewriter.getListener());
     Value non_broadcast_op =
-        Adaptor::CreateOp(op, result_type, lhs, rhs, if_eq_shapes_builder);
+        Adaptor::CreateOp(op, result_type, {lhs, rhs}, if_eq_shapes_builder);
     if_eq_shapes_builder.create<scf::YieldOp>(loc, non_broadcast_op);
 
     // If shapes do not have exactly one element, nor are equal