[IndVarSimplify] Rewrite loop exit values with their initial values from loop preheader

Summary:
This patch adds support to check if a loop has loop invariant conditions which lead to loop exits. If so, we know that if the exit path is taken, it is at the first loop iteration. If there is an induction variable used in that exit path whose value has not been updated, it will keep its initial value passing from loop preheader. We can therefore rewrite the exit value with
its initial value. This will help remove phis created by LCSSA and enable other optimizations like loop unswitch.


Reviewers: sanjoy

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D13974

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@251839 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Transforms/Scalar/IndVarSimplify.cpp

@@ -132,6 +132,7 @@ private:
 
   bool canLoopBeDeleted(Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet);
   void rewriteLoopExitValues(Loop *L, SCEVExpander &Rewriter);
+  void rewriteFirstIterationLoopExitValues(Loop *L);
 
   Value *linearFunctionTestReplace(Loop *L, const SCEV *BackedgeTakenCount,
                                    PHINode *IndVar, SCEVExpander &Rewriter);
@@ -700,6 +701,73 @@ void IndVarSimplify::rewriteLoopExitValues(Loop *L, SCEVExpander &Rewriter) {
   Rewriter.clearInsertPoint();
 }
 
+//===---------------------------------------------------------------------===//
+// rewriteFirstIterationLoopExitValues: Rewrite loop exit values if we know
+// they will exit at the first iteration.
+//===---------------------------------------------------------------------===//
+
+/// Check to see if this loop has loop invariant conditions which lead to loop
+/// exits. If so, we know that if the exit path is taken, it is at the first
+/// loop iteration. This lets us predict exit values of PHI nodes that live in
+/// loop header.
+void IndVarSimplify::rewriteFirstIterationLoopExitValues(Loop *L) {
+  // Verify the input to the pass is already in LCSSA form.
+  assert(L->isLCSSAForm(*DT));
+
+  SmallVector<BasicBlock *, 8> ExitBlocks;
+  L->getUniqueExitBlocks(ExitBlocks);
+
+  for (auto *ExitBB : ExitBlocks) {
+    BasicBlock::iterator begin = ExitBB->begin();
+    // If there are no more PHI nodes in this exit block, then no more
+    // values defined inside the loop are used on this path.
+    while (auto *PN = dyn_cast<PHINode>(begin++)) {
+      for (unsigned IncomingValIdx = 0, e = PN->getNumIncomingValues();
+          IncomingValIdx != e; ++IncomingValIdx) {
+        auto *IncomingBB = PN->getIncomingBlock(IncomingValIdx);
+        if (!L->contains(IncomingBB))
+          continue;
+
+        // Get condition that leads to the exit path.
+        auto *TermInst = IncomingBB->getTerminator();
+
+        Value *Cond = nullptr;
+        if (auto *BI = dyn_cast<BranchInst>(TermInst)) {
+          // Must be a conditional branch, otherwise the block
+          // should not be in the loop.
+          Cond = BI->getCondition();
+        } else if (auto *SI = dyn_cast<SwitchInst>(TermInst))
+          Cond = SI->getCondition();
+        else
+          continue;
+
+        // All non-instructions are loop-invariant.
+        if (isa<Instruction>(Cond) && !L->isLoopInvariant(Cond))
+            continue;
+
+        auto *ExitVal =
+            dyn_cast<PHINode>(PN->getIncomingValue(IncomingValIdx));
+
+        // Only deal with PHIs.
+        if (!ExitVal)
+          continue;
+
+        // If ExitVal is a PHI on the loop header, then we know its
+        // value along this exit because the exit can only be taken
+        // on the first iteration.
+        auto *LoopPreheader = L->getLoopPreheader();
+        assert(LoopPreheader && "Invalid loop");
+        if (ExitVal->getBasicBlockIndex(LoopPreheader) != -1) {
+          assert(ExitVal->getParent() == L->getHeader() &&
+              "ExitVal must be in loop header");
+          PN->setIncomingValue(IncomingValIdx,
+              ExitVal->getIncomingValueForBlock(LoopPreheader));
+        }
+      }
+    }
+  }
+}
+
 /// Check whether it is possible to delete the loop after rewriting exit
 /// value. If it is possible, ignore ReplaceExitValue and do rewriting
 /// aggressively.
