Julia custom alloc fixups (rust-lang#780)

* Julia custom alloc fixups

* Fix tape usage in rev

* fix

Author: wsmoses

enzyme/Enzyme/CApi.cpp

@@ -588,11 +588,22 @@ const char *EnzymeGradientUtilsInvertedPointersToString(GradientUtils *gutils,
 
 void EnzymeStringFree(const char *cstr) { delete[] cstr; }
 
-void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2) {
+void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2,
+                      LLVMBuilderRef B) {
   Instruction *I1 = cast<Instruction>(unwrap(inst1));
   Instruction *I2 = cast<Instruction>(unwrap(inst2));
-  if (I1 != I2)
+  if (I1 != I2) {
+    if (B != nullptr) {
+      IRBuilder<> &BR = *unwrap(B);
+      if (I1->getIterator() == BR.GetInsertPoint()) {
+        if (I2->getNextNode() == nullptr)
+          BR.SetInsertPoint(I1->getParent());
+        else
+          BR.SetInsertPoint(I1->getNextNode());
+      }
+    }
     I1->moveBefore(I2);
+  }
 }
 
 void EnzymeSetMustCache(LLVMValueRef inst1) {

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2312,10 +2312,6 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
     report_fatal_error("function failed verification (2)");
   }
 
-  StructType *sty = cast<StructType>(gutils->newFunc->getReturnType());
-  SmallVector<Type *, 4> RetTypes(sty->elements().begin(),
-                                  sty->elements().end());
-
   SmallVector<Type *, 4> MallocTypes;
 
   for (auto a : gutils->getTapeValues()) {
@@ -2338,6 +2334,30 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       forceAnonymousTape;
   bool noTape = MallocTypes.size() == 0 && !forceAnonymousTape;
 
+  StructType *sty = cast<StructType>(gutils->newFunc->getReturnType());
+  SmallVector<Type *, 4> RetTypes(sty->elements().begin(),
+                                  sty->elements().end());
+  if (!noTape) {
+    if (recursive && !omp) {
+      auto size =
+          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();
+      }
+    }
+  }
+
   int oldretIdx = -1;
   if (returnMapping.find(AugmentedStruct::Return) != returnMapping.end()) {
     oldretIdx = returnMapping[AugmentedStruct::Return];
@@ -2367,8 +2387,6 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
     }
     RetTypes.erase(RetTypes.begin() + tidx);
   } else if (recursive) {
-    assert(RetTypes[returnMapping.find(AugmentedStruct::Tape)->second] ==
-           Type::getInt8PtrTy(nf->getContext()));
   } else {
     RetTypes[returnMapping.find(AugmentedStruct::Tape)->second] = tapeType;
   }
@@ -2440,7 +2458,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
 
   IRBuilder<> ib(NewF->getEntryBlock().getFirstNonPHI());
 
-  Value *ret = noReturn ? nullptr : ib.CreateAlloca(RetType);
+  AllocaInst *ret = noReturn ? nullptr : ib.CreateAlloca(RetType);
   if (!noReturn && EnzymeZeroCache) {
     ib.CreateStore(Constant::getNullValue(RetType), ret);
   }
