Add nvidia reciprocal (#1586)

EnzymeAD · Dec 17, 2023 · a4f429f · a4f429f
1 parent d22dd49
commit a4f429f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td
@@ -257,6 +257,17 @@ def DiffeRet : DiffeRetIndex<[-1]>;
 def Shadow : Operation</*primal*/0, /*shadow*/1> {
 }
 
+//  1 / y -> -1/y^2 dy
+def : CallPattern<(Op $y),
+                  ["__nv_frcp_rd", "__nv_frcp_rn", "__nv_frcp_ru", "__nv_frcp_rz",
+                   "__nv_drcp_rd", "__nv_drcp_rn", "__nv_drcp_ru", "__nv_drcp_rz"],
+                    [
+                      (FNeg (CheckedDiv (DiffeRet), (FMul $y, $y)))
+                    ],
+                    (ForwardFromSummedReverse),
+                    [ReadNone, NoUnwind]
+                  >;
+
 def : CallPattern<(Op $x),
                   ["atan", "atanf", "atanl", "__fd_atan_1"],
                   [(CheckedDiv (DiffeRet), (FAdd (FMul $x, $x), (ConstantFP<"1.0"> $x)))],

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
@@ -93,6 +93,14 @@ const llvm::StringMap<llvm::Intrinsic::ID> LIBM_FUNCTIONS = {
     {"sin", Intrinsic::sin},
     {"tan", Intrinsic::not_intrinsic},
     {"acos", Intrinsic::not_intrinsic},
+    {"__nv_frcp_rd", Intrinsic::not_intrinsic},
+    {"__nv_frcp_rn", Intrinsic::not_intrinsic},
+    {"__nv_frcp_ru", Intrinsic::not_intrinsic},
+    {"__nv_frcp_rz", Intrinsic::not_intrinsic},
+    {"__nv_drcp_rd", Intrinsic::not_intrinsic},
+    {"__nv_drcp_rn", Intrinsic::not_intrinsic},
+    {"__nv_drcp_ru", Intrinsic::not_intrinsic},
+    {"__nv_drcp_rz", Intrinsic::not_intrinsic},
     {"__nv_isnand", Intrinsic::not_intrinsic},
     {"__nv_isnanf", Intrinsic::not_intrinsic},
     {"__nv_isinfd", Intrinsic::not_intrinsic},