[NFCI][LoopUnrollAndJam] Minor changes.

[Whitney]

Summary

1. Add assertions.
2. Use SmallVector instead of std::vector
3. Verify more analyses.

These changes are moved out of https://reviews.llvm.org/D73129 to simplify that review.

[fhahn]

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

525–527

Could you provide context on why not using DTU is preferred here? Is it because the existing copy/cloning code in the pass changes the DT in place? I just had a quick look and it looks like the DT is not accessed in between cloning until simplifying the loops after unrolling.

DTU has the advantage that we can preserve the PDT also for almost no extra cost (just pass a PDT to the DTU additionally). Also, batch updates can be much faster, given that the update set is large enough.

I think unless there's a strong reason, we should try to replace the existing uses of in-place updates with DTU update.

[Whitney]

Addressed Dave's comment.

[Whitney]

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

525–527

ok, I will revert my change here. The main reason I did that is it is easier to understand.

[Whitney]

Addressed Florian's comment.

[dmgreen]

Looks like some nice cleanups.

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

384

Why does the unroller have it's own version of this function, and how is it different to RemapInstruction?

If this is no longer needs to be shared with the unroller, it can be removed from UnrollLoop.h and made static in LoopUnroll.cpp.

[Whitney]

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

384

Looks to me that RemapInstruction is a superset of ::remapInstruction. I tried locally changing remapInstruction to call RemapInstruction(I, VMap, RF_NoModuleLevelChanges | RF_IgnoreMissingLocals); and all lit tests passed. Should be safe to remove ::remapInstruction. Should I do that in this patch?

[dmgreen]

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

384

There is a possibility that the loop unrollers remapInstruction is a subset for (compiler time) performance reasons. It may know that it needs to do less, so only handles what is required.

But there is a chance that it's just very old and no-one has noticed they are the same in the past though. Maybe look through git history to see if anything looks obvious?

If it does look OK to remove, we should probably change the Unroller and UnrollAndJam in the same patch. I imagine if there is a problem, the same conditions would apply to both.

[Whitney]

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

384

Created https://reviews.llvm.org/D73277 to address this.

[dmgreen]

Nice one. What's left here looks good to me.

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

DidSomething |=
         simplifyLoop(L, &DT, &LI, &SE, &AC, nullptr, false /* PreserveLCSSA */);
     DidSomething |= formLCSSARecursively(*L, DT, &LI, &SE);
   }
 
   SmallPriorityWorklist<Loop *, 4> Worklist;
   internal::appendLoopsToWorklist(reverse(LI), Worklist);
   while (!Worklist.empty()) {
     Loop *L = Worklist.pop_back_val();
-    formLCSSA(*L, DT, &LI, &SE);
     LoopUnrollResult Result =
         tryToUnrollAndJamLoop(L, DT, &LI, SE, TTI, AC, DI, ORE, OptLevel);
     if (Result != LoopUnrollResult::Unmodified)
       DidSomething = true;
   }
 
   return DidSomething;
 }

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

LoopUnrollResult llvm::UnrollAndJamLoop(
   SubLoopBlocksLast.push_back(SubLoop->getExitingBlock());
   AftBlocksFirst.push_back(SubLoop->getExitBlock());
   AftBlocksLast.push_back(L->getExitingBlock());
   // Maps Blocks[0] -> Blocks[It]
   ValueToValueMapTy LastValueMap;
 
   // Move any instructions from fore phi operands from AftBlocks into Fore.
   moveHeaderPhiOperandsToForeBlocks(
-      Header, LatchBlock, SubLoop->getLoopPreheader()->getTerminator(),
-      AftBlocks);
+      Header, LatchBlock, ForeBlocksLast[0]->getTerminator(), AftBlocks);
 
   // The current on-the-fly SSA update requires blocks to be processed in
   // reverse postorder so that LastValueMap contains the correct value at each
   // exit.
   LoopBlocksDFS DFS(L);
   DFS.perform(LI);
   // Stash the DFS iterators before adding blocks to the loop.
   LoopBlocksDFS::RPOIterator BlockBegin = DFS.beginRPO();
for (BasicBlock *BB : L->getBlocks())
             else
               LLVM_DEBUG(dbgs()
                          << "Failed to create new discriminator: "
                          << DIL->getFilename() << " Line: " << DIL->getLine());
           }
 
   // Copy all blocks
   for (unsigned It = 1; It != Count; ++It) {
-    std::vector<BasicBlock *> NewBlocks;
+    SmallVector<BasicBlock *, 4> NewBlocks;
     // Maps Blocks[It] -> Blocks[It-1]
     DenseMap<Value *, Value *> PrevItValueMap;
 
     for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
       ValueToValueMapTy VMap;
       BasicBlock *New = CloneBasicBlock(*BB, VMap, "." + Twine(It));
       Header->getParent()->getBasicBlockList().push_back(New);
 
for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
         assert(OriginalBBIDom);
         assert(LastValueMap[cast<Value>(OriginalBBIDom)]);
         DT->addNewBlock(
             New, cast<BasicBlock>(LastValueMap[cast<Value>(OriginalBBIDom)]));
       }
     }
 
     // Remap all instructions in the most recent iteration
