Fix enzyme_inactive removal on return (#1469)

EnzymeAD · Oct 9, 2023 · 3194829 · 3194829
1 parent a57a98f
commit 3194829
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -1740,6 +1740,18 @@ void clearFunctionAttributes(Function *f) {
     if (f->hasAttribute(llvm::AttributeList::ReturnIndex, attr)) {
       f->removeAttribute(llvm::AttributeList::ReturnIndex, attr);
     }
+#endif
+  }
+  for (auto attr : {"enzyme_inactive"}) {
+#if LLVM_VERSION_MAJOR >= 14
+    if (f->getAttributes().hasRetAttribute(attr)) {
+      f->removeRetAttr(attr);
+    }
+#else
+    if (f->getAttributes().hasAttribute(llvm::AttributeList::ReturnIndex,
+                                        attr)) {
+      f->removeAttribute(llvm::AttributeList::ReturnIndex, attr);
+    }
 #endif
   }
 }
@@ -2564,6 +2576,18 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
     }
 #endif
   }
+  for (auto attr : {"enzyme_inactive"}) {
+#if LLVM_VERSION_MAJOR >= 14
+    if (gutils->newFunc->getAttributes().hasRetAttribute(attr)) {
+      gutils->newFunc->removeRetAttr(attr);
+    }
+#else
+    if (gutils->newFunc->getAttributes().hasAttribute(
+            llvm::AttributeList::ReturnIndex, attr)) {
+      gutils->newFunc->removeAttribute(llvm::AttributeList::ReturnIndex, attr);
+    }
+#endif
+  }
 
   gutils->eraseFictiousPHIs();