Set formemset (rust-lang#881)

Author: wsmoses

enzyme/Enzyme/CApi.cpp

@@ -675,6 +675,11 @@ void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2,
   }
 }
 
+void EnzymeSetForMemSet(LLVMValueRef inst1) {
+  Instruction *I1 = cast<Instruction>(unwrap(inst1));
+  I1->setMetadata("enzyme_formemset", MDNode::get(I1->getContext(), {}));
+}
+
 void EnzymeSetMustCache(LLVMValueRef inst1) {
   Instruction *I1 = cast<Instruction>(unwrap(inst1));
   I1->setMetadata("enzyme_mustcache", MDNode::get(I1->getContext(), {}));

enzyme/Enzyme/CacheUtility.cpp

@@ -901,8 +901,13 @@ AllocaInst *CacheUtility::createCacheForScope(LimitContext ctx, Type *T,
           }
         }
 
-        if (ZeroInst)
+        if (ZeroInst) {
+          if (ZeroInst->getOperand(0) != malloccall) {
+            scopeInstructions[alloc].push_back(
+                cast<Instruction>(ZeroInst->getOperand(0)));
+          }
           scopeInstructions[alloc].push_back(ZeroInst);
+        }
         storealloc = allocationBuilder.CreateStore(firstallocation, storeInto);
 
         scopeAllocs[alloc].push_back(malloccall);

enzyme/Enzyme/GradientUtils.cpp

@@ -2214,6 +2214,8 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
     toadd = scopeAllocs[found2->second.first][0];
     for (auto u : toadd->users()) {
       if (auto ci = dyn_cast<CastInst>(u)) {
+        if (hasMetadata(ci, "enzyme_formemset"))
+          continue;
         toadd = ci;
       }
     }