[Mlir-commits] [mlir] [mlir][complex] Require fastmath for the add/sub and exp/log folds (PR #221384)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Fri Sep 11 07:36:16 PDT 2026
https://github.com/Tim096 updated https://github.com/llvm/llvm-project/pull/221384
>From 378e96fe4dfd7ebb1c7ccac632def90bcc082e01 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 11:24:45 +0800
Subject: [PATCH 1/2] [mlir][complex] Require fastmath for the add/sub and
exp/log folds
`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))`. Neither identity holds in
floating point: the intermediate result is rounded, so
(1.0 - 1.0e30) + 1.0e30 is 0.0, and (-0.0 - b) + b is +0.0. LLVM only
simplifies `(X - Y) + Y` and `(X + Y) - Y` with `reassoc` and `nsz` on
the outer operation, and only rewrites `log(exp(x))` when both calls
are fast.
Gate the add/sub folds on `reassoc` and `nsz` of the folded op and the
exp/log fold on `afn` of both ops, and add tests without the flags.
---
mlir/lib/Dialect/Complex/IR/ComplexOps.cpp | 35 +++++++++-----
mlir/test/Dialect/Complex/canonicalize.mlir | 53 +++++++++++++++++++--
2 files changed, 74 insertions(+), 14 deletions(-)
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>
}
>From eb99a882d843816888cb5228f63406dd1cfe62e4 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Fri, 11 Sep 2026 22:36:02 +0800
Subject: [PATCH 2/2] Use reassoc for the exp/log fold and take the test
operands as arguments
---
mlir/lib/Dialect/Complex/IR/ComplexOps.cpp | 9 +-
mlir/test/Dialect/Complex/canonicalize.mlir | 98 +++++++++------------
2 files changed, 48 insertions(+), 59 deletions(-)
diff --git a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
index 5f7c009e64585..e99de97ae105b 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
@@ -322,11 +322,10 @@ 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>())
- if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::afn) &&
- arith::bitEnumContainsAll(logOp.getFastmath(),
- arith::FastMathFlags::afn))
+ // exp(log(a)) is only an approximation of a, and log(a) has no finite
+ // result for a zero or non-finite a, so this needs `reassoc`.
+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::reassoc))
+ if (auto logOp = getOperand().getDefiningOp<LogOp>())
return logOp.getOperand();
return {};
diff --git a/mlir/test/Dialect/Complex/canonicalize.mlir b/mlir/test/Dialect/Complex/canonicalize.mlir
index 82d2d28fb4a9f..28fcf3a62e87d 100644
--- a/mlir/test/Dialect/Complex/canonicalize.mlir
+++ b/mlir/test/Dialect/Complex/canonicalize.mlir
@@ -64,48 +64,42 @@ 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 fastmath<reassoc,nsz> : 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 : 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
-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-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]] : complex<f32>
+ // CHECK: %[[ADD:.*]] = complex.add %[[SUB]], %[[B]] 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>
+ %sub = complex.sub %a, %b : complex<f32>
+ %add = complex.add %sub, %b fastmath<reassoc> : 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 fastmath<reassoc,nsz> : 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 : 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
-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-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]] : complex<f32>
+ // CHECK: %[[ADD:.*]] = complex.add %[[B]], %[[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>
+ %sub = complex.sub %a, %b : complex<f32>
+ %add = complex.add %b, %sub fastmath<nsz> : complex<f32>
return %add : complex<f32>
}
@@ -131,23 +125,22 @@ 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 fastmath<afn> : complex<f32>
- %exp = complex.exp %log fastmath<afn> : 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 : complex<f32>
+ %exp = complex.exp %log fastmath<reassoc> : 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-SAME: (%[[A:.*]]: complex<f32>)
+func.func @complex_exp_log_without_fast_math(%a: complex<f32>) -> complex<f32> {
+ // CHECK: %[[LOG:.*]] = complex.log %[[A]] fastmath<reassoc> : complex<f32>
+ // CHECK: %[[EXP:.*]] = complex.exp %[[LOG]] : complex<f32>
// CHECK-NEXT: return %[[EXP]] : complex<f32>
- %log = complex.log %complex1 : complex<f32>
- %exp = complex.exp %log fastmath<afn> : complex<f32>
+ %log = complex.log %a fastmath<reassoc> : complex<f32>
+ %exp = complex.exp %log : complex<f32>
return %exp : complex<f32>
}
@@ -172,25 +165,22 @@ 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 fastmath<reassoc,nsz> : 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 : 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
-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-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]] : complex<f32>
+ // CHECK: %[[SUB:.*]] = complex.sub %[[ADD]], %[[B]] : complex<f32>
// CHECK-NEXT: return %[[SUB]] : complex<f32>
- %add = complex.add %complex1, %complex2 : complex<f32>
- %sub = complex.sub %add, %complex2 : complex<f32>
+ %add = complex.add %a, %b : complex<f32>
+ %sub = complex.sub %add, %b : complex<f32>
return %sub : complex<f32>
}
More information about the Mlir-commits
mailing list