ActivityAnalysis: consider atomicrmw (#2205)

* ActivityAnalysis: consider atomicrmw * fix * fix
EnzymeAD · Dec 23, 2024 · 1ca3ceb · 1ca3ceb
1 parent 79ac406
commit 1ca3ceb
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -1554,6 +1554,26 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
             ReEvaluateValueIfInactiveValue[II->getOperand(0)].insert(TmpOrig);
           }
         }
+      } else if (auto RMW = dyn_cast<AtomicRMWInst>(TmpOrig)) {
+        if (directions == UP) {
+          if (isConstantValue(TR, RMW->getPointerOperand())) {
+            InsertConstantValue(TR, Val);
+            return true;
+          }
+        } else {
+          if (UpHypothesis->isConstantValue(TR, RMW->getPointerOperand())) {
+            InsertConstantValue(TR, Val);
+            insertConstantsFrom(TR, *UpHypothesis);
+            return true;
+          }
+        }
+        if (EnzymeEnableRecursiveHypotheses) {
+          ReEvaluateValueIfInactiveValue[RMW->getPointerOperand()].insert(Val);
+          if (TmpOrig != Val) {
+            ReEvaluateValueIfInactiveValue[RMW->getPointerOperand()].insert(
+                TmpOrig);
+          }
+        }
       } else if (auto op = dyn_cast<CallInst>(TmpOrig)) {
         if (isInactiveCall(*op) || op->hasFnAttr("enzyme_inactive_val") ||
             op->getAttributes().hasAttribute(llvm::AttributeList::ReturnIndex,
@@ -1940,7 +1960,7 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
           isRefSet(AARes)) {
         if (EnzymePrintActivity)
           llvm::errs() << "potential active load: " << *I << "\n";
-        if (isa<LoadInst>(I) || isNVLoad(I)) {
+        if (isa<LoadInst>(I) || isNVLoad(I) || isa<AtomicRMWInst>(I)) {
           // If the ref'ing value is a load check if the loaded value is
           // active
           if (!Hypothesis->isConstantValue(TR, I)) {

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -3822,6 +3822,9 @@ bool GradientUtils::legalRecompute(const Value *val,
     }
   }
 
+  if (isa<AtomicRMWInst>(val))
+    return false;
+
   if (auto phi = dyn_cast<PHINode>(val)) {
     if (auto uiv = hasUninverted(val)) {
       if (auto dli = dyn_cast_or_null<LoadInst>(uiv)) {
@@ -3835,6 +3838,13 @@ bool GradientUtils::legalRecompute(const Value *val,
       }
     }
 
+    auto found = fictiousPHIs.find(const_cast<llvm::PHINode *>(phi));
+    if (found != fictiousPHIs.end()) {
+      auto orig = found->second;
+      if (isa<AtomicRMWInst>(orig))
+        return false;
+    }
+
     if (phi->getNumIncomingValues() == 0) {
       llvm::errs() << *oldFunc << "\n";
       llvm::errs() << *newFunc << "\n";