Cache only the last-iteration value of loop PHIs used after loop exit In split mode (augmented forward + reverse), a loop header PHI chosen to be cached is stored once per iteration of its loop. When the loop is not rotated, a running reduction (e.g. `s` in `for (...) s += x[i]*x[i]; return sqrt(s);`) is the header PHI itself and is used after the loop exits. The reverse pass then only ever reads the entry of the final iteration, yet the tape holds trip-count + 1 values (per iteration of any enclosing loop). For a mag() call inside an O(n) loop this made the tape O(n^2) (2 GB at n = 16000). A use outside the PHI's loop can only observe the value of the PHI's last execution before the loop was exited (by dominance the use follows the loop within the same iteration of every enclosing loop, and the loop cannot be re-entered without a new iteration of the enclosing loop), independent of the number of exits. So, if no reverse-pass computation inside the loop needs the PHI, it suffices to keep one slot per iteration of the enclosing loops that is overwritten on every iteration; the slot holds the exit value. computeMinCache (augmented pass only) records such PHIs in cacheOnlyLastIteration: cached (knownRecomputeHeuristic == false), non-pointer header PHIs, not needed for loop bounds, all of whose in-loop users are non-terminator, non-memory, non-pointer instructions that neither use the PHI directly in their adjoint nor are needed in the reverse pass. For these, cacheForReverse uses the loop preheader as cache scope (a plain tape value if the loop is outermost). The decision is stored in AugmentedReturn and reused by the reverse pass and the split forward-mode derivative, so all tape users agree on the layout. Anything else keeps the previous behaviour. Combined mode is unaffected (it already uses fixLCSSA's manual_lcssa phi). On the reproducer (normalize/mag with only SROA before Enzyme, n = 16000): peak RSS 1956 MB -> 12 MB, gradient/primal time 5.9x -> 1.2x, identical gradient. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019RNZGgCdK1fbT2WPQpR7Br
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp index b259d3c..07d3475 100644 --- a/enzyme/Enzyme/EnzymeLogic.cpp +++ b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -2174,13 +2174,15 @@ else bb.CreateRet(cal); - return insert_or_assign<AugmentedCacheKey, AugmentedReturn>( - AugmentedCachedFunctions, tup, - AugmentedReturn(NewF, aug.tapeType, aug.tapeIndices, - aug.returns, aug.overwritten_args_map, - aug.can_modref_map, next_constant_args, - shadowReturnUsed)) - ->second; + auto &res = + insert_or_assign<AugmentedCacheKey, AugmentedReturn>( + AugmentedCachedFunctions, tup, + AugmentedReturn(NewF, aug.tapeType, aug.tapeIndices, aug.returns, + aug.overwritten_args_map, aug.can_modref_map, + next_constant_args, shadowReturnUsed)) + ->second; + res.lastIterationCached = aug.lastIterationCached; + return res; } if (foundcalled->hasStructRetAttr() && !todiff->hasStructRetAttr()) { @@ -2519,6 +2521,9 @@ AugmentedReturn(gutils->newFunc, nullptr, {}, returnMapping, overwritten_args_map, can_modref_map, constant_args, shadowReturnUsed)); + AugmentedCachedFunctions.find(tup)->second.lastIterationCached.insert( + gutils->cacheOnlyLastIteration.begin(), + gutils->cacheOnlyLastIteration.end()); auto getIndex = [&](Instruction *I, CacheType u, IRBuilder<> &B) -> unsigned { return gutils->getIndex( @@ -4331,6 +4336,10 @@ // requires is_value_needed_in_reverse, that needs unnecessaryValues // sets knownRecomputeHeuristic gutils->computeMinCache(); + if (augmenteddata) + gutils->cacheOnlyLastIteration.insert( + augmenteddata->lastIterationCached.begin(), + augmenteddata->lastIterationCached.end()); // Requires knownRecomputeCache to be set as call to getContext // itself calls createCacheForScope @@ -5003,6 +5012,10 @@ // requires is_value_needed_in_reverse, that needs unnecessaryValues // sets knownRecomputeHeuristic gutils->computeMinCache(); + if (augmenteddata) + gutils->cacheOnlyLastIteration.insert( + augmenteddata->lastIterationCached.begin(), + augmenteddata->lastIterationCached.end()); // Requires knownRecomputeCache to be set as call to getContext // itself calls createCacheForScope
diff --git a/enzyme/Enzyme/EnzymeLogic.h b/enzyme/Enzyme/EnzymeLogic.h index 5491e0b..b30045a 100644 --- a/enzyme/Enzyme/EnzymeLogic.h +++ b/enzyme/Enzyme/EnzymeLogic.h
@@ -119,6 +119,12 @@ std::set<ssize_t> tapeIndiciesToFree; + //! Loop header PHIs (of the original function) whose tape entry only holds + //! the value of the final iteration of their loop (see + //! GradientUtils::cacheOnlyLastIteration). Users of the tape must use the + //! same cache layout. + std::set<const llvm::Instruction *> lastIterationCached; + const std::vector<DIFFE_TYPE> constant_args; bool shadowReturnUsed;
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index 1b2838c..18660b4 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp
@@ -2696,6 +2696,28 @@ return lcssaPHI; } +/// If the value `malloc` (in the new function) only needs its value from the +/// final iteration of its loop to be cached (see cacheOnlyLastIteration), +/// return the block whose loop nest should define the cache instead of the +/// defining block: the preheader of the value's loop. This keeps one value per +/// iteration of the enclosing loops, overwritten on every iteration of the +/// value's own loop. Otherwise return null (use the default scope). +BasicBlock *GradientUtils::getCacheScopeForReverse(Value *malloc) { + auto inst = dyn_cast<Instruction>(malloc); + if (!inst) + return nullptr; + if (cacheOnlyLastIteration.size() && isa<PHINode>(inst)) { + auto orig = isOriginal(inst); + if (orig && cacheOnlyLastIteration.count(orig)) { + LoopContext lc; + if (getContext(inst->getParent(), lc) && lc.header == inst->getParent() && + isOriginalBlock(*lc.preheader)) + return lc.preheader; + } + } + return nullptr; +} + Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc, int idx, bool replace) { assert(malloc); @@ -2776,15 +2798,18 @@ return UndefValue::get(retType); } + BasicBlock *cacheScope = getCacheScopeForReverse(malloc); LimitContext ctx(/*ReverseLimit*/ reverseBlocks.size() > 0, BuilderQ.GetInsertBlock()); if (auto inst = dyn_cast<Instruction>(malloc)) ctx = LimitContext(/*ReverseLimit*/ reverseBlocks.size() > 0, - inst->getParent()); + cacheScope ? cacheScope : inst->getParent()); if (auto found = findInMap(scopeMap, malloc)) { ctx = found->second; } assert(isOriginalBlock(*ctx.Block)); + if (!cacheScope) + cacheScope = BuilderQ.GetInsertBlock(); bool inLoop; if (ctx.ForceSingleIteration) { @@ -2825,7 +2850,7 @@ /*inForwardPass*/ true, nullptr, LimitContext( /*ReverseLimit*/ reverseBlocks.size() > 0, - BuilderQ.GetInsertBlock())) + cacheScope)) .size(); i < limit; ++i) { if (!isa<PointerType>(innerType)) { @@ -2893,8 +2918,7 @@ } } - LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0, - BuilderQ.GetInsertBlock()); + LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0, cacheScope); AllocaInst *cache = createCacheForScope(lctx, innerType, "mdyncache_fromtape", ((DiffeGradientUtils *)this)->FreeMemory, false); @@ -3100,11 +3124,12 @@ return malloc; } + BasicBlock *cacheScope = getCacheScopeForReverse(malloc); LimitContext ctx(/*ReverseLimit*/ reverseBlocks.size() > 0, BuilderQ.GetInsertBlock()); if (auto inst = dyn_cast<Instruction>(malloc)) ctx = LimitContext(/*ReverseLimit*/ reverseBlocks.size() > 0, - inst->getParent()); + cacheScope ? cacheScope : inst->getParent()); if (auto found = findInMap(scopeMap, malloc)) { ctx = found->second; } @@ -3144,7 +3169,7 @@ ensureLookupCached( cast<Instruction>(malloc), /*shouldFree=*/reverseBlocks.size() > 0, - /*scope*/ nullptr, + /*scope*/ cacheScope, cast<Instruction>(malloc)->getMetadata(LLVMContext::MD_tbaa)); auto found2 = scopeMap.find(malloc); assert(found2 != scopeMap.end()); @@ -3171,7 +3196,8 @@ /*inForwardPass*/ true, nullptr, LimitContext( /*ReverseLimit*/ reverseBlocks.size() > 0, - BuilderQ.GetInsertBlock())) + cacheScope ? cacheScope + : BuilderQ.GetInsertBlock())) .size(); i < limit; ++i) { innerType = innerType->getPointerElementType(); @@ -8871,6 +8897,84 @@ } } } + + // A loop header PHI that is cached (rather than recomputed) is by default + // cached once per iteration of its loop. If however the only reverse-pass + // uses of the PHI are outside of its loop (e.g. the running sum of an + // unrotated reduction loop, used after the loop exits), those uses can + // only ever observe the value of the last execution of the PHI before the + // loop was exited. By SSA dominance such a use executes after the loop + // within the same iteration of every enclosing loop, and the loop cannot + // be re-entered without starting a new iteration of the enclosing loop. + // Thus it suffices to cache a single value per iteration of the + // enclosing loops, overwritten on each iteration of the PHI's loop. This + // holds independently of the number of exits of the loop. + // + // This is only legal if no reverse-pass computation within the loop needs + // the PHI: every user inside the loop must neither use the PHI directly + // in its adjoint, nor be needed (and thus potentially recomputed from the + // PHI) in the reverse pass, nor be control flow / memory / pointer + // related. This decision is made once when creating the augmented + // forward pass, and recorded in its AugmentedReturn such that all users + // of the tape (the reverse pass and the split forward-mode derivative) + // use the same cache layout. The combined mode is not affected, as it + // already looks such values up through an LCSSA phi (see fixLCSSA). + if (mode == DerivativeMode::ReverseModePrimal) { + std::map<UsageKey, bool> KnownSeen = populateSeenFromKnownRecompute(); + for (BasicBlock &BB : *oldFunc) { + if (notForAnalysis.count(&BB)) + continue; + Loop *L = OrigLI->getLoopFor(&BB); + if (!L || L->getHeader() != &BB) + continue; + for (PHINode &PN : BB.phis()) { + auto found = knownRecomputeHeuristic.find(&PN); + if (found == knownRecomputeHeuristic.end() || found->second) + continue; + if (PN.getType()->isPtrOrPtrVectorTy() || + !PN.getType()->isSingleValueType()) + continue; + if (NewLoopBoundReq.count(&PN)) + continue; + bool legal = true; + for (User *U : PN.users()) { + if (U == &PN) + continue; + auto UI = dyn_cast<Instruction>(U); + if (!UI) { + legal = false; + break; + } + if (notForAnalysis.count(UI->getParent())) + continue; + if (!L->contains(UI->getParent())) + continue; + if (UI->isTerminator() || UI->mayReadOrWriteMemory() || + UI->getType()->isPtrOrPtrVectorTy() || + !(UI->getType()->isSingleValueType()) || + NewLoopBoundReq.count(UI)) { + legal = false; + break; + } + if (DifferentialUseAnalysis::is_use_directly_needed_in_reverse( + this, &PN, minCutMode, UI, notForAnalysis, + QueryType::Primal)) { + legal = false; + break; + } + std::map<UsageKey, bool> Seen = KnownSeen; + if (DifferentialUseAnalysis::is_value_needed_in_reverse< + QueryType::Primal>(this, UI, minCutMode, Seen, + notForAnalysis)) { + legal = false; + break; + } + } + if (legal) + cacheOnlyLastIteration.insert(&PN); + } + } + } } if (rematerializableAllocations.size()) {
diff --git a/enzyme/Enzyme/GradientUtils.h b/enzyme/Enzyme/GradientUtils.h index f023b8f..d517f7d 100644 --- a/enzyme/Enzyme/GradientUtils.h +++ b/enzyme/Enzyme/GradientUtils.h
@@ -185,6 +185,16 @@ llvm::SmallPtrSet<llvm::Instruction *, 4> unnecessaryIntermediates; + /// Loop header PHIs (of the original function) that are cached for the + /// reverse pass, but whose value is only required in the reverse pass by + /// users outside of their loop. Such users can only observe the value of + /// the final iteration, so a single cache slot per iteration of the + /// enclosing loops (overwritten on every iteration of the PHI's own loop) + /// suffices, instead of a cache with one entry per iteration. Computed when + /// creating the augmented forward pass and copied from its AugmentedReturn + /// for the passes that consume its tape. + llvm::SmallPtrSet<const llvm::Instruction *, 4> cacheOnlyLastIteration; + const std::map<llvm::Instruction *, bool> *can_modref_map; const std::map<llvm::CallInst *, std::pair<bool, const std::vector<bool>>> *overwritten_args_map_ptr; @@ -382,6 +392,8 @@ std::map<std::pair<llvm::Instruction *, CacheType>, int> &mapping, llvm::IRBuilder<> &); + llvm::BasicBlock *getCacheScopeForReverse(llvm::Value *malloc); + llvm::Value *cacheForReverse(llvm::IRBuilder<> &BuilderQ, llvm::Value *malloc, int idx, bool replace = true);
diff --git a/enzyme/test/Enzyme/ReverseMode/loopexitcache.ll b/enzyme/test/Enzyme/ReverseMode/loopexitcache.ll new file mode 100644 index 0000000..e0fe37d --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/loopexitcache.ll
@@ -0,0 +1,142 @@ +; RUN: if [ %llvmver -ge 16 ]; then %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -enzyme-preopt=false -S | FileCheck %s; fi + +; The loops in @mag and @nested are not rotated: the running sum %s is a loop +; header phi that is only used after the loop exits. Only its value from the +; final iteration is needed in the reverse pass, so the augmented forward pass +; should cache a single value (per iteration of any enclosing loop) rather than +; one value per iteration of the loop defining %s. + +declare double @llvm.sqrt.f64(double) + +define double @mag(ptr noalias %x, i64 %n) { +entry: + br label %loop + +loop: + %i = phi i64 [ 0, %entry ], [ %inc, %body ] + %s = phi double [ 0.000000e+00, %entry ], [ %add, %body ] + %cmp = icmp slt i64 %i, %n + br i1 %cmp, label %body, label %exit + +body: + %gep = getelementptr inbounds double, ptr %x, i64 %i + %ld = load double, ptr %gep, align 8 + %mul = fmul double %ld, %ld + %add = fadd double %s, %mul + %inc = add nuw nsw i64 %i, 1 + br label %loop + +exit: + %r = call double @llvm.sqrt.f64(double %s) + ret double %r +} + +; The inner loop has two exits, neither of which depends on %s. +define void @nested(ptr noalias %out, ptr noalias %x, i64 %n) { +entry: + br label %outer + +outer: + %j = phi i64 [ 0, %entry ], [ %j.next, %outer.latch ] + %xj.p = getelementptr inbounds double, ptr %x, i64 %j + %xj = load double, ptr %xj.p, align 8 + br label %inner + +inner: + %i = phi i64 [ 0, %outer ], [ %inc, %body ] + %s = phi double [ %xj, %outer ], [ %add, %body ] + %cmp = icmp slt i64 %i, %n + br i1 %cmp, label %check, label %outer.latch + +check: + %gep = getelementptr inbounds double, ptr %x, i64 %i + %ld = load double, ptr %gep, align 8 + %neg = fcmp olt double %ld, 0.000000e+00 + br i1 %neg, label %outer.latch, label %body + +body: + %mul = fmul double %ld, %ld + %add = fadd double %s, %mul + %inc = add nuw nsw i64 %i, 1 + br label %inner + +outer.latch: + %sq = call double @llvm.sqrt.f64(double %s) + %res = fmul double %xj, %sq + %o.p = getelementptr inbounds double, ptr %out, i64 %j + store double %res, ptr %o.p, align 8 + %j.next = add nuw nsw i64 %j, 1 + %cmp2 = icmp slt i64 %j.next, %n + br i1 %cmp2, label %outer, label %exit + +exit: + ret void +} + +define void @normalize(ptr noalias %out2, ptr noalias %out, ptr noalias %x, i64 %n) { +entry: + br label %loop + +loop: + %j = phi i64 [ 0, %entry ], [ %j.next, %loop ] + %m = call double @mag(ptr %x, i64 %n) + call void @nested(ptr %out2, ptr %x, i64 %n) + %xj.p = getelementptr inbounds double, ptr %x, i64 %j + %xj = load double, ptr %xj.p, align 8 + %div = fdiv double %xj, %m + %o.p = getelementptr inbounds double, ptr %out, i64 %j + store double %div, ptr %o.p, align 8 + %j.next = add nuw nsw i64 %j, 1 + %cmp = icmp slt i64 %j.next, %n + br i1 %cmp, label %loop, label %exit + +exit: + ret void +} + +declare void @__enzyme_autodiff(ptr, ...) + +define void @dnormalize(ptr %out2, ptr %dout2, ptr %out, ptr %dout, ptr %x, ptr %dx, i64 %n) { +entry: + call void (ptr, ...) @__enzyme_autodiff(ptr @normalize, metadata !"enzyme_dup", ptr %out2, ptr %dout2, metadata !"enzyme_dup", ptr %out, ptr %dout, metadata !"enzyme_dup", ptr %x, ptr %dx, i64 %n) + ret void +} + +; For the nested loop, one value per iteration of the outer loop is cached +; (n doubles), instead of a growing per-inner-iteration buffer per outer +; iteration. +; CHECK: define internal ptr @augmented_nested(ptr noalias writeonly captures(none) %out, ptr captures(none) %"out'", ptr noalias readonly captures(none) %x, ptr captures(none) %"x'", i64 %n) +; CHECK-NEXT: entry: +; CHECK-NEXT: %smax = call i64 @llvm.smax.i64(i64 %n, i64 1) +; CHECK-NEXT: %mallocsize = mul nuw nsw i64 %smax, 8 +; CHECK-NEXT: %s_malloccache = tail call noalias nonnull ptr @malloc(i64 %mallocsize) +; CHECK-NEXT: br label %outer +; CHECK: inner: +; CHECK-NEXT: %iv1 = phi i64 [ %iv.next2, %body ], [ 0, %outer ] +; CHECK-NEXT: %s = phi double [ %xj, %outer ], [ %add, %body ] +; CHECK-NEXT: %[[gep:.+]] = getelementptr inbounds double, ptr %s_malloccache, i64 %iv +; CHECK-NEXT: store double %s, ptr %[[gep]] +; CHECK-NOT: realloc +; CHECK: ret ptr %s_malloccache + +; The tape of @mag is the final running sum itself (no per-iteration cache), +; overwritten in the loop header on every iteration. +; CHECK: define internal { double, double } @augmented_mag(ptr noalias readonly captures(none) %x, ptr captures(none) %"x'", i64 %n) +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = alloca { double, double } +; CHECK-NEXT: br label %loop +; CHECK: loop: +; CHECK-NEXT: %iv = phi i64 [ %iv.next, %body ], [ 0, %entry ] +; CHECK-NEXT: %s = phi double [ 0.000000e+00, %entry ], [ %add, %body ] +; CHECK-NEXT: store double %s, ptr %0 +; CHECK-NOT: malloc +; CHECK: exit: +; CHECK-NEXT: %r = call double @llvm.sqrt.f64(double %s) +; CHECK: ret { double, double } + +; CHECK: define internal void @diffemag(ptr noalias readonly captures(none) %x, ptr captures(none) %"x'", i64 %n, double %differeturn, double %s) +; CHECK-NOT: malloc +; CHECK-NOT: free +; CHECK: invertexit: +; CHECK-NEXT: %[[cmp:.+]] = fcmp fast ueq double %s, 0.000000e+00 +; CHECK-NEXT: %[[sq:.+]] = call fast double @llvm.sqrt.f64(double %s)
diff --git a/enzyme/test/Enzyme/ReverseMode/mincutoneneed.ll b/enzyme/test/Enzyme/ReverseMode/mincutoneneed.ll index 414ee1d..0e01bc4 100644 --- a/enzyme/test/Enzyme/ReverseMode/mincutoneneed.ll +++ b/enzyme/test/Enzyme/ReverseMode/mincutoneneed.ll
@@ -41,14 +41,14 @@ ret double %tmp60 } -; CHECK: define internal { double } @diffemid(double %x, double* %xp, double* %"xp'", double %differeturn, double* %tapeArg) +; Only the final value of %tmp24 is used (after the loop), so the tape holds +; just that value rather than one value per iteration. +; CHECK: define internal { double } @diffemid(double %x, double* %xp, double* %"xp'", double %differeturn, double %tmp24) ; CHECK-NEXT: entry: ; CHECK-NEXT: br label %loop ; CHECK: loop: ; preds = %loop, %entry ; CHECK-NEXT: %iv = phi i64 [ %iv.next, %loop ], [ 0, %entry ] -; CHECK-NEXT: %0 = getelementptr inbounds double, double* %tapeArg, i64 %iv -; CHECK-NEXT: %tmp24 = load double, double* %0, align 8, !invariant.group ! ; CHECK-NEXT: %iv.next = add nuw nsw i64 %iv, 1 ; CHECK-NEXT: %tmp27 = icmp eq i64 %iv, 16 ; CHECK-NEXT: br i1 %tmp27, label %invertexit, label %loop @@ -58,8 +58,6 @@ ; CHECK-NEXT: %[[a1:.+]] = fadd fast double %[[a0]], %[[a3:.+]] ; CHECK-NEXT: store double %[[a1]], double* %"xp'", align 8 ; CHECK-NEXT: %[[a2:.+]] = insertvalue { double } undef, double %[[m1diffex:.+]], 0 -; CHECK-NEXT: %4 = bitcast double* %tapeArg to i8* -; CHECK-NEXT: tail call void @free(i8* nonnull %4) ; CHECK-NEXT: ret { double } %[[a2]] ; CHECK: invertloop: ; preds = %invertexit, %incinvertloop