diff options
author | Kai Sasaki <lewuathe@gmail.com> | 2024-03-26 13:50:28 +0900 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-03-26 13:50:28 +0900 |
commit | 08a321e1b7ffdd72154ae808821d20fe511c83b6 (patch) | |
tree | 86f675d598c04b060a8823538669f8ba56444f58 /mlir | |
parent | 3e046ee0c944b5a55f30b3ce0153791d91dbb134 (diff) | |
download | llvm-08a321e1b7ffdd72154ae808821d20fe511c83b6.zip llvm-08a321e1b7ffdd72154ae808821d20fe511c83b6.tar.gz llvm-08a321e1b7ffdd72154ae808821d20fe511c83b6.tar.bz2 |
[mlir][complex] Canonicalize complex.div by one (#85513)
We can canonicalize the complex.div if the divisor is one (real = 1.0,
imag = 0.0) with the input number itself.
Ref: https://www.cuemath.com/numbers/division-of-complex-numbers/
Diffstat (limited to 'mlir')
-rw-r--r-- | mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td | 2 | ||||
-rw-r--r-- | mlir/lib/Dialect/Complex/IR/ComplexOps.cpp | 26 | ||||
-rw-r--r-- | mlir/test/Dialect/Complex/canonicalize.mlir | 56 |
3 files changed, 84 insertions, 0 deletions
diff --git a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td index e19d714..d660292 100644 --- a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td +++ b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td @@ -225,6 +225,8 @@ def DivOp : ComplexArithmeticOp<"div"> { %a = complex.div %b, %c : complex<f32> ``` }]; + + let hasFolder = 1; } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp index 5529dcc..1c81433 100644 --- a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp +++ b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp @@ -370,6 +370,32 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) { } //===----------------------------------------------------------------------===// +// DivOp +//===----------------------------------------------------------------------===// + +OpFoldResult DivOp::fold(FoldAdaptor adaptor) { + auto rhs = adaptor.getRhs(); + if (!rhs) + return {}; + + ArrayAttr arrayAttr = rhs.dyn_cast<ArrayAttr>(); + if (!arrayAttr || arrayAttr.size() != 2) + return {}; + + APFloat real = cast<FloatAttr>(arrayAttr[0]).getValue(); + APFloat imag = cast<FloatAttr>(arrayAttr[1]).getValue(); + + if (!imag.isZero()) + return {}; + + // complex.div(a, complex.constant<1.0, 0.0>) -> a + if (real == APFloat(real.getSemantics(), 1)) + return getLhs(); + + return {}; +} + +//===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/Complex/canonicalize.mlir b/mlir/test/Dialect/Complex/canonicalize.mlir index 64c3f31..16fa0fd 100644 --- a/mlir/test/Dialect/Complex/canonicalize.mlir +++ b/mlir/test/Dialect/Complex/canonicalize.mlir @@ -281,3 +281,59 @@ func.func @double_reverse_bitcast(%arg0 : complex<f32>) -> f64 { // CHECK: return %[[R0]] : f64 func.return %1 : f64 } + + +// CHECK-LABEL: func @div_one_f16 +// CHECK-SAME: (%[[ARG0:.*]]: f16, %[[ARG1:.*]]: f16) -> complex<f16> +func.func @div_one_f16(%arg0: f16, %arg1: f16) -> complex<f16> { + %create = complex.create %arg0, %arg1: complex<f16> + %one = complex.constant [1.0 : f16, 0.0 : f16] : complex<f16> + %div = complex.div %create, %one : complex<f16> + // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f16> + // CHECK-NEXT: return %[[CREATE]] + return %div : complex<f16> +} + +// CHECK-LABEL: func @div_one_f32 +// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32) -> complex<f32> +func.func @div_one_f32(%arg0: f32, %arg1: f32) -> complex<f32> { + %create = complex.create %arg0, %arg1: complex<f32> + %one = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32> + %div = complex.div %create, %one : complex<f32> + // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f32> + // CHECK-NEXT: return %[[CREATE]] + return %div : complex<f32> +} + +// CHECK-LABEL: func @div_one_f64 +// CHECK-SAME: (%[[ARG0:.*]]: f64, %[[ARG1:.*]]: f64) -> complex<f64> +func.func @div_one_f64(%arg0: f64, %arg1: f64) -> complex<f64> { + %create = complex.create %arg0, %arg1: complex<f64> + %one = complex.constant [1.0 : f64, 0.0 : f64] : complex<f64> + %div = complex.div %create, %one : complex<f64> + // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f64> + // CHECK-NEXT: return %[[CREATE]] + return %div : complex<f64> +} + +// CHECK-LABEL: func @div_one_f80 +// CHECK-SAME: (%[[ARG0:.*]]: f80, %[[ARG1:.*]]: f80) -> complex<f80> +func.func @div_one_f80(%arg0: f80, %arg1: f80) -> complex<f80> { + %create = complex.create %arg0, %arg1: complex<f80> + %one = complex.constant [1.0 : f80, 0.0 : f80] : complex<f80> + %div = complex.div %create, %one : complex<f80> + // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f80> + // CHECK-NEXT: return %[[CREATE]] + return %div : complex<f80> +} + +// CHECK-LABEL: func @div_one_f128 +// CHECK-SAME: (%[[ARG0:.*]]: f128, %[[ARG1:.*]]: f128) -> complex<f128> +func.func @div_one_f128(%arg0: f128, %arg1: f128) -> complex<f128> { + %create = complex.create %arg0, %arg1: complex<f128> + %one = complex.constant [1.0 : f128, 0.0 : f128] : complex<f128> + %div = complex.div %create, %one : complex<f128> + // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f128> + // CHECK-NEXT: return %[[CREATE]] + return %div : complex<f128> +} |