[Mlir-commits] [mlir] d424cf2 - [mlir][complex] Require fastmath for the add/sub and exp/log folds (#221384)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 8 05:19:34 PDT 2026
Author: Hung-Kuan Tseng
Date: 2026-10-08T14:19:27+02:00
New Revision: d424cf2a666554a3cc5581b3eadccf0d7b44d34b
URL: https://github.com/llvm/llvm-project/commit/d424cf2a666554a3cc5581b3eadccf0d7b44d34b
DIFF: https://github.com/llvm/llvm-project/commit/d424cf2a666554a3cc5581b3eadccf0d7b44d34b.diff
LOG: [mlir][complex] Require fastmath for the add/sub and exp/log folds (#221384)
`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(a)` has no finite result for a zero or
non-finite `a`, so `exp(log((NaN, 1.0)))` is `(NaN, NaN)` unfolded and
`(NaN, 1.0)` folded.
This gates the add/sub folds on `reassoc` and `nsz` of the folded op and
`reassoc` of the inner op, and the exp/log fold on `reassoc` of the
`exp`
and `reassoc`, `nnan`, `ninf` and `nsz` of the `log`. The LangRef treats
`reassoc` as a rewrite-based flag that every instruction in the
rewritten
expression must carry, so this is stricter than InstructionSimplify,
which only checks the outer instruction for the real
`(X - Y) + Y`, `(X + Y) - Y` and `exp(log(x))`. Each fold gets tests
with the flags missing, and the existing tests gain the flags and take
their operands as arguments.
Same direction as #212751 and #212781, which fixed the neighbouring
`add(a, 0)` and `log(exp(a))` folds.
Added:
Modified:
mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
mlir/test/Dialect/Complex/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
index 1d953ca81d751..65c92d0e7e0a0 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
@@ -252,14 +252,23 @@ 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 both ops need `reassoc`, and
+ // (-0.0 - b) + b is +0.0, so the add also needs `nsz`.
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::reassoc |
+ arith::FastMathFlags::nsz)) {
+ if (auto sub = getLhs().getDefiningOp<SubOp>())
+ if (getRhs() == sub.getRhs() &&
+ arith::bitEnumContainsAll(sub.getFastmath(),
+ arith::FastMathFlags::reassoc))
+ return sub.getLhs();
+
+ if (auto sub = getRhs().getDefiningOp<SubOp>())
+ if (getLhs() == sub.getRhs() &&
+ arith::bitEnumContainsAll(sub.getFastmath(),
+ arith::FastMathFlags::reassoc))
+ return sub.getLhs();
+ }
// complex.add(a, complex.constant<-0.0, -0.0>) -> a
if (auto constantOp = getRhs().getDefiningOp<ConstantOp>()) {
@@ -279,9 +288,15 @@ 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 both ops need `reassoc`, and
+ // (-0.0 + b) - b is +0.0, so the sub also needs `nsz`.
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::reassoc |
+ arith::FastMathFlags::nsz))
+ if (auto add = getLhs().getDefiningOp<AddOp>())
+ if (getRhs() == add.getRhs() &&
+ arith::bitEnumContainsAll(add.getFastmath(),
+ arith::FastMathFlags::reassoc))
+ return add.getLhs();
// complex.sub(a, complex.constant<0.0, 0.0>) -> a
if (auto constantOp = getRhs().getDefiningOp<ConstantOp>()) {
@@ -313,8 +328,16 @@ OpFoldResult NegOp::fold(FoldAdaptor adaptor) {
OpFoldResult ExpOp::fold(FoldAdaptor adaptor) {
// complex.exp(complex.log(a)) -> a
- if (auto logOp = getOperand().getDefiningOp<LogOp>())
- return logOp.getOperand();
+ // exp(log(a)) is only an approximation of a, so both ops need `reassoc`.
+ // log(a) is not finite for a zero or non-finite a, and its result depends
+ // on the sign of a zero, so the log also needs `nnan`, `ninf` and `nsz`.
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::reassoc))
+ if (auto logOp = getOperand().getDefiningOp<LogOp>())
+ if (arith::bitEnumContainsAll(
+ logOp.getFastmath(),
+ arith::FastMathFlags::reassoc | arith::FastMathFlags::nnan |
+ arith::FastMathFlags::ninf | arith::FastMathFlags::nsz))
+ return logOp.getOperand();
return {};
}
diff --git a/mlir/test/Dialect/Complex/canonicalize.mlir b/mlir/test/Dialect/Complex/canonicalize.mlir
index f41b0a6c2bc8b..c487b3ca7f06a 100644
--- a/mlir/test/Dialect/Complex/canonicalize.mlir
+++ b/mlir/test/Dialect/Complex/canonicalize.mlir
@@ -64,24 +64,64 @@ func.func @imag_of_create_op() -> 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>
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_add_sub_lhs(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK-NEXT: return %[[A]] : complex<f32>
+ %sub = complex.sub %a, %b fastmath<reassoc> : complex<f32>
+ %add = complex.add %sub, %b fastmath<reassoc,nsz> : complex<f32>
+ return %add : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_add_sub_lhs_without_fast_math
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_add_sub_lhs_without_fast_math(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK: %[[SUB:.*]] = complex.sub %[[A]], %[[B]] fastmath<reassoc> : complex<f32>
+ // CHECK: %[[ADD:.*]] = complex.add %[[SUB]], %[[B]] fastmath<reassoc> : complex<f32>
+ // CHECK-NEXT: return %[[ADD]] : complex<f32>
+ %sub = complex.sub %a, %b fastmath<reassoc> : complex<f32>
+ %add = complex.add %sub, %b fastmath<reassoc> : complex<f32>
+ return %add : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_add_sub_lhs_without_inner_reassoc
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_add_sub_lhs_without_inner_reassoc(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK: %[[SUB:.*]] = complex.sub %[[A]], %[[B]] : complex<f32>
+ // CHECK: %[[ADD:.*]] = complex.add %[[SUB]], %[[B]] fastmath<reassoc,nsz> : complex<f32>
+ // CHECK-NEXT: return %[[ADD]] : complex<f32>
+ %sub = complex.sub %a, %b : complex<f32>
+ %add = complex.add %sub, %b fastmath<reassoc,nsz> : 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>
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_add_sub_rhs(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK-NEXT: return %[[A]] : complex<f32>
+ %sub = complex.sub %a, %b fastmath<reassoc> : complex<f32>
+ %add = complex.add %b, %sub fastmath<reassoc,nsz> : complex<f32>
+ return %add : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_add_sub_rhs_without_fast_math
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_add_sub_rhs_without_fast_math(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK: %[[SUB:.*]] = complex.sub %[[A]], %[[B]] fastmath<reassoc> : complex<f32>
+ // CHECK: %[[ADD:.*]] = complex.add %[[B]], %[[SUB]] fastmath<nsz> : complex<f32>
+ // CHECK-NEXT: return %[[ADD]] : complex<f32>
+ %sub = complex.sub %a, %b fastmath<reassoc> : complex<f32>
+ %add = complex.add %b, %sub fastmath<nsz> : complex<f32>
+ return %add : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_add_sub_rhs_without_inner_reassoc
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_add_sub_rhs_without_inner_reassoc(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK: %[[SUB:.*]] = complex.sub %[[A]], %[[B]] : complex<f32>
+ // CHECK: %[[ADD:.*]] = complex.add %[[B]], %[[SUB]] fastmath<reassoc,nsz> : complex<f32>
+ // CHECK-NEXT: return %[[ADD]] : complex<f32>
+ %sub = complex.sub %a, %b : complex<f32>
+ %add = complex.add %b, %sub fastmath<reassoc,nsz> : complex<f32>
return %add : complex<f32>
}
@@ -107,15 +147,40 @@ func.func @complex_log_exp() -> complex<f32> {
}
// CHECK-LABEL: func @complex_exp_log
-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 : complex<f32>
+// CHECK-SAME: (%[[A:.*]]: complex<f32>)
+func.func @complex_exp_log(%a: complex<f32>) -> complex<f32> {
+ // CHECK-NEXT: return %[[A]] : complex<f32>
+ %log = complex.log %a fastmath<reassoc,nnan,ninf,nsz> : complex<f32>
+ %exp = complex.exp %log fastmath<reassoc> : complex<f32>
+ return %exp : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_exp_log_without_fast_math
+// CHECK-SAME: (%[[A:.*]]: complex<f32>)
+func.func @complex_exp_log_without_fast_math(%a: complex<f32>) -> complex<f32> {
+ // CHECK: %[[LOG:.*]] = complex.log %[[A]] fastmath<reassoc,nnan,ninf,nsz> : complex<f32>
+ // CHECK: %[[EXP:.*]] = complex.exp %[[LOG]] : complex<f32>
+ // CHECK-NEXT: return %[[EXP]] : complex<f32>
+ %log = complex.log %a fastmath<reassoc,nnan,ninf,nsz> : complex<f32>
%exp = complex.exp %log : complex<f32>
return %exp : complex<f32>
}
+// Each log is missing one of the flags the fold needs.
+// CHECK-LABEL: func @complex_exp_log_without_log_flags
+// CHECK-COUNT-4: complex.exp %{{.*}} fastmath<reassoc> : complex<f32>
+func.func @complex_exp_log_without_log_flags(%a: complex<f32>) -> (complex<f32>, complex<f32>, complex<f32>, complex<f32>) {
+ %log0 = complex.log %a fastmath<nnan,ninf,nsz> : complex<f32>
+ %exp0 = complex.exp %log0 fastmath<reassoc> : complex<f32>
+ %log1 = complex.log %a fastmath<reassoc,ninf,nsz> : complex<f32>
+ %exp1 = complex.exp %log1 fastmath<reassoc> : complex<f32>
+ %log2 = complex.log %a fastmath<reassoc,nnan,nsz> : complex<f32>
+ %exp2 = complex.exp %log2 fastmath<reassoc> : complex<f32>
+ %log3 = complex.log %a fastmath<reassoc,nnan,ninf> : complex<f32>
+ %exp3 = complex.exp %log3 fastmath<reassoc> : complex<f32>
+ return %exp0, %exp1, %exp2, %exp3 : complex<f32>, complex<f32>, complex<f32>, complex<f32>
+}
+
// CHECK-LABEL: func @complex_conj_conj
func.func @complex_conj_conj() -> complex<f32> {
%complex1 = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
@@ -137,13 +202,33 @@ func.func @complex_add_neg_zero() -> complex<f32> {
}
// CHECK-LABEL: func @complex_sub_add_lhs
-func.func @complex_sub_add_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>
- %add = complex.add %complex1, %complex2 : complex<f32>
- %sub = complex.sub %add, %complex2 : complex<f32>
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_sub_add_lhs(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK-NEXT: return %[[A]] : complex<f32>
+ %add = complex.add %a, %b fastmath<reassoc> : complex<f32>
+ %sub = complex.sub %add, %b fastmath<reassoc,nsz> : complex<f32>
+ return %sub : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_sub_add_lhs_without_fast_math
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_sub_add_lhs_without_fast_math(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK: %[[ADD:.*]] = complex.add %[[A]], %[[B]] fastmath<reassoc> : complex<f32>
+ // CHECK: %[[SUB:.*]] = complex.sub %[[ADD]], %[[B]] : complex<f32>
+ // CHECK-NEXT: return %[[SUB]] : complex<f32>
+ %add = complex.add %a, %b fastmath<reassoc> : complex<f32>
+ %sub = complex.sub %add, %b : complex<f32>
+ return %sub : complex<f32>
+}
+
+// CHECK-LABEL: func @complex_sub_add_lhs_without_inner_reassoc
+// CHECK-SAME: (%[[A:.*]]: complex<f32>, %[[B:.*]]: complex<f32>)
+func.func @complex_sub_add_lhs_without_inner_reassoc(%a: complex<f32>, %b: complex<f32>) -> complex<f32> {
+ // CHECK: %[[ADD:.*]] = complex.add %[[A]], %[[B]] : complex<f32>
+ // CHECK: %[[SUB:.*]] = complex.sub %[[ADD]], %[[B]] fastmath<reassoc,nsz> : complex<f32>
+ // CHECK-NEXT: return %[[SUB]] : complex<f32>
+ %add = complex.add %a, %b : complex<f32>
+ %sub = complex.sub %add, %b fastmath<reassoc,nsz> : complex<f32>
return %sub : complex<f32>
}
More information about the Mlir-commits
mailing list