[Mlir-commits] [mlir] [mlir][math] Fix `math.powf(a, b)` expansion for negative base (PR #193231)
Zmicier Prybysh
llvmlistbot at llvm.org
Tue Apr 21 07:27:48 PDT 2026
https://github.com/dimp-pl created https://github.com/llvm/llvm-project/pull/193231
Fixes #132202.
This PR is a small fix for an incorrectly applied rewriting of `powf` mathematical operation into a combination of `exp` and `log`.
The `-math-expand-ops` pass for `math.powf(a, b)` was rewriting the operation using the `exp(b * log(a))` formula unconditionally. That identity breaks for `a < 0` ( `log(a) = NaN` ), so the expansion returned `NaN` even when `b` is an integer, differring from `libm`'s `pow`, which would return a correct `non-NaN`) value.
With this fix, `convertPowfOp` now behaves as following:
1. The pre-existing algebraic shortcuts for constant exponents in `{0, +/-1, +/-0.5, +/-2, 3}` are considered first.
2. If the base is non-negative (constant \>= 0, or produced by `math.absf/sqrt/rsqrt/exp/exp2`), the cheap `exp(b * log(a))` form is used.
3. Otherwise, a generic branching form is used:
```
magnitude = exp(b * log|a|)
signedMag = bIsEven ? magnitude : -magnitude
negCase = bIsInt ? signedMag : NaN
result = (a < 0) ? negCase : magnitude
```
This implements `pow(a, b) = a^b` for `a >= 0`, `(+|a|^b)` for `a<0` with `b` an even integer, `(-|a|^b)` for `a<0` with `b` an odd integer, and `NaN` for `a<0` with `b` non-integer. Parity is tested via `floor(b/2)`.
Granted, there is a performance penalty associated with this change, as the generic version is slower than the current one. But at least it's correct, so I think that we will be able to convince maintainers that this should be merged.
>From dd5873a388cfdd9b609b092b918335eeae5a24ad Mon Sep 17 00:00:00 2001
From: Zmicier Prybysh <zprybysh at baylibre.com>
Date: Mon, 20 Apr 2026 12:15:26 +0200
Subject: [PATCH] [mlir][math] Fix `math.powf(a, b)` expansion for negative
base
Fixes #132202.
The -math-expand-ops for math.powf(a, b) was emitting
exp(b * log(a)) unconditionally. That identity breaks for a < 0
(log(a) = NaN), so the expansion returned NaN even when b is an
integer, differring from libm's `pow`, which would return a
correct (non-NaN) value.
convertPowfOp now behaves as following:
1. The pre-existing algebraic shortcuts for constant exponents
in {0, +/-1, +/-0.5, +/-2, 3} are considered first.
2. If the base is non-negative (constant >= 0, or produced by
math.absf/sqrt/rsqrt/exp/exp2), the cheap exp(b * log(a)) form
is used.
3. Otherwise, a generic branching form is used:
magnitude = exp(b * log|a|)
signedMag = bIsEven ? magnitude : -magnitude
negCase = bIsInt ? signedMag : NaN
result = (a < 0) ? negCase : magnitude
implementing pow(a, b) = a^b for a >= 0, (+|a|^b) for a<0
with b an even integer, (-|a|^b) for a<0 with b an odd
integer, and NaN for a<0 with b non-integer. Parity is
tested via floor(b/2) exactness.
---
.../lib/Dialect/Math/Transforms/ExpandOps.cpp | 77 +++++++++++++++--
mlir/test/Dialect/Math/expand-math.mlir | 85 ++++++++++++++++---
.../mlir-runner/test-expand-math-approx.mlir | 40 +++++++++
3 files changed, 184 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
index f76ddfae2a67a..7b889e54f829e 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
@@ -367,9 +367,24 @@ static LogicalResult convertFPowIOp(math::FPowIOp op,
return success();
}
-// Converts Powf(float a, float b) (meaning a^b) to exp^(b * ln(a))
-// Some special cases where b is constant are handled separately:
-// when b == 0, or |b| == 0.5, 1.0, or 2.0.
+static bool isKnownNonNegativeFloat(Value v) {
+ APFloat cst(0.0);
+ if (matchPattern(v, m_ConstantFloat(&cst)))
+ return !cst.isNegative() || cst.isZero();
+ Operation *defOp = v.getDefiningOp();
+ if (!defOp)
+ return false;
+ return isa<math::AbsFOp, math::SqrtOp, math::RsqrtOp, math::ExpOp,
+ math::Exp2Op>(defOp);
+}
+
+// Converts math.powf(a, b) (meaning a^b) using:
+// 1. Algebraic shortcuts when b is the constant 0, 1, -1, 0.5, -0.5, 2, -2, 3
+// 2. For a non-negative base, the classical identity a^b = exp(b * log(a))
+// 3. Otherwise, an IEEE-correct branching form (documented below) that
+// computes |a|^b (using the classical identity) and fixes the sign for
+// negative bases with integer exponents (producing NaN for negative bases
+// with non-integer exponents)
static LogicalResult convertPowfOp(math::PowFOp op, PatternRewriter &rewriter) {
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
Value operandA = op.getOperand(0);
@@ -433,10 +448,58 @@ static LogicalResult convertPowfOp(math::PowFOp op, PatternRewriter &rewriter) {
}
}
- Value logA = math::LogOp::create(b, operandA);
- Value mult = arith::MulFOp::create(b, operandB, logA);
- Value expResult = math::ExpOp::create(b, mult);
- rewriter.replaceOp(op, expResult);
+ if (isKnownNonNegativeFloat(operandA)) {
+ Value logA = math::LogOp::create(b, operandA);
+ Value mult = arith::MulFOp::create(b, operandB, logA);
+ rewriter.replaceOp(op, math::ExpOp::create(b, mult));
+ return success();
+ }
+
+ // Implements a^b for any base that may be negative by computing the
+ // magnitude through |a| (so log is always defined) and then considering
+ // the sign and NaN cases:
+ //
+ // | a^b if a >= 0 (== |a|^b)
+ // pow = | + |a|^b if a < 0 and b is an even integer
+ // | - |a|^b if a < 0 and b is an odd integer
+ // | NaN if a < 0 and b is not an integer
+ //
+ // Sequence of operations:
+ // magnitude = exp(b * log|a|) // == |a|^b
+ // bIsInt = (b == floor(b))
+ // bIsEven = (floor(b)/2 == floor(floor(b)/2)) // only makes sense if bIsInt
+ // signedMag = bIsEven ? magnitude : -magnitude
+ // negCase = bIsInt ? signedMag : NaN
+ // result = (a < 0) ? negCase : magnitude
+
+ Value absA = math::AbsFOp::create(b, operandA);
+ Value logAbsA = math::LogOp::create(b, absA);
+ Value mult = arith::MulFOp::create(b, operandB, logAbsA);
+ Value magnitude = math::ExpOp::create(b, mult);
+
+ // Integer / parity tests on b
+ Value floorB = math::FloorOp::create(b, operandB);
+ Value bIsInt =
+ arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, operandB, floorB);
+ Value half = createFloatConst(op->getLoc(), typeB, 0.5, rewriter);
+ Value halfB = arith::MulFOp::create(b, floorB, half);
+ Value halfBFloor = math::FloorOp::create(b, halfB);
+ Value bIsEven =
+ arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, halfB, halfBFloor);
+
+ // Sign fix for a < 0: odd integer -> flip sign, non-integer -> NaN
+ Value negMagnitude = arith::NegFOp::create(b, magnitude);
+ Value signedMag =
+ arith::SelectOp::create(b, bIsEven, magnitude, negMagnitude);
+ Value nan =
+ createFloatConst(op->getLoc(), typeA, APFloat::getQNaN(sem), rewriter);
+ Value negCase = arith::SelectOp::create(b, bIsInt, signedMag, nan);
+
+ // Dispatch on the sign of a
+ Value zero = createFloatConst(op->getLoc(), typeA, 0.0, rewriter);
+ Value aNeg =
+ arith::CmpFOp::create(b, arith::CmpFPredicate::OLT, operandA, zero);
+ rewriter.replaceOp(op, arith::SelectOp::create(b, aNeg, negCase, magnitude));
return success();
}
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 126270ca40130..212a794b1f2be 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -257,10 +257,67 @@ func.func @roundf_func(%a: f32) -> f32 {
// CHECK-LABEL: func @powf_func
// CHECK-SAME: (%[[ARG0:.+]]: f64, %[[ARG1:.+]]: f64) -> f64
func.func @powf_func(%a: f64, %b: f64) -> f64 {
- // CHECK: %[[LOGA:.+]] = math.log %[[ARG0]] : f64
- // CHECK: %[[MUL:.+]] = arith.mulf %[[ARG1]], %[[LOGA]] : f64
+ // CHECK: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f64
+ // CHECK: %[[NAN:.+]] = arith.constant 0x7FF8000000000000 : f64
+ // CHECK: %[[HALF:.+]] = arith.constant 5.000000e-01 : f64
+ // CHECK: %[[ABSA:.+]] = math.absf %[[ARG0]] : f64
+ // CHECK: %[[LOGABS:.+]] = math.log %[[ABSA]] : f64
+ // CHECK: %[[MUL:.+]] = arith.mulf %[[ARG1]], %[[LOGABS]] : f64
+ // CHECK: %[[MAG:.+]] = math.exp %[[MUL]] : f64
+ // CHECK: %[[FLOORB:.+]] = math.floor %[[ARG1]] : f64
+ // CHECK: %[[ISINT:.+]] = arith.cmpf oeq, %[[ARG1]], %[[FLOORB]] : f64
+ // CHECK: %[[HALFB:.+]] = arith.mulf %[[FLOORB]], %[[HALF]] : f64
+ // CHECK: %[[HALFBFLOOR:.+]] = math.floor %[[HALFB]] : f64
+ // CHECK: %[[ISEVEN:.+]] = arith.cmpf oeq, %[[HALFB]], %[[HALFBFLOOR]] : f64
+ // CHECK: %[[NEGMAG:.+]] = arith.negf %[[MAG]] : f64
+ // CHECK: %[[SIGNED:.+]] = arith.select %[[ISEVEN]], %[[MAG]], %[[NEGMAG]] : f64
+ // CHECK: %[[NEGCASE:.+]] = arith.select %[[ISINT]], %[[SIGNED]], %[[NAN]] : f64
+ // CHECK: %[[ANEG:.+]] = arith.cmpf olt, %[[ARG0]], %[[ZERO]] : f64
+ // CHECK: %[[RES:.+]] = arith.select %[[ANEG]], %[[NEGCASE]], %[[MAG]] : f64
+ // CHECK: return %[[RES]] : f64
+ %ret = math.powf %a, %b : f64
+ return %ret : f64
+}
+
+// CHECK-LABEL: func @powf_func_nonneg_base
+// CHECK-SAME: (%[[ARG0:.+]]: f64, %[[ARG1:.+]]: f64) -> f64
+func.func @powf_func_nonneg_base(%a: f64, %b: f64) -> f64 {
+ // CHECK: %[[ABS:.+]] = math.absf %[[ARG0]] : f64
+ // CHECK: %[[LOG:.+]] = math.log %[[ABS]] : f64
+ // CHECK: %[[MUL:.+]] = arith.mulf %[[ARG1]], %[[LOG]] : f64
+ // CHECK: %[[EXP:.+]] = math.exp %[[MUL]] : f64
+ // CHECK-NOT: arith.select
+ // CHECK: return %[[EXP]] : f64
+ %abs = math.absf %a : f64
+ %ret = math.powf %abs, %b : f64
+ return %ret : f64
+}
+
+// CHECK-LABEL: func @powf_func_const_nonneg_base
+// CHECK-SAME: (%[[ARG0:.+]]: f64)
+func.func @powf_func_const_nonneg_base(%b: f64) -> f64 {
+ // log(const) is gone because of constant folding, only mul + exp are left.
+ // CHECK: %[[LOGCST:.+]] = arith.constant {{.*}} : f64
+ // CHECK: %[[MUL:.+]] = arith.mulf %[[ARG0]], %[[LOGCST]] : f64
// CHECK: %[[EXP:.+]] = math.exp %[[MUL]] : f64
+ // CHECK-NOT: arith.select
+ // CHECK-NOT: math.absf
// CHECK: return %[[EXP]] : f64
+ %a = arith.constant 3.0 : f64
+ %ret = math.powf %a, %b : f64
+ return %ret : f64
+}
+
+// CHECK-LABEL: func @powf_func_const_neg_base
+func.func @powf_func_const_neg_base(%b: f64) -> f64 {
+ // Negative constant forces the general branching form.
+ // After constant folding, absf and the outer aNeg select should vanish,
+ // but the sign-fixing negf + select should stay.
+ // CHECK-NOT: math.absf
+ // CHECK-NOT: arith.cmpf olt
+ // CHECK: arith.negf
+ // CHECK: arith.select
+ %a = arith.constant -3.0 : f64
%ret = math.powf %a, %b : f64
return %ret : f64
}
@@ -728,10 +785,13 @@ func.func @math_fpowi_to_powf_tensor(%0 : tensor<8xf32>, %1: tensor<8xi32>) -> t
}
// CHECK-SAME: (%[[ARG0:.*]]: tensor<8xf32>, %[[ARG1:.*]]: tensor<8xi32>) -> tensor<8xf32> {
// CHECK: %[[TOFP:.*]] = arith.sitofp %[[ARG1]] : tensor<8xi32> to tensor<8xf32>
-// CHECK: %[[LOGA:.*]] = math.log %[[ARG0]] : tensor<8xf32>
-// CHECK: %[[MUL:.*]] = arith.mulf %[[TOFP]], %[[LOGA]] : tensor<8xf32>
-// CHECK: %[[EXP:.*]] = math.exp %[[MUL]] : tensor<8xf32>
-// CHECK: return %[[EXP]]
+// CHECK: %[[ABS:.*]] = math.absf %[[ARG0]] : tensor<8xf32>
+// CHECK: %[[LOGABS:.*]] = math.log %[[ABS]] : tensor<8xf32>
+// CHECK: %[[MUL:.*]] = arith.mulf %[[TOFP]], %[[LOGABS]] : tensor<8xf32>
+// CHECK: %[[MAG:.*]] = math.exp %[[MUL]] : tensor<8xf32>
+// CHECK: %[[ANEG:.*]] = arith.cmpf olt, %[[ARG0]]
+// CHECK: %[[RES:.*]] = arith.select %[[ANEG]]
+// CHECK: return %[[RES]]
// -----
// CHECK-LABEL: func.func @math_fpowi_to_powf_scalar
@@ -740,11 +800,14 @@ func.func @math_fpowi_to_powf_scalar(%0 : f32, %1: i64) -> f32 {
return %2 : f32
}
// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: i64) -> f32 {
-// CHECK: %[[TOFP:.*]] = arith.sitofp %[[ARG1]] : i64 to f32
-// CHECK: %[[LOGA:.*]] = math.log %[[ARG0]] : f32
-// CHECK: %[[MUL:.*]] = arith.mulf %[[TOFP]], %[[LOGA]] : f32
-// CHECK: %[[EXP:.*]] = math.exp %[[MUL]] : f32
-// CHECK: return %[[EXP]] : f32
+// CHECK: %[[TOFP:.*]] = arith.sitofp %[[ARG1]] : i64 to f32
+// CHECK: %[[ABS:.*]] = math.absf %[[ARG0]] : f32
+// CHECK: %[[LOGABS:.*]] = math.log %[[ABS]] : f32
+// CHECK: %[[MUL:.*]] = arith.mulf %[[TOFP]], %[[LOGABS]] : f32
+// CHECK: %[[MAG:.*]] = math.exp %[[MUL]] : f32
+// CHECK: %[[ANEG:.*]] = arith.cmpf olt, %[[ARG0]]
+// CHECK: %[[RES:.*]] = arith.select %[[ANEG]]
+// CHECK: return %[[RES]] : f32
// -----
diff --git a/mlir/test/mlir-runner/test-expand-math-approx.mlir b/mlir/test/mlir-runner/test-expand-math-approx.mlir
index 06b3171a2349e..51d48d72076be 100644
--- a/mlir/test/mlir-runner/test-expand-math-approx.mlir
+++ b/mlir/test/mlir-runner/test-expand-math-approx.mlir
@@ -259,6 +259,46 @@ func.func @powf() {
%l_p = arith.constant -2.0 : f32
%l_r = math.powf %k, %l_p : f32
vector.print %l_r : f32
+
+ // @func_powff32 is used so that both operands are SSA arguments
+ // — to force a branching fallback way of computing `powf`
+
+ // CHECK-NEXT: -0.125
+ %m = arith.constant -2.0 : f32
+ %m_p = arith.constant -3.0 : f32
+ call @func_powff32(%m, %m_p) : (f32, f32) -> ()
+
+ // CHECK-NEXT: 16
+ %n = arith.constant -2.0 : f32
+ %n_p = arith.constant 4.0 : f32
+ call @func_powff32(%n, %n_p) : (f32, f32) -> ()
+
+ // CHECK-NEXT: nan
+ %o = arith.constant -2.0 : f32
+ %o_p = arith.constant 3.5 : f32
+ call @func_powff32(%o, %o_p) : (f32, f32) -> ()
+
+ // CHECK-NEXT: -0
+ %p = arith.constant -94.0 : f32
+ %p_p = arith.constant -47.0 : f32
+ call @func_powff32(%p, %p_p) : (f32, f32) -> ()
+
+ // CHECK-NEXT: 0.125
+ %q = arith.constant -2.0 : f32
+ %q_abs = math.absf %q : f32
+ %q_p = arith.constant -3.0 : f32
+ %q_r = math.powf %q_abs, %q_p : f32
+ vector.print %q_r : f32
+
+ // CHECK-NEXT: 0
+ %r = arith.constant 0.0 : f32
+ %r_p = arith.constant 5.0 : f32
+ call @func_powff32(%r, %r_p) : (f32, f32) -> ()
+
+ // CHECK-NEXT: inf
+ %s = arith.constant 0.0 : f32
+ %s_p = arith.constant -5.0 : f32
+ call @func_powff32(%s, %s_p) : (f32, f32) -> ()
return
}
More information about the Mlir-commits
mailing list