Reuse intermediate loop variables on multi exit

EnzymeAD · Sep 19, 2023 · bf586a5 · bf586a5
1 parent 3870f89
commit bf586a5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -3686,6 +3686,7 @@ BasicBlock *GradientUtils::getReverseOrLatchMerge(BasicBlock *BB,
     IRBuilder<> tbuild(prevBlock);
 
     SmallVector<std::pair<Value *, Value *>, 1> lims;
+    ValueToValueMapTy available;
     for (auto I = exitingContexts.rbegin(), E = exitingContexts.rend(); I != E;
          I++) {
       auto &lc = *I;
@@ -3699,13 +3700,14 @@ BasicBlock *GradientUtils::getReverseOrLatchMerge(BasicBlock *BB,
         assert(/*ReverseLimit*/ reverseBlocks.size() > 0);
         LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                           lc.preheader);
-        lim = lookupValueFromCache(
-            lc.var->getType(),
-            /*forwardPass*/ false, tbuild, lctx, getDynamicLoopLimit(L),
-            /*isi1*/ false, /*available*/ ValueToValueMapTy());
+        lim = lookupValueFromCache(lc.var->getType(),
+                                   /*forwardPass*/ false, tbuild, lctx,
+                                   getDynamicLoopLimit(L),
+                                   /*isi1*/ false, available);
       } else {
-        lim = lookupM(lc.trueLimit, tbuild);
+        lim = lookupM(lc.trueLimit, tbuild, available);
       }
+      available[lc.var] = lim;
       lims.push_back(std::make_pair(lim, (Value *)lc.antivaralloc));
     }
 

diff --git a/enzyme/test/Enzyme/ReverseMode/multiloopexit.ll b/enzyme/test/Enzyme/ReverseMode/multiloopexit.ll
@@ -200,8 +200,7 @@ define float @c() {
 ; CHECK-NEXT:   store float %differeturn, float* %"sum.2'de", align 4
 ; CHECK-NEXT:   %[[i89:.+]] = load i64, i64* %loopLimit_cache, align 8
 ; CHECK-NEXT:   %[[i90:.+]] = load i64*, i64** %loopLimit_cache2, align 8
-; CHECK-NEXT:   %[[i91:.+]] = load i64, i64* %"iv'ac", align 4
-; CHECK-NEXT:   %[[i92:.+]] = getelementptr inbounds i64, i64* %[[i90]], i64 %[[i91]]
+; CHECK-NEXT:   %[[i92:.+]] = getelementptr inbounds i64, i64* %[[i90]], i64 %[[i89]]
 ; CHECK-NEXT:   %[[i93:.+]] = load i64, i64* %[[i92]], align 8
 ; CHECK-NEXT:   br label %mergeinvertloop2_exit