add user data to EnzymeRegisterCallHandler and add to header

EnzymeAD · Jan 16, 2025 · a801448 · a801448
1 parent 5651636
commit a801448
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/enzyme/Enzyme/CApi.cpp b/enzyme/Enzyme/CApi.cpp
@@ -340,9 +340,9 @@ void EnzymeRegisterAllocationHandler(char *Name, CustomShadowAlloc AHandle,
     };
 }
 
-void EnzymeRegisterCallHandler(char *Name,
+void EnzymeRegisterCallHandler(const char *Name,
                                CustomAugmentedFunctionForward FwdHandle,
-                               CustomFunctionReverse RevHandle) {
+                               CustomFunctionReverse RevHandle, void *data) {
   auto &pair = customCallHandlers[Name];
   pair.first = [=](IRBuilder<> &B, CallInst *CI, GradientUtils &gutils,
                    Value *&normalReturn, Value *&shadowReturn,
@@ -351,15 +351,15 @@ void EnzymeRegisterCallHandler(char *Name,
     LLVMValueRef shadowR = wrap(shadowReturn);
     LLVMValueRef tapeR = wrap(tape);
     uint8_t noMod =
-        FwdHandle(wrap(&B), wrap(CI), &gutils, &normalR, &shadowR, &tapeR);
+        FwdHandle(wrap(&B), wrap(CI), &gutils, &normalR, &shadowR, &tapeR, data);
     normalReturn = unwrap(normalR);
     shadowReturn = unwrap(shadowR);
     tape = unwrap(tapeR);
     return noMod != 0;
   };
   pair.second = [=](IRBuilder<> &B, CallInst *CI, DiffeGradientUtils &gutils,
                     Value *tape) {
-    RevHandle(wrap(&B), wrap(CI), &gutils, wrap(tape));
+    RevHandle(wrap(&B), wrap(CI), &gutils, wrap(tape), data);
   };
 }
 

diff --git a/enzyme/Enzyme/CApi.h b/enzyme/Enzyme/CApi.h
@@ -198,10 +198,11 @@ typedef uint8_t (*CustomAugmentedFunctionForward)(LLVMBuilderRef, LLVMValueRef,
                                                   GradientUtils *,
                                                   LLVMValueRef *,
                                                   LLVMValueRef *,
-                                                  LLVMValueRef *);
+                                                  LLVMValueRef *,
+                                                  void *);
 
 typedef void (*CustomFunctionReverse)(LLVMBuilderRef, LLVMValueRef,
-                                      DiffeGradientUtils *, LLVMValueRef);
+                                      DiffeGradientUtils *, LLVMValueRef, void *);
 
 LLVMValueRef EnzymeCreateForwardDiff(
     EnzymeLogicRef Logic, LLVMValueRef request_req, LLVMBuilderRef request_ip,
@@ -222,6 +223,14 @@ LLVMValueRef EnzymeCreatePrimalAndGradient(
     uint8_t *_overwritten_args, size_t overwritten_args_size,
     EnzymeAugmentedReturnPtr augmented, uint8_t AtomicAdd);
 
+void EnzymeRegisterCallHandler(const char *Name,
+                               CustomAugmentedFunctionForward FwdHandle,
+                               CustomFunctionReverse RevHandle,
+                               void *data);
+
+LLVMValueRef EnzymeGradientUtilsNewFromOriginal(GradientUtils *gutils,
+                                                LLVMValueRef val);
+
 #ifdef __cplusplus
 }
 #endif