[mlir] Add conversion from complex to standard dialect for EqualOp.

[akuegel]

Summary

This adds the straightforward conversion for EqualOp
(two complex numbers are equal if both the real and the imaginary part are equal).

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

matchAndRewrite(complex::AbsOp op, ArrayRef<Value> operands,
     Value realSqr = rewriter.create<MulFOp>(loc, real, real);
     Value imagSqr = rewriter.create<MulFOp>(loc, imag, imag);
     Value sqNorm = rewriter.create<AddFOp>(loc, realSqr, imagSqr);
 
     rewriter.replaceOpWithNewOp<math::SqrtOp>(op, sqNorm);
     return success();
   }
 };
+
+struct EqualOpConversion : public OpConversionPattern<complex::EqualOp> {
+  using OpConversionPattern<complex::EqualOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::EqualOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    complex::EqualOp::Adaptor transformed(operands);
+    auto loc = op.getLoc();
+    auto type =
+        transformed.lhs().getType().cast<ComplexType>().getElementType();
+
+    Value realLhs =
+        rewriter.create<complex::ReOp>(loc, type, transformed.lhs());
+    Value imagLhs =
+        rewriter.create<complex::ImOp>(loc, type, transformed.lhs());
+    Value realRhs =
+        rewriter.create<complex::ReOp>(loc, type, transformed.rhs());
+    Value imagRhs =
+        rewriter.create<complex::ImOp>(loc, type, transformed.rhs());
+    Value realEqual =
+        rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, realLhs, realRhs);
+    Value imagEqual =
+        rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, imagLhs, imagRhs);
+
+    rewriter.replaceOpWithNewOp<AndOp>(op, realEqual, imagEqual);
+    return success();
+  }
+};
 } // namespace
 
 void mlir::populateComplexToStandardConversionPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<AbsOpConversion>(patterns.getContext());
