[InstCombine] optimize exp(exp(x)) / exp(x) with fast-math

[zjaffal]

This patch enables us to replace exp(exp(x)) / exp(x) to exp(exp(x) - x) making us avoid the div instruction

Closes #65608

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Changes

This patch enables us to replace exp(exp(x)) / exp(x) to exp(exp(x) - x) making us avoid the div instruction

Closes #65608

Full diff: https://github.com/llvm/llvm-project/pull/66177.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp (+12)
• (added) llvm/test/Transforms/InstCombine/fdiv-exp.ll (+55)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 86cacf979839126..692fcf01455835a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -1748,6 +1748,18 @@ Instruction *InstCombinerImpl::visitFDiv(BinaryOperator &I) {
 
   if (Instruction *Mul = foldFDivPowDivisor(I, Builder))
     return Mul;
+  Value *ExpX;
+  // exp(exp(X)) / exp(X) -> exp(exp(X) - X)
+  if (match(Op0, m_Intrinsic<Intrinsic::exp>(m_Value(ExpX))) &&
+      match(Op1, m_Intrinsic<Intrinsic::exp>(m_Value(X))) &&
+      match(ExpX, m_Intrinsic<Intrinsic::exp>(m_Specific(X)))) {
+    // check that exp(x) is only used in the div expression.
+    if (Op1->hasNUses(2)) {
+      Value *XY = Builder.CreateFSubFMF(ExpX, X, &I);
+      auto *NewPow = Builder.CreateUnaryIntrinsic(Intrinsic::exp, XY, &I);
+      return replaceInstUsesWith(I, NewPow);
+    }
+  }
 
   // pow(X, Y) / X --> pow(X, Y-1)
   if (I.hasAllowReassoc() &&
diff --git a/llvm/test/Transforms/InstCombine/fdiv-exp.ll b/llvm/test/Transforms/InstCombine/fdiv-exp.ll
new file mode 100644
index 000000000000000..3ea281cd4bb3af5
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fdiv-exp.ll
@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -S -passes=instcombine < %s | FileCheck %s
+
+define double @fdiv_exp(double %x) {
+; CHECK-LABEL: define double @fdiv_exp
+; CHECK-SAME: (double [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[EXP_X:%.*]] = call fast double @llvm.exp.f64(double [[X]])
+; CHECK-NEXT:    [[TMP0:%.*]] = fsub fast double [[EXP_X]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = call fast double @llvm.exp.f64(double [[TMP0]])
+; CHECK-NEXT:    ret double [[DIV]]
+;
+entry:
+  %exp_x = call fast double @llvm.exp.f64(double %x)
+  %exp_exp_x = call fast double @llvm.exp.f64(double %exp_x)
+  %div = fdiv fast double %exp_exp_x, %exp_x
+  ret double %div
+}
+
+define double @fdiv_exp_multiple_uses(double %x) {
+; CHECK-LABEL: define double @fdiv_exp_multiple_uses
+; CHECK-SAME: (double [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[EXP_X:%.*]] = call fast double @llvm.exp.f64(double [[X]])
+; CHECK-NEXT:    [[EXP_EXP_X:%.*]] = call fast double @llvm.exp.f64(double [[EXP_X]])
+; CHECK-NEXT:    [[DIV:%.*]] = fdiv fast double [[EXP_EXP_X]], [[EXP_X]]
+; CHECK-NEXT:    call void @use(double [[EXP_X]])
+; CHECK-NEXT:    ret double [[DIV]]
+;
+entry:
+  %exp_x = call fast double @llvm.exp.f64(double %x)
+  %exp_exp_x = call fast double @llvm.exp.f64(double %exp_x)
+  %div = fdiv fast double %exp_exp_x, %exp_x
+  call void @use(double %exp_x)
+  ret double %div
+}
+
+define double @fdiv_exp_swapped(double %x) {
+; CHECK-LABEL: define double @fdiv_exp_swapped
+; CHECK-SAME: (double [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[EXP_X:%.*]] = call fast double @llvm.exp.f64(double [[X]])
+; CHECK-NEXT:    [[TMP0:%.*]] = fsub fast double [[X]], [[EXP_X]]
+; CHECK-NEXT:    [[DIV:%.*]] = call fast double @llvm.exp.f64(double [[TMP0]])
+; CHECK-NEXT:    ret double [[DIV]]
+;
+entry:
+  %exp_x = call fast double @llvm.exp.f64(double %x)
+  %exp_exp_x = call fast double @llvm.exp.f64(double %exp_x)
+  %div = fdiv fast double %exp_x, %exp_exp_x
+  ret double %div
+}
+
+declare double @llvm.exp.f64(double)
+declare void @use(double)

[nikic]

This pattern is too specific. I assume the motivation here is this special case where exp(x) has multiple uses, which means that the x / exp(y) to x * exp(-y) fold does not trigger. However, you should still be able to handle something more general, such as exp(x) / exp(y) to exp(x - y) where exp(x) has one use. This does not reduce the number of exp calls, but does reduce the number of division.

[nikic]

Not seeing any checks for FMF here? Please add negative tests that the transform is not performed without needed flags.

[zjaffal]

which means that the x / exp(y) to x * exp(-y) fold does not trigger

This example doesn't have an exp(x). We should be able to handle this form as well right? I guess we can port the optimisations from pow(x,y) to the exp(x)

However, you should still be able to handle something more general, such as exp(x) / exp(y) to exp(x - y) where exp(x) has one use

I guess my pattern should make sure that if x has multiple uses then all of them should be in the division operation subtree so this should be able to handle exp(x) / exp(y) and exp(exp(x)) / exp(x)

[zjaffal]

I think what we can do is similar to what is done for exp(x) * exp(y) by checking for

if (I.isOnlyUserOfAnyOperand())

to make sure that exp(x) has one user