[Mlir-commits] [mlir] [mlir][complex] Require fastmath for the add/sub and exp/log folds (PR #221384)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 4 20:25:35 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hung-Kuan Tseng (Tim096)
<details>
<summary>Changes</summary>
`complex.add(complex.sub(a, b), b)`, `complex.add(b, complex.sub(a, b))`
and `complex.sub(complex.add(a, b), b)` fold to `a` unconditionally, and
so does `complex.exp(complex.log(a))`:
```mlir
func.func @<!-- -->add_sub(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
%sub = complex.sub %a, %b : complex<f32>
%add = complex.add %sub, %b : complex<f32>
return %add : complex<f32>
}
// -canonicalize today
func.func @<!-- -->add_sub(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
return %a : complex<f32>
}
```
Neither identity holds in floating point. The intermediate result is
rounded, so with `a = (1.0, 0.0)` and `b = (1.0e30, 0.0)` the original
returns `(0.0, 0.0)`, and `(-0.0 - b) + b` is `+0.0`. `exp(log(a))` is an
approximation of `a` (and `log` of a zero is `-inf` before `exp` takes it
back). LLVM only simplifies `(X - Y) + Y` and `(X + Y) - Y` with
`reassoc` and `nsz` on the outer operation
(`InstructionSimplify.cpp`, `simplifyFAddInst`/`simplifyFSubInst`), and
only rewrites `log(exp(x))` when both calls are fast
(`SimplifyLibCalls.cpp`, `optimizeLog`).
This gates the add/sub folds on `reassoc` and `nsz` of the folded op and
the exp/log fold on `afn` of both ops, mirroring those rules, and adds a
test for each fold without the flags. The existing tests gain the flags.
Same direction as #<!-- -->212751 and #<!-- -->212781, which fixed the neighbouring
`add(a, 0)` and `log(exp(a))` folds.
---
Full diff: https://github.com/llvm/llvm-project/pull/221384.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Complex/IR/ComplexOps.cpp (+24-11)
- (modified) mlir/test/Dialect/Complex/canonicalize.mlir (+50-3)
``````````diff
diff --git a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
index 1d953ca81d751..5f7c009e64585 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
@@ -252,14 +252,19 @@ void ReOp::getCanonicalizationPatterns(RewritePatternSet &results,
OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
// 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();
+ // The intermediate result is rounded, so this needs `reassoc`, and
+ // (-0.0 - b) + b is +0.0, so it also needs `nsz`.
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::reassoc |
+ arith::FastMathFlags::nsz)) {
+ if (auto sub = getLhs().getDefiningOp<SubOp>())
+ if (getRhs() == sub.getRhs())
+ return sub.getLhs();
+
+ if (auto sub = getRhs().getDefiningOp<SubOp>())
+ if (getLhs() == sub.getRhs())
+ return sub.getLhs();
+ }
// complex.add(a, complex.constant<-0.0, -0.0>) -> a
if (auto constantOp = getRhs().getDefiningOp<ConstantOp>()) {
@@ -279,9 +284,13 @@ OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
// complex.sub(complex.add(a, b), b) -> a
- if (auto add = getLhs().getDefiningOp<AddOp>())
- if (getRhs() == add.getRhs())
- return add.getLhs();
+ // The intermediate result is rounded, so this needs `reassoc`, and
+ // (-0.0 + b) - b is +0.0, so it also needs `nsz`.
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::reassoc |
+ arith::FastMathFlags::nsz))
+ if (auto add = getLhs().getDefiningOp<AddOp>())
+ if (getRhs() == add.getRhs())
+ return add.getLhs();
// complex.sub(a, complex.constant<0.0, 0.0>) -> a
if (auto constantOp = getRhs().getDefiningOp<ConstantOp>()) {
@@ -313,8 +322,12 @@ OpFoldResult NegOp::fold(FoldAdaptor adaptor) {
OpFoldResult ExpOp::fold(FoldAdaptor adaptor) {
// complex.exp(complex.log(a)) -> a
+ // exp(log(a)) is only an approximation of a, so both ops need `afn`.
if (auto logOp = getOperand().getDefiningOp<LogOp>())
- return logOp.getOperand();
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::afn) &&
+ arith::bitEnumContainsAll(logOp.getFastmath(),
+ arith::FastMathFlags::afn))
+ return logOp.getOperand();
return {};
}
diff --git a/mlir/test/Dialect/Complex/canonicalize.mlir b/mlir/test/Dialect/Complex/canonicalize.mlir
index f41b0a6c2bc8b..82d2d28fb4a9f 100644
--- a/mlir/test/Dialect/Complex/canonicalize.mlir
+++ b/mlir/test/Dialect/Complex/canonicalize.mlir
@@ -70,7 +70,19 @@ func.func @complex_add_sub_lhs() -> 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>
+ %add = complex.add %sub, %complex2 fastmath<reassoc,nsz> : complex<f32>
+ return %add : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_add_sub_lhs_without_fast_math
+func.func @complex_add_sub_lhs_without_fast_math() -> 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: %[[SUB:.*]] = complex.sub
+ // CHECK: %[[ADD:.*]] = complex.add %[[SUB]], %{{.*}} fastmath<reassoc> : complex<f32>
+ // CHECK-NEXT: return %[[ADD]] : complex<f32>
+ %sub = complex.sub %complex1, %complex2 : complex<f32>
+ %add = complex.add %sub, %complex2 fastmath<reassoc> : complex<f32>
return %add : complex<f32>
}
@@ -81,7 +93,19 @@ func.func @complex_add_sub_rhs() -> 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>
+ %add = complex.add %complex2, %sub fastmath<reassoc,nsz> : complex<f32>
+ return %add : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_add_sub_rhs_without_fast_math
+func.func @complex_add_sub_rhs_without_fast_math() -> 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: %[[SUB:.*]] = complex.sub
+ // CHECK: %[[ADD:.*]] = complex.add %[[SUB]], %{{.*}} fastmath<nsz> : complex<f32>
+ // CHECK-NEXT: return %[[ADD]] : complex<f32>
+ %sub = complex.sub %complex1, %complex2 : complex<f32>
+ %add = complex.add %complex2, %sub fastmath<nsz> : complex<f32>
return %add : complex<f32>
}
@@ -111,8 +135,19 @@ func.func @complex_exp_log() -> complex<f32> {
%complex1 = complex.constant [1.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>
+ %log = complex.log %complex1 fastmath<afn> : complex<f32>
+ %exp = complex.exp %log fastmath<afn> : complex<f32>
+ return %exp : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_exp_log_without_fast_math
+func.func @complex_exp_log_without_fast_math() -> complex<f32> {
+ %complex1 = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
+ // CHECK: %[[LOG:.*]] = complex.log
+ // CHECK: %[[EXP:.*]] = complex.exp %[[LOG]] fastmath<afn> : complex<f32>
+ // CHECK-NEXT: return %[[EXP]] : complex<f32>
%log = complex.log %complex1 : complex<f32>
- %exp = complex.exp %log : complex<f32>
+ %exp = complex.exp %log fastmath<afn> : complex<f32>
return %exp : complex<f32>
}
@@ -143,6 +178,18 @@ func.func @complex_sub_add_lhs() -> complex<f32> {
// CHECK: %[[CPLX:.*]] = complex.constant [1.000000e+00 : f32, 0.000000e+00 : f32] : complex<f32>
// CHECK-NEXT: return %[[CPLX]] : complex<f32>
%add = complex.add %complex1, %complex2 : complex<f32>
+ %sub = complex.sub %add, %complex2 fastmath<reassoc,nsz> : complex<f32>
+ return %sub : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_sub_add_lhs_without_fast_math
+func.func @complex_sub_add_lhs_without_fast_math() -> 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: %[[ADD:.*]] = complex.add
+ // CHECK: %[[SUB:.*]] = complex.sub %[[ADD]], %{{.*}} : complex<f32>
+ // CHECK-NEXT: return %[[SUB]] : complex<f32>
+ %add = complex.add %complex1, %complex2 : complex<f32>
%sub = complex.sub %add, %complex2 : complex<f32>
return %sub : complex<f32>
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/221384
More information about the Mlir-commits
mailing list