Fix memtransfer of anything (#2152)

* Fix memtransfer of anything * fix * fmt * Fix --------- Co-authored-by: Paul Berg <[email protected]>
EnzymeAD · Nov 5, 2024 · 4025abd · 4025abd
1 parent f1f4d8e
commit 4025abd
Show file tree

Hide file tree

Showing 9 changed files with 51 additions and 33 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -3317,7 +3317,8 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
     } else {
       auto &DL = gutils->newFunc->getParent()->getDataLayout();
       auto vd = TR.query(orig_dst).Data0().ShiftIndices(DL, 0, size, 0);
-      vd |= TR.query(orig_src).Data0().ShiftIndices(DL, 0, size, 0);
+      vd |= TR.query(orig_src).Data0().PurgeAnything().ShiftIndices(DL, 0, size,
+                                                                    0);
       for (size_t i = 0; i < MTI.getNumOperands(); i++)
         if (MTI.getOperand(i) == orig_dst)
           if (MTI.getAttributes().hasParamAttr(i, "enzyme_type")) {

diff --git a/enzyme/Enzyme/DifferentialUseAnalysis.cpp b/enzyme/Enzyme/DifferentialUseAnalysis.cpp
@@ -873,7 +873,7 @@ void DifferentialUseAnalysis::minCut(const DataLayout &DL, LoopInfo &OrigLI,
   std::map<Node, Node> parent;
   bfs(G, Recomputes, parent);
 
-  std::deque<Value *> todo;
+  SetVector<Value *> todo;
 
   // Print all edges that are from a reachable vertex to
   // non-reachable vertex in the original graph
@@ -885,20 +885,36 @@ void DifferentialUseAnalysis::minCut(const DataLayout &DL, LoopInfo &OrigLI,
           assert(pair.first.V == N.V);
           MinReq.insert(N.V);
           if (Orig.find(Node(N.V, true)) != Orig.end()) {
-            todo.push_back(N.V);
+            todo.insert(N.V);
           }
         }
       }
   }
 
-  // When ambiguous, push to cache the last value in a computation chain
-  // This should be considered in a cost for the max flow
   while (todo.size()) {
     auto V = todo.front();
-    todo.pop_front();
+    todo.remove(V);
     auto found = Orig.find(Node(V, true));
     assert(found != Orig.end());
     const auto &mp = found->second;
+
+    assert(MinReq.count(V));
+
+    // Fix up non-cacheable calls to use their operand(s) instead
+    if (hasNoCache(V)) {
+      assert(!Required.count(V));
+      MinReq.remove(V);
+      for (auto &pair : Orig) {
+        if (pair.second.count(Node(V, false))) {
+          MinReq.insert(pair.first.V);
+          todo.insert(pair.first.V);
+        }
+      }
+      continue;
+    }
+
+    // When ambiguous, push to cache the last value in a computation chain
+    // This should be considered in a cost for the max flow
     if (mp.size() == 1 && !Required.count(V)) {
       bool potentiallyRecursive =
           isa<PHINode>((*mp.begin()).V) &&
@@ -967,21 +983,7 @@ void DifferentialUseAnalysis::minCut(const DataLayout &DL, LoopInfo &OrigLI,
         auto nnode = (*mp.begin()).V;
         MinReq.insert(nnode);
         if (Orig.find(Node(nnode, true)) != Orig.end())
-          todo.push_back(nnode);
-      }
-    }
-  }
-
-  // Fix up non-cacheable calls to use their operand(s) instead
-  for (auto V : Intermediates) {
-    if (!hasNoCache(V))
-      continue;
-    if (!MinReq.count(V))
-      continue;
-    MinReq.remove(V);
-    for (auto &pair : Orig) {
-      if (pair.second.count(Node(V, false))) {
-        MinReq.insert(pair.first.V);
+          todo.insert(nnode);
       }
     }
   }

diff --git a/enzyme/Enzyme/EnzymeLogic.h b/enzyme/Enzyme/EnzymeLogic.h
@@ -52,7 +52,6 @@
 
 extern "C" {
 extern llvm::cl::opt<bool> EnzymePrint;
-extern llvm::cl::opt<bool> EnzymeJuliaAddrLoad;
 }
 
 constexpr char EnzymeFPRTPrefix[] = "__enzyme_fprt_";

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -8177,6 +8177,9 @@ void GradientUtils::computeMinCache() {
                                    notForAnalysis);
             if (oneneed) {
               knownRecomputeHeuristic[&I] = false;
+
+              CountTrackedPointers T(I.getType());
+              assert(!T.derived);
             } else
               Recomputes.insert(&I);
           }
@@ -8267,8 +8270,14 @@ void GradientUtils::computeMinCache() {
         assert(legalRecompute(V, Available2, nullptr));
       }
       if (!NeedGraph.count(V)) {
+        assert(!MinReq.count(V));
         unnecessaryIntermediates.insert(cast<Instruction>(V));
       }
