Fix increment usage (rust-lang#468)

Author: wsmoses

enzyme/Enzyme/CacheUtility.cpp

@@ -192,7 +192,7 @@ std::pair<PHINode *, Instruction *> FindCanonicalIV(Loop *L, Type *Ty) {
 // Attempt to rewrite all phinode's in the loop in terms of the
 // induction variable
 void RemoveRedundantIVs(BasicBlock *Header, PHINode *CanonicalIV,
-                        MustExitScalarEvolution &SE,
+                        Instruction *Increment, MustExitScalarEvolution &SE,
                         std::function<void(Instruction *, Value *)> replacer,
                         std::function<void(Instruction *)> eraser) {
   assert(Header);
@@ -246,6 +246,37 @@ void RemoveRedundantIVs(BasicBlock *Header, PHINode *CanonicalIV,
     replacer(Tmp, NewIV);
     eraser(Tmp);
   }
+
+  // Replace existing increments with canonical Increment
+  Increment->moveAfter(CanonicalIV->getParent()->getFirstNonPHI());
+  SmallVector<Instruction *, 1> toErase;
+  for (auto use : CanonicalIV->users()) {
+    auto BO = dyn_cast<BinaryOperator>(use);
+    if (BO == nullptr)
+      continue;
+    if (BO->getOpcode() != BinaryOperator::Add)
+      continue;
+    if (use == Increment)
+      continue;
+
+    Value *toadd = nullptr;
+    if (BO->getOperand(0) == CanonicalIV) {
+      toadd = BO->getOperand(1);
+    } else {
+      assert(BO->getOperand(1) == CanonicalIV);
+      toadd = BO->getOperand(0);
+    }
+    if (auto CI = dyn_cast<ConstantInt>(toadd)) {
+      if (!CI->isOne())
+        continue;
+      BO->replaceAllUsesWith(Increment);
+      toErase.push_back(BO);
+    } else {
+      continue;
+    }
+  }
+  for (auto BO : toErase)
+    eraser(BO);
 }
 
 void CanonicalizeLatches(const Loop *L, BasicBlock *Header,
@@ -332,37 +363,6 @@ void CanonicalizeLatches(const Loop *L, BasicBlock *Header,
   // Replace previous increment usage with new increment value
   if (Increment) {
     Increment->moveAfter(CanonicalIV->getParent()->getFirstNonPHI());
-    std::vector<Instruction *> toerase;
-    // Replace existing increments with canonical Increment
-    for (auto use : CanonicalIV->users()) {
-      auto BO = dyn_cast<BinaryOperator>(use);
-      if (BO == nullptr)
-        continue;
-      if (BO->getOpcode() != BinaryOperator::Add)
-        continue;
-      if (use == Increment)
-        continue;
-
-      Value *toadd = nullptr;
-      if (BO->getOperand(0) == CanonicalIV) {
-        toadd = BO->getOperand(1);
-      } else {
-        assert(BO->getOperand(1) == CanonicalIV);
-        toadd = BO->getOperand(0);
-      }
-      if (auto CI = dyn_cast<ConstantInt>(toadd)) {
-        if (!CI->isOne())
-          continue;
-        BO->replaceAllUsesWith(Increment);
-        toerase.push_back(BO);
-      } else {
-        continue;
-      }
-    }
-    for (auto inst : toerase) {
-      gutils.erase(inst);
-    }
-    toerase.clear();
 
     if (latches.size() == 1 && isa<BranchInst>(latches[0]->getTerminator()) &&
         cast<BranchInst>(latches[0]->getTerminator())->isConditional())

enzyme/Enzyme/CacheUtility.h

@@ -390,7 +390,7 @@ InsertNewCanonicalIV(llvm::Loop *L, llvm::Type *Ty, std::string name = "iv");
 // induction variable
 void RemoveRedundantIVs(
     llvm::BasicBlock *Header, llvm::PHINode *CanonicalIV,
-    MustExitScalarEvolution &SE,
+    llvm::Instruction *Increment, MustExitScalarEvolution &SE,
     std::function<void(llvm::Instruction *, llvm::Value *)> replacer,
     std::function<void(llvm::Instruction *)> eraser);
 #endif

enzyme/Enzyme/EnzymeLogic.cpp

@@ -3476,11 +3476,11 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
         sucBB->removePredecessor(newBB);
       }
 
-      std::vector<Instruction *> toerase;
+      SmallVector<Instruction *, 2> toerase;
       for (auto &I : oBB) {
         toerase.push_back(&I);
       }
-      for (auto I : toerase) {
+      for (auto I : llvm::reverse(toerase)) {
         maker.eraseIfUnused(*I, /*erase*/ true,
                             /*check*/ key.mode ==
                                 DerivativeMode::ReverseModeCombined);

enzyme/Enzyme/FunctionUtils.cpp

@@ -808,7 +808,7 @@ void CanonicalizeLoops(Function *F, FunctionAnalysisManager &FAM) {
     PHINode *CanonicalIV = pair.first;
     assert(CanonicalIV);
     RemoveRedundantIVs(
-        L->getHeader(), CanonicalIV, SE,
+        L->getHeader(), CanonicalIV, pair.second, SE,
         [&](Instruction *I, Value *V) { I->replaceAllUsesWith(V); },
         [&](Instruction *I) { I->eraseFromParent(); });
   }

enzyme/test/Enzyme/ReverseMode/increplace.ll

@@ -0,0 +1,66 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -instsimplify -simplifycfg -loop-deletion -simplifycfg -S | FileCheck %s
+
+source_filename = "text"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
+target triple = "x86_64-linux-gnu"
+
+declare void @julia_throw_boundserror_1767(i64 %d)
+
+define double @sub(double* %i58) {
+bb:
+  br label %bb22
+
+bb22:                                             ; preds = %bb32, %bb
+  %i23 = phi i64 [ %i78, %bb32 ], [ 0, %bb ]
+  %i78 = add nuw nsw i64 %i23, 1
+  %i26 = icmp sle i64 %i23, 10
+  br i1 %i26, label %bb32, label %bb29
+
+bb29:                                             ; preds = %bb22
+  %i30 = add nuw nsw i64 %i23, 1
+  call void @julia_throw_boundserror_1767(i64 %i30)
+  unreachable
+
+bb32:                                             ; preds = %bb22
+  %c = icmp sle i64 %i23, 5
+  br i1 %c, label %bb22, label %exit
+
+exit:
+  %v = load double, double* %i58, align 8
+  ret double %v
+}
+
+
+declare dso_local double @__enzyme_autodiff(i8*, i64*, double*, double*)
+
+define void @main(double* %arg, double* %arg1) {
+bb:
+  %enzyme_dup = alloca i64, align 8
+  %i = tail call double @__enzyme_autodiff(i8* bitcast (double (double*)* @julia_arsum2_1761 to i8*), i64* %enzyme_dup, double* %arg, double* %arg1)
+  ret void
+}
+
+define double @julia_arsum2_1761(double* %arg) {
+bb:
+  %i23 = call double @sub(double* %arg)
+  store double 0.000000e+00, double* %arg
+  ret double %i23
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 2, !"Dwarf Version", i32 4}
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"jtbaa_data", !4, i64 0}
+!4 = !{!"jtbaa", !5, i64 0}
+!5 = !{!"jtbaa"}
+
+; Creating this should not segfault due to the use of the replaced iv in the unreachable block
+; CHECK: define internal void @diffesub(double* %i58, double* %"i58'", double %differeturn)
+; CHECK-NEXT: bb:
+; CHECK-NEXT:   %0 = load double, double* %"i58'", align 8
+; CHECK-NEXT:   %1 = fadd fast double %0, %differeturn
+; CHECK-NEXT:   store double %1, double* %"i58'", align 8
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }