From ac3be7e59da58e8e57d51d712b52270929986518 Mon Sep 17 00:00:00 2001
From: tyb0807 <son@brium.ai>
Date: Wed, 8 Jan 2025 22:52:25 +0100
Subject: [PATCH] Add derivative for LLVM:ExpOp (#2220)

Co-authored-by: BuildKite <buildkite@buildkite.com>
---
 enzyme/Enzyme/MLIR/Implementations/Common.td    |  4 ++++
 .../LLVMAutoDiffOpInterfaceImpl.cpp             |  1 +
 .../MLIR/Implementations/LLVMDerivatives.td     |  6 ++++++
 enzyme/test/MLIR/ForwardMode/llvm.mlir          | 17 +++++++++++++++++
 4 files changed, 28 insertions(+)

diff --git a/enzyme/Enzyme/MLIR/Implementations/Common.td b/enzyme/Enzyme/MLIR/Implementations/Common.td
index 9846f9ae6183..29f977c95a98 100644
--- a/enzyme/Enzyme/MLIR/Implementations/Common.td
+++ b/enzyme/Enzyme/MLIR/Implementations/Common.td
@@ -121,6 +121,7 @@ def TypeOf : Operation</*primal*/0, /*shadow*/0> {
 
 class ComplexInst<string m> : Inst<m, "complex">;
 class ArithInst<string m> : Inst<m, "arith">;
+class LlvmInst<string m> : Inst<m, "LLVM">;
 class MathInst<string m> : Inst<m, "math">;
 
 def AddF : ArithInst<"AddFOp">;
@@ -133,6 +134,9 @@ def RemF : ArithInst<"RemFOp">;
 def CheckedMulF : ArithInst<"MulFOp">;
 def CheckedDivF : ArithInst<"DivFOp">;
 
+def LlvmCheckedMulF : LlvmInst<"FMulOp">;
+def LlvmExpF : LlvmInst<"ExpOp">;
+
 def CosF : MathInst<"CosOp">;
 def SinF : MathInst<"SinOp">;
 def ExpF : MathInst<"ExpOp">;
diff --git a/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp
index b9e9ade7421f..11b191d7fdd6 100644
--- a/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp
+++ b/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp
@@ -15,6 +15,7 @@
 #include "Interfaces/AutoDiffOpInterface.h"
 #include "Interfaces/AutoDiffTypeInterface.h"
 #include "Interfaces/GradientUtils.h"
+#include "Interfaces/GradientUtilsReverse.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Support/LogicalResult.h"
diff --git a/enzyme/Enzyme/MLIR/Implementations/LLVMDerivatives.td b/enzyme/Enzyme/MLIR/Implementations/LLVMDerivatives.td
index e77e88aea47f..949eeb22e09b 100644
--- a/enzyme/Enzyme/MLIR/Implementations/LLVMDerivatives.td
+++ b/enzyme/Enzyme/MLIR/Implementations/LLVMDerivatives.td
@@ -26,3 +26,9 @@ def : ReadOnlyIdentityOp<"LLVM", "PtrToIntOp", [0]>;
 def : ReadOnlyIdentityOp<"LLVM", "IntToPtrOp", [0]>;
 
 def : AllocationOp<"LLVM", "AllocaOp">;
+
+def : MLIRDerivative<"LLVM", "ExpOp", (Op $x),
+                    [
+                      (LlvmCheckedMulF (DiffeRet), (LlvmExpF $x))
+                    ]
+                  >;
diff --git a/enzyme/test/MLIR/ForwardMode/llvm.mlir b/enzyme/test/MLIR/ForwardMode/llvm.mlir
index df7e572ce4ce..f6cd8c3a9c2c 100644
--- a/enzyme/test/MLIR/ForwardMode/llvm.mlir
+++ b/enzyme/test/MLIR/ForwardMode/llvm.mlir
@@ -13,6 +13,16 @@ module {
     %r = enzyme.fwddiff @square(%x, %dx) { activity=[#enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64) -> (f64)
     return %r : f64
   }
+
+  func.func @exp(%x: f32) -> f32 {
+    %0 = llvm.intr.exp(%x) : (f32) -> f32
+    return %0 : f32
+  }
+
+  func.func @dexp(%x: f32, %dx: f32) -> f32 {
+    %r = enzyme.fwddiff @exp(%x, %dx) { activity=[#enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f32, f32) -> f32
+    return %r : f32
+  }
 }
 
 // CHECK:   func.func private @fwddiffesquare(%[[arg0:.+]]: f64, %[[arg1:.+]]: f64) -> f64 {
@@ -29,3 +39,10 @@ module {
 // CHECK-NEXT:     %[[i7:.+]] = llvm.load %[[i1]] : !llvm.ptr -> f64
 // CHECK-NEXT:     return %[[i6]] : f64
 // CHECK-NEXT:   }
+
+// CHECK:  func.func private @fwddiffeexp(%[[arg0:.+]]: f32, %[[arg1:.+]]: f32) -> f32 {
+// CHECK-NEXT:    %[[der:.+]] = llvm.intr.exp(%[[arg0]]) : (f32) -> f32
+// CHECK-NEXT:    %[[res:.+]] = llvm.fmul %[[arg1]], %[[der]] : f32
+// CHECK-NEXT:    %[[exp:.+]] = llvm.intr.exp(%[[arg0]]) : (f32) -> f32
+// CHECK-NEXT:    return %[[res]] : f32
+// CHECK-NEXT:  }