Custom alloc zeroing (rust-lang#908)

* Custom zero

* Enable on differentials

* Cleanup

Author: wsmoses

enzyme/Enzyme/CApi.cpp

@@ -685,11 +685,6 @@ void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2,
   }
 }
 
-void EnzymeSetForMemSet(LLVMValueRef inst1) {
-  Instruction *I1 = cast<Instruction>(unwrap(inst1));
-  I1->setMetadata("enzyme_formemset", MDNode::get(I1->getContext(), {}));
-}
-
 void EnzymeSetMustCache(LLVMValueRef inst1) {
   Instruction *I1 = cast<Instruction>(unwrap(inst1));
   I1->setMetadata("enzyme_mustcache", MDNode::get(I1->getContext(), {}));

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2513,9 +2513,6 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
   IRBuilder<> ib(NewF->getEntryBlock().getFirstNonPHI());
 
   AllocaInst *ret = noReturn ? nullptr : ib.CreateAlloca(RetType);
-  if (!noReturn && EnzymeZeroCache) {
-    ib.CreateStore(Constant::getNullValue(RetType), ret);
-  }
 
   if (!noTape) {
     Value *tapeMemory;
@@ -2573,6 +2570,10 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
 #endif
         cast<GetElementPtrInst>(tapeMemory)->setIsInBounds(true);
       }
+      if (EnzymeZeroCache) {
+        ZeroMemory(ib, tapeType, tapeMemory,
+                   /*isTape*/ true);
+      }
     }
 
     unsigned i = 0;

enzyme/Enzyme/GradientUtils.cpp

@@ -2218,9 +2218,8 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
     toadd = scopeAllocs[found2->second.first][0];
     for (auto u : toadd->users()) {
       if (auto ci = dyn_cast<CastInst>(u)) {
-        if (hasMetadata(ci, "enzyme_formemset"))
-          continue;
         toadd = ci;
+        break;
       }
     }
 

enzyme/Enzyme/GradientUtils.h

@@ -1650,8 +1650,8 @@ class DiffeGradientUtils final : public GradientUtils {
 #else
       differentials[val]->setAlignment(Alignment);
 #endif
-      entryBuilder.CreateStore(Constant::getNullValue(type),
-                               differentials[val]);
+      ZeroMemory(entryBuilder, type, differentials[val],
+                 /*isTape*/ false);
     }
 #if LLVM_VERSION_MAJOR >= 15
     if (val->getContext().supportsTypedPointers()) {

enzyme/Enzyme/Utils.cpp

@@ -50,6 +50,8 @@ LLVMValueRef (*CustomAllocator)(LLVMBuilderRef, LLVMTypeRef,
                                 /*Count*/ LLVMValueRef,
                                 /*Align*/ LLVMValueRef, uint8_t,
                                 LLVMValueRef *) = nullptr;
+LLVMValueRef (*CustomZero)(LLVMBuilderRef, LLVMTypeRef,
+                           /*Ptr*/ LLVMValueRef, uint8_t) = nullptr;
 LLVMValueRef (*CustomDeallocator)(LLVMBuilderRef, LLVMValueRef) = nullptr;
 void (*CustomRuntimeInactiveError)(LLVMBuilderRef, LLVMValueRef,
                                    LLVMValueRef) = nullptr;
@@ -58,6 +60,15 @@ LLVMValueRef *(*EnzymePostCacheStore)(LLVMValueRef, LLVMBuilderRef,
 LLVMTypeRef (*EnzymeDefaultTapeType)(LLVMContextRef) = nullptr;
 }
 
+void ZeroMemory(llvm::IRBuilder<> &Builder, llvm::Type *T, llvm::Value *obj,
+                bool isTape) {
+  if (CustomZero) {
+    CustomZero(wrap(&Builder), wrap(T), wrap(obj), isTape);
+  } else {
+    Builder.CreateStore(Constant::getNullValue(T), obj);
+  }
+}
+
 llvm::SmallVector<llvm::Instruction *, 2> PostCacheStore(llvm::StoreInst *SI,
                                                          llvm::IRBuilder<> &B) {
   SmallVector<llvm::Instruction *, 2> res;

enzyme/Enzyme/Utils.h

@@ -92,6 +92,8 @@ llvm::Value *CreateAllocation(llvm::IRBuilder<> &B, llvm::Type *T,
                               llvm::Instruction **ZeroMem = nullptr,
                               bool isDefault = false);
 llvm::CallInst *CreateDealloc(llvm::IRBuilder<> &B, llvm::Value *ToFree);
+void ZeroMemory(llvm::IRBuilder<> &Builder, llvm::Type *T, llvm::Value *obj,
+                bool isTape);
 
 llvm::Value *CreateReAllocation(llvm::IRBuilder<> &B, llvm::Value *prev,
                                 llvm::Type *T, llvm::Value *OuterCount,