More powerful C API

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -3017,17 +3017,22 @@ class AdjointGenerator
       Value* tape = nullptr;
 
       if (Mode == DerivativeMode::ReverseModePrimal ||
-          Mode == DerivativeMode::ReverseModeCombined)
+          Mode == DerivativeMode::ReverseModeCombined) {
         found->second.first(BuilderZ, orig, *gutils, normalReturn, invertedReturn, tape);
+        if (tape)
+          gutils->cacheForReverse(BuilderZ, tape, getIndex(orig, CacheType::Tape));
+      }
 
       if (Mode == DerivativeMode::ReverseModeGradient ||
-          Mode == DerivativeMode::ReverseModeCombined)
+          Mode == DerivativeMode::ReverseModeCombined) {
+        if (Mode == DerivativeMode::ReverseModeGradient &&
+            augmentedReturn->tapeIndices.find(std::make_pair(orig, CacheType::Tape)) != augmentedReturn->tapeIndices.end()) {
+          tape = Builder2.CreatePHI(Type::getInt32Ty(orig->getContext()), 0);
+          tape = gutils->cacheForReverse(Builder2, (Value*)0x01, getIndex(orig, CacheType::Tape), /*ignoreType*/true);
+        }
+        if (tape)
+          tape = gutils->lookupM(tape, Builder2);
         found->second.second(Builder2, orig, *(DiffeGradientUtils*)gutils, tape);
-
-      assert(!tape && "Tape mechanism not implemented for custom yet");
-
-      if (Mode == DerivativeMode::ReverseModePrimal && tape) {
-        gutils->cacheForReverse(BuilderZ, tape, getIndex(orig, CacheType::Tape));
       }
 
       if (gutils->invertedPointers.count(orig)) {

enzyme/Enzyme/GradientUtils.cpp

@@ -1157,7 +1157,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
 }
 
 Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
-                                      int idx) {
+                                      int idx, bool ignoreType) {
   assert(malloc);
   assert(BuilderQ.GetInsertBlock()->getParent() == newFunc);
   if (mode == DerivativeMode::ReverseModeCombined) {
@@ -1268,7 +1268,7 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
           innerType != ret->getType()) {
         assert(innerType == Type::getInt8Ty(malloc->getContext()));
       } else {
-        if (innerType != malloc->getType()) {
+        if (!ignoreType && innerType != malloc->getType()) {
           llvm::errs() << *oldFunc << "\n";
           llvm::errs() << *newFunc << "\n";
           llvm::errs() << "innerType: " << *innerType << "\n";
@@ -1289,7 +1289,7 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
 
       auto v = lookupValueFromCache(/*forwardPass*/ true, BuilderQ, lctx, cache,
                                     isi1);
-      if (malloc) {
+      if (!ignoreType && malloc) {
         assert(v->getType() == malloc->getType());
       }
       insert_or_assign(scopeMap, v, std::make_pair(cache, ctx));
@@ -1462,7 +1462,8 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
         }
       }
       // llvm::errs() << "replacing " << *malloc << " with " << *ret << "\n";
-      cast<Instruction>(malloc)->replaceAllUsesWith(ret);
+      if (!ignoreType)
+        cast<Instruction>(malloc)->replaceAllUsesWith(ret);
       std::string n = malloc->getName().str();
       erase(cast<Instruction>(malloc));
       ret->setName(n);

enzyme/Enzyme/GradientUtils.h

@@ -723,7 +723,7 @@ class GradientUtils : public CacheUtility {
     }
   }
 
-  Value *cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc, int idx);
+  Value *cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc, int idx, bool ignoreType=false);
 
   const SmallVectorImpl<Value *> &getTapeValues() const {
     return addedTapeVals;