@@ -2173,6 +2241,11 @@ bool IndVarSimplify::runOnLoop(Loop *L, LPPassManager &LPM) {
   // loop may be sunk below the loop to reduce register pressure.
   sinkUnusedInvariants(L);
 
+  // rewriteFirstIterationLoopExitValues does not rely on the computation of
+  // trip count and therefore can further simplify exit values in addition to
+  // rewriteLoopExitValues.
+  rewriteFirstIterationLoopExitValues(L);
+
   // Clean up dead instructions.
   Changed |= DeleteDeadPHIs(L->getHeader(), TLI);
   // Check a post-condition.

test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll

@@ -0,0 +1,75 @@
+; RUN: opt -indvars -instcombine -S < %s | FileCheck %s
+
+;; Test that loop's exit value is rewritten to its initial
+;; value from loop preheader
+define i32 @test1(i32* %var) {
+; CHECK-LABEL: @test1
+entry:
+ %cond = icmp eq i32* %var, null 
+ br label %header
+
+header:
+ %phi_indvar = phi i32 [0, %entry], [%indvar, %loop]
+ br i1 %cond, label %loop, label %exit
+
+loop:
+ %indvar = add i32 %phi_indvar, 1
+ br label %header
+
+exit:
+; CHECK: ret i32 0
+ ret i32 %phi_indvar
+}
+
+
+;; Test that inner loop's exit value is first rewritten to outer
+;; loop's induction variable, and then further rewritten to a 
+;; constant when process outer loop.
+define i32 @test2(i32* %var1, i32* %var2) {
+; CHECK-LABEL: @test2
+entry:
+ %cond1 = icmp eq i32* %var1, null 
+ %cond2 = icmp eq i32* %var2, null 
+ br label %outer_header
+
+outer_header:
+ %phi_outer = phi i32 [0, %entry], [%indvar_outer, %inner_exit]
+ br label %inner_header
+ 
+inner_header:
+ %phi_inner = phi i32 [%phi_outer, %outer_header], [%indvar_inner, %loop]
+ br i1 %cond1, label %loop, label %exit
+
+loop:
+ %indvar_inner = add i32 %phi_inner, 1
+ br i1 %cond2, label %inner_header, label %inner_exit
+ 
+inner_exit:
+ %indvar_outer = add i32 %phi_outer, 1
+ br label %outer_header
+
+exit:
+;; %phi_inner is first rewritten to %phi_outer
+;; and then %phi_outer is rewritten to 0
+ %ret_val = add i32 %phi_inner, %phi_outer
+; CHECK: ret i32 0
+ ret i32 %ret_val
+}
+
+;; Test that we can not rewrite loop exit value if it's not
+;; a phi node (%indvar is an add instruction in this test).
+define i32 @test3(i32* %var) {
+; CHECK-LABEL: @test3
+entry:
+ %cond = icmp eq i32* %var, null 
+ br label %header
+
+header:
+ %phi_indvar = phi i32 [0, %entry], [%indvar, %header]
+ %indvar = add i32 %phi_indvar, 1
+ br i1 %cond, label %header, label %exit
+
+exit:
+; CHECK: ret i32 %indvar
+ ret i32 %indvar
+}

test/Transforms/LoopUnroll/scevunroll.ll

@@ -184,7 +184,7 @@ for.body87:
 ; CHECK: for.body:
 ; CHECK: %b.03 = phi i32 [ 0, %entry ], [ %add, %for.cond ]
 ; CHECK: return:
-; CHECK: %b.03.lcssa = phi i32 [ %b.03, %for.body ], [ %b.03, %for.cond ]
+; CHECK: %b.03.lcssa = phi i32 [ %b.03, %for.body ], [ 0, %for.cond ]
 define void @nsw_latch(i32* %a) nounwind {
 entry:
   br label %for.body