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

[Whitney]

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.

[Meinersbur]

Could you add the example in the summary as regression test?

[Whitney]

Could you add the example in the summary as regression test?

There is an example in summary which would fail before and pass now. Do you mean write it as complete LLVMIR?

[Meinersbur]

In D74382#1879821, @Whitney wrote:

There is an example in summary which would fail before and pass now. Do you mean write it as complete LLVMIR?

Yes, I though about a regression test in either llvm/test or llvm/unittests.

[Whitney]

In D74382#1879841, @Meinersbur wrote:

In D74382#1879821, @Whitney wrote:

There is an example in summary which would fail before and pass now. Do you mean write it as complete LLVMIR?

Yes, I though about a regression test in either llvm/test or llvm/unittests.

OK, I will change the example to complete LLVMIR. I though of a regression test in llvm/unittests too, however, I cannot find an easy way to manipulate the blocks list.

[Whitney]

@Meinersbur I have updated the example in the summary with complete LLVMIR, please let me know if that's sufficient. Thanks!

[Meinersbur]

In D74382#1885245, @Whitney wrote:

@Meinersbur I have updated the example in the summary with complete LLVMIR, please let me know if that's sufficient. Thanks!

What I meant is to update this patch to include the example that triggers this assertion as a regression test so running ninja check-llvm confirms that the assertion is not triggered anymore.

[Whitney]

In D74382#1886689, @Meinersbur wrote:

In D74382#1885245, @Whitney wrote:

@Meinersbur I have updated the example in the summary with complete LLVMIR, please let me know if that's sufficient. Thanks!

What I meant is to update this patch to include the example that triggers this assertion as a regression test so running ninja check-llvm confirms that the assertion is not triggered anymore.

The assertion only get triggered if the blocks returned from getBlocks() of loop header.i is in the order of: header.i, latch.k, header.k, header.j, latch.j, latch.i, when cloneLoopWithPreheader is called with loop header.i.
The only way I see of changing the block order is using reverseBlock(unsigned from). Do you know of other ways I can modify the order of the blocks stored in a loop?

[Meinersbur]

Sorry, I was assuming that getBlocks() follows the order in the read input files (so I did not understand why it would not be testable), but it is probably an order derived from the DFS traversal.

LGTM.

llvm/lib/Transforms/Utils/CloneFunction.cpp

for (BasicBlock *BB : OrigLoop->getBlocks()) {
     Loop *&NewLoop = LMap[CurLoop];
     assert(NewLoop && "Expecting new loop to be allocated");
 
     BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
     VMap[BB] = NewBB;
 
     // Update LoopInfo.
     NewLoop->addBasicBlockToLoop(NewBB, *LI);
-    if (BB == CurLoop->getHeader())
-      NewLoop->moveToHeader(NewBB);
 
     // Add DominatorTree node. After seeing all blocks, update to correct
     // IDom.
     DT->addNewBlock(NewBB, NewPH);
 
     Blocks.push_back(NewBB);
   }
 
   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]),
                                  cast<BasicBlock>(VMap[IDomBB]));
   }
 
   // Move them physically from the end of the block list.
   F->getBasicBlockList().splice(Before->getIterator(), F->getBasicBlockList(),