From e9b8c6266cee82693603638d34df1455fbaa1b50 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Sun, 8 Oct 2023 16:44:28 -0400
Subject: [PATCH] Fix enzyme_inactive removal on return

---
 enzyme/Enzyme/EnzymeLogic.cpp | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
index 51c4aa300f36..c51bb58625d6 100644
--- 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();