Adapt to upstream

EnzymeAD · Feb 6, 2025 · 73e48de · 73e48de
1 parent df197be
commit 73e48de
Show file tree

Hide file tree

Showing 11 changed files with 142 additions and 85 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -161,16 +161,14 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
     AttributeList AL;
     AL = AL.addParamAttribute(DT->getContext(), 0,
                               Attribute::AttrKind::ReadOnly);
-    AL = AL.addParamAttribute(DT->getContext(), 0,
-                              Attribute::AttrKind::NoCapture);
+    AL = addFunctionNoCapture(DT->getContext(), AL, 0);
     AL =
         AL.addParamAttribute(DT->getContext(), 0, Attribute::AttrKind::NoAlias);
     AL =
         AL.addParamAttribute(DT->getContext(), 0, Attribute::AttrKind::NonNull);
     AL = AL.addParamAttribute(DT->getContext(), 1,
                               Attribute::AttrKind::WriteOnly);
-    AL = AL.addParamAttribute(DT->getContext(), 1,
-                              Attribute::AttrKind::NoCapture);
+    AL = addFunctionNoCapture(DT->getContext(), AL, 1);
     AL =
         AL.addParamAttribute(DT->getContext(), 1, Attribute::AttrKind::NoAlias);
     AL =
@@ -208,11 +206,11 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
     auto alloc = IRBuilder<>(gutils->inversionAllocs).CreateAlloca(rankTy);
     AttributeList AL;
     AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::ReadOnly);
-    AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NoCapture);
+    AL = addFunctionNoCapture(context, AL, 0);
     AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NoAlias);
     AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NonNull);
     AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::WriteOnly);
-    AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NoCapture);
+    AL = addFunctionNoCapture(context, AL, 1);
     AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NoAlias);
     AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NonNull);
     AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
@@ -241,11 +239,11 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
     auto alloc = IRBuilder<>(gutils->inversionAllocs).CreateAlloca(rankTy);
     AttributeList AL;
     AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::ReadOnly);
-    AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NoCapture);
+    AL = addFunctionNoCapture(context, AL, 0);
     AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NoAlias);
     AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NonNull);
     AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::WriteOnly);
-    AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NoCapture);
+    AL = addFunctionNoCapture(context, AL, 1);
     AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NoAlias);
     AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NonNull);
     AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,

diff --git a/enzyme/Enzyme/CacheUtility.cpp b/enzyme/Enzyme/CacheUtility.cpp
@@ -194,8 +194,8 @@ std::pair<PHINode *, Instruction *> FindCanonicalIV(Loop *L, Type *Ty) {
       continue;
     if (!Inc)
       continue;
-    if (Inc != Header->getFirstNonPHIOrDbg())
-      Inc->moveBefore(Header->getFirstNonPHIOrDbg());
+    if (Inc != getFirstNonPHIOrDbg(Header))
+      Inc->moveBefore(getFirstNonPHIOrDbg(Header));
     return std::make_pair(PN, Inc);
   }
   llvm::errs() << *Header << "\n";

diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp
@@ -125,7 +125,7 @@ bool attributeKnownFunctions(llvm::Function &F) {
   if (F.getName() == "fprintf") {
     for (auto &arg : F.args()) {
       if (arg.getType()->isPointerTy()) {
-        arg.addAttr(Attribute::NoCapture);
+        addFunctionNoCapture(&F, arg.getArgNo());
         changed = true;
       }
     }
@@ -148,7 +148,7 @@ bool attributeKnownFunctions(llvm::Function &F) {
       for (auto &arg : F.args()) {
         if (arg.getType()->isPointerTy()) {
           arg.addAttr(Attribute::ReadNone);
-          arg.addAttr(Attribute::NoCapture);
+          addFunctionNoCapture(&F, arg.getArgNo());
         }
       }
   }
@@ -168,7 +168,7 @@ bool attributeKnownFunctions(llvm::Function &F) {
     F.addFnAttr(Attribute::NoSync);
     for (int i = 0; i < 2; i++)
       if (F.getFunctionType()->getParamType(i)->isPointerTy()) {
-        F.addParamAttr(i, Attribute::NoCapture);
+        addFunctionNoCapture(&F, i);
         F.addParamAttr(i, Attribute::WriteOnly);
       }
   }
@@ -192,7 +192,7 @@ bool attributeKnownFunctions(llvm::Function &F) {
     F.addFnAttr(Attribute::NoSync);
     F.addParamAttr(0, Attribute::WriteOnly);
     if (F.getFunctionType()->getParamType(2)->isPointerTy()) {
-      F.addParamAttr(2, Attribute::NoCapture);
+      addFunctionNoCapture(&F, 2);
       F.addParamAttr(2, Attribute::WriteOnly);
     }
     F.addParamAttr(6, Attribute::WriteOnly);
@@ -211,7 +211,7 @@ bool attributeKnownFunctions(llvm::Function &F) {
     F.addFnAttr(Attribute::NoSync);
     F.addParamAttr(0, Attribute::ReadOnly);
     if (F.getFunctionType()->getParamType(2)->isPointerTy()) {
-      F.addParamAttr(2, Attribute::NoCapture);
+      addFunctionNoCapture(&F, 2);
       F.addParamAttr(2, Attribute::ReadOnly);
     }
     F.addParamAttr(6, Attribute::WriteOnly);
@@ -231,12 +231,12 @@ bool attributeKnownFunctions(llvm::Function &F) {
     F.addFnAttr(Attribute::NoSync);
 
     if (F.getFunctionType()->getParamType(0)->isPointerTy()) {
-      F.addParamAttr(0, Attribute::NoCapture);
+      addFunctionNoCapture(&F, 0);
       F.addParamAttr(0, Attribute::ReadOnly);
     }
     if (F.getFunctionType()->getParamType(1)->isPointerTy()) {
       F.addParamAttr(1, Attribute::WriteOnly);
-      F.addParamAttr(1, Attribute::NoCapture);
+      addFunctionNoCapture(&F, 1);
     }
   }
   if (F.getName() == "MPI_Wait" || F.getName() == "PMPI_Wait") {
@@ -246,9 +246,9 @@ bool attributeKnownFunctions(llvm::Function &F) {
     F.addFnAttr(Attribute::WillReturn);
     F.addFnAttr(Attribute::NoFree);
     F.addFnAttr(Attribute::NoSync);
-    F.addParamAttr(0, Attribute::NoCapture);
+    addFunctionNoCapture(&F, 0);
     F.addParamAttr(1, Attribute::WriteOnly);
-    F.addParamAttr(1, Attribute::NoCapture);
+    addFunctionNoCapture(&F, 1);
   }
   if (F.getName() == "MPI_Waitall" || F.getName() == "PMPI_Waitall") {
     changed = true;
@@ -257,9 +257,9 @@ bool attributeKnownFunctions(llvm::Function &F) {
     F.addFnAttr(Attribute::WillReturn);
     F.addFnAttr(Attribute::NoFree);
     F.addFnAttr(Attribute::NoSync);
-    F.addParamAttr(1, Attribute::NoCapture);
+    addFunctionNoCapture(&F, 1);
     F.addParamAttr(2, Attribute::WriteOnly);
-    F.addParamAttr(2, Attribute::NoCapture);
+    addFunctionNoCapture(&F, 2);
   }
   // Map of MPI function name to the arg index of its type argument
   std::map<std::string, int> MPI_TYPE_ARGS = {
@@ -2347,7 +2347,7 @@ class EnzymeBase {
           for (size_t i = 0; i < num_args; ++i) {
             if (CI->getArgOperand(i)->getType()->isPointerTy()) {
               CI->addParamAttr(i, Attribute::ReadNone);
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2361,7 +2361,7 @@ class EnzymeBase {
           for (size_t i = 0; i < num_args; ++i) {
             if (CI->getArgOperand(i)->getType()->isPointerTy()) {
               CI->addParamAttr(i, Attribute::ReadNone);
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2375,7 +2375,7 @@ class EnzymeBase {
           for (size_t i = 0; i < num_args; ++i) {
             if (CI->getArgOperand(i)->getType()->isPointerTy()) {
               CI->addParamAttr(i, Attribute::ReadNone);
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2389,7 +2389,7 @@ class EnzymeBase {
           for (size_t i = 0; i < num_args; ++i) {
             if (CI->getArgOperand(i)->getType()->isPointerTy()) {
               CI->addParamAttr(i, Attribute::ReadNone);
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2439,9 +2439,9 @@ class EnzymeBase {
           CI->addAttribute(AttributeList::FunctionIndex, Attribute::ReadOnly);
 #endif
           CI->addParamAttr(1, Attribute::ReadOnly);
-          CI->addParamAttr(1, Attribute::NoCapture);
+          addCallSiteNoCapture(CI, 1);
           CI->addParamAttr(3, Attribute::ReadOnly);
-          CI->addParamAttr(3, Attribute::NoCapture);
+          addCallSiteNoCapture(CI, 3);
         }
         if (Fn->getName() == "frexp" || Fn->getName() == "frexpf" ||
             Fn->getName() == "frexpl") {
@@ -2502,7 +2502,7 @@ class EnzymeBase {
           for (size_t i : {0, 1}) {
             if (i < num_args &&
                 CI->getArgOperand(i)->getType()->isPointerTy()) {
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2527,7 +2527,7 @@ class EnzymeBase {
           for (size_t i : {0, 2}) {
             if (i < num_args &&
                 CI->getArgOperand(i)->getType()->isPointerTy()) {
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2553,7 +2553,7 @@ class EnzymeBase {
           for (size_t i : {0, 1, 2, 3}) {
             if (i < num_args &&
                 CI->getArgOperand(i)->getType()->isPointerTy()) {
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2579,7 +2579,7 @@ class EnzymeBase {
           for (size_t i : {0}) {
             if (i < num_args &&
                 CI->getArgOperand(i)->getType()->isPointerTy()) {
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }
@@ -2601,7 +2601,7 @@ class EnzymeBase {
           for (size_t i = 0; i < num_args; ++i) {
             if (CI->getArgOperand(i)->getType()->isPointerTy()) {
               CI->addParamAttr(i, Attribute::ReadOnly);
-              CI->addParamAttr(i, Attribute::NoCapture);
+              addCallSiteNoCapture(CI, i);
             }
           }
         }

diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -2713,9 +2713,17 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
   auto i = nf->arg_begin(), j = NewF->arg_begin();
   while (i != nf->arg_end()) {
     VMap[i] = j;
-    if (nf->hasParamAttribute(attrIndex, Attribute::NoCapture)) {
-      NewF->addParamAttr(attrIndex, Attribute::NoCapture);
+#if LLVM_VERSION_MAJOR > 20
+    if (nf->hasParamAttribute(attrIndex, Attribute::Captures)) {
+      NewF->addParamAttr(attrIndex,
+                         nf->getParamAttribute(attrIndex, Attribute::Captures));
     }
+#else
+    if (nf->hasParamAttribute(attrIndex, Attribute::NoCaptures)) {
+      NewF->addParamAttr(
+          attrIndex, nf->getParamAttribute(attrIndex, Attribute::NoCapture));
+    }
+#endif
     if (nf->hasParamAttribute(attrIndex, Attribute::NoAlias)) {
       NewF->addParamAttr(attrIndex, Attribute::NoAlias);
     }

diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp
@@ -698,7 +698,7 @@ OldAllocationSize(Value *Ptr, CallInst *Loc, Function *NewF, IntegerType *T,
     AttributeList list;
     list = list.addFnAttribute(NewF->getContext(), Attribute::ReadOnly);
     list = list.addParamAttribute(NewF->getContext(), 0, Attribute::ReadNone);
-    list = list.addParamAttribute(NewF->getContext(), 0, Attribute::NoCapture);
+    addFunctionNoCapture(NewF->getContext(), list, 0);
     auto allocSize = NewF->getParent()->getOrInsertFunction(
         allocName,
         FunctionType::get(
@@ -1109,7 +1109,7 @@ static void SimplifyMPIQueries(Function &NewF, FunctionAnalysisManager &FAM) {
         B.SetInsertPoint(Bound->getNextNode());
       }
       B.CreateStore(B.CreateLoad(AI2->getAllocatedType(), AI2), AI);
-      Bound->addParamAttr(i, Attribute::NoCapture);
+      addCallSiteNoCapture(Bound, i);
     }
   }
   PreservedAnalyses PA;

diff --git a/enzyme/Enzyme/TraceGenerator.cpp b/enzyme/Enzyme/TraceGenerator.cpp
@@ -63,7 +63,7 @@ void TraceGenerator::visitFunction(Function &F) {
     return;
 
   auto fn = tutils->newFunc;
-  auto entry = fn->getEntryBlock().getFirstNonPHIOrDbgOrLifetime();
+  auto entry = getFirstNonPHIOrDbgOrLifetime(&fn->getEntryBlock());
 
   while (isa<AllocaInst>(entry) && entry->getNextNode()) {
     entry = entry->getNextNode();

diff --git a/enzyme/Enzyme/TraceInterface.cpp b/enzyme/Enzyme/TraceInterface.cpp
@@ -311,7 +311,7 @@ DynamicTraceInterface::DynamicTraceInterface(Value *dynamicInterface,
   assert(dynamicInterface);
 
   auto &M = *F->getParent();
-  IRBuilder<> Builder(F->getEntryBlock().getFirstNonPHIOrDbg());
+  IRBuilder<> Builder(getFirstNonPHIOrDbg(&F->getEntryBlock()));
 
   getTraceFunction = MaterializeInterfaceFunction(
       Builder, dynamicInterface, getTraceTy(), 0, M, "get_trace");

diff --git a/enzyme/Enzyme/TraceUtils.cpp b/enzyme/Enzyme/TraceUtils.cpp
@@ -196,11 +196,8 @@ TraceUtils::ValueToVoidPtrAndSize(IRBuilder<> &Builder, Value *val,
         Builder.CreateIntToPtr(cast, getInt8PtrTy(cast->getContext()));
     return {retval, ConstantInt::get(size_type, valsize / 8)};
   } else {
-    auto insertPoint = Builder.GetInsertBlock()
-                           ->getParent()
-                           ->getEntryBlock()
-                           .getFirstNonPHIOrDbgOrLifetime();
-    IRBuilder<> AllocaBuilder(insertPoint);
+    IRBuilder<> AllocaBuilder(getFirstNonPHIOrDbgOrLifetime(
+        &Builder.GetInsertBlock()->getParent()->getEntryBlock()));
     auto alloca = AllocaBuilder.CreateAlloca(val->getType(), nullptr,
                                              val->getName() + ".ptr");
     Builder.CreateStore(val, alloca);
@@ -248,7 +245,8 @@ CallInst *TraceUtils::InsertChoice(IRBuilder<> &Builder, Value *address,
   auto call = Builder.CreateCall(interface->insertChoiceTy(),
                                  interface->insertChoice(Builder), args);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+
+  addCallSiteNoCapture(call, 1);
   return call;
 }
 
@@ -259,7 +257,7 @@ CallInst *TraceUtils::InsertCall(IRBuilder<> &Builder, Value *address,
   auto call = Builder.CreateCall(interface->insertCallTy(),
                                  interface->insertCall(Builder), args);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
 #if LLVM_VERSION_MAJOR >= 14
   call->addAttributeAtIndex(
       AttributeList::FunctionIndex,
@@ -283,7 +281,7 @@ CallInst *TraceUtils::InsertArgument(IRBuilder<> &Builder, Value *name,
   auto call = Builder.CreateCall(interface->insertArgumentTy(),
                                  interface->insertArgument(Builder), args);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return call;
 }
 
@@ -322,7 +320,7 @@ CallInst *TraceUtils::InsertChoiceGradient(IRBuilder<> &Builder,
 
   auto call = Builder.CreateCall(interface_type, interface_function, args);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return call;
 }
 
@@ -339,7 +337,7 @@ CallInst *TraceUtils::InsertArgumentGradient(IRBuilder<> &Builder,
 
   auto call = Builder.CreateCall(interface_type, interface_function, args);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return call;
 }
 
@@ -352,16 +350,14 @@ CallInst *TraceUtils::GetTrace(IRBuilder<> &Builder, Value *address,
   auto call = Builder.CreateCall(interface->getTraceTy(),
                                  interface->getTrace(Builder), args, Name);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return call;
 }
 
 Instruction *TraceUtils::GetChoice(IRBuilder<> &Builder, Value *address,
                                    Type *choiceType, const Twine &Name) {
-  IRBuilder<> AllocaBuilder(Builder.GetInsertBlock()
-                                ->getParent()
-                                ->getEntryBlock()
-                                .getFirstNonPHIOrDbgOrLifetime());
+  IRBuilder<> AllocaBuilder(getFirstNonPHIOrDbgOrLifetime(
+      &Builder.GetInsertBlock()->getParent()->getEntryBlock()));
   AllocaInst *store_dest =
       AllocaBuilder.CreateAlloca(choiceType, nullptr, Name + ".ptr");
   auto preallocated_size = choiceType->getPrimitiveSizeInBits() / 8;
@@ -385,7 +381,7 @@ Instruction *TraceUtils::GetChoice(IRBuilder<> &Builder, Value *address,
                      Attribute::get(call->getContext(), "enzyme_inactive"));
 #endif
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return Builder.CreateLoad(choiceType, store_dest, "from.trace." + Name);
 }
 
@@ -396,7 +392,7 @@ Instruction *TraceUtils::HasChoice(IRBuilder<> &Builder, Value *address,
   auto call = Builder.CreateCall(interface->hasChoiceTy(),
                                  interface->hasChoice(Builder), args, Name);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return call;
 }
 
@@ -407,7 +403,7 @@ Instruction *TraceUtils::HasCall(IRBuilder<> &Builder, Value *address,
   auto call = Builder.CreateCall(interface->hasCallTy(),
                                  interface->hasCall(Builder), args, Name);
   call->addParamAttr(1, Attribute::ReadOnly);
-  call->addParamAttr(1, Attribute::NoCapture);
+  addCallSiteNoCapture(call, 1);
   return call;
 }