+    remapInstructionsInBlocks(NewBlocks, LastValueMap);
     for (BasicBlock *NewBlock : NewBlocks) {
       for (Instruction &I : *NewBlock) {
-        ::remapInstruction(&I, LastValueMap);
         if (auto *II = dyn_cast<IntrinsicInst>(&I))
           if (II->getIntrinsicID() == Intrinsic::assume)
             AC->registerAssumption(II);
       }
     }
 
     // Alter the ForeBlocks phi's, pointing them at the latest version of the
     // value from the previous iteration's phis
LoopUnrollResult llvm::UnrollAndJamLoop(
 
   // Update the PHI values outside the loop to point to the last block
   updatePHIBlocksAndValues(LoopExit, AftBlocksLast[0], AftBlocksLast.back(),
                            LastValueMap);
 
   // Update ForeBlocks successors and phi nodes
   BranchInst *ForeTerm =
       cast<BranchInst>(ForeBlocksLast.back()->getTerminator());
-  BasicBlock *Dest = SubLoopBlocksFirst[0];
-  ForeTerm->setSuccessor(0, Dest);
+  assert(ForeTerm->getNumSuccessors() == 1 && "Expecting one successor");
+  ForeTerm->setSuccessor(0, SubLoopBlocksFirst[0]);
 
   if (CompletelyUnroll) {
     while (PHINode *Phi = dyn_cast<PHINode>(ForeBlocksFirst[0]->begin())) {
       Phi->replaceAllUsesWith(Phi->getIncomingValueForBlock(Preheader));
       Phi->getParent()->getInstList().erase(Phi);
     }
   } else {
     // Update the PHI values to point to the last aft block
     updatePHIBlocksAndValues(ForeBlocksFirst[0], AftBlocksLast[0],
                              AftBlocksLast.back(), LastValueMap);
   }
 
   for (unsigned It = 1; It != Count; It++) {
     // Remap ForeBlock successors from previous iteration to this
     BranchInst *ForeTerm =
         cast<BranchInst>(ForeBlocksLast[It - 1]->getTerminator());
-    BasicBlock *Dest = ForeBlocksFirst[It];
-    ForeTerm->setSuccessor(0, Dest);
+    assert(ForeTerm->getNumSuccessors() == 1 && "Expecting one successor");
+    ForeTerm->setSuccessor(0, ForeBlocksFirst[It]);
   }
 
   // Subloop successors and phis
   BranchInst *SubTerm =
       cast<BranchInst>(SubLoopBlocksLast.back()->getTerminator());
   SubTerm->setSuccessor(!SubLoopContinueOnTrue, SubLoopBlocksFirst[0]);
   SubTerm->setSuccessor(SubLoopContinueOnTrue, AftBlocksFirst[0]);
   updatePHIBlocks(SubLoopBlocksFirst[0], ForeBlocksLast[0],
for (unsigned It = 1; It != Count; It++) {
     updatePHIBlocks(SubLoopBlocksFirst[It], ForeBlocksLast[It],
                     ForeBlocksLast.back());
     updatePHIBlocks(SubLoopBlocksFirst[It], SubLoopBlocksLast[It],
                     SubLoopBlocksLast.back());
     movePHIs(SubLoopBlocksFirst[It], SubLoopBlocksFirst[0]);
   }
 
   // Aft blocks successors and phis
-  BranchInst *Term = cast<BranchInst>(AftBlocksLast.back()->getTerminator());
+  BranchInst *AftTerm = cast<BranchInst>(AftBlocksLast.back()->getTerminator());
   if (CompletelyUnroll) {
-    BranchInst::Create(LoopExit, Term);
-    Term->eraseFromParent();
+    BranchInst::Create(LoopExit, AftTerm);
+    AftTerm->eraseFromParent();
   } else {
-    Term->setSuccessor(!ContinueOnTrue, ForeBlocksFirst[0]);
+    AftTerm->setSuccessor(!ContinueOnTrue, ForeBlocksFirst[0]);
+    assert(AftTerm->getSuccessor(ContinueOnTrue) == LoopExit &&
+           "Expecting the ContinueOnTrue successor of AftTerm to be LoopExit");
   }
   updatePHIBlocks(AftBlocksFirst[0], SubLoopBlocksLast[0],
                   SubLoopBlocksLast.back());
 
   for (unsigned It = 1; It != Count; It++) {
     // Replace the conditional branch of the previous iteration subloop with an
     // unconditional one to this one
     BranchInst *AftTerm =
         cast<BranchInst>(AftBlocksLast[It - 1]->getTerminator());
     BranchInst::Create(AftBlocksFirst[It], AftTerm);
     AftTerm->eraseFromParent();
 
     updatePHIBlocks(AftBlocksFirst[It], SubLoopBlocksLast[It],
                     SubLoopBlocksLast.back());
     movePHIs(AftBlocksFirst[It], AftBlocksFirst[0]);
   }
 
-  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
   // Dominator Tree. Remove the old links between Fore, Sub and Aft, adding the
   // new ones required.
   if (Count != 1) {
-    SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
-    DTUpdates.emplace_back(DominatorTree::UpdateKind::Delete, ForeBlocksLast[0],
-                           SubLoopBlocksFirst[0]);
-    DTUpdates.emplace_back(DominatorTree::UpdateKind::Delete,
-                           SubLoopBlocksLast[0], AftBlocksFirst[0]);
-
-    DTUpdates.emplace_back(DominatorTree::UpdateKind::Insert,
-                           ForeBlocksLast.back(), SubLoopBlocksFirst[0]);
-    DTUpdates.emplace_back(DominatorTree::UpdateKind::Insert,
-                           SubLoopBlocksLast.back(), AftBlocksFirst[0]);
-    DTU.applyUpdatesPermissive(DTUpdates);
+    DT->changeImmediateDominator(SubLoopBlocksFirst[0], ForeBlocksLast.back());
+    DT->changeImmediateDominator(AftBlocksFirst[0], SubLoopBlocksLast.back());
+    DT->changeImmediateDominator(LoopExit, AftBlocksLast.back());
   }
 
   // Merge adjacent basic blocks, if possible.
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
   SmallPtrSet<BasicBlock *, 16> MergeBlocks;
   MergeBlocks.insert(ForeBlocksLast.begin(), ForeBlocksLast.end());
   MergeBlocks.insert(SubLoopBlocksLast.begin(), SubLoopBlocksLast.end());
   MergeBlocks.insert(AftBlocksLast.begin(), AftBlocksLast.end());
   while (!MergeBlocks.empty()) {
     BasicBlock *BB = *MergeBlocks.begin();
     BranchInst *Term = dyn_cast<BranchInst>(BB->getTerminator());
     if (Term && Term->isUnconditional() && L->contains(Term->getSuccessor(0))) {
LoopUnrollResult llvm::UnrollAndJamLoop(
   // inserted code, doing constant propagation and dead code elimination as we
   // go.
   simplifyLoopAfterUnroll(SubLoop, true, LI, SE, DT, AC);
   simplifyLoopAfterUnroll(L, !CompletelyUnroll && Count > 1, LI, SE, DT, AC);
 
   NumCompletelyUnrolledAndJammed += CompletelyUnroll;
   ++NumUnrolledAndJammed;
 
+  // Update LoopInfo if the loop is completely removed.
+  if (CompletelyUnroll)
+    LI->erase(L);
+
 #ifndef NDEBUG
   // We shouldn't have done anything to break loop simplify form or LCSSA.
-  Loop *OuterL = L->getParentLoop();
-  Loop *OutestLoop = OuterL ? OuterL : (!CompletelyUnroll ? L : SubLoop);
+  Loop *OutestLoop = SubLoop->getParentLoop()
+                         ? SubLoop->getParentLoop()->getParentLoop()
+                               ? SubLoop->getParentLoop()->getParentLoop()
+                               : SubLoop->getParentLoop()
+                         : SubLoop;
+  assert(DT->verify());
+  LI->verify(*DT);
   assert(OutestLoop->isRecursivelyLCSSAForm(*DT, *LI));
   if (!CompletelyUnroll)
     assert(L->isLoopSimplifyForm());
   assert(SubLoop->isLoopSimplifyForm());
-  assert(DT->verify());
+  SE->verify();
 #endif
 
-  // Update LoopInfo if the loop is completely removed.
-  if (CompletelyUnroll)
-    LI->erase(L);
-
   return CompletelyUnroll ? LoopUnrollResult::FullyUnrolled
                           : LoopUnrollResult::PartiallyUnrolled;
 }
 
 static bool getLoadsAndStores(BasicBlockSet &Blocks,
                               SmallVector<Value *, 4> &MemInstr) {
   // Scan the BBs and collect legal loads and stores.
   // Returns false if non-simple loads/stores are found.