Loop-level rematerialization (rust-lang#516)

* Add primitive loop rematerialization

* Rematerialize loops

* Add and fix test

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -767,7 +767,7 @@ class AdjointGenerator
     for (auto pair : gutils->rematerializableAllocations) {
       if (is_value_needed_in_reverse<ValueType::Primal>(
               TR, gutils, pair.first, Mode, Seen, oldUnreachable)) {
-        if (pair.second.second.count(&SI)) {
+        if (pair.second.stores.count(&SI)) {
           return;
         }
       }
@@ -2539,7 +2539,7 @@ class AdjointGenerator
     for (auto pair : gutils->rematerializableAllocations) {
       if (is_value_needed_in_reverse<ValueType::Primal>(
               TR, gutils, pair.first, Mode, Seen, oldUnreachable)) {
-        if (pair.second.second.count(&MS)) {
+        if (pair.second.stores.count(&MS)) {
           rematerialized = true;
           break;
         }
@@ -5398,8 +5398,12 @@ class AdjointGenerator
                                        Builder2, /*lookup*/ true));
         cal->setCallingConv(dwait->getCallingConv());
         cal->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
+#if LLVM_VERSION_MAJOR >= 14
+        cal->addFnAttr(Attribute::AlwaysInline);
+#else
         cal->addAttribute(AttributeList::FunctionIndex,
                           Attribute::AlwaysInline);
+#endif
         Builder2.CreateBr(endBlock);
 
         Builder2.SetInsertPoint(endBlock);
@@ -5536,8 +5540,12 @@ class AdjointGenerator
                                        Builder2, /*lookup*/ true));
         cal->setCallingConv(dwait->getCallingConv());
         cal->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
+#if LLVM_VERSION_MAJOR >= 14
+        cal->addFnAttr(Attribute::AlwaysInline);
+#else
         cal->addAttribute(AttributeList::FunctionIndex,
                           Attribute::AlwaysInline);
+#endif
         Builder2.CreateBr(eloopBlock);
 
         Builder2.SetInsertPoint(eloopBlock);
@@ -8060,7 +8068,7 @@ class AdjointGenerator
               Seen[UsageKey(pair.first, ValueType::Primal)] = false;
           bool rematerializedPrimal = false;
           for (auto pair : gutils->rematerializableAllocations) {
-            if (pair.second.second.count(orig) &&
+            if (pair.second.stores.count(orig) &&
                 is_value_needed_in_reverse<ValueType::Primal>(
                     TR, gutils, pair.first, Mode, Seen, oldUnreachable)) {
               rematerializedPrimal = true;
@@ -8818,7 +8826,7 @@ class AdjointGenerator
 
       // If a rematerializable allocation.
       for (auto rmat : gutils->rematerializableAllocations) {
-        if (rmat.second.second.count(orig)) {
+        if (rmat.second.stores.count(orig)) {
           // Leave the original free behavior since this won't be used
           // in the reverse pass in split mode
           if (Mode == DerivativeMode::ReverseModePrimal) {

enzyme/Enzyme/DifferentialUseAnalysis.h

@@ -364,8 +364,8 @@ static inline bool is_value_needed_in_reverse(
           // we'll set it to unused, then check the gep, then here we'll
           // directly say unused by induction instead of checking the final
           // loads.
-          if (pair.second.second.count(SI))
-            for (LoadInst *L : pair.second.first)
+          if (pair.second.stores.count(SI))
+            for (LoadInst *L : pair.second.loads)
               if (is_value_needed_in_reverse<VT>(TR, gutils, L, mode, seen,
                                                  oldUnreachable)) {
                 return seen[idx] = true;
@@ -535,21 +535,20 @@ static inline int cmpLoopNest(Loop *prev, Loop *next) {
   return -1;
 }
 
-static inline void
-minCut(const DataLayout &DL, LoopInfo &OrigLI,
-       const SmallPtrSetImpl<Value *> &Recomputes,
-       const SmallPtrSetImpl<Value *> &Intermediates,
-       SmallPtrSetImpl<Value *> &Required, SmallPtrSetImpl<Value *> &MinReq,
-       const ValueMap<Value *, std::pair<SmallPtrSet<LoadInst *, 1>,
-                                         SmallPtrSet<Instruction *, 1>>>
-           &rematerializableAllocations) {
+static inline void minCut(const DataLayout &DL, LoopInfo &OrigLI,
+                          const SmallPtrSetImpl<Value *> &Recomputes,
+                          const SmallPtrSetImpl<Value *> &Intermediates,
+                          SmallPtrSetImpl<Value *> &Required,
+                          SmallPtrSetImpl<Value *> &MinReq,
+                          const ValueMap<Value *, GradientUtils::Rematerializer>
+                              &rematerializableAllocations) {
   Graph G;
   for (auto V : Intermediates) {
     G[Node(V, false)].insert(Node(V, true));
     for (auto U : V->users()) {
       if (auto I = dyn_cast<Instruction>(U)) {
         for (auto pair : rematerializableAllocations) {
-          if (Intermediates.count(pair.first) && pair.second.second.count(I))
+          if (Intermediates.count(pair.first) && pair.second.stores.count(I))
             G[Node(V, true)].insert(Node(pair.first, false));
         }
       }
@@ -560,7 +559,7 @@ minCut(const DataLayout &DL, LoopInfo &OrigLI,
   }
   for (auto pair : rematerializableAllocations) {
     if (Intermediates.count(pair.first)) {
-      for (LoadInst *L : pair.second.first) {
+      for (LoadInst *L : pair.second.loads) {
         if (Intermediates.count(L)) {
           G[Node(pair.first, true)].insert(Node(L, false));
         }

enzyme/Enzyme/EnzymeLogic.cpp

@@ -105,8 +105,7 @@ struct CacheAnalysis {
 
   const ValueMap<const CallInst *, SmallPtrSet<const CallInst *, 1>>
       &allocationsWithGuaranteedFree;
-  const ValueMap<Value *, std::pair<SmallPtrSet<LoadInst *, 1>,
-                                    SmallPtrSet<Instruction *, 1>>>
+  const ValueMap<Value *, GradientUtils::Rematerializer>
       &rematerializableAllocations;
   TypeResults &TR;
   AAResults &AA;
@@ -123,8 +122,7 @@ struct CacheAnalysis {
   CacheAnalysis(
       const ValueMap<const CallInst *, SmallPtrSet<const CallInst *, 1>>
           &allocationsWithGuaranteedFree,
-      const ValueMap<Value *, std::pair<SmallPtrSet<LoadInst *, 1>,
-                                        SmallPtrSet<Instruction *, 1>>>
+      const ValueMap<Value *, GradientUtils::Rematerializer>
           &rematerializableAllocations,
       TypeResults &TR, AAResults &AA, Function *oldFunc, ScalarEvolution &SE,
       LoopInfo &OrigLI, DominatorTree &OrigDT, TargetLibraryInfo &TLI,
@@ -269,6 +267,13 @@ struct CacheAnalysis {
         }
       }
 
+    // Any load from a rematerializable allocation is definitionally
+    // reloadable. Notably we don't need to perform the allFollowers
+    // of check as the loop scope caching should allow us to ignore
+    // such stores.
+    if (rematerializableAllocations.count(obj))
+      return false;
+
     // If not running combined, check if pointer operand is overwritten
     // by a subsequent call (i.e. not this function).
     bool can_modref = false;