[LSR] Optimize unused IVs to final values in the exit block

Loop Strength Reduce sometimes optimizes away all uses of an induction variable
from a loop but leaves the IV increments. When the only remaining use of the IV
is the PHI in the exit block, this patch will call rewriteLoopExitValues to
replace the exit block PHI with the final value of the IV to skip the updates
in each loop iteration.

Differential Revision: https://reviews.llvm.org/D118808

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

@@ -113,6 +113,7 @@
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include <algorithm>
 #include <cassert>
@@ -5600,6 +5601,27 @@ void LSRInstance::Rewrite(const LSRUse &LU, const LSRFixup &LF,
     DeadInsts.emplace_back(OperandIsInstr);
 }
 
+// Check if there are any loop exit values which are only used once within the
+// loop which may potentially be optimized with a call to rewriteLoopExitValue.
+static bool LoopExitValHasSingleUse(Loop *L) {
+  BasicBlock *ExitBB = L->getExitBlock();
+  if (!ExitBB)
+    return false;
+
+  for (PHINode &ExitPhi : ExitBB->phis()) {
+    if (ExitPhi.getNumIncomingValues() != 1)
+      break;
+
+    BasicBlock *Pred = ExitPhi.getIncomingBlock(0);
+    Value *IVNext = ExitPhi.getIncomingValueForBlock(Pred);
+    // One use would be the exit phi node, and there should be only one other
+    // use for this to be considered.
+    if (IVNext->getNumUses() == 2)
+      return true;
+  }
+  return false;
+}
+
 /// Rewrite all the fixup locations with new values, following the chosen
 /// solution.
 void LSRInstance::ImplementSolution(
@@ -6374,6 +6396,24 @@ static bool ReduceLoopStrength(Loop *L, IVUsers &IU, ScalarEvolution &SE,
       DeleteDeadPHIs(L->getHeader(), &TLI, MSSAU.get());
     }
   }
+  // LSR may at times remove all uses of an induction variable from a loop.
+  // The only remaining use is the PHI in the exit block.
+  // When this is the case, if the exit value of the IV can be calculated using
+  // SCEV, we can replace the exit block PHI with the final value of the IV and
+  // skip the updates in each loop iteration.
+  if (L->isRecursivelyLCSSAForm(DT, LI) && LoopExitValHasSingleUse(L)) {
+    SmallVector<WeakTrackingVH, 16> DeadInsts;
+    const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
+    SCEVExpander Rewriter(SE, DL, "lsr", false);
+    int Rewrites = rewriteLoopExitValues(L, &LI, &TLI, &SE, &TTI, Rewriter, &DT,
+                                         OnlyCheapRepl, DeadInsts);
+    if (Rewrites) {
+      Changed = true;
+      RecursivelyDeleteTriviallyDeadInstructionsPermissive(DeadInsts, &TLI,
+                                                           MSSAU.get());
+      DeleteDeadPHIs(L->getHeader(), &TLI, MSSAU.get());
+    }
+  }
 
   if (SalvageableDVI.empty())
     return Changed;

llvm/test/Transforms/LoopStrengthReduce/remove_scev_indvars.ll

@@ -0,0 +1,57 @@
+; RUN: opt < %s -S -loop-reduce | FileCheck %s
+
+define void @testIVNext(i64* nocapture %a, i64 signext %m, i64 signext %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %indvars.iv.prol = phi i64 [ %indvars.iv.next.prol, %for.body ], [ %m, %entry ]
+  %i = phi i64 [ %i.next, %for.body ], [ 0, %entry ]
+  %uglygep138 = getelementptr i64, i64* %a, i64 %i
+  store i64 55, i64* %uglygep138, align 4
+  %indvars.iv.next.prol = add nuw nsw i64 %indvars.iv.prol, 1
+  %i.next = add i64 %i, 1
+  %i.cmp.not = icmp eq i64 %i.next, %n
+  br i1 %i.cmp.not, label %for.exit, label %for.body
+
+; CHECK: entry:
+; CHECK: %0 = add i64 %n, %m
+; CHECK-NOT : %indvars.iv.next.prol
+; CHECK-NOT: %indvars.iv.prol
+; CHECK: %indvars.iv.unr = phi i64 [ %0, %for.exit ]
+for.exit:
+  %indvars.iv.next.prol.lcssa = phi i64 [ %indvars.iv.next.prol, %for.body ]
+  br label %exit
+
+exit:
+  %indvars.iv.unr = phi i64 [ %indvars.iv.next.prol.lcssa, %for.exit ]
+  ret void
+}
+
+define void @testIV(i64* nocapture %a, i64 signext %m, i64 signext %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %iv.prol = phi i64 [ %iv.next.prol, %for.body ], [ %m, %entry ]
+  %i = phi i64 [ %i.next, %for.body ], [ 0, %entry ]
+  %uglygep138 = getelementptr i64, i64* %a, i64 %i
+  store i64 55, i64* %uglygep138, align 4
+  %iv.next.prol = add nuw nsw i64 %iv.prol, 1
+  %i.next = add i64 %i, 1
+  %i.cmp.not = icmp eq i64 %i.next, %n
+  br i1 %i.cmp.not, label %for.exit, label %for.body
+
+; CHECK: entry:
+; CHECK: %0 = add i64 %n, %m
+; CHECK: %1 = add i64 %0, -1
+; CHECK-NOT: %iv.next.prol
+; CHECK-NOT: %iv.prol
+; CHECK: %indvars.iv.unr = phi i64 [ %1, %for.exit ]
+for.exit:
+  %iv.prol.lcssa = phi i64 [ %iv.prol, %for.body ]
+  br label %exit
+exit:
+  %indvars.iv.unr = phi i64 [%iv.prol.lcssa, %for.exit]
+  ret void
+}