+
+      if (NeedGraph.count(V) && MinReq.count(V)) {
+        CountTrackedPointers T(V->getType());
+        assert(!T.derived);
+      }
     }
   }
 }

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
@@ -108,11 +108,6 @@ const llvm::StringMap<llvm::Intrinsic::ID> LIBM_FUNCTIONS = {
     {"__nv_drcp_rn", Intrinsic::not_intrinsic},
     {"__nv_drcp_ru", Intrinsic::not_intrinsic},
     {"__nv_drcp_rz", Intrinsic::not_intrinsic},
-    {"__nv_isnand", Intrinsic::not_intrinsic},
-    {"__nv_isnanf", Intrinsic::not_intrinsic},
-    {"__nv_isinfd", Intrinsic::not_intrinsic},
-    {"__nv_isinff", Intrinsic::not_intrinsic},
-    {"__nv_acos", Intrinsic::not_intrinsic},
     {"asin", Intrinsic::not_intrinsic},
     {"__nv_asin", Intrinsic::not_intrinsic},
     {"atan", Intrinsic::not_intrinsic},

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h
@@ -60,6 +60,7 @@ extern const llvm::StringMap<llvm::Intrinsic::ID> LIBM_FUNCTIONS;
 
 static inline bool isMemFreeLibMFunction(llvm::StringRef str,
                                          llvm::Intrinsic::ID *ID = nullptr) {
+  llvm::StringRef ogstr = str;
   if (startsWith(str, "__") && endsWith(str, "_finite")) {
     str = str.substr(2, str.size() - 2 - 7);
   } else if (startsWith(str, "__fd_") && endsWith(str, "_1")) {
@@ -72,7 +73,8 @@ static inline bool isMemFreeLibMFunction(llvm::StringRef str,
       *ID = LIBM_FUNCTIONS.find(str.str())->second;
     return true;
   }
-  if (endsWith(str, "f") || endsWith(str, "l")) {
+  if (endsWith(str, "f") || endsWith(str, "l") ||
+      (startsWith(ogstr, "__nv_") && endsWith(str, "d"))) {
     if (LIBM_FUNCTIONS.find(str.substr(0, str.size() - 1).str()) !=
         LIBM_FUNCTIONS.end()) {
       if (ID)

diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h
@@ -94,6 +94,7 @@ extern llvm::cl::opt<bool> EnzymePrintPerf;
 extern llvm::cl::opt<bool> EnzymeStrongZero;
 extern llvm::cl::opt<bool> EnzymeBlasCopy;
 extern llvm::cl::opt<bool> EnzymeLapackCopy;
+extern llvm::cl::opt<bool> EnzymeJuliaAddrLoad;
 extern LLVMValueRef (*CustomErrorHandler)(const char *, LLVMValueRef, ErrorType,
                                           const void *, LLVMValueRef,
                                           LLVMBuilderRef);
@@ -1184,6 +1185,15 @@ static inline bool hasNoCache(llvm::Value *op) {
   if (auto I = dyn_cast<Instruction>(op))
     if (hasMetadata(I, "enzyme_nocache"))
       return true;
+
+  if (EnzymeJuliaAddrLoad) {
+    if (auto PT = dyn_cast<PointerType>(op->getType())) {
+      if (PT->getAddressSpace() == 11 || PT->getAddressSpace() == 13) {
+        if (isa<CastInst>(op) || isa<GetElementPtrInst>(op))
+          return true;
+      }
+    }
+  }
   return false;
 }
 

diff --git a/enzyme/test/Enzyme/ForwardMode/memcpyanyflt.ll b/enzyme/test/Enzyme/ForwardMode/memcpyanyflt.ll
@@ -26,9 +26,7 @@ attributes #0 = { argmemonly nounwind }
 
 ; CHECK: define internal void @fwddiffememcpy_ptr(i8* nocapture %dst, i8* nocapture %"dst'", i8* nocapture readonly %src, i8* nocapture %"src'", i64 %num)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %"dst'", i8* bitcast ({ i64, double }* @_j_const2 to i8*), i64 8, i1 false)
-; CHECK-NEXT:   %0 = getelementptr inbounds i8, i8* %"dst'", i64 8
-; CHECK-NEXT:   tail call void @llvm.memset.p0i8.i64(i8* align 1 %0, i8 0, i64 8, i1 true)
+; CHECK-NEXT:   tail call void @llvm.memset.p0i8.i64(i8* align 1 %"dst'", i8 0, i64 16, i1 true) 
 ; CHECK-NEXT:   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %dst, i8* bitcast ({ i64, double }* @_j_const2 to i8*), i64 16, i1 false)
 ; CHECK-NEXT:   ret void
 ; CHECK-NEXT: }
diff --git a/enzyme/test/Enzyme/ReverseMode/nullcp.ll b/enzyme/test/Enzyme/ReverseMode/nullcp.ll
@@ -1,5 +1,5 @@
-; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -enzyme -mem2reg -instsimplify -adce -loop-deletion -correlated-propagation -simplifycfg -S | FileCheck %s; fi
-; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,adce,loop(loop-deletion),correlated-propagation,%simplifycfg)" -S | FileCheck %s
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -enzyme -mem2reg -instsimplify -adce -loop-deletion -correlated-propagation -simplifycfg  -enzyme-runtime-error -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,adce,loop(loop-deletion),correlated-propagation,%simplifycfg)"  -enzyme-runtime-error -S | FileCheck %s
 
 source_filename = "nullcp.c"
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
@@ -130,6 +130,8 @@ attributes #6 = { nounwind }
 ; CHECK-NEXT:   %0 = bitcast double* %dst to i8*
 ; CHECK-NEXT:   %1 = bitcast double* %src to i8*
 ; CHECK-NEXT:   %mul = shl i64 %n, 3
+; CHECK-NEXT:   %2 = call i32 @puts(i8* getelementptr inbounds ([126 x i8], [126 x i8]* @.str.1, i32 0, i32 0))
+; CHECK-NEXT:   call void @exit(i32 1)
 ; CHECK-NEXT:   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %"'ipc", i8* align 8 %1, i64 %mul, i1 false)
 ; CHECK-NEXT:   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %0, i8* align 8 %1, i64 %mul, i1 false)
 ; CHECK-NEXT:   br label %invertif.end