Don't erase custom call for primal (rust-lang#247)

* Don't erase custom call for primal

* Error on token type storage

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -3942,7 +3942,19 @@ class AdjointGenerator
         }
       }
 
-      if (subretused) {
+      bool primalNeededInReverse;
+
+      if (gutils->knownRecomputeHeuristic.count(orig)) {
+        primalNeededInReverse = !gutils->knownRecomputeHeuristic[orig];
+      } else {
+        std::map<UsageKey, bool> Seen;
+        for (auto pair : gutils->knownRecomputeHeuristic)
+          Seen[UsageKey(pair.first, ValueType::Primal)] = false;
+        primalNeededInReverse = is_value_needed_in_reverse<ValueType::Primal>(
+            TR, gutils, orig, Mode, Seen, oldUnreachable);
+      }
+
+      if (subretused && primalNeededInReverse) {
         if (normalReturn != newCall) {
           assert(normalReturn->getType() == newCall->getType());
           gutils->replaceAWithB(newCall, normalReturn);
@@ -3952,7 +3964,9 @@ class AdjointGenerator
         normalReturn = gutils->cacheForReverse(BuilderZ, normalReturn,
                                                getIndex(orig, CacheType::Self));
       } else {
-        eraseIfUnused(*orig, /*erase*/ true, /*check*/ false);
+        if (!orig->mayWriteToMemory() ||
+            Mode == DerivativeMode::ReverseModeGradient)
+          eraseIfUnused(*orig, /*erase*/ true, /*check*/ false);
       }
       return;
     }
@@ -4439,17 +4453,16 @@ class AdjointGenerator
             return;
           }
           SmallVector<Value *, 1> iargs;
-          IRBuilder<> Builder2(call.getParent());
-          getReverseBuilder(Builder2);
+          IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&call));
           for (size_t i = 0, end = orig->getNumArgOperands(); i < end; ++i) {
             auto arg = orig->getArgOperand(i);
             if (!gutils->isConstantValue(arg)) {
-              Value *ptrshadow = gutils->invertPointerM(arg, Builder2);
+              Value *ptrshadow = gutils->invertPointerM(arg, BuilderZ);
               iargs.push_back(ptrshadow);
             }
           }
           if (iargs.size()) {
-            Builder2.CreateCall(called, iargs);
+            BuilderZ.CreateCall(called, iargs);
           }
           return;
         }

enzyme/Enzyme/DifferentialUseAnalysis.h

@@ -349,6 +349,11 @@ static inline bool is_value_needed_in_reverse(
     if (!direct)
       continue;
 
+    if (inst->getType()->isTokenTy()) {
+      llvm::errs() << " need " << *inst << " via " << *user << "\n";
+    }
+    assert(!inst->getType()->isTokenTy());
+
     return seen[idx] = true;
   }
   return false;

enzyme/Enzyme/EnzymeLogic.cpp

@@ -658,8 +658,9 @@ struct CacheAnalysis {
 
         // We do not need uncacheable args for intrinsic functions. So skip such
         // callsites.
-        if (isa<IntrinsicInst>(&inst)) {
-          continue;
+        if (auto II = dyn_cast<IntrinsicInst>(&inst)) {
+          if (!II->getCalledFunction()->getName().startswith("llvm.julia"))
+            continue;
         }
 
         // For all other calls, we compute the uncacheable args for this

enzyme/Enzyme/GradientUtils.cpp

@@ -1349,6 +1349,13 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
     }
   }
 
+  if (malloc->getType()->isTokenTy()) {
+    llvm::errs() << " oldFunc: " << *oldFunc << "\n";
+    llvm::errs() << " newFunc: " << *newFunc << "\n";
+    llvm::errs() << " malloc: " << *malloc << "\n";
+  }
+  assert(!malloc->getType()->isTokenTy());
+
   if (tape) {
     if (idx >= 0 && !tape->getType()->isStructTy()) {
       llvm::errs() << "cacheForReverse incorrect tape type: " << *tape

enzyme/Enzyme/GradientUtils.h

@@ -700,12 +700,6 @@ class GradientUtils : public CacheUtility {
         bb.CreateCall(oldFunc->getParent()->getOrInsertFunction(
                           "julia.write_barrier", FT),
                       anti);
-        if (mode != DerivativeMode::ReverseModeCombined) {
-          EmitFailure("SplitGCAllocation", orig->getDebugLoc(), orig,
-                      "Not handling Julia shadow GC allocation in split mode ",
-                      *orig);
-          return anti;
-        }
       }
 
       if (orig->getCalledFunction()->getName() == "swift_allocObject") {