[mlir] Fold complex.create(complex.re(op), complex.im(op))

[akuegel]

Summary

If both operands of create decompose the same complex value, it can be folded.

[akuegel]

Sync to HEAD.

[akuegel]

Fix CHECK-LABEL line in test to check the full function name.

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

TypesMatchWith<"complex element type matches real operand type",
                     "complex", "real",
                     "$_self.cast<ComplexType>().getElementType()">,
      TypesMatchWith<"complex element type matches imaginary operand type",
                     "complex", "imaginary",
                     "$_self.cast<ComplexType>().getElementType()">]> {
 
   let summary = "complex number creation operation";
   let description = [{
-    The `complex.complex` operation creates a complex number from two
+    The `complex.create` operation creates a complex number from two
     floating-point operands, the real and the imaginary part.
 
     Example:
 
     ```mlir
     %a = complex.create %b, %c : complex<f32>
     ```
   }];
 
   let arguments = (ins AnyFloat:$real, AnyFloat:$imaginary);
   let results = (outs Complex<AnyFloat>:$complex);
 
   let assemblyFormat = "$real `,` $imaginary attr-dict `:` type($complex)";
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DivOp
 //===----------------------------------------------------------------------===//
 
 def DivOp : ComplexArithmeticOp<"div"> {
   let summary = "complex division";

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

 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"
 
-OpFoldResult ReOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult CreateOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "binary op takes two operands");
+  // Fold complex.create(complex.re(op), complex.im(op)).
+  if (auto reOp = getOperand(0).getDefiningOp<ReOp>()) {
+    if (auto imOp = getOperand(1).getDefiningOp<ImOp>()) {
+      if (reOp.getOperand() == imOp.getOperand()) {
+        return reOp.getOperand();
+      }
+    }
+  }
+  return {};
+}
+
+OpFoldResult ImOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 1 && "unary op takes 1 operand");
   ArrayAttr arrayAttr = operands[0].dyn_cast_or_null<ArrayAttr>();
   if (arrayAttr && arrayAttr.size() == 2)
-    return arrayAttr[0];
+    return arrayAttr[1];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
-    return createOp.getOperand(0);
+    return createOp.getOperand(1);
   return {};
 }
 
-OpFoldResult ImOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult ReOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 1 && "unary op takes 1 operand");
   ArrayAttr arrayAttr = operands[0].dyn_cast_or_null<ArrayAttr>();
   if (arrayAttr && arrayAttr.size() == 2)
-    return arrayAttr[1];
+    return arrayAttr[0];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
-    return createOp.getOperand(1);
+    return createOp.getOperand(0);
   return {};
 }

mlir/test/Dialect/Complex/canonicalize.mlir

 // RUN: mlir-opt %s -canonicalize | FileCheck %s
 
+// CHECK-LABEL: func @create_of_real_and_imag
+// CHECK-SAME: (%[[CPLX:.*]]: complex<f32>)
+func @create_of_real_and_imag(%cplx: complex<f32>) -> complex<f32> {
+  // CHECK-NEXT: return %[[CPLX]] : complex<f32>
+  %real = complex.re %cplx : complex<f32>
+  %imag = complex.im %cplx : complex<f32>
+  %complex = complex.create %real, %imag : complex<f32>
+  return %complex : complex<f32>
+}
+
+// CHECK-LABEL: func @create_of_real_and_imag_different_operand
+// CHECK-SAME: (%[[CPLX:.*]]: complex<f32>, %[[CPLX2:.*]]: complex<f32>)
+func @create_of_real_and_imag_different_operand(
+    %cplx: complex<f32>, %cplx2 : complex<f32>) -> complex<f32> {
+  // CHECK-NEXT: %[[REAL:.*]] = complex.re %[[CPLX]] : complex<f32>
+  // CHECK-NEXT: %[[IMAG:.*]] = complex.im %[[CPLX2]] : complex<f32>
+  // CHECK-NEXT: %[[COMPLEX:.*]] = complex.create %[[REAL]], %[[IMAG]] : complex<f32>
+  %real = complex.re %cplx : complex<f32>
+  %imag = complex.im %cplx2 : complex<f32>
+  %complex = complex.create %real, %imag : complex<f32>
+  return %complex: complex<f32>
+}
+
 // CHECK-LABEL: func @real_of_const(
 func @real_of_const() -> f32 {
   // CHECK: %[[CST:.*]] = constant 1.000000e+00 : f32
   // CHECK-NEXT: return %[[CST]] : f32
   %complex = constant [1.0 : f32, 0.0 : f32] : complex<f32>
   %1 = complex.re %complex : complex<f32>
   return %1 : f32
 }