From d2a49ce9fc46c7cba6b807130ad632d0e3f73ed4 Mon Sep 17 00:00:00 2001
From: Philip Reames <listmail@philipreames.com>
Date: Tue, 14 May 2019 17:20:10 +0000
Subject: [PATCH] [IndVars] Extend reasoning about loop invariant exits to
 non-header blocks

Noticed while glancing through the code for other reasons.  The extension is trivial enough, decided to just do it.



git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@360694 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Scalar/IndVarSimplify.cpp      | 16 +++--
 .../IndVarSimplify/rewrite-loop-exit-value.ll | 65 +++++++++++++++++++
 test/Transforms/LoopUnroll/scevunroll.ll      |  2 +-
 3 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/lib/Transforms/Scalar/IndVarSimplify.cpp b/lib/Transforms/Scalar/IndVarSimplify.cpp
index ecf0bef13df..333de75c244 100644
--- a/lib/Transforms/Scalar/IndVarSimplify.cpp
+++ b/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -728,11 +728,13 @@ bool IndVarSimplify::rewriteFirstIterationLoopExitValues(Loop *L) {
            IncomingValIdx != E; ++IncomingValIdx) {
         auto *IncomingBB = PN.getIncomingBlock(IncomingValIdx);
 
-        // We currently only support loop exits from loop header. If the
-        // incoming block is not loop header, we need to recursively check
-        // all conditions starting from loop header are loop invariants.
-        // Additional support might be added in the future.
-        if (IncomingBB != LoopHeader)
+        // Can we prove that the exit must run on the first iteration if it
+        // runs at all?  (i.e. early exits are fine for our purposes, but
+        // traces which lead to this exit being taken on the 2nd iteration
+        // aren't.)  Note that this is about whether the exit branch is
+        // executed, not about whether it is taken.
+        if (!L->getLoopLatch() ||
+            !DT->dominates(IncomingBB, L->getLoopLatch()))
           continue;
 
         // Get condition that leads to the exit path.
@@ -753,8 +755,8 @@ bool IndVarSimplify::rewriteFirstIterationLoopExitValues(Loop *L) {
 
         auto *ExitVal = dyn_cast<PHINode>(PN.getIncomingValue(IncomingValIdx));
 
-        // Only deal with PHIs.
-        if (!ExitVal)
+        // Only deal with PHIs in the loop header.
+        if (!ExitVal || ExitVal->getParent() != L->getHeader())
           continue;
 
         // If ExitVal is a PHI on the loop header, then we know its
diff --git a/test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll b/test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll
index c8a0c5d3c73..a8f713bbf12 100644
--- a/test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll
+++ b/test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll
@@ -94,3 +94,68 @@ exit:
   ret i32 %phi_indvar
 }
 
+
+; Multiple exits dominating latch
+define i32 @test4(i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test4(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[HEADER:%.*]]
+; CHECK:       header:
+; CHECK-NEXT:    br i1 [[COND1:%.*]], label [[LOOP:%.*]], label [[EXIT:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 [[COND2:%.*]], label [[HEADER]], label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  br label %header
+
+header:
+  %phi_indvar = phi i32 [0, %entry], [%indvar, %loop]
+  br i1 %cond1, label %loop, label %exit
+
+loop:
+  %indvar = add i32 %phi_indvar, 1
+  br i1 %cond2, label %header, label %exit
+
+exit:
+  ret i32 %phi_indvar
+}
+
+; A conditionally executed exit.
+define i32 @test5(i1* %addr, i1 %cond2) {
+; CHECK-LABEL: @test5(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[HEADER:%.*]]
+; CHECK:       header:
+; CHECK-NEXT:    [[PHI_INDVAR:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INDVAR:%.*]], [[LOOP:%.*]] ]
+; CHECK-NEXT:    [[COND1:%.*]] = load volatile i1, i1* [[ADDR:%.*]], align 1
+; CHECK-NEXT:    br i1 [[COND1]], label [[LOOP]], label [[MAYBE:%.*]]
+; CHECK:       maybe:
+; CHECK-NEXT:    br i1 [[COND2:%.*]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDVAR]] = add i32 [[PHI_INDVAR]], 1
+; CHECK-NEXT:    br label [[HEADER]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i32 [[PHI_INDVAR]]
+;
+entry:
+  br label %header
+
+header:
+  %phi_indvar = phi i32 [0, %entry], [%indvar, %loop]
+  %cond1 = load volatile i1, i1* %addr
+  br i1 %cond1, label %loop, label %maybe
+
+maybe:
+  br i1 %cond2, label %loop, label %exit
+
+loop:
+  %indvar = add i32 %phi_indvar, 1
+  br label %header
+
+exit:
+  ret i32 %phi_indvar
+}
+
+
diff --git a/test/Transforms/LoopUnroll/scevunroll.ll b/test/Transforms/LoopUnroll/scevunroll.ll
index a5c9a6efacf..afee41e11aa 100644
--- a/test/Transforms/LoopUnroll/scevunroll.ll
+++ b/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
-- 
2.50.1