add enzyme_noatomic attribute

enzyme/Enzyme/DiffeGradientUtils.cpp

@@ -51,6 +51,26 @@
 
 using namespace llvm;
 
+namespace {
+bool noAtomicsForContext(const Instruction *orig, const Value *origptr) {
+  if (orig) {
+    if (const Function *F = orig->getFunction()) {
+      if (F->hasFnAttribute("enzyme_noatomic")) {
+        return true;
+      }
+    }
+  }
+  const Value *base = getBaseObject(origptr);
+  if (auto *arg = dyn_cast<Argument>(base)) {
+    if (const Function *F = arg->getParent()) {
+      return F->getAttributes().hasParamAttr(arg->getArgNo(),
+                                             "enzyme_noatomic");
+    }
+  }
+  return false;
+}
+} // namespace
+
 DiffeGradientUtils::DiffeGradientUtils(
     EnzymeLogic &Logic, Function *newFunc_, Function *oldFunc_,
     TargetLibraryInfo &TLI, TypeAnalysis &TA, TypeResults TR,
@@ -918,6 +938,8 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig,
   // all additional parallelism in this function is outlined.
   if (backwardsOnlyShadows.find(TmpOrig) != backwardsOnlyShadows.end())
     Atomic = false;
+  if (Atomic && noAtomicsForContext(orig, origptr))
+    Atomic = false;
 
   if (Atomic) {
     // For amdgcn constant AS is 4 and if the primal is in it we need to cast

enzyme/Enzyme/PreserveNVVM.cpp

@@ -372,6 +372,15 @@ bool preserveNVVM(bool Begin, Module &M) {
               continue;
             }
 
+            if (AS == "enzyme_noatomic" && Func) {
+              Func->addAttribute(
+                  AttributeList::FunctionIndex,
+                  Attribute::get(Func->getContext(), "enzyme_noatomic"));
+              changed = true;
+              replacements.push_back(Constant::getNullValue(CAOp->getType()));
+              continue;
+            }
+
             if (AS == "enzyme_shouldrecompute" && Func) {
               Func->addAttribute(
                   AttributeList::FunctionIndex,