Permit Zero Initialization of Cache (rust-lang#386)

* Zero Outermost Cache

* WIP: Silly billy

* Zero cache option

Co-authored-by: Valentin Churavy <v.churavy@gmail.com>

Author: wsmoses

enzyme/Enzyme/CacheUtility.cpp

@@ -35,6 +35,10 @@ llvm::cl::opt<bool>
     EfficientBoolCache("enzyme-smallbool", cl::init(false), cl::Hidden,
                        cl::desc("Place 8 bools together in a single byte"));
 
+llvm::cl::opt<bool> EnzymeZeroCache("enzyme-zero-cache", cl::init(false),
+                                    cl::Hidden,
+                                    cl::desc("Zero initialize the cache"));
+
 llvm::cl::opt<bool>
     EnzymePrintPerf("enzyme-print-perf", cl::init(false), cl::Hidden,
                     cl::desc("Enable Enzyme to print performance info"));
@@ -640,6 +644,9 @@ AllocaInst *CacheUtility::createCacheForScope(LimitContext ctx, Type *T,
 #endif
     }
   }
+  if (EnzymeZeroCache && sublimits.size() == 0)
+    scopeInstructions[alloc].push_back(
+        entryBuilder.CreateStore(Constant::getNullValue(types.back()), alloc));
 
   Type *BPTy = Type::getInt8PtrTy(T->getContext());
 
@@ -699,6 +706,20 @@ AllocaInst *CacheUtility::createCacheForScope(LimitContext ctx, Type *T,
               cast<CallInst>(cast<Instruction>(firstallocation)->getOperand(0));
         }
 
+        if (EnzymeZeroCache && i == 0) {
+          Value *args[] = {
+              malloccall,
+              ConstantInt::get(Type::getInt8Ty(malloccall->getContext()), 0),
+              malloccall->getArgOperand(0),
+              ConstantInt::getFalse(malloccall->getContext())};
+          Type *tys[] = {args[0]->getType(), args[2]->getType()};
+
+          scopeInstructions[alloc].push_back(allocationBuilder.CreateCall(
+              Intrinsic::getDeclaration(newFunc->getParent(), Intrinsic::memset,
+                                        tys),
+              args));
+        }
+
         // Assert computation of size of array doesn't wrap
         if (auto BI = dyn_cast<BinaryOperator>(malloccall->getArgOperand(0))) {
           if ((BI->getOperand(0) == byteSizeOfType &&
@@ -788,8 +809,9 @@ AllocaInst *CacheUtility::createCacheForScope(LimitContext ctx, Type *T,
         Value *realloccall = nullptr;
 
         realloccall = build.CreateCall(
-            getOrInsertExponentialAllocator(*newFunc->getParent()), idxs,
-            name + "_realloccache");
+            getOrInsertExponentialAllocator(*newFunc->getParent(),
+                                            EnzymeZeroCache && i == 0),
+            idxs, name + "_realloccache");
         scopeAllocs[alloc].push_back(cast<CallInst>(realloccall));
         allocation = build.CreateBitCast(realloccall, allocation->getType());
         storealloc = build.CreateStore(allocation, storeInto);

enzyme/Enzyme/CacheUtility.h

@@ -43,6 +43,8 @@
 extern "C" {
 /// Pack 8 bools together in a single byte
 extern llvm::cl::opt<bool> EfficientBoolCache;
+
+extern llvm::cl::opt<bool> EnzymeZeroCache;
 }
 
 /// Container for all loop information to synthesize gradients

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2144,6 +2144,9 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
   IRBuilder<> ib(NewF->getEntryBlock().getFirstNonPHI());
 
   Value *ret = noReturn ? nullptr : ib.CreateAlloca(RetType);
+  if (!noReturn && EnzymeZeroCache) {
+    ib.CreateStore(Constant::getNullValue(RetType), ret);
+  }
 
   if (!noTape) {
     Value *tapeMemory;
@@ -2169,6 +2172,19 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
         malloccall->addAttribute(AttributeList::ReturnIndex, attr);
 #endif
       }
+      if (EnzymeZeroCache) {
+        IRBuilder<> B(malloccall->getNextNode());
+        Value *args[] = {
+            malloccall,
+            ConstantInt::get(Type::getInt8Ty(malloccall->getContext()), 0),
+            malloccall->getArgOperand(0),
+            ConstantInt::getFalse(malloccall->getContext())};
+        Type *tys[] = {args[0]->getType(), args[2]->getType()};
+
+        B.CreateCall(Intrinsic::getDeclaration(NewF->getParent(),
+                                               Intrinsic::memset, tys),
+                     args);
+      }
 #if LLVM_VERSION_MAJOR >= 14
       malloccall->addDereferenceableRetAttr(size->getLimitedValue());
       AttrBuilder B;

enzyme/Enzyme/Utils.cpp

@@ -531,11 +531,14 @@ llvm::Value *getOrInsertOpFloatSum(llvm::Module &M, llvm::Type *OpPtr,
   return B2.CreateLoad(GV);
 }
 
-Function *getOrInsertExponentialAllocator(Module &M) {
+Function *getOrInsertExponentialAllocator(Module &M, bool ZeroInit) {
   Type *BPTy = Type::getInt8PtrTy(M.getContext());
   Type *types[] = {BPTy, Type::getInt64Ty(M.getContext()),
                    Type::getInt64Ty(M.getContext())};
   std::string name = "__enzyme_exponentialallocation";
+  if (ZeroInit)
+    name += "zero";
+  assert(!ZeroInit && "Zero initialization within reallocation not handled");
   FunctionType *FT =
       FunctionType::get(Type::getInt8PtrTy(M.getContext()), types, false);
 

enzyme/Enzyme/Utils.h

@@ -833,7 +833,7 @@ enum class MPI_CallType {
 llvm::Value *getOrInsertOpFloatSum(llvm::Module &M, llvm::Type *OpPtr,
                                    ConcreteType CT, llvm::Type *intType,
                                    llvm::IRBuilder<> &B2);
-llvm::Function *getOrInsertExponentialAllocator(llvm::Module &M);
+llvm::Function *getOrInsertExponentialAllocator(llvm::Module &M, bool ZeroInit);
 
 class AssertingReplacingVH : public llvm::CallbackVH {
 public: