Reverse-mode derivatives for math.atan, math.absf, arith.select

Author: pengmai

enzyme/Enzyme/MLIR/Implementations/ArithDerivatives.td

@@ -31,6 +31,13 @@ def : MLIRDerivative<"arith", "DivFOp", (Op $x, $y),
                     ],
                     (CheckedDivF (SubF (SelectIfActive $x, (MulF (Shadow $x), $y), (ConstantFP<"0","arith", "ConstantOp"> $x)), (SelectIfActive $y, (MulF (Shadow $y), $x), (ConstantFP<"0","arith","ConstantOp"> $y))), (MulF $y, $y))
                   >;
+def : MLIRDerivative<"arith", "SelectOp", (Op $pred, $x, $y),
+                    [
+                      (AssertingInactiveArg),
+                      (Arith_Select $pred, (DiffeRet), (ConstantFP<"0","arith","ConstantOp"> $x)),
+                      (Arith_Select $pred, (ConstantFP<"0","arith","ConstantOp"> $x), (DiffeRet)),
+                    ]
+                  >;
 
 def ExtF : ArithInst<"ExtFOp">;
 def TruncF : ArithInst<"TruncFOp">;

enzyme/Enzyme/MLIR/Implementations/Common.td

@@ -155,6 +155,7 @@ def CmpF : ArithInst<"CmpFOp">;
 def Arith_Select : ArithInst<"SelectOp">;
 
 def Arith_OEQ : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "arith::CmpFPredicate::OEQ">;
+def Arith_OGE : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "arith::CmpFPredicate::OGE">;
 
 def CheckedMulF : SubRoutine<(Op $diffret, $x),
                   (

enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td

@@ -20,3 +20,14 @@ def : MLIRDerivative<"math", "SqrtOp", (Op $x),
                       (Arith_Select (CmpF (Arith_OEQ), $x, (ConstantFP<"0","arith","ConstantOp"> $x):$zero), $zero, (DivF (DiffeRet), (MulF (ConstantFP<"2.0","arith","ConstantOp"> $x), (SqrtF $x))))
                     ]
                   >;
+def : MLIRDerivative<"math", "AtanOp", (Op $x),
+                    [
+                      (CheckedMulF (DiffeRet), (DivF (ConstantFP<"1.0","arith","ConstantOp"> $x):$one, (AddF (MulF $x, $x), $one)))
+                    ]
+                  >;
+def : MLIRDerivative<"math", "AbsFOp", (Op $x),
+                    [
+                      // TODO: handle complex
+                      (Arith_Select (CmpF (Arith_OGE), $x, (ConstantFP<"0","arith","ConstantOp"> $x)), (DiffeRet), (NegF (DiffeRet)))
+                    ]
+                  >;