Attribute::Memory fixes (#7)

but still getting zeros instead of derivatives

Author: stepasite

enzyme/Enzyme/Enzyme.cpp

@@ -2017,7 +2017,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,
@@ -2040,7 +2042,10 @@ class EnzymeBase {
         if (Fn->getName() == "frexp" || Fn->getName() == "frexpf" ||
             Fn->getName() == "frexpl") {
 #if LLVM_VERSION_MAJOR >= 16
-          CI->addAttribute(AttributeList::FunctionIndex, Attribute::Memory);
+          CI->addAttribute(AttributeList::FunctionIndex,
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::argMemOnly()));
 #else
           CI->addAttribute(AttributeList::FunctionIndex, Attribute::ArgMemOnly);
 #endif
@@ -2061,7 +2066,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,
@@ -2072,7 +2079,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleOrArgMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,
@@ -2096,7 +2105,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleOrArgMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,
@@ -2123,7 +2134,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleOrArgMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,
@@ -2150,7 +2163,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleOrArgMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,
@@ -2182,7 +2197,9 @@ class EnzymeBase {
 #if LLVM_VERSION_MAJOR >= 16
           Fn->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
           CI->addAttribute(AttributeList::FunctionIndex,
-                           Attribute::Memory);
+                           Attribute::getWithMemoryEffects(
+                            CI->getContext(),
+                            MemoryEffects::inaccessibleOrArgMemOnly()));
 #else
           Fn->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
           CI->addAttribute(AttributeList::FunctionIndex,

enzyme/Enzyme/FunctionUtils.cpp

@@ -927,9 +927,7 @@ Function *CreateMPIWrapper(Function *F) {
     Attribute::NoFree,
     Attribute::NoSync,
 #endif
-#if LLVM_VERSION_MAJOR >= 16
-    Attribute::Memory
-#else
+#if LLVM_VERSION_MAJOR < 16
     Attribute::InaccessibleMemOnly
 #endif
   };
@@ -940,6 +938,11 @@ Function *CreateMPIWrapper(Function *F) {
     W->addAttribute(AttributeList::FunctionIndex, attr);
 #endif
   }
+#if LLVM_VERSION_MAJOR >= 16
+    W->addFnAttr(Attribute::getWithMemoryEffects(
+                        F->getContext(),
+                        MemoryEffects::inaccessibleOrArgMemOnly()));
+#endif
 #if LLVM_VERSION_MAJOR >= 14
   W->addFnAttr(Attribute::get(F->getContext(), "enzyme_inactive"));
 #else