diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index ef6799c79ba4..8901f782a88c 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp @@ -3686,6 +3686,7 @@ BasicBlock *GradientUtils::getReverseOrLatchMerge(BasicBlock *BB, IRBuilder<> tbuild(prevBlock); SmallVector, 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 index 463247f0bc7d..b1e71bcb0a69 100644 --- 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