[LoopUnroll] Simplify loops created by unrolling.

[mzolotukhin]

Summary

Currently loop-unrolling doesn't preserve loop-simplified form. This patch
fixes it by resimplifying affected loops.

[chandlerc]

lib/Transforms/Utils/LoopUnroll.cpp

386–387 ↗

(On Diff #66747)

So, this is done because we leave the first copy of the loop alone when unrolling?

If not, shouldn't we handle this below so that we don't add inner loops that won't actually exist after unrolling?

692 ↗

(On Diff #66747)

Related to the above TODO comment, the ways in which this can break loop simplified form seem fairly constrained. If this is a compile time issue, its likely a more targeted approach could be used... Happy to just leave a comment here if you don't yet know whether we need to consider the compile time here.

[mzolotukhin]

lib/Transforms/Utils/LoopUnroll.cpp

386–387 ↗

(On Diff #66747)

this is done because we leave the first copy of the loop alone when unrolling?

Yes, exactly. When we unroll something like this:

outer_loop {
  inner_loop
}

we could get something like this:

outer_loop {
  inner_loop
  inner_loop2
}

or even

inner_loop
inner_loop2

in the case of complete unrolling.

inner_loop2 is a new loop, created by loop-unrolling (see lines 417-418). inner_loop existed in the original IR and remains (almost) untouched. It still needs to be simplified, as it'll be using the same exit block, as the cloned loops.

692 ↗

(On Diff #66747)

the ways in which this can break loop simplified form seem fairly constrained

This is true, that's why I left the TODO there. It should be definitely possible to fix only broken parts (i.e. insert preheaders and split exit-edges). On the other hand, we've been calling simplifyLoop on OuterL without any issues before, which also simplifies all nested loops - that was my motivation for keeping it simple for now. My idea was that comment from above relates to this spot as well, just didn't want to replicate the same TODO.

[chandlerc]

Patch LGTM.

Would it make sense to only fill in LoopsToSimplify if OuterL is null? Either way, feel free to land.

[mzolotukhin]

Thanks, committed in r278038.

Would it make sense to only fill in LoopsToSimplify if OuterL is null?

I think the checks will clutter the code for not so big benefit, so I decided to go without them.

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

bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
   LoopBlocksDFS DFS(L);
   DFS.perform(LI);
 
   // Stash the DFS iterators before adding blocks to the loop.
   LoopBlocksDFS::RPOIterator BlockBegin = DFS.beginRPO();
   LoopBlocksDFS::RPOIterator BlockEnd = DFS.endRPO();
 
   std::vector<BasicBlock*> UnrolledLoopBlocks = L->getBlocks();
+
+  // Loop Unrolling might create new loops. While we do preserve LoopInfo, we
+  // might break loop-simplified form for these loops (as they, e.g., would
+  // share the same exit blocks). We'll keep track of loops for which we can
+  // break this so that later we can re-simplify them.
+  SmallSetVector<Loop *, 4> LoopsToSimplify;
+  for (Loop *SubLoop : *L)
+    LoopsToSimplify.insert(SubLoop);
+
   for (unsigned It = 1; It != Count; ++It) {
     std::vector<BasicBlock*> NewBlocks;
     SmallDenseMap<const Loop *, Loop *, 4> NewLoops;
     NewLoops[L] = L;
 
     for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
       ValueToValueMapTy VMap;
       BasicBlock *New = CloneBasicBlock(*BB, VMap, "." + Twine(It));
for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
           assert(*BB == OldLoop->getHeader() &&
                  "Header should be first in RPO");
 
           Loop *NewLoopParent = NewLoops.lookup(OldLoop->getParentLoop());
           assert(NewLoopParent &&
                  "Expected parent loop before sub-loop in RPO");
           NewLoop = new Loop;
           NewLoopParent->addChildLoop(NewLoop);
+          LoopsToSimplify.insert(NewLoop);
 
           // Forget the old loop, since its inputs may have changed.
           if (SE)
             SE->forgetLoop(OldLoop);
         }
         NewLoop->addBasicBlockToLoop(New, *LI);
       }
 
bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
   // If we have a pass and a DominatorTree we should re-simplify impacted loops
   // to ensure subsequent analyses can rely on this form. We want to simplify
   // at least one layer outside of the loop that was unrolled so that any
   // changes to the parent loop exposed by the unrolling are considered.
   if (DT) {
     if (!OuterL && !CompletelyUnroll)
       OuterL = L;
     if (OuterL) {
+      // OuterL includes all loops for which we can break loop-simplify, so
+      // it's sufficient to simplify only it (it'll recursively simplify inner
+      // loops too).
+      // TODO: That potentially might be compile-time expensive. We should try
+      // to fix the loop-simplified form incrementally.
       simplifyLoop(OuterL, DT, LI, SE, AC, PreserveLCSSA);
 
       // LCSSA must be performed on the outermost affected loop. The unrolled
       // loop's last loop latch is guaranteed to be in the outermost loop after
       // LoopInfo's been updated by markAsRemoved.
       Loop *LatchLoop = LI->getLoopFor(Latches.back());
       if (!OuterL->contains(LatchLoop))
         while (OuterL->getParentLoop() != LatchLoop)
           OuterL = OuterL->getParentLoop();
 
       if (NeedToFixLCSSA)
         formLCSSARecursively(*OuterL, *DT, LI, SE);
       else
         assert(OuterL->isLCSSAForm(*DT) &&
                "Loops should be in LCSSA form after loop-unroll.");
+    } else {
+      // Simplify loops for which we might've broken loop-simplify form.
+      for (Loop *SubLoop : LoopsToSimplify)
+        simplifyLoop(SubLoop, DT, LI, SE, AC, PreserveLCSSA);
     }
   }
 
   return true;
 }
 
 /// Given an llvm.loop loop id metadata node, returns the loop hint metadata
 /// node with the given name (for example, "llvm.loop.unroll.count"). If no