[Mlir-commits] [mlir] [mlir][math] Fold exp(a) / exp(b) into exp(a - b) (PR #220010)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 31 08:44:23 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Md Abdullah Shahneous Bari (mshahneo)
<details>
<summary>Changes</summary>
Add ExpQuotientStrengthReduction to the Math algebraic simplification patterns, rewriting `exp(a) / exp(b)` into `exp(a - b)` (and likewise for `exp2`). This trades a division and an exponential for a subtraction, and shows up in online-softmax kernels, where the running-max correction term is computed as a quotient of two exponentials.
The rewrite is not value-safe, so it is gated on `arcp` and `reassoc` on the division: it is the fused form of the two steps LLVM's InstCombine uses to reach the same result, `Z / exp(Y) --> Z * exp(-Y)` (needs `arcp`) followed by `exp(X) * exp(Y) --> exp(X + Y)` (needs `reassoc`). Besides rounding, it also changes the intermediate overflow behaviour: for a large `a == b` the original expression is `inf / inf`, i.e. NaN, while the folded one is `exp(0.0)`, i.e. 1.0.
Because a new exponential is created, the pattern additionally requires that at least one of the two exponentials it feeds on dies with the division, so the number of exponentials never grows. This is the fused equivalent of the `isOnlyUserOfAnyOperand()` check InstCombine applies to `exp(X) * exp(Y)`; it is marginally stronger than InstCombine, which requires the divisor specifically to have a single use.
The new `exp` only carries the fast-math flags common to both exponentials it replaces, so its accuracy contract cannot be loosened by flags that only the division had.
The pattern is rooted on `arith.divf` rather than being an `arith` canonicalization, since `arith` cannot depend on `math`.
---
Full diff: https://github.com/llvm/llvm-project/pull/220010.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp (+62)
- (modified) mlir/test/Dialect/Math/algebraic-simplification.mlir (+111)
``````````diff
diff --git a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
index 7ce7b2153ffb4..9993e30d35d09 100644
--- a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
@@ -242,6 +242,65 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
return success();
}
+//----------------------------------------------------------------------------//
+// ExpOp/Exp2Op quotient strength reduction.
+//----------------------------------------------------------------------------//
+
+namespace {
+/// Replaces `exp(a) / exp(b)` with `exp(a - b)`, and likewise for `exp2`,
+/// trading a division and an exponential for a subtraction.
+template <typename ExpOpTy>
+struct ExpQuotientStrengthReduction : public OpRewritePattern<arith::DivFOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(arith::DivFOp op,
+ PatternRewriter &rewriter) const final {
+ auto numerator = op.getLhs().getDefiningOp<ExpOpTy>();
+ auto denominator = op.getRhs().getDefiningOp<ExpOpTy>();
+ if (!numerator || !denominator)
+ return failure();
+
+ // The rewrite is only valid when the division may be turned into a
+ // reciprocal multiplication and then reassociated with the exponentials:
+ // exp(a) / exp(b) --> exp(a) * exp(-b) --> exp(a + -b)
+ // This mirrors LLVM's InstCombine, which reaches the same result with
+ // `arcp` for the first step and `reassoc` for the second. Note that the
+ // rewrite also changes the overflow behaviour: for a large `a == b` the
+ // original expression is `inf / inf`, i.e. NaN, while the folded one is
+ // `exp(0.0)`, i.e. 1.0.
+ arith::FastMathFlags fmf = op.getFastmath();
+ if (!bitEnumContainsAll(fmf, arith::FastMathFlags::arcp |
+ arith::FastMathFlags::reassoc))
+ return failure();
+
+ // The rewrite introduces a new exponential, so it is only profitable if at
+ // least one of the two it feeds on dies with the division; the exponential
+ // count then never grows while a division is traded for a subtraction.
+ // This is the fused equivalent of the `isOnlyUserOfAnyOperand()` check
+ // LLVM's InstCombine applies to `exp(X) * exp(Y) --> exp(X + Y)`.
+ Operation *divOp = op;
+ auto diesWithDivision = [divOp](Operation *exp) {
+ return llvm::all_of(exp->getUsers(),
+ [divOp](Operation *user) { return user == divOp; });
+ };
+ if (!diesWithDivision(numerator) && !diesWithDivision(denominator))
+ return failure();
+
+ // Do not loosen the accuracy contract of the exponentials being replaced:
+ // the new one may only carry the flags that both of them carry.
+ arith::FastMathFlags expFmf =
+ numerator.getFastmath() & denominator.getFastmath();
+
+ Value exponent =
+ arith::SubFOp::create(rewriter, op.getLoc(), numerator.getOperand(),
+ denominator.getOperand(), fmf);
+ rewriter.replaceOpWithNewOp<ExpOpTy>(op, exponent, expFmf);
+ return success();
+ }
+};
+} // namespace
+
//----------------------------------------------------------------------------//
void mlir::populateMathAlgebraicSimplificationPatterns(
@@ -252,4 +311,7 @@ void mlir::populateMathAlgebraicSimplificationPatterns(
PowIStrengthReduction<math::FPowIOp, arith::DivFOp, arith::MulFOp>,
PowIStrengthReduction<complex::PowiOp, complex::DivOp, complex::MulOp>>(
patterns.getContext(), /*exponentThreshold=*/8);
+ patterns.add<ExpQuotientStrengthReduction<math::ExpOp>,
+ ExpQuotientStrengthReduction<math::Exp2Op>>(
+ patterns.getContext());
}
diff --git a/mlir/test/Dialect/Math/algebraic-simplification.mlir b/mlir/test/Dialect/Math/algebraic-simplification.mlir
index 7342600748967..2b712fd42a7cd 100644
--- a/mlir/test/Dialect/Math/algebraic-simplification.mlir
+++ b/mlir/test/Dialect/Math/algebraic-simplification.mlir
@@ -349,3 +349,114 @@ func.func @fpowi_exp_three(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf
%3 = math.fpowi %arg1, %vm1 : vector<4xf32>, vector<4xi32>
return %0, %1, %2, %3 : f32, vector<4xf32>, f32, vector<4xf32>
}
+
+// CHECK-LABEL: @exp_quotient(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32,
+// CHECK-SAME: %[[ARG2:.+]]: vector<4xf32>, %[[ARG3:.+]]: vector<4xf32>
+func.func @exp_quotient(%arg0: f32, %arg1: f32, %arg2: vector<4xf32>,
+ %arg3: vector<4xf32>) -> (f32, vector<4xf32>) {
+ // CHECK: %[[SSUB:.*]] = arith.subf %[[ARG0]], %[[ARG1]] fastmath<fast> : f32
+ // CHECK: %[[SCALAR:.*]] = math.exp %[[SSUB]] fastmath<fast> : f32
+ // CHECK: %[[VSUB:.*]] = arith.subf %[[ARG2]], %[[ARG3]] fastmath<fast> : vector<4xf32>
+ // CHECK: %[[VECTOR:.*]] = math.exp2 %[[VSUB]] fastmath<fast> : vector<4xf32>
+ // CHECK: return %[[SCALAR]], %[[VECTOR]]
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = math.exp %arg1 fastmath<fast> : f32
+ %2 = arith.divf %0, %1 fastmath<fast> : f32
+ %3 = math.exp2 %arg2 fastmath<fast> : vector<4xf32>
+ %4 = math.exp2 %arg3 fastmath<fast> : vector<4xf32>
+ %5 = arith.divf %3, %4 fastmath<fast> : vector<4xf32>
+ return %2, %5 : f32, vector<4xf32>
+}
+
+// The rewrite only needs `arcp` and `reassoc` on the division; the flags of the
+// new `exp` are the ones shared by both exponentials it replaces.
+// CHECK-LABEL: @exp_quotient_minimal_fastmath(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32
+func.func @exp_quotient_minimal_fastmath(%arg0: f32, %arg1: f32) -> f32 {
+ // CHECK: %[[SUB:.*]] = arith.subf %[[ARG0]], %[[ARG1]] fastmath<reassoc,arcp> : f32
+ // CHECK: %[[EXP:.*]] = math.exp %[[SUB]] fastmath<afn> : f32
+ // CHECK: return %[[EXP]]
+ %0 = math.exp %arg0 fastmath<afn,ninf> : f32
+ %1 = math.exp %arg1 fastmath<afn> : f32
+ %2 = arith.divf %0, %1 fastmath<reassoc,arcp> : f32
+ return %2 : f32
+}
+
+// Negative test - the numerator is still needed, but folding remains
+// profitable because it removes the division.
+// CHECK-LABEL: @exp_quotient_numerator_multiple_uses(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32
+func.func @exp_quotient_numerator_multiple_uses(%arg0: f32, %arg1: f32) -> (f32, f32) {
+ // CHECK: %[[NUM:.*]] = math.exp %[[ARG0]] fastmath<fast> : f32
+ // CHECK: %[[SUB:.*]] = arith.subf %[[ARG0]], %[[ARG1]] fastmath<fast> : f32
+ // CHECK: %[[EXP:.*]] = math.exp %[[SUB]] fastmath<fast> : f32
+ // CHECK: return %[[EXP]], %[[NUM]]
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = math.exp %arg1 fastmath<fast> : f32
+ %2 = arith.divf %0, %1 fastmath<fast> : f32
+ return %2, %0 : f32, f32
+}
+
+// The divisor is still needed, but folding remains profitable because the
+// numerator dies and the division becomes a subtraction.
+// CHECK-LABEL: @exp_quotient_denominator_multiple_uses(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32
+func.func @exp_quotient_denominator_multiple_uses(%arg0: f32, %arg1: f32) -> (f32, f32) {
+ // CHECK: %[[DEN:.*]] = math.exp %[[ARG1]] fastmath<fast> : f32
+ // CHECK: %[[SUB:.*]] = arith.subf %[[ARG0]], %[[ARG1]] fastmath<fast> : f32
+ // CHECK: %[[EXP:.*]] = math.exp %[[SUB]] fastmath<fast> : f32
+ // CHECK: return %[[EXP]], %[[DEN]]
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = math.exp %arg1 fastmath<fast> : f32
+ %2 = arith.divf %0, %1 fastmath<fast> : f32
+ return %2, %1 : f32, f32
+}
+
+// Negative test - neither exponential dies, so the fold would add a third one.
+// CHECK-LABEL: @exp_quotient_both_multiple_uses(
+func.func @exp_quotient_both_multiple_uses(%arg0: f32, %arg1: f32) -> (f32, f32, f32) {
+ // CHECK-COUNT-2: math.exp
+ // CHECK-NOT: math.exp
+ // CHECK: arith.divf
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = math.exp %arg1 fastmath<fast> : f32
+ %2 = arith.divf %0, %1 fastmath<fast> : f32
+ return %2, %0, %1 : f32, f32, f32
+}
+
+// Both operands come from the same exponential, which therefore dies.
+// CHECK-LABEL: @exp_quotient_same_exponential(
+// CHECK-SAME: %[[ARG0:.+]]: f32
+func.func @exp_quotient_same_exponential(%arg0: f32) -> f32 {
+ // CHECK: %[[SUB:.*]] = arith.subf %[[ARG0]], %[[ARG0]] fastmath<fast> : f32
+ // CHECK: %[[EXP:.*]] = math.exp %[[SUB]] fastmath<fast> : f32
+ // CHECK: return %[[EXP]]
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = arith.divf %0, %0 fastmath<fast> : f32
+ return %1 : f32
+}
+
+// Negative test - not enough fastmath flags on the division.
+// CHECK-LABEL: @exp_quotient_not_enough_fastmath(
+func.func @exp_quotient_not_enough_fastmath(%arg0: f32, %arg1: f32, %arg2: f32,
+ %arg3: f32) -> (f32, f32) {
+ // CHECK-COUNT-2: arith.divf
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = math.exp %arg1 fastmath<fast> : f32
+ %2 = arith.divf %0, %1 fastmath<reassoc> : f32
+ %3 = math.exp %arg2 fastmath<fast> : f32
+ %4 = math.exp %arg3 fastmath<fast> : f32
+ %5 = arith.divf %3, %4 fastmath<arcp> : f32
+ return %2, %5 : f32, f32
+}
+
+// Negative test - mismatched exponential bases.
+// CHECK-LABEL: @exp_quotient_mixed_bases(
+func.func @exp_quotient_mixed_bases(%arg0: f32, %arg1: f32) -> f32 {
+ // CHECK: arith.divf
+ %0 = math.exp %arg0 fastmath<fast> : f32
+ %1 = math.exp2 %arg1 fastmath<fast> : f32
+ %2 = arith.divf %0, %1 fastmath<fast> : f32
+ return %2 : f32
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/220010
More information about the Mlir-commits
mailing list