Cleanup use of LI when needed OrigLI (#2755)

* Fix LoopAnalysis crash on incomplete new_func

* Fix LoopAnalysis crash on incomplete new_func using isOriginal

* Cleanup use of LI when needed OrigLI

* Fix LoopAnalysis crash on incomplete new_func and update manual check

* Fix LoopAnalysis crash on incomplete new_func and fix inverted condition

* Fix LoopAnalysis crash and fix manual addition logic for non-loop values

* Fix LoopAnalysis crash and include non-instructions in manual induction check

* fix

* fix

* fix

* f

* fix

* fix

* fix

* Fix clang-format indentation in GradientUtils.cpp (#2756)

* Initial plan

* Fix formatting in GradientUtils.cpp

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

---------

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

Author: wsmoses

enzyme/Enzyme/GradientUtils.cpp

@@ -1611,16 +1611,19 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
 
     // Don't attempt to unroll a loop induction variable in other
     // circumstances
-    auto &LLI = Logic.PPC.FAM.getResult<LoopAnalysis>(*parent->getParent());
     std::set<BasicBlock *> prevIteration;
-    if (LLI.isLoopHeader(parent)) {
+    BasicBlock *origParent = isOriginal(parent);
+    assert(origParent);
+    if (OrigLI->isLoopHeader(origParent)) {
       if (phi->getNumIncomingValues() != 2) {
         assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
         goto endCheck;
       }
-      auto L = LLI.getLoopFor(parent);
+      auto OrigL = OrigLI->getLoopFor(origParent);
       for (auto PH : predecessors(parent)) {
-        if (L->contains(PH))
+        BasicBlock *origPH = isOriginal(PH);
+        assert(origPH);
+        if (OrigL->contains(origPH))
           prevIteration.insert(PH);
       }
       if (prevIteration.size() && !legalRecompute(phi, available, &BuilderM)) {
@@ -1629,16 +1632,33 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
       }
     }
     for (auto &val : phi->incoming_values()) {
-      if (isPotentialLastLoopValue(val, parent, LLI)) {
-        if (unwrapMode == UnwrapMode::LegalFullUnwrap) {
-          llvm::errs() << " module: " << *newFunc->getParent() << "\n";
-          llvm::errs() << " newFunc: " << *newFunc << "\n";
-          llvm::errs() << " parent: " << *parent << "\n";
-          llvm::errs() << " val: " << *val << "\n";
+      auto inst = dyn_cast<Instruction>(val);
+      if (!inst)
+        continue;
+      auto origInstParent = isOriginal(inst->getParent());
+      assert(origInstParent);
+      const llvm::Loop *InstLoop = OrigLI->getLoopFor(origInstParent);
+      if (!InstLoop) {
+        continue;
+      }
+      bool isParentLoop = false;
+      for (const llvm::Loop *L = OrigLI->getLoopFor(origParent); L;
+           L = L->getParentLoop()) {
+        if (L == InstLoop) {
+          isParentLoop = true;
+          break;
         }
-        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
-        goto endCheck;
       }
+      if (isParentLoop)
+        continue;
+      if (unwrapMode == UnwrapMode::LegalFullUnwrap) {
+        llvm::errs() << " module: " << *newFunc->getParent() << "\n";
+        llvm::errs() << " newFunc: " << *newFunc << "\n";
+        llvm::errs() << " parent: " << *parent << "\n";
+        llvm::errs() << " val: " << *val << "\n";
+      }
+      assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
+      goto endCheck;
     }
 
     if (phi->getNumIncomingValues() == 1) {

enzyme/test/Enzyme/ReverseMode/enzyme_loop_analysis_crash.ll

@@ -0,0 +1,33 @@
+; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -S | FileCheck %s
+
+define hidden void @_Z10entry_wrapRN6enzyme6tensorIfJLm2ELm3EEEES2_RKS1_(i8* %out_0, i8* %out_1, i8* %in_0) #0 {
+entry:
+  %out_0.addr = alloca i8*, align 8
+  %out_1.addr = alloca i8*, align 8
+  %in_0.addr = alloca i8*, align 8
+  store i8* %out_0, i8** %out_0.addr, align 8
+  store i8* %out_1, i8** %out_1.addr, align 8
+  store i8* %in_0, i8** %in_0.addr, align 8
+  %0 = load i8*, i8** %out_0.addr, align 8
+  %1 = load i8*, i8** %out_1.addr, align 8
+  %2 = load i8*, i8** %in_0.addr, align 8
+  call void @_Z4myfnILm2ELm3EEvRN6enzyme6tensorIfJXT_EXT0_EEEES3_RKS2_(i8* %0, i8* %1, i8* %2)
+  ret void
+}
+
+define hidden void @_Z4myfnILm2ELm3EEvRN6enzyme6tensorIfJXT_EXT0_EEEES3_RKS2_(i8* %0, i8* %1, i8* %2) #0 {
+entry:
+  ret void
+}
+
+declare void @__enzyme_autodiff(...)
+
+define void @test_derivative(i8* %out_0, i8* %out_0_d, i8* %out_1, i8* %out_1_d, i8* %in_0, i8* %in_0_d) {
+entry:
+  call void (...) @__enzyme_autodiff(i8* bitcast (void (i8*, i8*, i8*)* @_Z10entry_wrapRN6enzyme6tensorIfJLm2ELm3EEEES2_RKS1_ to i8*), metadata !"enzyme_dup", i8* %out_0, i8* %out_0_d, metadata !"enzyme_dup", i8* %out_1, i8* %out_1_d, metadata !"enzyme_dup", i8* %in_0, i8* %in_0_d)
+  ret void
+}
+
+attributes #0 = { mustprogress nounwind }
+
+; CHECK: define internal void @diffe_Z10entry_wrap