add tanh support for llvm19+

EnzymeAD · Oct 2, 2024 · a1d7334 · a1d7334
1 parent 7425c37
commit a1d7334
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 6 deletions.
diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp
@@ -2812,6 +2812,11 @@ bool guaranteedDataDependent(Value *z) {
     case Intrinsic::sqrt:
     case Intrinsic::sin:
     case Intrinsic::cos:
+#if LLVM_VERSION_MAJOR >= 19
+    case Intrinsic::sinh:
+    case Intrinsic::cosh:
+    case Intrinsic::tanh:
+#endif
       return guaranteedDataDependent(II->getArgOperand(0));
     default:
       break;

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -4180,6 +4180,11 @@ bool GradientUtils::shouldRecompute(const Value *val,
     case Intrinsic::sin:
     case Intrinsic::cos:
     case Intrinsic::exp:
+#if LLVM_VERSION_MAJOR >= 19
+    case Intrinsic::tanh:
+    case Intrinsic::cosh:
+    case Intrinsic::sinh:
+#endif
     case Intrinsic::log:
     case Intrinsic::nvvm_ldu_global_i:
     case Intrinsic::nvvm_ldu_global_p:

diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td
@@ -335,12 +335,6 @@ def : CallPattern<(Op $x, $y),
                   [ReadNone, NoUnwind]
                   >;
 
-def : CallPattern<(Op $x),
-                  ["tanh"],
-                  [(FDiv (DiffeRet), (FMul(Call<(SameTypesFunc<"cosh">), [ReadNone,NoUnwind]> $x):$c, $c))],
-                  (ForwardFromSummedReverse),
-                  [ReadNone, NoUnwind]
-                  >;
 def : CallPattern<(Op $x),
                   ["tanhf"],
                   [(FDiv (DiffeRet), (FMul(Call<(SameTypesFunc<"coshf">), [ReadNone,NoUnwind]> $x):$c, $c))],
@@ -872,6 +866,12 @@ def : CallPattern<(Op (Op $x, $y):$z),
                   [ReadNone, NoUnwind]
                   >;
 
+def : IntrPattern<(Op $x),
+                  [["tanh"]],
+                  [(FDiv (DiffeRet), (FMul(Intrinsic<"cosh"> $x):$c, $c))],
+                  (ForwardFromSummedReverse)
+                  >;
+
 def : IntrPattern<(Op $x),
                   [["sin"]],
                   [(FMul (DiffeRet), (Intrinsic<"cos"> $x))]  ,

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
@@ -118,9 +118,15 @@ const llvm::StringMap<llvm::Intrinsic::ID> LIBM_FUNCTIONS = {
     {"atan", Intrinsic::not_intrinsic},
     {"atan2", Intrinsic::not_intrinsic},
     {"__nv_atan2", Intrinsic::not_intrinsic},
+#if LLVM_VERSION_MAJOR >= 19
+    {"cosh", Intrinsic::cosh},
+    {"sinh", Intrinsic::sinh},
+    {"tanh", Intrinsic::tanh},
+#else
     {"cosh", Intrinsic::not_intrinsic},
     {"sinh", Intrinsic::not_intrinsic},
     {"tanh", Intrinsic::not_intrinsic},
+#endif
     {"acosh", Intrinsic::not_intrinsic},
     {"asinh", Intrinsic::not_intrinsic},
     {"atanh", Intrinsic::not_intrinsic},
@@ -3849,6 +3855,11 @@ void TypeAnalyzer::visitIntrinsicInst(llvm::IntrinsicInst &I) {
   case Intrinsic::exp2:
   case Intrinsic::sin:
   case Intrinsic::cos:
+#if LLVM_VERSION_MAJOR >= 19
+  case Intrinsic::sinh:
+  case Intrinsic::cosh:
+  case Intrinsic::tanh:
+#endif
   case Intrinsic::floor:
   case Intrinsic::ceil:
   case Intrinsic::trunc:

diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h
@@ -1693,6 +1693,11 @@ static inline bool isNoEscapingAllocation(const llvm::Function *F) {
   case Intrinsic::exp:
   case Intrinsic::cos:
   case Intrinsic::sin:
+#if LLVM_VERSION_MAJOR >= 19
+  case Intrinsic::tanh:
+  case Intrinsic::cosh:
+  case Intrinsic::sinh:
+#endif
   case Intrinsic::copysign:
   case Intrinsic::fabs:
     return true;

diff --git a/enzyme/test/Enzyme/ReverseMode/tanh19.ll b/enzyme/test/Enzyme/ReverseMode/tanh19.ll
@@ -0,0 +1,29 @@
+; RUN: if [ %llvmver -ge 19 ]; then %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s; fi
+
+; Function Attrs: mustprogress nocallback nofree nosync nounwind speculatable willreturn memory(none)
+declare double @llvm.tanh.f64(double) #14
+
+; Function Attrs: nounwind readnone uwtable
+define double @tester(double %x) {
+entry:
+  %0 = call double @llvm.tanh.f64(double %x)
+  ret double %0
+}
+
+define double @test_derivative(double %x) {
+entry:
+  %0 = tail call double (double (double)*, ...) @__enzyme_autodiff(double (double)* nonnull @tester, double %x)
+  ret double %0
+}
+
+; Function Attrs: nounwind
+declare double @__enzyme_autodiff(double (double)*, ...)
+
+; CHECK: define internal { double } @diffetester(double %x, double %differeturn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:  %0 = call fast double @llvm.cosh.f64(double %x)
+; CHECK-NEXT:  %1 = fmul fast double %0, %0
+; CHECK-NEXT:  %2 = fdiv fast double %differeturn, %1
+; CHECK-NEXT:  %3 = insertvalue { double } undef, double %2, 0
+; CHECK-NEXT:  ret { double } %3
+; CHECK-NEXT: }