+  patterns.add<AbsOpConversion, EqualOpConversion>(patterns.getContext());
 }
 
 namespace {
 struct ConvertComplexToStandardPass
     : public ConvertComplexToStandardBase<ConvertComplexToStandardPass> {
   void runOnFunction() override;
 };
 
 void ConvertComplexToStandardPass::runOnFunction() {
   auto function = getFunction();
 
   // Convert to the Standard dialect using the converter defined above.
   RewritePatternSet patterns(&getContext());
   populateComplexToStandardConversionPatterns(patterns);
 
   ConversionTarget target(getContext());
   target.addLegalDialect<StandardOpsDialect, math::MathDialect,
                          complex::ComplexDialect>();
-  target.addIllegalOp<complex::AbsOp>();
+  target.addIllegalOp<complex::AbsOp, complex::EqualOp>();
   if (failed(applyPartialConversion(function, target, std::move(patterns))))
     signalPassFailure();
 }
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createConvertComplexToStandardPass() {
   return std::make_unique<ConvertComplexToStandardPass>();
 }

mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir

 // RUN: mlir-opt %s -convert-complex-to-standard | FileCheck %s
 
 // CHECK-LABEL: func @complex_abs
 // CHECK-SAME: %[[ARG:.*]]: complex<f32>
 func @complex_abs(%arg: complex<f32>) -> f32 {
   %abs = complex.abs %arg: complex<f32>
   return %abs : f32
 }
 // CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
 // CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
 // CHECK-DAG: %[[REAL_SQ:.*]] = mulf %[[REAL]], %[[REAL]] : f32
 // CHECK-DAG: %[[IMAG_SQ:.*]] = mulf %[[IMAG]], %[[IMAG]] : f32
 // CHECK: %[[SQ_NORM:.*]] = addf %[[REAL_SQ]], %[[IMAG_SQ]] : f32
 // CHECK: %[[NORM:.*]] = math.sqrt %[[SQ_NORM]] : f32
 // CHECK: return %[[NORM]] : f32
 
+// CHECK-LABEL: func @complex_eq
+// CHECK-SAME: %[[LHS:.*]]: complex<f32>, %[[RHS:.*]]: complex<f32>
+func @complex_eq(%lhs: complex<f32>, %rhs: complex<f32>) -> i1 {
+  %eq = complex.eq %lhs, %rhs: complex<f32>
+  return %eq : i1
+}
+// CHECK: %[[REAL_LHS:.*]] = complex.re %[[LHS]] : complex<f32>
+// CHECK: %[[IMAG_LHS:.*]] = complex.im %[[LHS]] : complex<f32>
+// CHECK: %[[REAL_RHS:.*]] = complex.re %[[RHS]] : complex<f32>
+// CHECK: %[[IMAG_RHS:.*]] = complex.im %[[RHS]] : complex<f32>
+// CHECK-DAG: %[[REAL_EQUAL:.*]] = cmpf oeq, %[[REAL_LHS]], %[[REAL_RHS]] : f32
+// CHECK-DAG: %[[IMAG_EQUAL:.*]] = cmpf oeq, %[[IMAG_LHS]], %[[IMAG_RHS]] : f32
+// CHECK: %[[EQUAL:.*]] = and %[[REAL_EQUAL]], %[[IMAG_EQUAL]] : i1
+// CHECK: return %[[EQUAL]] : i1

mlir/test/Conversion/ComplexToStandard/full-conversion.mlir

 // RUN: mlir-opt %s -convert-complex-to-standard -convert-complex-to-llvm -convert-std-to-llvm | FileCheck %s
 
 // CHECK-LABEL: llvm.func @complex_abs
 // CHECK-SAME: %[[ARG:.*]]: ![[C_TY:.*]])
 func @complex_abs(%arg: complex<f32>) -> f32 {
   %abs = complex.abs %arg: complex<f32>
   return %abs : f32
 }
 // CHECK: %[[REAL:.*]] = llvm.extractvalue %[[ARG]][0] : ![[C_TY]]
 // CHECK: %[[IMAG:.*]] = llvm.extractvalue %[[ARG]][1] : ![[C_TY]]
 // CHECK-DAG: %[[REAL_SQ:.*]] = llvm.fmul %[[REAL]], %[[REAL]]  : f32
 // CHECK-DAG: %[[IMAG_SQ:.*]] = llvm.fmul %[[IMAG]], %[[IMAG]]  : f32
 // CHECK: %[[SQ_NORM:.*]] = llvm.fadd %[[REAL_SQ]], %[[IMAG_SQ]]  : f32
 // CHECK: %[[NORM:.*]] = "llvm.intr.sqrt"(%[[SQ_NORM]]) : (f32) -> f32
 // CHECK: llvm.return %[[NORM]] : f32
 
+// CHECK-LABEL: llvm.func @complex_eq
+// CHECK-SAME: %[[LHS:.*]]: ![[C_TY:.*]], %[[RHS:.*]]: ![[C_TY:.*]])
+func @complex_eq(%lhs: complex<f32>, %rhs: complex<f32>) -> i1 {
+  %eq = complex.eq %lhs, %rhs: complex<f32>
+  return %eq : i1
+}
+// CHECK: %[[REAL_LHS:.*]] = llvm.extractvalue %[[LHS]][0] : ![[C_TY]]
+// CHECK: %[[IMAG_LHS:.*]] = llvm.extractvalue %[[LHS]][1] : ![[C_TY]]
+// CHECK: %[[REAL_RHS:.*]] = llvm.extractvalue %[[RHS]][0] : ![[C_TY]]
+// CHECK: %[[IMAG_RHS:.*]] = llvm.extractvalue %[[RHS]][1] : ![[C_TY]]
+// CHECK-DAG: %[[REAL_EQUAL:.*]] = llvm.fcmp "oeq" %[[REAL_LHS]], %[[REAL_RHS]]  : f32
+// CHECK-DAG: %[[IMAG_EQUAL:.*]] = llvm.fcmp "oeq" %[[IMAG_LHS]], %[[IMAG_RHS]]  : f32
+// CHECK: %[[EQUAL:.*]] = llvm.and %[[REAL_EQUAL]], %[[IMAG_EQUAL]] : i1
+// CHECK: llvm.return %[[EQUAL]] : i1