Add complex sqrt

EnzymeAD · Mar 5, 2024 · 53f7e8c · 53f7e8c
1 parent 0b62188
commit 53f7e8c
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td
@@ -934,6 +934,13 @@ def : IntrPattern<(Op $x),
  (ForwardFromSummedReverse)
  >;
 
+
+def : IntrPattern<(Op (Op $x, $y)$z),
+ [["cmplx_sqrt"]],
+ [(Select (And (FCmpUEQ $x, (ConstantFP<"0"> $x)), (FCmpUEQ $y, (ConstantFP<"0"> $y))), (ConstantCFP<"0", "0"> $z), (Conj (CFDiv (Conj (DiffeRet)), (CFMul (ConstantCFP<"2", "0"> $z), (Call<(SameFunc), [ReadNone,NoUnwind]> $z)))))],
+ (Select (And (FCmpUEQ $x, (ConstantFP<"0"> $x)), (FCmpUEQ $y, (ConstantFP<"0"> $y))), (ConstantCFP<"0", "0"> $z), (CFDiv (Shadow $z), (CFMul (ConstantCFP<"2", "0"> $z), (Call<(SameFunc), [ReadNone,NoUnwind]> $z))))
+ >;
+
 def : IntrPattern<(Op $x, $y),
  [["pow"]],
  [