Add addr replacement handling of memtransfer (rust-lang#879)

* Add addr replacement handling of memtransfer

* Update test

Author: wsmoses

enzyme/Enzyme/FunctionUtils.cpp

@@ -350,10 +350,35 @@ void RecursivelyReplaceAddressSpace(Value *AI, Value *rep, bool legal) {
           Intrinsic::getDeclaration(MS->getParent()->getParent()->getParent(),
                                     Intrinsic::memset, tys),
           nargs));
-      nMS->copyIRFlags(MS);
+      nMS->copyMetadata(*MS);
+      nMS->setAttributes(MS->getAttributes());
       toErase.push_back(MS);
       continue;
     }
+    if (auto MTI = dyn_cast<MemTransferInst>(inst)) {
+      IRBuilder<> B(MTI);
+
+      Value *nargs[4] = {MTI->getArgOperand(0), MTI->getArgOperand(1),
+                         MTI->getArgOperand(2), MTI->getArgOperand(3)};
+
+      if (nargs[0] == prev)
+        nargs[0] = rep;
+
+      if (nargs[1] == prev)
+        nargs[1] = rep;
+
+      Type *tys[] = {nargs[0]->getType(), nargs[1]->getType(),
+                     nargs[2]->getType()};
+
+      auto nMTI = cast<CallInst>(B.CreateCall(
+          Intrinsic::getDeclaration(MTI->getParent()->getParent()->getParent(),
+                                    MTI->getIntrinsicID(), tys),
+          nargs));
+      nMTI->copyMetadata(*MTI);
+      nMTI->setAttributes(MTI->getAttributes());
+      toErase.push_back(MTI);
+      continue;
+    }
     if (auto CI = dyn_cast<CallInst>(inst)) {
       if (auto F = CI->getCalledFunction()) {
         if (F->getName() == "julia.write_barrier" && legal) {

enzyme/Enzyme/Utils.cpp

@@ -63,7 +63,6 @@ llvm::SmallVector<llvm::Instruction *, 2> PostCacheStore(llvm::StoreInst *SI,
   SmallVector<llvm::Instruction *, 2> res;
   if (EnzymePostCacheStore) {
     uint64_t size = 0;
-    LLVMValueRef V2 = nullptr;
     auto ptr = EnzymePostCacheStore(wrap(SI), wrap(&B), &size);
     for (size_t i = 0; i < size; i++) {
       res.push_back(cast<Instruction>(unwrap(ptr[i])));

enzyme/test/Enzyme/ReverseMode/writeonlyretcjl.ll

@@ -40,7 +40,7 @@ entry:
 ; CHECK-NEXT:   %"r'ai" = alloca double, i64 1, align 8
 ; CHECK-NEXT:   %0 = bitcast double* %"r'ai" to {}*
 ; CHECK-NEXT:   %1 = bitcast {}* %0 to i8*
-; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* %1, i8 0, i64 8, i1 false)
+; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* nonnull dereferenceable(8) dereferenceable_or_null(8) %1, i8 0, i64 8, i1 false)
 ; CHECK-NEXT:   %2 = load double, double* %"r'ai", align 8
 ; CHECK-NEXT:   %3 = fadd fast double %2, %differeturn
 ; CHECK-NEXT:   store double %3, double* %"r'ai", align 8,

enzyme/test/Enzyme/ReverseMode/writeonlyretjl.ll

@@ -72,7 +72,7 @@ entry:
 ; CHECK-NEXT:   %"r'ai" = alloca double, i64 1, align 8
 ; CHECK-NEXT:   %0 = bitcast double* %"r'ai" to {}*
 ; CHECK-NEXT:   %1 = bitcast {}* %0 to i8*
-; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* %1, i8 0, i64 8, i1 false)
+; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* nonnull dereferenceable(8) dereferenceable_or_null(8) %1, i8 0, i64 8, i1 false)
 ; CHECK-NEXT:   %2 = load double, double* %"r'ai", align 8
 ; CHECK-NEXT:   %3 = fadd fast double %2, %differeturn
 ; CHECK-NEXT:   store double %3, double* %"r'ai", align 8

enzyme/test/Enzyme/ReverseMode/writeonlyretjlptr.ll

@@ -64,7 +64,7 @@ define double @dsquare({} addrspace(10)* %x, {} addrspace(10)* %dx) {
 ; CHECK-NEXT:   %"r'ai" = alloca {} addrspace(10)*, i64 1, align 8
 ; CHECK-NEXT:   %1 = bitcast {} addrspace(10)** %"r'ai" to {}*
 ; CHECK-NEXT:   %2 = bitcast {}* %1 to i8*
-; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* %2, i8 0, i64 8, i1 false)
+; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* nonnull dereferenceable(8) dereferenceable_or_null(8) %2, i8 0, i64 8, i1 false)
 ; CHECK-NEXT:   call void @augmented_subsq({} addrspace(10)** %r, {} addrspace(10)** %"r'ai", {} addrspace(10)* %x, {} addrspace(10)* %"x'")
 ; CHECK-NEXT:   %l = load {} addrspace(10)*, {} addrspace(10)** %r, align 8
 ; CHECK-NEXT:   %bc = bitcast {} addrspace(10)* %l to double addrspace(10)*
@@ -78,7 +78,7 @@ define double @dsquare({} addrspace(10)* %x, {} addrspace(10)* %dx) {
 ; CHECK-NEXT:   %"r'ai" = alloca {} addrspace(10)*, i64 1, align 8
 ; CHECK-NEXT:   %0 = bitcast {} addrspace(10)** %"r'ai" to {}*
 ; CHECK-NEXT:   %1 = bitcast {}* %0 to i8*
-; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* %1, i8 0, i64 8, i1 false)
+; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* nonnull dereferenceable(8) dereferenceable_or_null(8) %1, i8 0, i64 8, i1 false)
 ; CHECK-NEXT:   %"l'ipl" = load {} addrspace(10)*, {} addrspace(10)** %"r'ai", align 8
 ; CHECK-NEXT:   %"bc'ipc" = bitcast {} addrspace(10)* %"l'ipl" to double addrspace(10)*
 ; CHECK-NEXT:   %2 = load double, double addrspace(10)* %"bc'ipc"