[mlir][complex] Canonicalization for consecutive complex.add and sub

Add basic canonicalization for consecutive complex.add and sub operations.

Reviewed By: pifon2a

Differential Revision: https://reviews.llvm.org/D128702
This commit is contained in:
lewuathe 2022-06-28 11:33:02 +02:00 committed by Alexander Belyaev
parent 5ae9b42efb
commit 036a699675
3 changed files with 45 additions and 0 deletions

View File

@ -73,6 +73,8 @@ def AddOp : ComplexArithmeticOp<"add"> {
%a = complex.add %b, %c : complex<f32>
```
}];
let hasFolder = 1;
}
//===----------------------------------------------------------------------===//

View File

@ -8,6 +8,7 @@
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
using namespace mlir;
using namespace mlir::complex;
@ -103,6 +104,26 @@ OpFoldResult ReOp::fold(ArrayRef<Attribute> operands) {
return {};
}
//===----------------------------------------------------------------------===//
// AddOp
//===----------------------------------------------------------------------===//
OpFoldResult AddOp::fold(ArrayRef<Attribute> operands) {
assert(operands.size() == 2 && "binary op takes 2 operands");
// complex.add(complex.sub(a, b), b) -> a
if (auto sub = getLhs().getDefiningOp<SubOp>())
if (getRhs() == sub.getRhs())
return sub.getLhs();
// complex.add(b, complex.sub(a, b)) -> a
if (auto sub = getRhs().getDefiningOp<SubOp>())
if (getLhs() == sub.getRhs())
return sub.getLhs();
return {};
}
//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//

View File

@ -62,3 +62,25 @@ func.func @imag_of_create_op() -> f32 {
%1 = complex.im %complex : complex<f32>
return %1 : f32
}
// CHECK-LABEL: func @complex_add_sub_lhs
func.func @complex_add_sub_lhs() -> complex<f32> {
%complex1 = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
%complex2 = complex.constant [0.0 : f32, 2.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>
%add = complex.add %sub, %complex2 : complex<f32>
return %add : complex<f32>
}
// CHECK-LABEL: func @complex_add_sub_rhs
func.func @complex_add_sub_rhs() -> complex<f32> {
%complex1 = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
%complex2 = complex.constant [0.0 : f32, 2.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>
%add = complex.add %complex2, %sub : complex<f32>
return %add : complex<f32>
}