Keep the explicit cache decision for a loopy reduction's start value computeMinCache runs a min cut over Intermediates, then marks everything in that set as recomputed or cached according to which side of the cut it landed on: knownRecomputeHeuristic[V] = !MinReq.count(V); Membership of Intermediates therefore carries an invariant, established where the set is built: a value only joins it if it is legal to recompute, since anything not selected by the cut will be. Values which are not legal to recompute are excluded there, having already been given an explicit knownRecomputeHeuristic[V] = false by the scan above. pushLoopyPHIPreheader broke that invariant. Added in 04b80f67 to pull a loopy reduction's start value into the need graph, it inserted straight into Intermediates -- unconditionally, and after the cut had already been computed, so the value could never be selected for caching. For a start value that is not legal to recompute, that overwrote a correct "cache this" decision with "recompute this", and computeMinCache then caught its own inconsistency: Assertion `legalRecompute(V, Available2, nullptr)' failed. In the reported case the start value is a load whose memory may be written before the reverse pass runs, so it must be cached. Have pushLoopyPHIPreheader collect the values instead, and let computeMinCache apply the same legality gate the main worklist uses. A value that passes joins Intermediates as before; one that does not keeps the decision already made for it. Either way it is pushed onto the worklist, so the need graph is unchanged. Fixes #3040 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01P65yn8LELKWU1dq6AQvA5f
diff --git a/enzyme/Enzyme/DifferentialUseAnalysis.cpp b/enzyme/Enzyme/DifferentialUseAnalysis.cpp index 7510cf5..9beefe5 100644 --- a/enzyme/Enzyme/DifferentialUseAnalysis.cpp +++ b/enzyme/Enzyme/DifferentialUseAnalysis.cpp
@@ -1285,9 +1285,9 @@ return true; } -void pushLoopyPHIPreheader(const GradientUtils *gutils, llvm::Value *V, - llvm::SetVector<llvm::Value *> &Intermediates, - std::deque<llvm::Value *> &todo) { +void pushLoopyPHIPreheader( + const GradientUtils *gutils, llvm::Value *V, + llvm::SmallVectorImpl<llvm::Value *> &preheaderVals) { using namespace llvm; if (auto P0 = dyn_cast<PHINode>(V)) { if (!gutils->OrigLI) @@ -1348,8 +1348,7 @@ // instructions). if (!isa<Instruction>(Pstart)) break; - Intermediates.insert(Pstart); - todo.push_back(Pstart); + preheaderVals.push_back(Pstart); if (auto phi = dyn_cast<PHINode>(Pstart)) { if (phi->getNumIncomingValues() == 1) Pstart = phi->getIncomingValue(0);
diff --git a/enzyme/Enzyme/DifferentialUseAnalysis.h b/enzyme/Enzyme/DifferentialUseAnalysis.h index 17fb4a8..dae104f 100644 --- a/enzyme/Enzyme/DifferentialUseAnalysis.h +++ b/enzyme/Enzyme/DifferentialUseAnalysis.h
@@ -84,9 +84,10 @@ const llvm::PHINode *P0, const llvm::Value *incomingVal); +/// Collect the preheader-incoming values of V, if V is the header phi of a +/// loopy reduction whose start value is needed in the reverse pass. void pushLoopyPHIPreheader(const GradientUtils *gutils, llvm::Value *V, - llvm::SetVector<llvm::Value *> &Intermediates, - std::deque<llvm::Value *> &todo); + llvm::SmallVectorImpl<llvm::Value *> &preheaderVals); template <QueryType VT, bool OneLevel = false> inline bool is_value_needed_in_reverse(
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index 9fec75e..d847ed0 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp
@@ -8745,10 +8745,36 @@ Required, MinReq, this, TLI); SmallPtrSet<Value *, 5> NeedGraph; + // The start value of a loopy reduction is needed in the reverse pass, but + // is only discovered now that the min cut is known. Since the cut has + // already been computed, such a value can no longer be selected for + // caching -- everything placed into Intermediates from here on will be + // marked as recomputed below. Only add it if that is legal; if it is not, + // it has already been given an explicit cache decision by the scan above, + // which must not be overwritten. Either way push it onto the worklist so + // the need graph stays complete. + auto pushLoopyPHIPreheader = [&](Value *V) { + SmallVector<Value *, 2> preheaderVals; + DifferentialUseAnalysis::pushLoopyPHIPreheader(this, V, preheaderVals); + for (Value *PV : preheaderVals) { + ValueToValueMapTy Available2; + for (auto a : Available) + Available2[a.first] = a.second; + for (Loop *L = OrigLI->getLoopFor(cast<Instruction>(PV)->getParent()); + L != nullptr; L = L->getParentLoop()) { + for (auto v : LoopAvail[L]) { + Available2[v] = v; + } + } + if (legalRecompute(PV, Available2, nullptr)) + Intermediates.insert(PV); + todo.push_back(PV); + } + }; + for (Value *V : MinReq) { NeedGraph.insert(V); - DifferentialUseAnalysis::pushLoopyPHIPreheader(this, V, Intermediates, - todo); + pushLoopyPHIPreheader(V); } for (Value *V : Required) { todo.push_back(V); @@ -8759,8 +8785,7 @@ if (NeedGraph.count(V)) continue; NeedGraph.insert(V); - DifferentialUseAnalysis::pushLoopyPHIPreheader(this, V, Intermediates, - todo); + pushLoopyPHIPreheader(V); auto I = dyn_cast<Instruction>(V); if (!I) continue;
diff --git a/enzyme/test/Enzyme/ReverseMode/loopy-reduction-cached-preheader.ll b/enzyme/test/Enzyme/ReverseMode/loopy-reduction-cached-preheader.ll new file mode 100644 index 0000000..af0ed89 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/loopy-reduction-cached-preheader.ll
@@ -0,0 +1,52 @@ +; RUN: if [ %llvmver -ge 15 ]; then %opt < %s %OPnewLoadEnzyme -passes="enzyme" -enzyme-preopt=false -S | FileCheck %s; fi + +; The start value of a loopy reduction is needed in the reverse pass, but is +; only discovered after the min cut has been computed. Here that start value +; (%seed) is a load which is not legal to recompute, so it must keep the +; explicit cache decision made for it rather than being marked as recomputed. + +target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128" + +@enzyme_primal_return = external global ptr +@enzyme_dup = external global ptr +@enzyme_const = external global ptr + +define hidden double @wrap(ptr nocapture readonly %diag, ptr nocapture writeonly %out) { +start: + %seed = load double, ptr %diag, align 8 + br label %loop + +loop: + %i = phi i64 [ 0, %start ], [ %i.next, %loop ] + %gu = phi double [ %seed, %start ], [ %gu.next, %loop ] + %i.next = add nuw nsw i64 %i, 1 + %p = getelementptr inbounds double, ptr %diag, i64 %i + %v = load double, ptr %p, align 8 + %cmp = fcmp ogt double %gu, %v + %vplus = fadd double %v, 0.000000e+00 + %gu.next = select i1 %cmp, double %gu, double %vplus + %done = icmp eq i64 %i.next, 8 + br i1 %done, label %exit, label %loop + +exit: + %gu.final = phi double [ %gu.next, %loop ] + %isneg = fcmp ogt double %gu.final, 0.000000e+00 + %sigma = select i1 %isneg, double 0.000000e+00, double %gu.final + %slot = getelementptr inbounds i8, ptr %out, i64 56 + store double %sigma, ptr %slot, align 8 + ret double 0.000000e+00 +} + +define double @entry(ptr %diag, ptr %ddiag, ptr %out) { + %r = tail call double (...) @__enzyme_autodiff(ptr @wrap, ptr @enzyme_primal_return, ptr @enzyme_dup, ptr %diag, ptr %ddiag, ptr @enzyme_const, ptr %out) + ret double %r +} + +declare double @__enzyme_autodiff(...) + +; CHECK: define internal { double } @diffewrap(ptr {{.*}}%diag, ptr {{.*}}%"diag'", ptr {{.*}}%out, double %differeturn) +; The start value is loaded once in the primal, and its adjoint accumulated in +; the reverse; it is never recomputed there. +; CHECK: %seed = load double, ptr %diag +; CHECK: invertstart: +; CHECK-NOT: load double, ptr %diag