Preserve domtree and loop-simplify for runtime unrolling.

[efriedma]

Summary

Mostly straightforward changes; we just didn't do the computation before. One sort of interesting change in LoopUnroll.cpp: we weren't handling dominance for children of the loop latch correctly, but foldBlockIntoPredecessor hid the problem for complete unrolling.

Currently punting on loop peeling; made some minor changes to isolate that problem to LoopUnrollPeel.cpp.

Adds a flag -unroll-verify-domtree; this is on by default for +Asserts builds.

[mkuper]

lib/Transforms/Utils/LoopUnroll.cpp

55 ↗

(On Diff #82393)

We want to enforce this being on in new tests too, I assume?

619 ↗

(On Diff #82393)

Isn't it the case that when the latch ends with an unconditional branch, that branch is towards the header block?
If so, we should not have any children outside the loop.

[efriedma]

lib/Transforms/Utils/LoopUnroll.cpp

55 ↗

(On Diff #82393)

We probably want this on in new tests, yes, but I have no idea how we would enforce it.

619 ↗

(On Diff #82393)

Maybe this is more clear? "The latch is special because we can emit an unconditional branch in the unrolled loop even if the original latch block ends in a conditional branch."

[mkuper]

lib/Transforms/Utils/LoopUnroll.cpp

55 ↗

(On Diff #82393)

In code review, I meant. :-)

619 ↗

(On Diff #82393)

Ah, ok, this makes sense.
This still looks a bit weird to me, though. Any chance you can give an example of when this ends up different than the nearest common dominator with Latches[0]?

(Feel free to ignore me - I'm not sure I understand this enough to LGTM it anyway. :-) )

[mzolotukhin]

lib/Transforms/Utils/LoopUnroll.cpp

55 ↗

(On Diff #82393)

Can we make the default value true for debug builds (it will match the existing behavior)?

619 ↗

(On Diff #82393)

I'm also curious about an example.

Also, this part is not clear to me: "Since the latch is always at the bottom of the loop, new dominator must also be a latch." Why can't a mid-block from the body be a dominator of a cloned latch?

639–642 ↗

(On Diff #82393)

Thanks for fixing this!

644–646 ↗

(On Diff #82393)

Do I understand it correctly that using -verify-dom-info wasn't sufficient because of foldBlockIntoPredecessor, and that's why we're now introducing unroll-verify-domtree?

700–701 ↗

(On Diff #82393)

Why can we remove this?

[efriedma]

lib/Transforms/Utils/LoopUnroll.cpp

55 ↗

(On Diff #82393)

Okay. (I initially didn't do this because it's O(N^2) in theory, but it turns out it doesn't really matter that much.)

619 ↗

(On Diff #82393)

This shows up for example, if you completely unroll a single-block loop with a constant trip count; the immediate dominator of the exit is the last iteration of the loop, not the first iteration. (foldBlockIntoPredecessor hides this problem because the block in question gets folded away when a loop is completely unrolled.)

Why can't a mid-block from the body be a dominator of a cloned latch?

I'm not sure what you're asking here. ChildrenToUpdate is the set of loop exit blocks which are dominated by BB, so this code only touches the idom for blocks outside the loop. The latch's idom isn't relevant.

644–646 ↗

(On Diff #82393)

Yes... we want to verify at this point to try and catch as many mistakes as possible.

700–701 ↗

(On Diff #82393)

There are three possibilities for unrolling: we could be completely unrolling a loop, we could be runtime unrolling a loop, or we could be peeling a loop. If we're completely unrolling, these lines are a no-op. If we're runtime unrolling, we now correctly preserve LoopSimplify for L (and we don't break it for any loop outside of L). If we're peeling, we explicitly recreate LoopSimplify earlier.

[mzolotukhin]

lib/Transforms/Utils/LoopUnroll.cpp

55 ↗

(On Diff #82393)

Thanks! Also, if you change it to true, then changes in the tests are no longer needed.

619 ↗

(On Diff #82393)

I've just understood what the problem is, thanks! Yes, it all makes sense now.

625–631 ↗

(On Diff #82393)

I think we should be able to tell which latch ends with a conditional. It should be either the first one (if all latches end with conditional branches), or the last one. Does it make sense? If so, can we remove this loop?

700–701 ↗

(On Diff #82393)

But isn't it also used for LCSSA? Maybe I'm missing something, but it's still not obvious to me that we can remove it.

[efriedma]

lib/Transforms/Utils/LoopUnroll.cpp

700–701 ↗

(On Diff #82393)

!CompletelyUnroll implies !NeedToFixLCSSA.

[mzolotukhin]

lib/Transforms/Utils/LoopUnroll.cpp

700–701 ↗

(On Diff #82393)

Ah, right!

[efriedma]

lib/Transforms/Utils/LoopUnroll.cpp

625–631 ↗

(On Diff #82393)

The actual conditional is complicated... if we're completely unrolling, it's either the first one or the last one, depending on PreserveCondBr. If we're using a remainder loop, it's always the last one. If we're partially unrolling without a remainder loop, it might not be the first or last one; there's a complicated conditional involving BreakingTrip and TripMultiple. Anyway, this is much more straightforward than trying to duplicate the NeedConditional logic.

[efriedma]

Verify domtree by default in +Asserts mode. Clarify comments.

[mzolotukhin]

Looks good to me, thanks!

Michael

lib/Transforms/Utils/LoopUnroll.cpp

625–631 ↗

(On Diff #82393)

Makes sense.

llvm/trunk/include/llvm/Transforms/Utils/UnrollLoop.h

bool UnrollRuntimeLoopRemainder(Loop *L, unsigned Count,
                                 bool UseEpilogRemainder, LoopInfo *LI,
                                 ScalarEvolution *SE, DominatorTree *DT,
                                 bool PreserveLCSSA);
 
 void computePeelCount(Loop *L, unsigned LoopSize,
                       TargetTransformInfo::UnrollingPreferences &UP);
 
 bool peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI, ScalarEvolution *SE,
-              DominatorTree *DT, bool PreserveLCSSA);
+              DominatorTree *DT, AssumptionCache *AC, bool PreserveLCSSA);
 
 MDNode *GetUnrollMetadata(MDNode *LoopID, StringRef Name);
 }
 
 #endif

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

 STATISTIC(NumCompletelyUnrolled, "Number of loops completely unrolled");
 STATISTIC(NumUnrolled, "Number of loops unrolled (completely or otherwise)");
 
 static cl::opt<bool>
 UnrollRuntimeEpilog("unroll-runtime-epilog", cl::init(false), cl::Hidden,
                     cl::desc("Allow runtime unrolled loops to be unrolled "
                              "with epilog instead of prolog."));
 
+static cl::opt<bool>
+UnrollVerifyDomtree("unroll-verify-domtree", cl::Hidden,
+                    cl::desc("Verify domtree after unrolling"),
+#ifdef NDEBUG
+    cl::init(false)
+#else
+    cl::init(true)
+#endif
+                    );
+
 /// Convert the instruction operands from referencing the current values into
 /// those specified by VMap.
 static inline void remapInstruction(Instruction *I,
                                     ValueToValueMapTy &VMap) {
   for (unsigned op = 0, E = I->getNumOperands(); op != E; ++op) {
     Value *Op = I->getOperand(op);
     ValueToValueMapTy::iterator It = VMap.find(Op);
     if (It != VMap.end())
bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
   // flag is specified.
   bool RuntimeTripCount = (TripCount == 0 && Count > 0 && AllowRuntime);
 
   assert((!RuntimeTripCount || !PeelCount) &&
          "Did not expect runtime trip-count unrolling "
          "and peeling for the same loop");
 
   if (PeelCount)
-    peelLoop(L, PeelCount, LI, SE, DT, PreserveLCSSA);
+    peelLoop(L, PeelCount, LI, SE, DT, AC, PreserveLCSSA);
 
   // Loops containing convergent instructions must have a count that divides
   // their TripMultiple.
   DEBUG(
       {
         bool HasConvergent = false;
         for (auto &BB : L->blocks())
           for (auto &I : *BB)
if (NeedConditional) {
           }
         }
       }
       // Replace the conditional branch with an unconditional one.
       BranchInst::Create(Dest, Term);
       Term->eraseFromParent();
     }
   }
+
   // Update dominators of blocks we might reach through exits.
   // Immediate dominator of such block might change, because we add more
   // routes which can lead to the exit: we can now reach it from the copied
-  // iterations too. Thus, the new idom of the block will be the nearest
-  // common dominator of the previous idom and common dominator of all copies of
-  // the previous idom. This is equivalent to the nearest common dominator of
-  // the previous idom and the first latch, which dominates all copies of the
-  // previous idom.
+  // iterations too.
   if (DT && Count > 1) {
     for (auto *BB : OriginalLoopBlocks) {
       auto *BBDomNode = DT->getNode(BB);
       SmallVector<BasicBlock *, 16> ChildrenToUpdate;
       for (auto *ChildDomNode : BBDomNode->getChildren()) {
         auto *ChildBB = ChildDomNode->getBlock();
         if (!L->contains(ChildBB))
           ChildrenToUpdate.push_back(ChildBB);
       }
-      BasicBlock *NewIDom = DT->findNearestCommonDominator(BB, Latches[0]);
+      BasicBlock *NewIDom;
+      if (BB == LatchBlock) {
+        // The latch is special because we emit unconditional branches in
+        // some cases where the original loop contained a conditional branch.
+        // Since the latch is always at the bottom of the loop, if the latch
+        // dominated an exit before unrolling, the new dominator of that exit
+        // must also be a latch.  Specifically, the dominator is the first
+        // latch which ends in a conditional branch, or the last latch if
+        // there is no such latch.
+        NewIDom = Latches.back();
+        for (BasicBlock *IterLatch : Latches) {
+          TerminatorInst *Term = IterLatch->getTerminator();
+          if (isa<BranchInst>(Term) && cast<BranchInst>(Term)->isConditional()) {
+            NewIDom = IterLatch;
+            break;
+          }
+        }
+      } else {
+        // The new idom of the block will be the nearest common dominator
+        // of all copies of the previous idom. This is equivalent to the
+        // nearest common dominator of the previous idom and the first latch,
+        // which dominates all copies of the previous idom.
+        NewIDom = DT->findNearestCommonDominator(BB, LatchBlock);
+      }
       for (auto *ChildBB : ChildrenToUpdate)
         DT->changeImmediateDominator(ChildBB, NewIDom);
     }
   }
 
+  if (DT && UnrollVerifyDomtree)
+    DT->verifyDomTree();
+
   // Merge adjacent basic blocks, if possible.
   SmallPtrSet<Loop *, 4> ForgottenLoops;
   for (BasicBlock *Latch : Latches) {
     BranchInst *Term = cast<BranchInst>(Latch->getTerminator());
     if (Term->isUnconditional()) {
       BasicBlock *Dest = Term->getSuccessor(0);
       if (BasicBlock *Fold =
               foldBlockIntoPredecessor(Dest, LI, SE, ForgottenLoops, DT)) {
         // Dest has been folded into Fold. Update our worklists accordingly.
         std::replace(Latches.begin(), Latches.end(), Dest, Fold);
         UnrolledLoopBlocks.erase(std::remove(UnrolledLoopBlocks.begin(),
                                              UnrolledLoopBlocks.end(), Dest),
                                  UnrolledLoopBlocks.end());
       }
     }
   }
 
-  // FIXME: We only preserve DT info for complete unrolling now. Incrementally
-  // updating domtree after partial loop unrolling should also be easy.
-  if (DT && !CompletelyUnroll)
-    DT->recalculate(*L->getHeader()->getParent());
-  else if (DT)
-    DEBUG(DT->verifyDomTree());
-
   // Simplify any new induction variables in the partially unrolled loop.
   if (SE && !CompletelyUnroll && Count > 1) {
     SmallVector<WeakVH, 16> DeadInsts;
     simplifyLoopIVs(L, SE, DT, LI, DeadInsts);
 
     // Aggressively clean up dead instructions that simplifyLoopIVs already
     // identified. Any remaining should be cleaned up below.
     while (!DeadInsts.empty())
bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
   if (PreserveLCSSA && OuterL && CompletelyUnroll && !NeedToFixLCSSA)
     NeedToFixLCSSA |= ::needToInsertPhisForLCSSA(OuterL, UnrolledLoopBlocks, LI);
 
   // 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);
 

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

 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/UnrollLoop.h"
 #include <algorithm>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-unroll"
 STATISTIC(NumPeeled, "Number of loops peeled");

 /// Rather, each iteration is peeled off separately, and needs to check the
 /// exit condition.
 /// For loops that dynamically execute \p PeelCount iterations or less
 /// this provides a benefit, since the peeled off iterations, which account
 /// for the bulk of dynamic execution, can be further simplified by scalar
 /// optimizations.
 bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
                     ScalarEvolution *SE, DominatorTree *DT,
-                    bool PreserveLCSSA) {
+                    AssumptionCache *AC, bool PreserveLCSSA) {
   if (!canPeel(L))
     return false;
 
   LoopBlocksDFS LoopBlocks(L);
   LoopBlocks.perform(LI);
 
   BasicBlock *Header = L->getHeader();
   BasicBlock *PreHeader = L->getLoopPreheader();
else
       BackEdgeWeight = 1;
     MDBuilder MDB(LatchBR->getContext());
     MDNode *WeightNode =
         HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)
                   : MDB.createBranchWeights(BackEdgeWeight, ExitWeight);
     LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
   }
 
+  // FIXME: Incrementally update domtree.
+  DT->recalculate(*L->getHeader()->getParent());
+
   // If the loop is nested, we changed the parent loop, update SE.
-  if (Loop *ParentLoop = L->getParentLoop())
+  if (Loop *ParentLoop = L->getParentLoop()) {
     SE->forgetLoop(ParentLoop);
 
+    // FIXME: Incrementally update loop-simplify
+    simplifyLoop(ParentLoop, DT, LI, SE, AC, PreserveLCSSA);
+  } else {
+    // FIXME: Incrementally update loop-simplify
+    simplifyLoop(L, DT, LI, SE, AC, PreserveLCSSA);
+  }
+
   NumPeeled++;
 
   return true;
 }

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

static void ConnectProlog(Loop *L, Value *BECount, unsigned Count,
   assert(Exit && "Loop must have a single exit block only");
   // Split the exit to maintain loop canonicalization guarantees
   SmallVector<BasicBlock*, 4> Preds(predecessors(Exit));
   SplitBlockPredecessors(Exit, Preds, ".unr-lcssa", DT, LI,
                          PreserveLCSSA);
   // Add the branch to the exit block (around the unrolled loop)
   B.CreateCondBr(BrLoopExit, Exit, NewPreHeader);
   InsertPt->eraseFromParent();
+  if (DT)
+    DT->changeImmediateDominator(Exit, PrologExit);
 }
 
 /// Connect the unrolling epilog code to the original loop.
 /// The unrolling epilog code contains code to execute the
 /// 'extra' iterations if the run-time trip count modulo the
 /// unroll count is non-zero.
 ///
 /// This function performs the following:
for (Instruction &BBI : *Succ) {
       VPN->setIncomingValue(VPN->getBasicBlockIndex(EpilogPreHeader), NewPN);
     }
   }
 
   Instruction *InsertPt = NewExit->getTerminator();
   IRBuilder<> B(InsertPt);
   Value *BrLoopExit = B.CreateIsNotNull(ModVal, "lcmp.mod");
   assert(Exit && "Loop must have a single exit block only");
-  // Split the exit to maintain loop canonicalization guarantees
+  // Split the epilogue exit to maintain loop canonicalization guarantees
   SmallVector<BasicBlock*, 4> Preds(predecessors(Exit));
   SplitBlockPredecessors(Exit, Preds, ".epilog-lcssa", DT, LI,
                          PreserveLCSSA);
   // Add the branch to the exit block (around the unrolling loop)
   B.CreateCondBr(BrLoopExit, EpilogPreHeader, Exit);
   InsertPt->eraseFromParent();
+  if (DT)
+    DT->changeImmediateDominator(Exit, NewExit);
+
+  // Split the main loop exit to maintain canonicalization guarantees.
+  SmallVector<BasicBlock*, 4> NewExitPreds{Latch};
+  SplitBlockPredecessors(NewExit, NewExitPreds, ".loopexit", DT, LI,
+                         PreserveLCSSA);
 }
 
 /// Create a clone of the blocks in a loop and connect them together.
 /// If CreateRemainderLoop is false, loop structure will not be cloned,
 /// otherwise a new loop will be created including all cloned blocks, and the
 /// iterator of it switches to count NewIter down to 0.
 /// The cloned blocks should be inserted between InsertTop and InsertBot.
 /// If loop structure is cloned InsertTop should be new preheader, InsertBot
 /// new loop exit.
 ///
 static void CloneLoopBlocks(Loop *L, Value *NewIter,
                             const bool CreateRemainderLoop,
                             const bool UseEpilogRemainder,
                             BasicBlock *InsertTop, BasicBlock *InsertBot,
                             BasicBlock *Preheader,
                             std::vector<BasicBlock *> &NewBlocks,
                             LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap,
-                            LoopInfo *LI) {
+                            DominatorTree *DT, 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();
for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
 
     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 (DT) {
+      if (Header == *BB) {
+        // The header is dominated by the preheader.
+        DT->addNewBlock(NewBB, InsertTop);
+      } else {
+        // Copy information from original loop to unrolled loop.
+        BasicBlock *IDomBB = DT->getNode(*BB)->getIDom()->getBlock();
+        DT->addNewBlock(NewBB, cast<BasicBlock>(VMap[IDomBB]));
+      }
+    }
+
     if (Latch == *BB) {
       // For the last block, if CreateRemainderLoop is false, create a direct
       // jump to InsertBot. If not, create a loop back to cloned head.
       VMap.erase((*BB)->getTerminator());
       BasicBlock *FirstLoopBB = cast<BasicBlock>(VMap[Header]);
       BranchInst *LatchBR = cast<BranchInst>(NewBB->getTerminator());
       IRBuilder<> Builder(LatchBR);
       if (!CreateRemainderLoop) {
Value *BranchVal =
                                            ConstantInt::get(BECount->getType(),
                                                             Count - 1)) :
                            B.CreateIsNotNull(ModVal, "lcmp.mod");
   BasicBlock *RemainderLoop = UseEpilogRemainder ? NewExit : PrologPreHeader;
   BasicBlock *UnrollingLoop = UseEpilogRemainder ? NewPreHeader : PrologExit;
   // Branch to either remainder (extra iterations) loop or unrolling loop.
   B.CreateCondBr(BranchVal, RemainderLoop, UnrollingLoop);
   PreHeaderBR->eraseFromParent();
+  if (DT) {
+    if (UseEpilogRemainder)
+      DT->changeImmediateDominator(NewExit, PreHeader);
+    else
+      DT->changeImmediateDominator(PrologExit, PreHeader);
+  }
   Function *F = Header->getParent();
   // Get an ordered list of blocks in the loop to help with the ordering of the
   // cloned blocks in the prolog/epilog code
   LoopBlocksDFS LoopBlocks(L);
   LoopBlocks.perform(LI);
 
   //
   // For each extra loop iteration, create a copy of the loop's basic blocks
   // and generate a condition that branches to the copy depending on the
   // number of 'left over' iterations.
   //
   std::vector<BasicBlock *> NewBlocks;
   ValueToValueMapTy VMap;
 
   // For unroll factor 2 remainder loop will have 1 iterations.
   // Do not create 1 iteration loop.
   bool CreateRemainderLoop = (Count != 2);
 
   // Clone all the basic blocks in the loop. If Count is 2, we don't clone
   // the loop, otherwise we create a cloned loop to execute the extra
   // iterations. This function adds the appropriate CFG connections.
   BasicBlock *InsertBot = UseEpilogRemainder ? Exit : PrologExit;
   BasicBlock *InsertTop = UseEpilogRemainder ? EpilogPreHeader : PrologPreHeader;
   CloneLoopBlocks(L, ModVal, CreateRemainderLoop, UseEpilogRemainder, InsertTop,
-                  InsertBot, NewPreHeader, NewBlocks, LoopBlocks, VMap, LI);
+                  InsertBot, NewPreHeader, NewBlocks, LoopBlocks, VMap, DT, LI);
 
   // Insert the cloned blocks into the function.
   F->getBasicBlockList().splice(InsertBot->getIterator(),
                                 F->getBasicBlockList(),
                                 NewBlocks[0]->getIterator(),
                                 F->end());
 
   // Loop structure should be the following: