Add custom anonymous type for tape (rust-lang#859)

* Add custom anonymous type for tape

* Fix cast

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -12063,7 +12063,9 @@ class AdjointGenerator
           shouldFree()) {
         assert(tape);
         auto tapep = BuilderZ.CreatePointerCast(
-            tape, PointerType::getUnqual(fnandtapetype->tapeType));
+            tape, PointerType::get(
+                      fnandtapetype->tapeType,
+                      cast<PointerType>(tape->getType())->getAddressSpace()));
 #if LLVM_VERSION_MAJOR > 7
         auto truetape =
             BuilderZ.CreateLoad(fnandtapetype->tapeType, tapep, "tapeld");

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2355,17 +2355,8 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
           gutils->newFunc->getParent()->getDataLayout().getTypeAllocSizeInBits(
               tapeType);
       if (size != 0) {
-        auto i64 = Type::getInt64Ty(gutils->newFunc->getContext());
-        BasicBlock *BB = BasicBlock::Create(gutils->newFunc->getContext(),
-                                            "entry", gutils->newFunc);
-        IRBuilder<> B(BB);
-
-        CallInst *malloccall;
-        CreateAllocation(B, tapeType, ConstantInt::get(i64, 1), "tapemem",
-                         &malloccall, nullptr);
         RetTypes[returnMapping.find(AugmentedStruct::Tape)->second] =
-            malloccall->getType();
-        BB->eraseFromParent();
+            getDefaultAnonymousTapeType(gutils->newFunc->getContext());
       }
     }
   }
@@ -2485,13 +2476,13 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       if (size != 0) {
         CallInst *malloccall = nullptr;
         Instruction *zero = nullptr;
-        tapeMemory =
-            CreateAllocation(ib, tapeType, ConstantInt::get(i64, 1), "tapemem",
-                             &malloccall, EnzymeZeroCache ? &zero : nullptr);
+        tapeMemory = CreateAllocation(
+            ib, tapeType, ConstantInt::get(i64, 1), "tapemem", &malloccall,
+            EnzymeZeroCache ? &zero : nullptr, /*isDefault*/ true);
         memory = malloccall;
       } else {
-        memory =
-            ConstantPointerNull::get(Type::getInt8PtrTy(NewF->getContext()));
+        memory = ConstantPointerNull::get(
+            getDefaultAnonymousTapeType(NewF->getContext()));
       }
       Value *Idxs[] = {
           ib.getInt32(0),

enzyme/Enzyme/FunctionUtils.cpp

@@ -1942,7 +1942,7 @@ FunctionType *getFunctionTypeForClone(
       returnValue == ReturnType::TapeAndReturn ||
       returnValue == ReturnType::Tape) {
     RetTypes.clear();
-    RetTypes.push_back(Type::getInt8PtrTy(FTy->getContext()));
+    RetTypes.push_back(getDefaultAnonymousTapeType(FTy->getContext()));
     if (returnValue == ReturnType::TapeAndTwoReturns) {
       RetTypes.push_back(FTy->getReturnType());
       RetTypes.push_back(

enzyme/Enzyme/Utils.cpp

@@ -48,12 +48,13 @@ void (*CustomErrorHandler)(const char *, LLVMValueRef, ErrorType,
                            void *) = nullptr;
 LLVMValueRef (*CustomAllocator)(LLVMBuilderRef, LLVMTypeRef,
                                 /*Count*/ LLVMValueRef,
-                                /*Align*/ LLVMValueRef) = nullptr;
+                                /*Align*/ LLVMValueRef, uint8_t) = nullptr;
 LLVMValueRef (*CustomDeallocator)(LLVMBuilderRef, LLVMValueRef) = nullptr;
 void (*CustomRuntimeInactiveError)(LLVMBuilderRef, LLVMValueRef,
                                    LLVMValueRef) = nullptr;
 LLVMValueRef (*EnzymePostCacheStore)(LLVMValueRef, LLVMBuilderRef,
                                      LLVMValueRef *) = nullptr;
+LLVMTypeRef (*EnzymeDefaultTapeType)(LLVMContextRef) = nullptr;
 }
 
 llvm::SmallVector<llvm::Instruction *, 2> PostCacheStore(llvm::StoreInst *SI,
@@ -70,6 +71,12 @@ llvm::SmallVector<llvm::Instruction *, 2> PostCacheStore(llvm::StoreInst *SI,
   return res;
 }
 
+llvm::PointerType *getDefaultAnonymousTapeType(llvm::LLVMContext &C) {
+  if (EnzymeDefaultTapeType)
+    return cast<PointerType>(unwrap(EnzymeDefaultTapeType(wrap(&C))));
+  return Type::getInt8PtrTy(C);
+}
+
 Function *getOrInsertExponentialAllocator(Module &M, Function *newFunc,
                                           bool ZeroInit, llvm::Type *RT) {
   bool custom = true;
@@ -239,15 +246,16 @@ llvm::Value *CreateReAllocation(llvm::IRBuilder<> &B, llvm::Value *prev,
 }
 
 Value *CreateAllocation(IRBuilder<> &Builder, llvm::Type *T, Value *Count,
-                        Twine Name, CallInst **caller, Instruction **ZeroMem) {
+                        Twine Name, CallInst **caller, Instruction **ZeroMem,
+                        bool isDefault) {
   Value *res;
   auto &M = *Builder.GetInsertBlock()->getParent()->getParent();
   auto AlignI = M.getDataLayout().getTypeAllocSizeInBits(T) / 8;
   auto Align = ConstantInt::get(Count->getType(), AlignI);
   CallInst *malloccall = nullptr;
   if (CustomAllocator) {
-    res = unwrap(
-        CustomAllocator(wrap(&Builder), wrap(T), wrap(Count), wrap(Align)));
+    res = unwrap(CustomAllocator(wrap(&Builder), wrap(T), wrap(Count),
+                                 wrap(Align), isDefault));
     if (auto I = dyn_cast<Instruction>(res))
       I->setName(Name);
 

enzyme/Enzyme/Utils.h

@@ -86,7 +86,8 @@ llvm::SmallVector<llvm::Instruction *, 2> PostCacheStore(llvm::StoreInst *SI,
 llvm::Value *CreateAllocation(llvm::IRBuilder<> &B, llvm::Type *T,
                               llvm::Value *Count, llvm::Twine Name = "",
                               llvm::CallInst **caller = nullptr,
-                              llvm::Instruction **ZeroMem = nullptr);
+                              llvm::Instruction **ZeroMem = nullptr,
+                              bool isDefault = false);
 llvm::CallInst *CreateDealloc(llvm::IRBuilder<> &B, llvm::Value *ToFree);
 
 llvm::Value *CreateReAllocation(llvm::IRBuilder<> &B, llvm::Value *prev,
@@ -95,6 +96,8 @@ llvm::Value *CreateReAllocation(llvm::IRBuilder<> &B, llvm::Value *prev,
                                 llvm::CallInst **caller = nullptr,
                                 bool ZeroMem = false);
 
+llvm::PointerType *getDefaultAnonymousTapeType(llvm::LLVMContext &C);
+
 extern std::map<std::string, std::function<llvm::Value *(
                                  llvm::IRBuilder<> &, llvm::CallInst *,
                                  llvm::ArrayRef<llvm::Value *>)>>