[mlir] Add conversion from Complex to Standard dialect for NotEqualOp.

[akuegel]

Summary

We can use a common templatized conversion for EqualOp and NotEqualOp.
This allows to reuse most of the logic for the conversion of EqualOp.

[pifon2a]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

72

Would it make sense to get rid of this if statement and just pass the replacement op type, like

ComparisonOpConversion<complex::EqualOp, CmpFPredicate::OEQ, AndOp>?

or even better, make AndOp the function of typename ComparisonOp.

[akuegel]

Derive result combination op type via std::conditional_t from comparison op type.

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

72

Thanks, learned now about std::conditional_t
Done.

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

matchAndRewrite(complex::AbsOp op, ArrayRef<Value> operands,
     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;
+template <typename ComparisonOp, CmpFPredicate p>
+struct ComparisonOpConversion : public OpConversionPattern<ComparisonOp> {
+  using OpConversionPattern<ComparisonOp>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(complex::EqualOp op, ArrayRef<Value> operands,
+  matchAndRewrite(ComparisonOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    complex::EqualOp::Adaptor transformed(operands);
+    typename ComparisonOp::Adaptor transformed(operands);
     auto loc = op.getLoc();
-    auto type =
-        transformed.lhs().getType().cast<ComplexType>().getElementType();
+    auto type = transformed.lhs()
+                    .getType()
+                    .template 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);
+    Value realComparison = rewriter.create<CmpFOp>(loc, p, realLhs, realRhs);
+    Value imagComparison = rewriter.create<CmpFOp>(loc, p, imagLhs, imagRhs);
 
-    rewriter.replaceOpWithNewOp<AndOp>(op, realEqual, imagEqual);
+    if (p == CmpFPredicate::OEQ)
+      rewriter.replaceOpWithNewOp<AndOp>(op, realComparison, imagComparison);
+    else if (p == CmpFPredicate::UNE)
+      rewriter.replaceOpWithNewOp<OrOp>(op, realComparison, imagComparison);
+    else
+      return failure();
     return success();
   }
 };
 } // namespace
 
 void mlir::populateComplexToStandardConversionPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<AbsOpConversion, EqualOpConversion>(patterns.getContext());
+  patterns.add<AbsOpConversion,
+               ComparisonOpConversion<complex::EqualOp, CmpFPredicate::OEQ>,
+               ComparisonOpConversion<complex::NotEqualOp, CmpFPredicate::UNE>>(
+      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, complex::EqualOp>();
+  target.addIllegalOp<complex::AbsOp, complex::EqualOp, complex::NotEqualOp>();
   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

 // 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
+
+// CHECK-LABEL: func @complex_neq
+// CHECK-SAME: %[[LHS:.*]]: complex<f32>, %[[RHS:.*]]: complex<f32>
+func @complex_neq(%lhs: complex<f32>, %rhs: complex<f32>) -> i1 {
+  %neq = complex.neq %lhs, %rhs: complex<f32>
+  return %neq : 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_NOT_EQUAL:.*]] = cmpf une, %[[REAL_LHS]], %[[REAL_RHS]] : f32
+// CHECK-DAG: %[[IMAG_NOT_EQUAL:.*]] = cmpf une, %[[IMAG_LHS]], %[[IMAG_RHS]] : f32
+// CHECK: %[[NOT_EQUAL:.*]] = or %[[REAL_NOT_EQUAL]], %[[IMAG_NOT_EQUAL]] : i1
+// CHECK: return %[[NOT_EQUAL]] : i1

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

 // 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
+
+// CHECK-LABEL: llvm.func @complex_neq
+// CHECK-SAME: %[[LHS:.*]]: ![[C_TY:.*]], %[[RHS:.*]]: ![[C_TY:.*]])
+func @complex_neq(%lhs: complex<f32>, %rhs: complex<f32>) -> i1 {
+  %neq = complex.neq %lhs, %rhs: complex<f32>
+  return %neq : 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_NOT_EQUAL:.*]] = llvm.fcmp "une" %[[REAL_LHS]], %[[REAL_RHS]]  : f32
+// CHECK-DAG: %[[IMAG_NOT_EQUAL:.*]] = llvm.fcmp "une" %[[IMAG_LHS]], %[[IMAG_RHS]]  : f32
+// CHECK: %[[NOT_EQUAL:.*]] = llvm.or %[[REAL_NOT_EQUAL]], %[[IMAG_NOT_EQUAL]] : i1
+// CHECK: llvm.return %[[NOT_EQUAL]] : i1