Simplify enzymejl_needs_restoration (#2187)

EnzymeAD · Dec 6, 2024 · a9cc3c9 · a9cc3c9
1 parent 57b718b
commit a9cc3c9
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp
@@ -92,7 +92,9 @@ void getFunction(const Twine &curIndent, raw_ostream &os, StringRef callval,
       os << curIndent << "auto " << callval
          << " = gutils->oldFunc->getParent()->getOrInsertFunction(";
       os << Def->getValueInit("name")->getAsString();
-      os << ", " << FT << ", called->getAttributes()).getCallee();\n";
+      os << ", " << FT
+         << ", called->getAttributes().removeFnAttribute(called->getContext(), "
+            "\"enzymejl_needs_restoration\")).getCallee();\n";
       os << curIndent << "auto " << cconv << " = cast<CallInst>(&" << origName
          << ")->getCallingConv();\n";
       return;
@@ -118,7 +120,9 @@ void getFunction(const Twine &curIndent, raw_ostream &os, StringRef callval,
       os << curIndent << "auto " << callval
          << " = gutils->oldFunc->getParent()->getOrInsertFunction(";
       os << Def->getValueInit("name")->getAsString();
-      os << ", " << FT << ", called->getAttributes()).getCallee();\n";
+      os << ", " << FT
+         << ", called->getAttributes().removeFnAttribute(called->getContext(), "
+            "\"enzymejl_needs_restoration\")).getCallee();\n";
       os << curIndent << "auto " << cconv << " = cast<CallInst>(&" << origName
          << ")->getCallingConv();\n";
       return;
@@ -133,7 +137,9 @@ void getFunction(const Twine &curIndent, raw_ostream &os, StringRef callval,
       os << curIndent << "auto " << callval
          << " = gutils->oldFunc->getParent()->getOrInsertFunction(";
       os << Def->getValueInit("name")->getAsString();
-      os << ", " << FT << ", called->getAttributes()).getCallee();\n";
+      os << ", " << FT
+         << ", called->getAttributes().removeFnAttribute(called->getContext(), "
+            "\"enzymejl_needs_restoration\")).getCallee();\n";
       os << curIndent << "auto " << cconv << " = cast<CallInst>(&" << origName
          << ")->getCallingConv();\n";
       return;