From 7a43b382ce9da2c5a4196d6cf4184d2983906400 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Thu, 14 Jul 2022 11:45:35 +0200
Subject: [PATCH] [IndVars] Make sure header phi simplification preserves LCSSA
 form

When simplifying instructions, make sure that the replacement
preserves LCSSA form. This fixes the issue reported at:
https://reviews.llvm.org/D129293#3650851
---
 llvm/lib/Transforms/Scalar/IndVarSimplify.cpp      |  9 ++---
 .../IndVarSimplify/lcssa-preservation.ll           | 39 ++++++++++++++++++++--
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
index 7a574f1..a9ca0bd 100644
--- a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
+++ b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -1307,7 +1307,7 @@ static void foldExit(const Loop *L, BasicBlock *ExitingBB, bool IsTaken,
 }
 
 static void replaceLoopPHINodesWithPreheaderValues(
-    Loop *L, SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
+    LoopInfo *LI, Loop *L, SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
   assert(L->isLoopSimplifyForm() && "Should only do it in simplify form!");
   auto *LoopPreheader = L->getLoopPreheader();
   auto *LoopHeader = L->getHeader();
@@ -1332,7 +1332,8 @@ static void replaceLoopPHINodesWithPreheaderValues(
     if (!L->contains(I))
       continue;
 
-    if (Value *Res = simplifyInstruction(I, I->getModule()->getDataLayout())) {
+    Value *Res = simplifyInstruction(I, I->getModule()->getDataLayout());
+    if (Res && LI->replacementPreservesLCSSAForm(I, Res)) {
       for (User *U : I->users())
         Worklist.push_back(cast<Instruction>(U));
       I->replaceAllUsesWith(Res);
@@ -1586,7 +1587,7 @@ bool IndVarSimplify::optimizeLoopExits(Loop *L, SCEVExpander &Rewriter) {
       // unconditional exit, we can still replace header phis with their
       // preheader value.
       if (!L->contains(BI->getSuccessor(CI->isNullValue())))
-        replaceLoopPHINodesWithPreheaderValues(L, DeadInsts);
+        replaceLoopPHINodesWithPreheaderValues(LI, L, DeadInsts);
       return true;
     }
 
@@ -1673,7 +1674,7 @@ bool IndVarSimplify::optimizeLoopExits(Loop *L, SCEVExpander &Rewriter) {
     // the header PHIs with values coming from the preheader.
     if (ExitCount->isZero()) {
       foldExit(L, ExitingBB, true, DeadInsts);
-      replaceLoopPHINodesWithPreheaderValues(L, DeadInsts);
+      replaceLoopPHINodesWithPreheaderValues(LI, L, DeadInsts);
       Changed = true;
       continue;
     }
diff --git a/llvm/test/Transforms/IndVarSimplify/lcssa-preservation.ll b/llvm/test/Transforms/IndVarSimplify/lcssa-preservation.ll
index 3da159e3..5e7b5ad 100644
--- a/llvm/test/Transforms/IndVarSimplify/lcssa-preservation.ll
+++ b/llvm/test/Transforms/IndVarSimplify/lcssa-preservation.ll
@@ -45,9 +45,44 @@ inner.latch:
 outer.latch:
   br i1 undef, label %outer.header, label %exit
 
-
-
 exit:
   %exit.phi = phi i32 [ %inc, %inner.latch ], [ undef, %outer.latch ]
   ret void
 }
+
+define i64 @unconditional_exit_simplification(i64 %arg) {
+; CHECK-LABEL: @unconditional_exit_simplification(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP1:%.*]]
+; CHECK:       loop1:
+; CHECK-NEXT:    br label [[LOOP2:%.*]]
+; CHECK:       loop2:
+; CHECK-NEXT:    [[IV2:%.*]] = phi i64 [ 0, [[LOOP1]] ], [ 1, [[LOOP2]] ]
+; CHECK-NEXT:    br i1 true, label [[LOOP2]], label [[LOOP1_LATCH:%.*]]
+; CHECK:       loop1.latch:
+; CHECK-NEXT:    [[RES_LCSSA:%.*]] = phi i64 [ [[IV2]], [[LOOP2]] ]
+; CHECK-NEXT:    br i1 false, label [[LOOP1]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[RES_LCSSA2:%.*]] = phi i64 [ [[RES_LCSSA]], [[LOOP1_LATCH]] ]
+; CHECK-NEXT:    ret i64 [[RES_LCSSA2]]
+;
+entry:
+  br label %loop1
+
+loop1:
+  %iv1 = phi i64 [ 0, %entry ], [ 1, %loop1.latch ]
+  br label %loop2
+
+loop2:
+  %iv2 = phi i64 [ 0, %loop1 ], [ 1, %loop2 ]
+  %res = add nuw nsw i64 %iv1, %iv2
+  br i1 true, label %loop2, label %loop1.latch
+
+loop1.latch:
+  %res.lcssa = phi i64 [ %res, %loop2 ]
+  br i1 false, label %loop1, label %exit
+
+exit:
+  %res.lcssa2 = phi i64 [ %res.lcssa, %loop1.latch ]
+  ret i64 %res.lcssa2
+}
-- 
2.7.4