mlir: implement more Complex derivatives (#2691)

* mlir: implement more Complex derivatives

* mul per differet

Author: Pangoraw

enzyme/Enzyme/MLIR/Implementations/Common.td

@@ -180,9 +180,14 @@ def CheckedDivF : SubRoutine<(Op $diffret, $x),
 def LlvmCheckedMulF : LlvmInst<"FMulOp">;
 def LlvmExpF : LlvmInst<"ExpOp">;
 
+def ComplexCreate : ComplexInst<"CreateOp">;
+def ComplexRe : ComplexInst<"ReOp">;
+def ComplexIm : ComplexInst<"ImOp">;
+
 def CosF : MathInst<"CosOp">;
 def SinF : MathInst<"SinOp">;
 def ExpF : MathInst<"ExpOp">;
 def SqrtF : MathInst<"SqrtOp">;
+def AbsF : MathInst<"AbsFOp">;
 
 #endif // ENZYME_MLIR_IMPLEMENTATIONS_COMMON

enzyme/Enzyme/MLIR/Implementations/ComplexDerivatives.td

@@ -15,3 +15,33 @@ def : MLIRDerivative<"complex", "MulOp", (Op $x, $y),
                       (CMul (DiffeRet), $x)
                     ]
                   >;
+
+def : MLIRDerivative<"complex", "ImOp", (Op $x),
+                    [
+                      (ComplexCreate
+                        (TypeOf $x),
+                        (ConstantFP<"0", "arith", "ConstantOp">),
+                        (NegF (DiffeRet))
+                      )
+                    ],
+                    (ComplexIm (Shadow $x))
+                  >;
+
+def : MLIRDerivative<"complex", "ReOp", (Op $x),
+                    [
+                      (ComplexCreate
+                        (TypeOf $x),
+                        (DiffeRet),
+                        (ConstantFP<"0", "arith", "ConstantOp">)
+                      )
+                    ],
+                    (ComplexRe (Shadow $x))
+                  >;
+
+def : MLIRDerivative<"complex", "CreateOp", (Op $re, $im),
+                    [
+                      (ComplexRe (DiffeRet)),
+                      (ComplexIm (DiffeRet))
+                    ]
+                  >;
+

enzyme/Enzyme/MLIR/Implementations/MathAutoDiffOpInterfaceImpl.cpp

@@ -16,6 +16,7 @@
 #include "Interfaces/GradientUtils.h"
 #include "Interfaces/GradientUtilsReverse.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Support/LogicalResult.h"

enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td

@@ -10,7 +10,7 @@ def : MLIRDerivative<"math", "ExpOp", (Op $x),
                       (CheckedMulF (DiffeRet), (ExpF $x))
                     ]
                   >;
-def : MLIRDerivative<"math", "SinOp", (Op $x), 
+def : MLIRDerivative<"math", "SinOp", (Op $x),
                     [
                       (CheckedMulF (DiffeRet), (CosF $x))
                     ]
@@ -27,7 +27,13 @@ def : MLIRDerivative<"math", "AtanOp", (Op $x),
                   >;
 def : MLIRDerivative<"math", "AbsFOp", (Op $x),
                     [
-                      // TODO: handle complex
-                      (Arith_Select (CmpF (Arith_OGE), $x, (ConstantFP<"0","arith","ConstantOp"> $x)), (DiffeRet), (NegF (DiffeRet)))
+                      (SelectIfComplex $x,
+                        (ComplexCreate
+                          (TypeOf $x),
+                          (CheckedMulF (DiffeRet), (DivF (ComplexRe $x), (AbsF $x))),
+                          (CheckedMulF (DiffeRet), (NegF (DivF (ComplexIm $x), (AbsF $x))))
+                        ),
+                        (Arith_Select (CmpF (Arith_OGE), $x, (ConstantFP<"0","arith","ConstantOp"> $x)), (DiffeRet), (NegF (DiffeRet)))
+                      )
                     ]
                   >;

enzyme/test/MLIR/ReverseMode/complex.mlir

@@ -0,0 +1,24 @@
+// RUN: %eopt %s --enzyme-wrap="infn=main outfn= argTys=enzyme_active retTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s
+
+module {
+  func.func @main(%x: complex<f32>) -> (f32, f32) {
+    %0 = complex.re %x : complex<f32>
+    %1 = complex.im %x : complex<f32>
+    return %0, %1 : f32, f32
+  }
+}
+
+// CHECK:  func.func @main(%arg0: complex<f32>, %arg1: f32, %arg2: f32) -> complex<f32> {
+// CHECK-NEXT:    %cst = arith.constant 0.000000e+00 : f32
+// CHECK-NEXT:    %cst_0 = complex.constant [0.000000e+00 : f32, 0.000000e+00 : f32] : complex<f32>
+// CHECK-NEXT:    %0 = arith.addf %arg1, %cst : f32
+// CHECK-NEXT:    %1 = arith.addf %arg2, %cst : f32
+// CHECK-NEXT:    %2 = arith.negf %1 : f32
+// CHECK-NEXT:    %3 = complex.create %cst, %2 : complex<f32>
+// CHECK-NEXT:    %4 = complex.conj %3 : complex<f32>
+// CHECK-NEXT:    %5 = complex.add %cst_0, %4 : complex<f32>
+// CHECK-NEXT:    %6 = complex.create %0, %cst : complex<f32>
+// CHECK-NEXT:    %7 = complex.conj %6 : complex<f32>
+// CHECK-NEXT:    %8 = complex.add %5, %7 : complex<f32>
+// CHECK-NEXT:    return %8 : complex<f32>
+// CHECK-NEXT:  }