[llvm] 0a70edd - [CloneFunction] Update loop headers after cloning all blocks in loop.

[Whitney Tsang via llvm-commits]

Author: Whitney Tsang
Date: 2020-02-21T22:18:24Z
New Revision: 0a70edd696264ba69914f4d4ecb28d559cbe5444

URL: https://github.com/llvm/llvm-project/commit/0a70edd696264ba69914f4d4ecb28d559cbe5444
DIFF: https://github.com/llvm/llvm-project/commit/0a70edd696264ba69914f4d4ecb28d559cbe5444.diff

LOG: [CloneFunction] Update loop headers after cloning all blocks in loop.

Summary:
Blocks in a loop can be in any order as long as the loop header is the
first block in Blocks.
With some order of Blocks, cloneLoopWithPreheader would trigger the
assertion in addBasicBlockToLoop.

Example:

define void @test(i64 %N) {
preheader.i:
  br label %header.i

header.i:
  %i = phi i64 [ 0, %preheader.i ], [ %inc.i, %latch.i ]
  br label %header.j

header.j:
  %j = phi i64 [ 0, %header.i ], [ %inc.j, %latch.j ]
  br label %header.k

header.k:
  %k = phi i64 [ 0, %header.j ], [ %inc.k, %latch.k ]
  call void @baz(i64 %i, i64 %j, i64 %k)
  br label %latch.k

latch.k:
  %inc.k = add nsw i64 %k, 1
  %cmp.k = icmp slt i64 %inc.k, %N
  br i1 %cmp.k, label %header.k, label %latch.j

latch.j:
  %inc.j = add nsw i64 %j, 1
  %cmp.j = icmp slt i64 %inc.j, %N
  br i1 %cmp.j, label %header.j, label %latch.i

latch.i:
  %inc.i = add nsw i64 %i, 1
  %cmp.i = icmp slt i64 %inc.i, %N
  br i1 %cmp.i, label %header.i, label %exit.i

exit.i:
  ret void
}
declare void @baz(i64, i64, i64)
If the blocks of loop-i is in the order: header.i, latch.k, header.k,
header.j, latch.j, latch.i,
then cloneLoopWithPreheader would trigger the assertion in
addBasicBlockToLoop
assert(contains(SameHeader) && getHeader() == SameHeader->getHeader() &&
"Incorrect LI specified for this loop!");

As latch.k is in both loop-j and loop-k, it would be set as the header
of both loops after adding latch.k.
If we update loop headers during cloning blocks, then after adding
header.k,
the header of loop-k would be updated with header.k,
while the header of loop-j stays as latch.k.

When adding header.j, SameHeader is loop-k, SameHeader->getHeader() is
header.k, but getHeader() is latch.k, which trigger the assertion.
Reviewer: jdoerfert, Meinersbur, fhahn, kbarton, hfinkel, bmahjour,
etiotto
Reviewed By: Meinersbur
Subscribers: hiraditya, llvm-commits
Tag: LLVM
Differential Revision: https://reviews.llvm.org/D74382

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/CloneFunction.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index 75e8963303c2..1e020db21db1 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp
@@ -804,8 +804,6 @@ Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
 
     // Update LoopInfo.
     NewLoop->addBasicBlockToLoop(NewBB, *LI);
-    if (BB == CurLoop->getHeader())
-      NewLoop->moveToHeader(NewBB);
 
     // Add DominatorTree node. After seeing all blocks, update to correct
     // IDom.
@@ -815,6 +813,11 @@ Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
   }
 
   for (BasicBlock *BB : OrigLoop->getBlocks()) {
+    // Update loop headers.
+    Loop *CurLoop = LI->getLoopFor(BB);
+    if (BB == CurLoop->getHeader())
+      LMap[CurLoop]->moveToHeader(cast<BasicBlock>(VMap[BB]));
+
     // Update DominatorTree.
     BasicBlock *IDomBB = DT->getNode(BB)->getIDom()->getBlock();
     DT->changeImmediateDominator(cast<BasicBlock>(VMap[BB]),