Fix cuda realloc (#2091)

* Fix cuda realloc

* fix

enzyme/Enzyme/FunctionUtils.cpp

@@ -656,11 +656,18 @@ OldAllocationSize(Value *Ptr, CallInst *Loc, Function *NewF, IntegerType *T,
  }
  if (success)
  continue;
+
+ auto v2 = simplifyLoad(LI);
+ if (v2) {
+ todo.push_back({v2, next.second});
+ continue;
+ }
  }
 
  EmitFailure("DynamicReallocSize", Loc->getDebugLoc(), Loc,
  "could not statically determine size of realloc ", *Loc,
  " - because of - ", *next.first);
+ return AI;
 
  std::string allocName;
  switch (llvm::Triple(NewF->getParent()->getTargetTriple()).getOS()) {
@@ -833,6 +840,12 @@ void PreProcessCache::ReplaceReallocs(Function *NewF, bool mem2reg) {
  if (mem2reg) {
  auto PA = PromotePass().run(*NewF, FAM);
  FAM.invalidate(*NewF, PA);
+#if !defined(FLANG)
+ PA = GVNPass().run(*NewF, FAM);
+#else
+ PA = GVN().run(*NewF, FAM);
+#endif
+ FAM.invalidate(*NewF, PA);
  }
 
  SmallVector<CallInst *, 4> ToConvert;

enzyme/Enzyme/Utils.cpp

@@ -45,6 +45,12 @@
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Verifier.h"
 
+#if LLVM_VERSION_MAJOR >= 16
+#include "llvm/TargetParser/Triple.h"
+#else
+#include "llvm/ADT/Triple.h"
+#endif
+
 #include "llvm-c/Core.h"
 
 #include "LibraryFuncs.h"
@@ -203,7 +209,10 @@ Function *getOrInsertExponentialAllocator(Module &M, Function *newFunc,
  ConstantInt::get(next->getType(), 0),
  B.CreateLShr(next, ConstantInt::get(next->getType(), 1)));
 
- if (!custom) {
+ auto Arch = llvm::Triple(M.getTargetTriple()).getArch();
+ bool forceMalloc = Arch == Triple::nvptx || Arch == Triple::nvptx64;
+
+ if (!custom && !forceMalloc) {
  auto reallocF = M.getOrInsertFunction("realloc", allocType, allocType,
  Type::getInt64Ty(M.getContext()));
 

enzyme/test/Enzyme/ReverseMode/sharedcachefwd.ll

@@ -197,13 +197,11 @@ attributes #6 = { nounwind }
 ; CHECK-NEXT: call void @llvm.nvvm.barrier0()
 ; CHECK-NEXT: %[[v9:.+]] = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.y()
 ; CHECK-NEXT: %mul = shl i32 %[[v9]], 4
-; CHECK-NEXT: %[[v10:.+]] = call i32 @llvm.nvvm.read.ptx.sreg.tid.y()
-; CHECK-NEXT: %add = add i32 %mul, %[[v10]]
+; CHECK-NEXT: %add = add i32 %mul, %1
 ; CHECK-NEXT: %conv = zext i32 %add to i64
 ; CHECK-NEXT: %[[v11:.+]] = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
 ; CHECK-NEXT: %mul3 = shl i32 %[[v11]], 4
-; CHECK-NEXT: %[[v12:.+]] = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
-; CHECK-NEXT: %add5 = add i32 %mul3, %[[v12]]
+; CHECK-NEXT: %add5 = add i32 %mul3, %0
 ; CHECK-NEXT: %conv6 = zext i32 %add5 to i64
 ; CHECK-NEXT: %[[v13:.+]] = call i32 @llvm.nvvm.read.ptx.sreg.nctaid.x()
 ; CHECK-NEXT: %conv8 = zext i32 %[[v13]] to i64
@@ -212,10 +210,10 @@ attributes #6 = { nounwind }
 
 ; CHECK: for.body.lr.ph:
 ; CHECK-NEXT: %mul9 = mul i64 %conv, %n
-; CHECK-NEXT: %conv13 = zext i32 %[[v12]] to i64
+; CHECK-NEXT: %conv13 = zext i32 %0 to i64
 ; CHECK-NEXT: %add11 = add i64 %mul9, %conv13
 ; CHECK-NEXT: %mul15 = mul i64 %n, %n
-; CHECK-NEXT: %idxprom = zext i32 %[[v10]] to i64
+; CHECK-NEXT: %idxprom = zext i32 %1 to i64
 ; CHECK-NEXT: %arrayidx2195 = getelementptr inbounds [16 x [16 x float]], [16 x [16 x float]] addrspace(3)* @_ZZ22gpu_square_matrix_multPfS_S_mE6tile_a, i64 0, i64 %idxprom, i64 %conv13
 ; CHECK-NEXT: %arrayidx21 = addrspacecast float addrspace(3)* %arrayidx2195 to float*
 ; CHECK-NEXT: %arrayidx4097 = getelementptr inbounds [16 x [16 x float]], [16 x [16 x float]] addrspace(3)* @_ZZ22gpu_square_matrix_multPfS_S_mE6tile_b, i64 0, i64 %idxprom, i64 %conv13
@@ -403,11 +401,11 @@ attributes #6 = { nounwind }
 ; CHECK: invertfor.body44_phimerge_phimerge:
 ; CHECK-NEXT: %[[v33:.+]] = phi {{(fast )?}}float [ %[[_unwrap51]], %invertfor.body44_phimerge_phirc ], [ 0.000000e+00, %invertfor.body44_phimerge ]
 ; CHECK-NEXT: %[[m1diffe:.+]] = fmul fast float %"add56'de.1", %[[v33]]
-; CHECK-NEXT: %[[conv13_unwrap53]] = zext i32 %[[v12]] to i64
+; CHECK-NEXT: %[[conv13_unwrap53]] = zext i32 %0 to i64
 ; CHECK-NEXT: %"arrayidx54101'ipg_unwrap" = getelementptr inbounds [16 x [16 x float]], [16 x [16 x float]] addrspace(3)* @_ZZ22gpu_square_matrix_multPfS_S_mE6tile_b_shadow, i64 0, i64 %[[idxprom_unwrap20]], i64 %[[conv13_unwrap53]]
 ; CHECK-NEXT: %"arrayidx54'ipc_unwrap" = addrspacecast float addrspace(3)* %"arrayidx54101'ipg_unwrap" to float*
 ; CHECK-NEXT: %{{.+}} = atomicrmw fadd float* %"arrayidx54'ipc_unwrap", float %[[m1diffe]] monotonic
-; CHECK-NEXT: %[[idxprom_unwrap54]] = zext i32 %[[v10]] to i64
+; CHECK-NEXT: %[[idxprom_unwrap54]] = zext i32 %1 to i64
 ; CHECK-NEXT: %"arrayidx4999'ipg_unwrap" = getelementptr inbounds [16 x [16 x float]], [16 x [16 x float]] addrspace(3)* @_ZZ22gpu_square_matrix_multPfS_S_mE6tile_a_shadow, i64 0, i64 %[[idxprom_unwrap54]], i64 %[[idxprom_unwrap20]]
 ; CHECK-NEXT: %"arrayidx49'ipc_unwrap" = addrspacecast float addrspace(3)* %"arrayidx4999'ipg_unwrap" to float*
 ; CHECK-NEXT: %{{.+}} = atomicrmw fadd float* %"arrayidx49'ipc_unwrap", float %[[m0diffe]] monotonic

enzyme/test/Enzyme/ReverseMode/sharedmem.ll

@@ -107,12 +107,10 @@ attributes #4 = { nounwind }
 ; CHECK: invertbb: ; preds = %shblock, %bb
 ; CHECK-NEXT: call void @llvm.nvvm.barrier0()
 ; CHECK-NEXT: %tmp = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.y() #{{.*}}, !range !12
-; CHECK-NEXT: %tmp4 = tail call i32 @llvm.nvvm.read.ptx.sreg.tid.y() #{{.*}}, !range !13
-; CHECK-NEXT: %tmp5 = add nuw nsw i32 %tmp4, %tmp
+; CHECK-NEXT: %tmp5 = add nuw nsw i32 %1, %tmp
 ; CHECK-NEXT: %tmp6 = zext i32 %tmp5 to i64
 ; CHECK-NEXT: %tmp7 = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() #{{.*}}, !range !14
-; CHECK-NEXT: %tmp8 = tail call i32 @llvm.nvvm.read.ptx.sreg.tid.x() #{{.*}}, !range !13
-; CHECK-NEXT: %tmp9 = add nuw i32 %tmp8, %tmp7
+; CHECK-NEXT: %tmp9 = add nuw i32 %0, %tmp7
 ; CHECK-NEXT: %tmp10 = zext i32 %tmp9 to i64
 ; CHECK-NEXT: %tmp11 = mul i64 %tmp6, %arg3
 ; CHECK-NEXT: %tmp12 = add i64 %tmp11, %tmp10

enzyme/test/Enzyme/ReverseModeVector/vecloadatomic.ll

@@ -71,39 +71,37 @@ entry:
 ; CHECK-NEXT: tail call void @free(i8* nonnull %tapeArg)
 ; CHECK-NEXT: %"i3'de" = alloca [2 x i64]
 ; CHECK-NEXT: store [2 x i64] zeroinitializer, [2 x i64]* %"i3'de"
-; CHECK-NEXT: %0 = extractvalue [2 x i64 addrspace(13)*] %"i7'", 0
-; CHECK-NEXT: %1 = load i64, i64 addrspace(13)* %0
-; CHECK-NEXT: %2 = extractvalue [2 x i64 addrspace(13)*] %"i7'", 1
-; CHECK-NEXT: %3 = load i64, i64 addrspace(13)* %2
-; CHECK-NEXT: %4 = extractvalue [2 x i64 addrspace(13)*] %"i7'", 0
-; CHECK-NEXT: store i64 0, i64 addrspace(13)* %4
-; CHECK-NEXT: %5 = extractvalue [2 x i64 addrspace(13)*] %"i7'", 1
-; CHECK-NEXT: store i64 0, i64 addrspace(13)* %5
-; CHECK-NEXT: %6 = getelementptr inbounds [2 x i64], [2 x i64]* %"i3'de", i32 0, i32 0
-; CHECK-NEXT: %7 = load i64, i64* %6
-; CHECK-NEXT: %8 = bitcast i64 %7 to double
-; CHECK-NEXT: %9 = bitcast i64 %1 to double
-; CHECK-NEXT: %10 = fadd fast double %8, %9
-; CHECK-NEXT: %11 = bitcast double %10 to i64
-; CHECK-NEXT: store i64 %11, i64* %6
-; CHECK-NEXT: %12 = getelementptr inbounds [2 x i64], [2 x i64]* %"i3'de", i32 0, i32 1
-; CHECK-NEXT: %13 = load i64, i64* %12
-; CHECK-NEXT: %14 = bitcast i64 %13 to double
-; CHECK-NEXT: %15 = bitcast i64 %3 to double
-; CHECK-NEXT: %16 = fadd fast double %14, %15
-; CHECK-NEXT: %17 = bitcast double %16 to i64
-; CHECK-NEXT: store i64 %17, i64* %12
-; CHECK-NEXT: %18 = load [2 x i64], [2 x i64]* %"i3'de"
+; CHECK-NEXT: %[[i0:.+]] = extractvalue [2 x i64 addrspace(13)*] %"i7'", 0
+; CHECK-NEXT: %[[i1:.+]] = load i64, i64 addrspace(13)* %[[i0]]
+; CHECK-NEXT: %[[i2:.+]] = extractvalue [2 x i64 addrspace(13)*] %"i7'", 1
+; CHECK-NEXT: %[[i3:.+]] = load i64, i64 addrspace(13)* %[[i2]]
+; CHECK-NEXT: store i64 0, i64 addrspace(13)* %[[i0]]
+; CHECK-NEXT: store i64 0, i64 addrspace(13)* %[[i2]]
+; CHECK-NEXT: %[[i6:.+]] = getelementptr inbounds [2 x i64], [2 x i64]* %"i3'de", i32 0, i32 0
+; CHECK-NEXT: %[[i7:.+]] = load i64, i64* %[[i6]]
+; CHECK-NEXT: %[[i8:.+]] = bitcast i64 %[[i7]] to double
+; CHECK-NEXT: %[[i9:.+]] = bitcast i64 %[[i1]] to double
+; CHECK-NEXT: %[[i10:.+]] = fadd fast double %[[i8]], %[[i9]]
+; CHECK-NEXT: %[[i11:.+]] = bitcast double %[[i10]] to i64
+; CHECK-NEXT: store i64 %[[i11]], i64* %[[i6]]
+; CHECK-NEXT: %[[i12:.+]] = getelementptr inbounds [2 x i64], [2 x i64]* %"i3'de", i32 0, i32 1
+; CHECK-NEXT: %[[i13:.+]] = load i64, i64* %[[i12]]
+; CHECK-NEXT: %[[i14:.+]] = bitcast i64 %[[i13]] to double
+; CHECK-NEXT: %[[i15:.+]] = bitcast i64 %[[i3]] to double
+; CHECK-NEXT: %[[i16:.+]] = fadd fast double %[[i14]], %[[i15]]
+; CHECK-NEXT: %[[i17:.+]] = bitcast double %[[i16]] to i64
+; CHECK-NEXT: store i64 %[[i17]], i64* %[[i12]]
+; CHECK-NEXT: %[[i18:.+]] = load [2 x i64], [2 x i64]* %"i3'de"
 ; CHECK-NEXT: store [2 x i64] zeroinitializer, [2 x i64]* %"i3'de"
-; CHECK-NEXT: %19 = extractvalue [2 x i64 addrspace(12)*] %"i2'", 0
-; CHECK-NEXT: %20 = bitcast i64 addrspace(12)* %19 to double addrspace(12)*
-; CHECK-NEXT: %21 = extractvalue [2 x i64 addrspace(12)*] %"i2'", 1
-; CHECK-NEXT: %22 = bitcast i64 addrspace(12)* %21 to double addrspace(12)*
-; CHECK-NEXT: %23 = extractvalue [2 x i64] %18, 0
-; CHECK-NEXT: %24 = bitcast i64 %23 to double
-; CHECK-NEXT: %25 = extractvalue [2 x i64] %18, 1
-; CHECK-NEXT: %26 = bitcast i64 %25 to double
-; CHECK-NEXT: %27 = atomicrmw fadd double addrspace(12)* %20, double %24 monotonic
-; CHECK-NEXT: %28 = atomicrmw fadd double addrspace(12)* %22, double %26 monotonic
+; CHECK-NEXT: %[[i19:.+]] = extractvalue [2 x i64 addrspace(12)*] %"i2'", 0
+; CHECK-NEXT: %[[i20:.+]] = bitcast i64 addrspace(12)* %[[i19]] to double addrspace(12)*
+; CHECK-NEXT: %[[i21:.+]] = extractvalue [2 x i64 addrspace(12)*] %"i2'", 1
+; CHECK-NEXT: %[[i22:.+]] = bitcast i64 addrspace(12)* %[[i21]] to double addrspace(12)*
+; CHECK-NEXT: %[[i23:.+]] = extractvalue [2 x i64] %[[i18]], 0
+; CHECK-NEXT: %[[i24:.+]] = bitcast i64 %[[i23]] to double
+; CHECK-NEXT: %[[i25:.+]] = extractvalue [2 x i64] %[[i18]], 1
+; CHECK-NEXT: %[[i26:.+]] = bitcast i64 %[[i25]] to double
+; CHECK-NEXT: %[[i27:.+]] = atomicrmw fadd double addrspace(12)* %[[i20]], double %[[i24]] monotonic
+; CHECK-NEXT: %[[i28:.+]] = atomicrmw fadd double addrspace(12)* %[[i22]], double %[[i26]] monotonic
 ; CHECK-NEXT: ret void
 ; CHECK-NEXT: }