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

Differential Revision: https://reviews.llvm.org/D103148
This commit is contained in:
Adrian Kuegel 2021-05-26 12:28:14 +02:00
parent 8ac66d61ea
commit dee46d0829
3 changed files with 45 additions and 8 deletions

View File

@ -88,7 +88,7 @@ def CreateOp : Complex_Op<"create",
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:
@ -102,6 +102,7 @@ def CreateOp : Complex_Op<"create",
let results = (outs Complex<AnyFloat>:$complex);
let assemblyFormat = "$real `,` $imaginary attr-dict `:` type($complex)";
let hasFolder = 1;
}
//===----------------------------------------------------------------------===//

View File

@ -19,13 +19,16 @@ using namespace mlir::complex;
#define GET_OP_CLASSES
#include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"
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[0];
if (auto createOp = getOperand().getDefiningOp<CreateOp>())
return createOp.getOperand(0);
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 {};
}
@ -38,3 +41,13 @@ OpFoldResult ImOp::fold(ArrayRef<Attribute> operands) {
return createOp.getOperand(1);
return {};
}
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[0];
if (auto createOp = getOperand().getDefiningOp<CreateOp>())
return createOp.getOperand(0);
return {};
}

View File

@ -1,5 +1,28 @@
// 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