Implement tgamma derivative (#2140)

EnzymeAD · Oct 30, 2024 · d2d3649 · d2d3649
1 parent c0710f6
commit d2d3649
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td
@@ -559,6 +559,17 @@ def : CallPattern<(Op $x, $y),
                   [ReadNone, NoUnwind]
                   >;
 
+def : CallPattern<(Op $x),
+                  ["tgamma"],
+                  [
+                  (
+                   FMul (Call<(ArgAsRetTypesFunc<"digamma">), [ReadNone,NoUnwind]> $x), (DiffeRet)
+                  )
+                  ],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+
 def : CallPattern<(Op $x),
                   ["logabsgamma"],
                   [

diff --git a/enzyme/test/Enzyme/ReverseMode/tgamma.ll b/enzyme/test/Enzyme/ReverseMode/tgamma.ll
@@ -33,3 +33,12 @@ declare double @tgamma(double)
 ; CHECK-NEXT: }
 
 ; CHECK: define internal fastcc { double } @diffea5(double %arg, double %differeturn)
+; CHECK-NEXT: bb:
+; CHECK-NEXT:   %i8 = call double @tgamma(double %arg) #1
+; CHECK-NEXT:   %0 = fmul fast double %differeturn, %arg
+; CHECK-NEXT:   %1 = fmul fast double %i8, %differeturn
+; CHECK-NEXT:   %2 = call fast double @digamma(double %arg) #0
+; CHECK-NEXT:   %3 = fmul fast double %2, %0
+; CHECK-NEXT:   %4 = fadd fast double %1, %3
+; CHECK-NEXT:   %5 = insertvalue { double } undef, double %4, 0
+; CHECK-NEXT:   ret { double } %5