fix segfault

Author: wsmoses

enzyme/Enzyme/EnzymeLogic.cpp

@@ -398,6 +398,19 @@ std::pair<Function*,StructType*> CreateAugmentedPrimal(Function* todiff, AAResul
   if (gutils->newFunc->hasFnAttribute(Attribute::OptimizeNone))
     gutils->newFunc->removeFnAttr(Attribute::OptimizeNone);
 
+  if (auto bytes = gutils->newFunc->getDereferenceableBytes(llvm::AttributeList::ReturnIndex)) {
+    AttrBuilder ab;
+    ab.addDereferenceableAttr(bytes);
+    gutils->newFunc->removeAttributes(llvm::AttributeList::ReturnIndex, ab);
+  }
+
+  if (gutils->newFunc->hasAttribute(llvm::AttributeList::ReturnIndex, llvm::Attribute::NoAlias)) {
+    gutils->newFunc->removeAttribute(llvm::AttributeList::ReturnIndex, llvm::Attribute::NoAlias);
+  }
+  if (gutils->newFunc->hasAttribute(llvm::AttributeList::ReturnIndex, llvm::Attribute::ZExt)) {
+    gutils->newFunc->removeAttribute(llvm::AttributeList::ReturnIndex, llvm::Attribute::ZExt);
+  }
+
   if (llvm::verifyFunction(*gutils->newFunc, &llvm::errs())) {
       llvm::errs() << *gutils->oldFunc << "\n";
       llvm::errs() << *gutils->newFunc << "\n";

enzyme/Enzyme/GradientUtils.cpp

@@ -955,6 +955,7 @@ void GradientUtils::branchToCorrespondingTarget(BasicBlock* ctx, IRBuilder <>& B
 
       Value* phi = lookupValueFromCache(BuilderM, ctx, cache);
 
+
       if (replacePHIs == nullptr) {
           SwitchInst* swtch = BuilderM.CreateSwitch(phi, *done[std::make_pair(block, si->getDefaultDest())].begin());
           for (auto switchcase : si->cases()) {
@@ -974,6 +975,12 @@ void GradientUtils::branchToCorrespondingTarget(BasicBlock* ctx, IRBuilder <>& B
                   }
                   val = BuilderM.CreateNot(val);
               }
+              if (&*BuilderM.GetInsertPoint() == pair.second) {
+                if (pair.second->getNextNode())
+                  BuilderM.SetInsertPoint(pair.second->getNextNode());
+                else
+                  BuilderM.SetInsertPoint(pair.second->getParent());
+              }
               pair.second->replaceAllUsesWith(val);
               pair.second->eraseFromParent();
           }
@@ -1049,6 +1056,12 @@ void GradientUtils::branchToCorrespondingTarget(BasicBlock* ctx, IRBuilder <>& B
           } else {
               val = BuilderM.CreateICmpEQ(ConstantInt::get(T, i), which);
           }
+          if (&*BuilderM.GetInsertPoint() == found->second) {
+            if (found->second->getNextNode())
+              BuilderM.SetInsertPoint(found->second->getNextNode());
+            else
+              BuilderM.SetInsertPoint(found->second->getParent());
+          }
           found->second->replaceAllUsesWith(val);
           found->second->eraseFromParent();
       }