[mlir][complex] Canonicalize re/im(neg(create))

[antiagainst]

Summary

When can just convert this to arith.negf.

[kuhar]

LGTM.

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

let description = [{
 
     ```mlir
     %a = complex.im %b : complex<f32>
     ```
   }];
 
   let results = (outs AnyFloat:$imaginary);
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // LogOp
 //===----------------------------------------------------------------------===//
 
 def LogOp : ComplexUnaryOp<"log", [SameOperandsAndResultType]> {
   let summary = "computes natural logarithm of a complex number";
let description = [{
 
     ```mlir
     %a = complex.re %b : complex<f32>
     ```
   }];
 
   let results = (outs AnyFloat:$real);
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RsqrtOp
 //===----------------------------------------------------------------------===//
 
 def RsqrtOp : ComplexUnaryOp<"rsqrt", [SameOperandsAndResultType]> {
   let summary = "complex reciprocal of square root";

mlir/lib/Dialect/Complex/IR/ComplexOps.cpp

 //===- ComplexOps.cpp - MLIR Complex Operations ---------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 using namespace mlir::complex;
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
ArrayAttr arrayAttr =
       llvm::dyn_cast_if_present<ArrayAttr>(adaptor.getComplex());
   if (arrayAttr && arrayAttr.size() == 2)
     return arrayAttr[1];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
     return createOp.getOperand(1);
   return {};
 }
 
+namespace {
+template <typename OpKind, int ComponentIndex>
+struct FoldComponentNeg final : OpRewritePattern<OpKind> {
+  using OpRewritePattern<OpKind>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpKind op,
+                                PatternRewriter &rewriter) const override {
+    auto negOp = op.getOperand().template getDefiningOp<NegOp>();
+    if (!negOp)
+      return failure();
+
+    auto createOp = negOp.getComplex().template getDefiningOp<CreateOp>();
+    if (!createOp)
+      return failure();
+
+    Type elementType = createOp.getType().getElementType();
+    assert(isa<FloatType>(elementType));
+
+    rewriter.replaceOpWithNewOp<arith::NegFOp>(
+        op, elementType, createOp.getOperand(ComponentIndex));
+    return success();
+  }
+};
+} // namespace
+
+void ImOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                       MLIRContext *context) {
+  results.add<FoldComponentNeg<ImOp, 1>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ReOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ReOp::fold(FoldAdaptor adaptor) {
   ArrayAttr arrayAttr =
       llvm::dyn_cast_if_present<ArrayAttr>(adaptor.getComplex());
   if (arrayAttr && arrayAttr.size() == 2)
     return arrayAttr[0];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
     return createOp.getOperand(0);
   return {};
 }
 
+void ReOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                       MLIRContext *context) {
+  results.add<FoldComponentNeg<ReOp, 0>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AddOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
   // complex.add(complex.sub(a, b), b) -> a
   if (auto sub = getLhs().getDefiningOp<SubOp>())
     if (getRhs() == sub.getRhs())

mlir/test/Dialect/Complex/canonicalize.mlir

 func.func @complex_sub_zero() -> complex<f32> {
   %complex1 = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
   %complex2 = complex.constant [0.0 : f32, 0.0 : f32] : complex<f32>
   // CHECK: %[[CPLX:.*]] = complex.constant [1.000000e+00 : f32, 0.000000e+00 : f32] : complex<f32>
   // CHECK-NEXT: return %[[CPLX:.*]] : complex<f32>
   %sub = complex.sub %complex1, %complex2 : complex<f32>
   return %sub : complex<f32>
 }
+
+// CHECK-LABEL: func @re_neg
+//  CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32)
+func.func @re_neg(%arg0: f32, %arg1: f32) -> f32 {
+  %create = complex.create %arg0, %arg1: complex<f32>
+  // CHECK: %[[NEG:.*]] = arith.negf %[[ARG0]]
+  %neg = complex.neg %create : complex<f32>
+  %re = complex.re %neg : complex<f32>
+  // CHECK-NEXT: return %[[NEG]]
+  return %re : f32
+}
+
+// CHECK-LABEL: func @im_neg
+//  CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32)
+func.func @im_neg(%arg0: f32, %arg1: f32) -> f32 {
+  %create = complex.create %arg0, %arg1: complex<f32>
+  // CHECK: %[[NEG:.*]] = arith.negf %[[ARG1]]
+  %neg = complex.neg %create : complex<f32>
+  %im = complex.im %neg : complex<f32>
+  // CHECK-NEXT: return %[[NEG]]
+  return %im : f32
+}