[LoopUnroll] Properly update loop-info when cloning prologues and epilogues.

[mzolotukhin]

Summary

When cloning blocks for prologue/epilogue we need to replicate the loop
structure from the original loop. It wasn't a problem for the innermost
loops, but it led to an incorrect loop info when we unrolled a loop with
a child loop - in this case created prologue-loop had a child loop, but
loop info didn't reflect that.

This fixes PR28888.

[evstupac]

Hi Michael,

The fix is good. I have some suggestions in inline comments.

Thanks,
Evgeny

lib/Transforms/Utils/LoopUnrollRuntime.cpp

306 ↗

(On Diff #70213)

Since
line 632: LoopBlocksDFS LoopBlocks(L);
Is there a case when LI->getLoopFor(*BB) is not L?

309 ↗

(On Diff #70213)

Should be simplified to just:
if (CreateRemainderLoop) {

[evstupac]

lib/Transforms/Utils/LoopUnrollRuntime.cpp

306 ↗

(On Diff #70213)

I see when we clone loop nest.
Please ignore this.

[chandlerc]

Really nice. LGTM with straight forward comments / refactorings below.

lib/Transforms/Utils/LoopUnrollRuntime.cpp

307–309 ↗

(On Diff #70213)

Factoring this into a helper or a lambda would make it more readable IMO because you could use early returns.

[mzolotukhin]

Thanks for the review! I'll commit it soon.

Michael

lib/Transforms/Utils/LoopUnrollRuntime.cpp

307–309 ↗

(On Diff #70213)

That's a good idea. I'll see if it makes the code better before I commit.

309 ↗

(On Diff #70213)

I assume that this comment can also be ignored now, but I'll expand on this just in case.

If we unroll a loop with inner loops, then even if the remainder is not a loop itself, it will contain cloned inner loops. This check catches exactly this case: if the block BB belongs not to the L (but to some of its children), then we'll need to create cloned version of this inner loop in the remainder.

llvm/trunk/lib/Transforms/Utils/LoopUnrollRuntime.cpp

static void CloneLoopBlocks(Loop *L, Value *NewIter,
                             LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap,
                             LoopInfo *LI) {
   StringRef suffix = UseEpilogRemainder ? "epil" : "prol";
   BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   Function *F = Header->getParent();
   LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();
   LoopBlocksDFS::RPOIterator BlockEnd = LoopBlocks.endRPO();
-  Loop *NewLoop = nullptr;
   Loop *ParentLoop = L->getParentLoop();
-  if (CreateRemainderLoop) {
-    NewLoop = new Loop();
-    if (ParentLoop)
-      ParentLoop->addChildLoop(NewLoop);
-    else
-      LI->addTopLevelLoop(NewLoop);
-  }
+
+  // The map from original loops to their cloned copies.
+  SmallDenseMap<const Loop *, Loop *, 4> NewLoops;
 
   // For each block in the original loop, create a new copy,
   // and update the value map with the newly created values.
   for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
     BasicBlock *NewBB = CloneBasicBlock(*BB, VMap, "." + suffix, F);
     NewBlocks.push_back(NewBB);
 
+    // Figure out which loop NewBB is in.
+    auto findClonedLoop = [&](const Loop *OldLoop) {
+      Loop *&NewLoop = NewLoops[OldLoop];
+      // If we've encountered this loop before, return it right away.
+      if (NewLoop)
+        return NewLoop;
+
+      // If BB is from L, and we're not creating a remainder, the loop for
+      // NewBB will be ParentLoop, which might be null. Update NewLoops map and
+      // return ParentLoop.
+      if (OldLoop == L && !CreateRemainderLoop)
+        return (NewLoop = ParentLoop);
+
+      // Now we know that there should be a cloned counterpart for OldLoop, but
+      // we haven't seen it yet. Note that OldLoop might be L if we're
+      // generating a remainder loop, or it can be an inner loop of L - in this
+      // case we'll recreate the loop structure of L in its clone.
+
+      // This is a first block belonging to OldLoop encountered in our RPO
+      // traversal.
+      assert(*BB == OldLoop->getHeader() && "Header should be first in RPO");
+
+      NewLoop = new Loop;
+      Loop *OldLoopParent = OldLoop->getParentLoop();
+      // If OldLoop has a parent loop, we have two options:
+      //   1. ParentLoop is the parent of L. It won't be cloned, and it will
+      //      be a parent for NewLoop too.
+      //   2. ParentLoop is not a parent of L. In this case, it should be one
+      //      of the cloned loops and we should be able to find it in our map.
+      //
+      // If OldLoop doesn't have a parent, then NewLoop should be yet another
+      // top-level loop.
+      if (OldLoopParent) {
+        Loop *NewLoopParent = ParentLoop == OldLoopParent
+                                  ? ParentLoop
+                                  : NewLoops.lookup(OldLoopParent);
+        assert(NewLoopParent && "Expected parent loop before sub-loop in RPO");
+        NewLoopParent->addChildLoop(NewLoop);
+      } else
+        LI->addTopLevelLoop(NewLoop);
+      return NewLoop;
+    };
+
+    Loop *NewLoop = findClonedLoop(LI->getLoopFor(*BB));
+
+    assert(NewLoop ||
+           (!CreateRemainderLoop && !ParentLoop) &&
+               "NewLoop can only be null if we are cloning top-level loop "
+               "without creating a remainder loop.");
     if (NewLoop)
       NewLoop->addBasicBlockToLoop(NewBB, *LI);
-    else if (ParentLoop)
-      ParentLoop->addBasicBlockToLoop(NewBB, *LI);
 
     VMap[*BB] = NewBB;
     if (Header == *BB) {
       // For the first block, add a CFG connection to this newly
       // created block.
       InsertTop->getTerminator()->setSuccessor(0, NewBB);
     }
 
if (!CreateRemainderLoop) {
       BasicBlock *NewLatch = cast<BasicBlock>(VMap[Latch]);
       idx = NewPHI->getBasicBlockIndex(Latch);
       Value *InVal = NewPHI->getIncomingValue(idx);
       NewPHI->setIncomingBlock(idx, NewLatch);
       if (Value *V = VMap.lookup(InVal))
         NewPHI->setIncomingValue(idx, V);
     }
   }
-  if (NewLoop) {
+  if (CreateRemainderLoop) {
+    Loop *NewLoop = NewLoops[L];
     // Add unroll disable metadata to disable future unrolling for this loop.
     SmallVector<Metadata *, 4> MDs;
     // Reserve first location for self reference to the LoopID metadata node.
     MDs.push_back(nullptr);
     MDNode *LoopID = NewLoop->getLoopID();
     if (LoopID) {
       // First remove any existing loop unrolling metadata.
       for (unsigned i = 1, ie = LoopID->getNumOperands(); i < ie; ++i) {

llvm/trunk/test/Transforms/LoopUnroll/pr28888.ll

+; RUN: opt -loop-unroll -verify-loop-info -unroll-runtime-epilog=false -unroll-count=4 -S < %s | FileCheck %s -check-prefix=PROLOG
+; RUN: opt -loop-unroll -verify-loop-info -unroll-runtime-epilog=true  -unroll-count=4 -S < %s | FileCheck %s -check-prefix=EPILOG
+
+; PR28888
+; Check that loop info is correct if we unroll an outer loop, and thus the
+; remainder loop has a child loop.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; PROLOG-LABEL: @foo
+; EPILOG-LABEL: @foo
+define void @foo(i1 %x) #0 {
+bb:
+  br label %bb1
+
+bb1:
+  br label %bb2
+
+; PROLOG: bb2.prol:
+; EPILOG: bb2.epil:
+bb2:
+  %tmp = phi i64 [ 0, %bb1 ], [ %tmp2, %bb5 ]
+  br label %bb3
+
+bb3:
+  br label %bb4
+
+bb4:
+  br i1 %x, label %bb3, label %bb5
+
+; PROLOG: bb5.3:
+; EPILOG: bb5.3:
+bb5:
+  %tmp2 = add nuw nsw i64 %tmp, 1
+  %tmp3 = trunc i64 %tmp2 to i32
+  %tmp4 = icmp eq i32 %tmp3, undef
+  br i1 %tmp4, label %bb6, label %bb2
+
+bb6:
+  br label %bb1
+}
+
+attributes #0 = { "target-cpu"="x86-64" }