@@ -2455,10 +2473,9 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       if (size != 0) {
         CallInst *malloccall = nullptr;
         Instruction *zero = nullptr;
-        IRBuilder<> Builder(NewF->getEntryBlock().getFirstNonPHI());
-        tapeMemory = CreateAllocation(
-            Builder, tapeType, ConstantInt::get(i64, 1), "tapemem", &malloccall,
-            EnzymeZeroCache ? &zero : nullptr);
+        tapeMemory =
+            CreateAllocation(ib, tapeType, ConstantInt::get(i64, 1), "tapemem",
+                             &malloccall, EnzymeZeroCache ? &zero : nullptr);
         memory = malloccall;
       } else {
         memory =
@@ -3690,7 +3707,10 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
       IRBuilder<> BuilderZ(gutils->inversionAllocs);
       if (!augmenteddata->tapeType->isEmptyTy()) {
         auto tapep = BuilderZ.CreatePointerCast(
-            additionalValue, PointerType::getUnqual(augmenteddata->tapeType));
+            additionalValue,
+            PointerType::get(augmenteddata->tapeType,
+                             cast<PointerType>(additionalValue->getType())
+                                 ->getAddressSpace()));
 #if LLVM_VERSION_MAJOR > 7
         LoadInst *truetape =
             BuilderZ.CreateLoad(augmenteddata->tapeType, tapep, "truetape");
@@ -3701,7 +3721,7 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
                               MDNode::get(truetape->getContext(), {}));
 
         if (!omp && gutils->FreeMemory) {
-          CreateDealloc(BuilderZ, tapep);
+          CreateDealloc(BuilderZ, additionalValue);
         }
         additionalValue = truetape;
       } else {

enzyme/Enzyme/FunctionUtils.cpp

@@ -379,11 +379,14 @@ void RecursivelyReplaceAddressSpace(Value *AI, Value *rep, bool legal) {
 /// pass. Specifically if we're not topLevel all allocations must be upgraded
 /// Even if topLevel any allocations that aren't in the entry block (and
 /// therefore may not be reachable in the reverse pass) must be upgraded.
-static inline void UpgradeAllocasToMallocs(Function *NewF,
-                                           DerivativeMode mode) {
+static inline void
+UpgradeAllocasToMallocs(Function *NewF, DerivativeMode mode,
+                        SmallPtrSetImpl<llvm::BasicBlock *> &Unreachable) {
   SmallVector<AllocaInst *, 4> ToConvert;
 
   for (auto &BB : *NewF) {
+    if (Unreachable.count(&BB))
+      continue;
     for (auto &I : BB) {
       if (auto AI = dyn_cast<AllocaInst>(&I)) {
         bool UsableEverywhere = AI->getParent() == &NewF->getEntryBlock();
@@ -1653,7 +1656,8 @@ Function *PreProcessCache::preprocessForClone(Function *F,
       mode == DerivativeMode::ReverseModeCombined) {
     // For subfunction calls upgrade stack allocations to mallocs
     // to ensure availability in the reverse pass
-    UpgradeAllocasToMallocs(NewF, mode);
+    auto unreachable = getGuaranteedUnreachable(NewF);
+    UpgradeAllocasToMallocs(NewF, mode, unreachable);
   }
 
   CanonicalizeLoops(NewF, FAM);

enzyme/Enzyme/Utils.cpp

@@ -128,9 +128,14 @@ Value *CreateAllocation(IRBuilder<> &Builder, llvm::Type *T, Value *Count,
     *caller = malloccall;
   }
   if (ZeroMem) {
+    auto PT = cast<PointerType>(malloccall->getType());
+    Value *tozero = malloccall;
+    if (!PT->getElementType()->isIntegerTy(8))
+      tozero = Builder.CreatePointerCast(
+          tozero, PointerType::get(Type::getInt8Ty(PT->getContext()),
+                                   PT->getAddressSpace()));
     Value *args[] = {
-        malloccall,
-        ConstantInt::get(Type::getInt8Ty(malloccall->getContext()), 0),
+        tozero, ConstantInt::get(Type::getInt8Ty(malloccall->getContext()), 0),
         Builder.CreateMul(Align, Count, "", true, true),
         ConstantInt::getFalse(malloccall->getContext())};
     Type *tys[] = {args[0]->getType(), args[2]->getType()};

enzyme/test/Enzyme/ReverseMode/loadcall.ll

@@ -72,8 +72,7 @@ attributes #4 = { nounwind }
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   %0 = bitcast i8* %tapeArg to double**
 ; CHECK-NEXT:   %[[callp:.+]] = load double*, double** %0{{(, align 8)?}}, !enzyme_mustcache !
-; CHECK-NEXT:   %[[t0:.+]] = bitcast double** %0 to i8*
-; CHECK-NEXT:   tail call void @free(i8* nonnull %[[t0]])
+; CHECK-NEXT:   tail call void @free(i8* nonnull %tapeArg)
 ; CHECK-NEXT:   %"loaded'de" = alloca double
 ; CHECK-NEXT:   store double 0.000000e+00, double* %"loaded'de"
 ; CHECK-NEXT:   br label %invertentry