Also memset undef as sret

EnzymeAD · Sep 18, 2023 · 1ef5a56 · 1ef5a56
1 parent eb841ac
commit 1ef5a56
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -2754,7 +2754,12 @@ class AdjointGenerator
         if (CI->isZero()) {
           auto root = getBaseObject(MS.getOperand(0));
           bool writtenTo = false;
-          if (isa<AllocaInst>(root) || isAllocationCall(root, gutils->TLI)) {
+          bool undefMemory =
+              isa<AllocaInst>(root) || isAllocationCall(root, gutils->TLI);
+          if (auto arg = dyn_cast<Argument>(root))
+            if (arg->hasStructRetAttr())
+              undefMemory = true;
+          if (undefMemory) {
             Instruction *cur = MS.getPrevNode();
             while (cur) {
               if (cur == root)