[LoopBoundSplit] Update phi node in exit block

It fixes https://bugs.llvm.org/show_bug.cgi?id=51700

Differential Revision:

1 files changed, 15 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Scalar/LoopBoundSplit.cpp b/llvm/lib/Transforms/Scalar/LoopBoundSplit.cpp
index 993b154..a89f0aa 100644
--- a/llvm/lib/Transforms/Scalar/LoopBoundSplit.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopBoundSplit.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/LoopBoundSplit.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
@@ -398,6 +399,20 @@ static bool splitLoopBound(Loop &L, DominatorTree &DT, LoopInfo &LI,
   else
     ExitingCond.BI->setSuccessor(1, PostLoopPreHeader);
 
+  // Update phi node in exit block of post-loop.
+  for (PHINode &PN : PostLoop->getExitBlock()->phis()) {
+    for (auto i : seq<int>(0, PN.getNumOperands())) {
+      // Check incoming block is pre-loop's exiting block.
+      if (PN.getIncomingBlock(i) == L.getExitingBlock()) {
+        // Replace pre-loop's exiting block by post-loop's preheader.
+        PN.setIncomingBlock(i, PostLoopPreHeader);
+        // Add a new incoming value with post-loop's exiting block.
+        PN.addIncoming(VMap[PN.getIncomingValue(i)],
+                       PostLoop->getExitingBlock());
+      }
+    }
+  }
+
   // Update dominator tree.
   DT.changeImmediateDominator(PostLoopPreHeader, L.getExitingBlock());
   DT.changeImmediateDominator(PostLoop->getExitBlock(), PostLoopPreHeader);