Atomicrmw undef fix (#1900)

* Atomicrmw undef fix

* fixup

* earlyinc

* fix diffuse custreamsync

* Fix cusync

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -344,9 +344,14 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
                  Constant::getNullValue(gutils->getShadowType(inst.getType())),
                  Builder2);
     }
-    if (!inst.getType()->isVoidTy())
-      gutils->getNewFromOriginal(&inst)->replaceAllUsesWith(
-          UndefValue::get(inst.getType()));
+#if LLVM_VERSION_MAJOR >= 12
+    if (!inst.getType()->isVoidTy()) {
+      for (auto &U :
+           make_early_inc_range(gutils->getNewFromOriginal(&inst)->uses())) {
+        U.set(UndefValue::get(inst.getType()));
+      }
+    }
+#endif
     eraseIfUnused(inst, /*erase*/ true, /*check*/ false);
     return;
   }
@@ -920,8 +925,14 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
         setDiffe(&I, Constant::getNullValue(gutils->getShadowType(I.getType())),
                  BuilderZ);
     }
-    gutils->replaceAWithB(gutils->getNewFromOriginal(&I),
-                          UndefValue::get(I.getType()));
+#if LLVM_VERSION_MAJOR >= 12
+    if (!I.getType()->isVoidTy()) {
+      for (auto &U :
+           make_early_inc_range(gutils->getNewFromOriginal(&I)->uses())) {
+        U.set(UndefValue::get(I.getType()));
+      }
+    }
+#endif
     eraseIfUnused(I, /*erase*/ true, /*check*/ false);
     return;
   }

enzyme/Enzyme/DifferentialUseAnalysis.cpp

@@ -525,6 +525,15 @@ bool DifferentialUseAnalysis::is_use_directly_needed_in_reverse(
 
     if (!shadow) {
 
+      // Need the primal request in reverse.
+      if (funcName == "cuStreamSynchronize")
+        if (val == CI->getArgOperand(0)) {
+          if (EnzymePrintDiffUse)
+            llvm::errs() << " Need: primal(" << to_string(qtype) << ") of "
+                         << *val << " in reverse for cuda sync " << *CI << "\n";
+          return true;
+        }
+
       // Only need the primal request.
       if (funcName == "MPI_Wait" || funcName == "PMPI_Wait")
         if (val != CI->getArgOperand(0))

enzyme/Enzyme/GradientUtils.cpp

@@ -8978,7 +8978,7 @@ void GradientUtils::replaceAWithB(Value *A, Value *B, bool storeInCache) {
   // Check that the replacement doesn't already exist in the mapping
   // thereby resulting in a conflict.
 #ifndef NDEBUG
-  {
+  if (!isa<UndefValue>(B)) {
     auto found = newToOriginalFn.find(A);
     if (found != newToOriginalFn.end()) {
       auto foundB = newToOriginalFn.find(B);

enzyme/test/Enzyme/ReverseMode/custreamsync.ll

@@ -0,0 +1,34 @@
+; RUN: if [ %llvmver -lt 16 ] && [ %llvmver -ge 12 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -S | FileCheck %s; fi
+; RUN: if [ %llvmver -ge 12 ]; then %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme" -S | FileCheck %s; fi
+
+declare void @__enzyme_reverse(...) 
+
+declare i64 @foo() "enzyme_inactive" "enzyme_nofree" "enzyme_no_escaping_allocation"
+
+declare void @cuStreamSynchronize(i64)
+
+define void @square() {
+entry:
+  %z = call i64 @foo()
+  call void @cuStreamSynchronize(i64 %z)
+  ret void
+}
+
+define void @dsquare() {
+entry:
+  tail call void (...) @__enzyme_reverse(void ()* nonnull @square, i8* null)
+  ret void
+}
+
+; CHECK: define internal void @diffesquare(i8* %tapeArg)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = bitcast i8* %tapeArg to i64*
+; CHECK-NEXT:   %z = load i64, i64* %0
+; CHECK-NEXT:   tail call void @free(i8* nonnull %tapeArg)
+; CHECK-NEXT:   br label %invertentry
+
+; CHECK: invertentry:                                      ; preds = %entry
+; CHECK-NEXT:   call void @cuStreamSynchronize(i64 %z)
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }
+