Fix readonly or throw to have local notions (#2416)

wsmoses · web-flow · commit 4a3c72685643 · 2025-08-18T06:17:25.000+02:00
* Fix readonly or throw to have local notions

* also loop cache to compile time error rather than assertion

* fmt

* fix

* fmt

* fmt

* change to local

* fix

* fix
diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -1961,7 +1961,10 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
         if (CB->onlyAccessesInaccessibleMemory())
           AARes = ModRefInfo::NoModRef;
 
-        bool ReadOnly = isReadOnlyOrThrow(CB);
+        bool ReadOnly = isLocalReadOnlyOrThrow(CB);
+        if (CB->hasStructRetAttr() &&
+            getBaseObject(CB->getArgOperand(0)) == getBaseObject(Val))
+          ReadOnly = false;
 
         bool WriteOnly = isWriteOnly(CB);
 
@@ -3008,7 +3011,10 @@ bool ActivityAnalyzer::isValueInactiveFromUsers(TypeResults const &TR,
 
         mayCapture |= !NoCapture;
 
-        bool ReadOnly = isReadOnlyOrThrow(call) || isReadOnly(call, idx);
+        bool ReadOnly = isReadOnly(call, idx);
+        if (!ReadOnly && isLocalReadOnlyOrThrow(call) && idx != 0 &&
+            call->hasStructRetAttr())
+          ReadOnly = true;
 
         mayWrite |= !ReadOnly;
 
@@ -3392,8 +3398,9 @@ bool ActivityAnalyzer::isValueActivelyStoredOrReturned(TypeResults const &TR,
 
     if (auto inst = dyn_cast<Instruction>(a)) {
       if (!inst->mayWriteToMemory() ||
-          (isa<CallInst>(inst) && (AA.onlyReadsMemory(cast<CallInst>(inst)) ||
-                                   isReadOnlyOrThrow(cast<CallInst>(inst))))) {
+          (isa<CallInst>(inst) &&
+           (AA.onlyReadsMemory(cast<CallInst>(inst)) ||
+            isLocalReadOnlyOrThrow(cast<CallInst>(inst))))) {
         // if not written to memory and returning a known constant, this
         // cannot be actively returned/stored
         if (inst->getParent()->getParent() == TR.getFunction() &&
diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp
@@ -1524,7 +1524,7 @@ void SplitPHIs(llvm::Function &F) {
 // returns if newly legal, subject to the pending calls
 bool DetectReadonlyOrThrowFn(llvm::Function &F,
                              SmallPtrSetImpl<Function *> &calls_todo,
-                             llvm::TargetLibraryInfo &TLI) {
+                             llvm::TargetLibraryInfo &TLI, bool &local) {
   if (isReadOnlyOrThrow(&F))
     return false;
   if (F.empty())
@@ -1539,8 +1539,10 @@ bool DetectReadonlyOrThrowFn(llvm::Function &F,
         continue;
       if (hasMetadata(&I, "enzyme_ReadOnlyOrThrow"))
         continue;
+      if (hasMetadata(&I, "enzyme_LocalReadOnlyOrThrow"))
+        continue;
       if (auto CI = dyn_cast<CallBase>(&I)) {
-        if (isReadOnlyOrThrow(CI)) {
+        if (isLocalReadOnlyOrThrow(CI)) {
           continue;
         }
         if (isAllocationCall(CI, TLI)) {
@@ -1574,26 +1576,80 @@ bool DetectReadonlyOrThrowFn(llvm::Function &F,
         // seen outside the function. Note, even if one stored into x =
         // malloc(..), and stored x into a global/arg pointer, that second store
         // would trigger not readonly.
-        if (isa<AllocaInst>(Obj) || isAllocationCall(Obj, TLI))
+        if (isa<AllocaInst>(Obj))
+          continue;
+        if (isAllocationCall(Obj, TLI)) {
+          if (local)
+            continue;
+          if (notCaptured(Obj))
+            continue;
+          local = true;
           continue;
+        }
+        if (auto arg = dyn_cast<Argument>(Obj)) {
+          if (arg->hasStructRetAttr() ||
+              arg->getParent()
+                  ->getAttribute(arg->getArgNo() + AttributeList::FirstArgIndex,
+                                 "enzymejl_returnRoots")
+                  .isValid()) {
+            local = true;
+            continue;
+          }
+        }
       }
       if (auto MTI = dyn_cast<MemTransferInst>(&I)) {
         auto Obj = getBaseObject(MTI->getOperand(0));
         // Storing into local memory is fine since it definitionally will not be
         // seen outside the function. Note, even if one stored into x =
         // malloc(..), and stored x into a global/arg pointer, that second store
         // would trigger not readonly.
-        if (isa<AllocaInst>(Obj) || isAllocationCall(Obj, TLI))
+        if (isa<AllocaInst>(Obj))
+          continue;
+        if (isAllocationCall(Obj, TLI)) {
+          if (local)
+            continue;
+          if (notCaptured(Obj))
+            continue;
+          local = true;
           continue;
+        }
+        if (auto arg = dyn_cast<Argument>(Obj)) {
+          if (arg->hasStructRetAttr() ||
+              arg->getParent()
+                  ->getAttribute(arg->getArgNo() + AttributeList::FirstArgIndex,
+                                 "enzymejl_returnRoots")
+                  .isValid()) {
+            local = true;
+            continue;
+          }
+        }
       }
       if (auto MSI = dyn_cast<MemSetInst>(&I)) {
         auto Obj = getBaseObject(MSI->getOperand(0));
         // Storing into local memory is fine since it definitionally will not be
         // seen outside the function. Note, even if one stored into x =
         // malloc(..), and stored x into a global/arg pointer, that second store
         // would trigger not readonly.
-        if (isa<AllocaInst>(Obj) || isAllocationCall(Obj, TLI))
+        if (isa<AllocaInst>(Obj))
+          continue;
+        if (isAllocationCall(Obj, TLI)) {
+          if (local)
+            continue;
+          if (notCaptured(Obj))
+            continue;
+          local = true;
           continue;
+        }
+        if (auto arg = dyn_cast<Argument>(Obj)) {
+          if (arg->hasStructRetAttr() ||
+              arg->getParent()
+                  ->getAttribute(arg->getArgNo() + AttributeList::FirstArgIndex,
+                                 "enzymejl_returnRoots")
+                  .isValid()) {
+            local = true;
+            continue;
+          }
+        }
       }
       // ignore atomic load impacts
       if (isa<LoadInst>(&I))
@@ -1620,7 +1676,10 @@ bool DetectReadonlyOrThrowFn(llvm::Function &F,
   }
 
   if (calls_todo.size() == 0) {
-    F.addFnAttr("enzyme_ReadOnlyOrThrow");
+    if (local)
+      F.addFnAttr("enzyme_LocalReadOnlyOrThrow");
+    else
+      F.addFnAttr("enzyme_ReadOnlyOrThrow");
   }
   return true;
 }
@@ -1629,6 +1688,17 @@ bool DetectReadonlyOrThrow(Module &M) {
 
   bool changed = false;
 
+  PassBuilder PB;
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
+  PB.registerModuleAnalyses(MAM);
+  PB.registerFunctionAnalyses(FAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.registerCGSCCAnalyses(CGAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+
   // Set of functions newly deduced readonlyorthrow by this pass
   SmallVector<llvm::Function *> todo;
 
@@ -1640,21 +1710,15 @@ bool DetectReadonlyOrThrow(Module &M) {
   // prerequisite for being readonly. Inverse of `todo_map`
   DenseMap<llvm::Function *, SmallPtrSet<Function *, 1>> inverse_todo_map;
 
-  PassBuilder PB;
-  LoopAnalysisManager LAM;
-  FunctionAnalysisManager FAM;
-  CGSCCAnalysisManager CGAM;
-  ModuleAnalysisManager MAM;
-  PB.registerModuleAnalyses(MAM);
-  PB.registerFunctionAnalyses(FAM);
-  PB.registerLoopAnalyses(LAM);
-  PB.registerCGSCCAnalyses(CGAM);
-  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  SmallPtrSet<Function *, 1> LocalReadOnlyFunctions;
 
   for (Function &F : M) {
     SmallPtrSet<Function *, 1> calls_todo;
     auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
-    if (DetectReadonlyOrThrowFn(F, calls_todo, TLI)) {
+    bool local = false;
+    if (DetectReadonlyOrThrowFn(F, calls_todo, TLI, local)) {
+      if (local)
+        LocalReadOnlyFunctions.insert(&F);
       if (calls_todo.size() == 0) {
         changed = true;
         todo.push_back(&F);
@@ -1681,7 +1745,10 @@ bool DetectReadonlyOrThrow(Module &M) {
       auto &fwd_set = found2->second;
       fwd_set.erase(cur);
       if (fwd_set.size() == 0) {
-        F2->addFnAttr("enzyme_ReadOnlyOrThrow");
+        if (LocalReadOnlyFunctions.contains(F2))
+          F2->addFnAttr("enzyme_LocalReadOnlyOrThrow");
+        else
+          F2->addFnAttr("enzyme_ReadOnlyOrThrow");
         todo.push_back(F2);
         todo_map.erase(F2);
       }
@@ -2340,8 +2407,9 @@ Function *PreProcessCache::preprocessForClone(Function *F,
 
   {
     SmallPtrSet<Function *, 1> calls_todo;
+    bool local = false;
     DetectReadonlyOrThrowFn(*NewF, calls_todo,
-                            FAM.getResult<TargetLibraryAnalysis>(*NewF));
+                            FAM.getResult<TargetLibraryAnalysis>(*NewF), local);
   }
 
   if (EnzymePrint)
diff --git a/enzyme/Enzyme/FunctionUtils.h b/enzyme/Enzyme/FunctionUtils.h
@@ -64,8 +64,13 @@ extern llvm::cl::opt<bool> EnzymeAlwaysInlineDiff;
 // Perform an analysis to detect functions which only write to visible memory
 // outside the function if an error is not throw. Such a function can touch
 // inaccessible memory [e.g. the insides of malloc/etc], and the only violation
-// is whether existing memory before the call is written to. In other words,
-// malloc, calloc, copy_array, and friends, are all considered
+// is whether existing memory before the call is written to.
+// If non-local, returning memory written to is a violation (since it writes to
+// externally visible memory).
+// If local, returning memory written to is fine (since existing memory before
+// the call remains unchanged).
+// In other words, malloc [local and non-local], calloc [local and non-local],
+// copy_array [local only], and friends, are all considered
 // readonly_or_throw, as they only either read externally visible state, throw
 // an error, or write to inaccesible memory.
 bool DetectReadonlyOrThrow(llvm::Module &M);
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -2750,14 +2750,27 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
         assert(innerType == Type::getInt8Ty(malloc->getContext()));
       } else {
         if (innerType != malloc->getType()) {
-          llvm::errs() << *oldFunc << "\n";
-          llvm::errs() << *newFunc << "\n";
-          llvm::errs() << "innerType: " << *innerType << "\n";
-          llvm::errs() << "malloc->getType(): " << *malloc->getType() << "\n";
-          llvm::errs() << "ret: " << *ret << " - " << *ret->getType() << "\n";
-          llvm::errs() << "malloc: " << *malloc << "\n";
-          assert(0 && "illegal loop cache type");
-          llvm_unreachable("illegal loop cache type");
+          std::string str;
+          raw_string_ostream ss(str);
+          ss << "Illegal loop cache type:\n";
+          ss << *oldFunc << "\n";
+          ss << *newFunc << "\n";
+          ss << "innerType: " << *innerType << "\n";
+          ss << "malloc->getType(): " << *malloc->getType() << "\n";
+          ss << "ret: " << *ret << " - " << *ret->getType() << "\n";
+          ss << "malloc: " << *malloc << "\n";
+          if (CustomErrorHandler) {
+            CustomErrorHandler(str.c_str(), wrap(malloc),
+                               ErrorType::InternalError, nullptr, nullptr,
+                               nullptr);
+          } else {
+            DebugLoc loc;
+            if (auto I = dyn_cast<Instruction>(malloc))
+              EmitFailure("LoopCache", I->getDebugLoc(), I, ss.str());
+            else
+              EmitFailure("LoopCache", DebugLoc(), newFunc, ss.str());
+          }
+          return UndefValue::get(malloc->getType());
         }
       }
 
diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp
@@ -3867,7 +3867,7 @@ bool notCapturedBefore(llvm::Value *V, Instruction *inst,
   else
     VI = VI->getNextNode();
   SmallPtrSet<BasicBlock *, 1> regionBetween;
-  {
+  if (inst) {
     SmallVector<BasicBlock *, 1> todo;
     todo.push_back(VI->getParent());
     while (todo.size()) {
@@ -3893,15 +3893,17 @@ bool notCapturedBefore(llvm::Value *V, Instruction *inst,
     auto UI = std::get<0>(pair);
     auto level = std::get<1>(pair);
     auto prev = std::get<2>(pair);
-    if (!regionBetween.count(UI->getParent()))
-      continue;
-    if (UI->getParent() == VI->getParent()) {
-      if (UI->comesBefore(VI))
+    if (inst) {
+      if (!regionBetween.count(UI->getParent()))
         continue;
+      if (UI->getParent() == VI->getParent()) {
+        if (UI->comesBefore(VI))
+          continue;
+      }
+      if (UI->getParent() == inst->getParent())
+        if (inst->comesBefore(UI))
+          continue;
     }
-    if (UI->getParent() == inst->getParent())
-      if (inst->comesBefore(UI))
-        continue;
 
     if (isPointerArithmeticInst(UI, /*includephi*/ true,
                                 /*includebin*/ true)) {
@@ -3961,6 +3963,8 @@ bool notCapturedBefore(llvm::Value *V, Instruction *inst,
   return true;
 }
 
+bool notCaptured(llvm::Value *V) { return notCapturedBefore(V, nullptr, 0); }
+
 // Return true if guaranteed not to alias
 // Return false if guaranteed to alias [with possible offset depending on flag].
 // Return {} if no information is given.
diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h
@@ -1646,6 +1646,56 @@ static inline bool isReadOnly(const llvm::CallBase *call, ssize_t arg = -1) {
   return false;
 }
 
+// Whether the function does not write to memory visible before the function in
+// all cases that it doesn't error. In other words, the legal operations here
+// are:
+//.  1) Throw [in which case any operation guaranteed to throw is valid]
+//.  2) Read from any memory
+//.  3) Write to memory which did not exist did not exist prior to the function
+// call. This means that one can write .     to memory whose allocation happened
+// within the call to F (including a local alloca, a malloc call, even if .
+// returned). This is also legal to write to an sret and/or returnroots
+// parameter (which must be an alloca).
+static inline bool isLocalReadOnlyOrThrow(const llvm::Function *F) {
+  if (isReadOnly(F))
+    return true;
+
+  if (F->hasFnAttribute("enzyme_LocalReadOnlyOrThrow") ||
+      F->hasFnAttribute("enzyme_ReadOnlyOrThrow"))
+    return true;
+
+  return false;
+}
+
+static inline bool isLocalReadOnlyOrThrow(const llvm::CallBase *call) {
+  if (isReadOnly(call))
+    return true;
+
+  if (call->hasFnAttr("enzyme_LocalReadOnlyOrThrow") ||
+      call->hasFnAttr("enzyme_ReadOnlyOrThrow"))
+    return true;
+
+  if (auto F = getFunctionFromCall(call)) {
+    // Do not use function attrs for if different calling conv, such as a julia
+    // call wrapping args into an array. This is because the wrapped array
+    // may be nocapure/readonly, but the actual arg (which will be put in the
+    // array) may not be.
+    if (F->getCallingConv() == call->getCallingConv())
+      if (isLocalReadOnlyOrThrow(F))
+        return true;
+  }
+  return false;
+}
+
+// Whether the function does not write to memory visible outside the function in
+// all cases that it doesn't error. In other words, the legal operations here
+// are:
+//.  1) Throw [in which case any operation guaranteed to throw is valid]
+//.  2) Read from any memory
+//.  3) Write to memory which did not exist did not exist prior to the function
+// call. This means that one can write .     to memory whose lifetime is
+// entirely contained within F (including a local alloca, a malloc call locally
+// freed, but not .     a returned malloc call).
 static inline bool isReadOnlyOrThrow(const llvm::Function *F) {
   if (isReadOnly(F))
     return true;
@@ -2265,6 +2315,9 @@ bool isNVLoad(const llvm::Value *V);
 bool notCapturedBefore(llvm::Value *V, llvm::Instruction *inst,
                        size_t checkLoadCaptured);
 
+//! Check if value if b captured
+bool notCaptured(llvm::Value *V);
+
 // Return true if guaranteed not to alias
 // Return false if guaranteed to alias [with possible offset depending on flag].
 // Return {} if no information is given.
