[CodeGenPrepare] Fix for using outdated/corrupt LoopInfo

[chill]

Summary

Some transformation in CodeGenPrepare pass may create and/or delete
basic block, but they don't update the LoopInfo, so the LoopInfo may
end up containing dangling pointers and sometimes reused basic blocks,
which leads to "interesting" non-deterministic behaviour.

These transformations do not seem to alter the loop structure of the
function, and updating the loop info is quite straighforward.

[efriedma]

We should probably have code somewhere that calls verify() on the LoopInfo, so we're confident these updates are correct, and to catch future issues. (If we do it once at the end of CodeGenPrepare, it should be cheap enough to just turn it on with +Asserts? Or we could stick it under EXPENSIVE_CHECKS, I guess.)

llvm/lib/CodeGen/CodeGenPrepare.cpp

2230

Wrong indentatation.

[chill]

There are multiple issues with this patch, I may go back to rebuilding the LoopInfo, wel'll see ...

[efriedma]

llvm/lib/CodeGen/CodeGenPrepare.cpp

594

Doing this for every basic block is going to get way too expensive, even for something labeled "expensive".

[chill]

Rebase/refresh, no changes.

[efriedma]

llvm/lib/CodeGen/CodeGenPrepare.cpp

659

Given the way the code is structured, isn't this always running just after you've just called releaseMemory()/analyze() on LI?

[chill]

llvm/lib/CodeGen/CodeGenPrepare.cpp

659

D'oh, indeed, after we exit the loop because of MadeChange being false it means all the code after LI->analyze didn't do anything, including changing the loop structure.

[chill]

Fix some more issues with maintaining DT/LI up to date:

If the DT was modified in a normal (not huge) function, the over-the-blocks loop
will exit with an outdated DT, which is then used by mergeSExts.
Thus the reset of the DT was moved immediately after returning from optimizeBlock.
This ensures that mergeSExts will work with an up to date DT and allows
to avoid unnecessary rebuilds of the DT by moving the initial reset outside the outer loop.
The latter necessitates updating the LI and DT in splitLargeGEPOffsets and
eliminateFallThrough (only if DT is already present).

[efriedma]

LGTM

llvm/lib/CodeGen/CodeGenPrepare.cpp

636

Weird indentation

6998

I think we have a SplitBasicBlock utility that handles DT/LI automatically, but I don't want to make any complex refactorings block this.

[nikic]

Doesn't this change essentially break the current model where DT is supposed to only be calculated lazily, if it is needed? Now we'll end up always calculating it on each iteration.

I think if you do this, you also need to start keeping DT up to date everywhere, instead of using the current crude invalidation approach.

[chill]

In D150384#4432424, @nikic wrote:

Doesn't this change essentially break the current model where DT is supposed to only be calculated lazily, if it is needed? Now we'll end up always calculating it on each iteration.

I think if you do this, you also need to start keeping DT up to date everywhere, instead of using the current crude invalidation approach.

In D150384#4432424, @nikic wrote:

Doesn't this change essentially break the current model where DT is supposed to only be calculated lazily, if it is needed? Now we'll end up always calculating it on each iteration.

I think if you do this, you also need to start keeping DT up to date everywhere, instead of using the current crude invalidation approach.

Now that we don't have (?) correctness/corruption issues, I'll see what I can do about updating the DT. (From a quick glance it does not seem like a terribly huge amount of work).

llvm/lib/CodeGen/CodeGenPrepare.cpp

 class CodeGenPrepare : public FunctionPass {
   const TargetMachine *TM = nullptr;
   const TargetSubtargetInfo *SubtargetInfo = nullptr;
   const TargetLowering *TLI = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
   const TargetTransformInfo *TTI = nullptr;
   const BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
   const TargetLibraryInfo *TLInfo = nullptr;
-  const LoopInfo *LI = nullptr;
+  LoopInfo *LI = nullptr;
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   ProfileSummaryInfo *PSI = nullptr;
 
   /// As we scan instructions optimizing them, this is the next instruction
   /// to optimize. Transforms that can invalidate this should update it.
   BasicBlock::iterator CurInstIterator;
 
private:
   DominatorTree &getDT(Function &F) {
     if (!DT)
       DT = std::make_unique<DominatorTree>(F);
     return *DT;
   }
 
   void removeAllAssertingVHReferences(Value *V);
   bool eliminateAssumptions(Function &F);
-  bool eliminateFallThrough(Function &F);
+  bool eliminateFallThrough(Function &F, DominatorTree *DT = nullptr);
   bool eliminateMostlyEmptyBlocks(Function &F);
   BasicBlock *findDestBlockOfMergeableEmptyBlock(BasicBlock *BB);
   bool canMergeBlocks(const BasicBlock *BB, const BasicBlock *DestBB) const;
   void eliminateMostlyEmptyBlock(BasicBlock *BB);
   bool isMergingEmptyBlockProfitable(BasicBlock *BB, BasicBlock *DestBB,
                                      bool isPreheader);
   bool makeBitReverse(Instruction &I);
   bool optimizeBlock(BasicBlock &BB, ModifyDT &ModifiedDT);
bool CodeGenPrepare::runOnFunction(Function &F) {
   // to help generate sane code for PHIs involving such edges.
   EverMadeChange |=
       SplitIndirectBrCriticalEdges(F, /*IgnoreBlocksWithoutPHI=*/true);
 
   // If we are optimzing huge function, we need to consider the build time.
   // Because the basic algorithm's complex is near O(N!).
   IsHugeFunc = F.size() > HugeFuncThresholdInCGPP;
 
+  // Transformations above may invalidate dominator tree and/or loop info.
+  DT.reset();
+  LI->releaseMemory();
+  LI->analyze(getDT(F));
+
   bool MadeChange = true;
   bool FuncIterated = false;
   while (MadeChange) {
     MadeChange = false;
-    DT.reset();
 
     for (BasicBlock &BB : llvm::make_early_inc_range(F)) {
       if (FuncIterated && !FreshBBs.contains(&BB))
         continue;
 
       ModifyDT ModifiedDTOnIteration = ModifyDT::NotModifyDT;
       bool Changed = optimizeBlock(BB, ModifiedDTOnIteration);
 
+      if (ModifiedDTOnIteration == ModifyDT::ModifyBBDT)
+        DT.reset();
+
       MadeChange |= Changed;
       if (IsHugeFunc) {
         // If the BB is updated, it may still has chance to be optimized.
         // This usually happen at sink optimization.
         // For example:
         //
         // bb0：
         // %and = and i32 %a, 4
         // %cmp = icmp eq i32 %and, 0
         //
         // If the %cmp sink to other BB, the %and will has chance to sink.
         if (Changed)
           FreshBBs.insert(&BB);
         else if (FuncIterated)
           FreshBBs.erase(&BB);
-
-        if (ModifiedDTOnIteration == ModifyDT::ModifyBBDT)
-          DT.reset();
       } else {
         // For small/normal functions, we restart BB iteration if the dominator
         // tree of the Function was changed.
         if (ModifiedDTOnIteration != ModifyDT::NotModifyDT)
           break;
       }
     }
     // We have iterated all the BB in the (only work for huge) function.
     FuncIterated = IsHugeFunc;
 
     if (EnableTypePromotionMerge && !ValToSExtendedUses.empty())
       MadeChange |= mergeSExts(F);
     if (!LargeOffsetGEPMap.empty())
       MadeChange |= splitLargeGEPOffsets();
     MadeChange |= optimizePhiTypes(F);
 
     if (MadeChange)
-      eliminateFallThrough(F);
+      eliminateFallThrough(F, DT.get());
+
+#ifndef NDEBUG
+    if (MadeChange && VerifyLoopInfo)
+      LI->verify(getDT(F));
+#endif
 
     // Really free removed instructions during promotion.
     for (Instruction *I : RemovedInsts)
       I->deleteValue();
 
     EverMadeChange |= MadeChange;
     SeenChainsForSExt.clear();
     ValToSExtendedUses.clear();
     RemovedInsts.clear();
     LargeOffsetGEPMap.clear();
     LargeOffsetGEPID.clear();
   }
 
   NewGEPBases.clear();
   SunkAddrs.clear();
 
   if (!DisableBranchOpts) {
     MadeChange = false;
     // Use a set vector to get deterministic iteration order. The order the
     // blocks are removed may affect whether or not PHI nodes in successors
     // are removed.
     SmallSetVector<BasicBlock *, 8> WorkList;
     for (BasicBlock &BB : F) {
       SmallVector<BasicBlock *, 2> Successors(successors(&BB));
       MadeChange |= ConstantFoldTerminator(&BB, true);
       if (!MadeChange)
         continue;
 
       for (BasicBlock *Succ : Successors)
void LLVM_ATTRIBUTE_UNUSED CodeGenPrepare::verifyBFIUpdates(Function &F) {
   BranchProbabilityInfo NewBPI(F, NewLI, TLInfo);
   BlockFrequencyInfo NewBFI(F, NewBPI, NewLI);
   NewBFI.verifyMatch(*BFI);
 }
 
 /// Merge basic blocks which are connected by a single edge, where one of the
 /// basic blocks has a single successor pointing to the other basic block,
 /// which has a single predecessor.
-bool CodeGenPrepare::eliminateFallThrough(Function &F) {
+bool CodeGenPrepare::eliminateFallThrough(Function &F, DominatorTree *DT) {
   bool Changed = false;
   // Scan all of the blocks in the function, except for the entry block.
   // Use a temporary array to avoid iterator being invalidated when
   // deleting blocks.
   SmallVector<WeakTrackingVH, 16> Blocks;
   for (auto &Block : llvm::drop_begin(F))
     Blocks.push_back(&Block);
 
   SmallSet<WeakTrackingVH, 16> Preds;
   for (auto &Block : Blocks) {
     auto *BB = cast_or_null<BasicBlock>(Block);
     if (!BB)
       continue;
     // If the destination block has a single pred, then this is a trivial
     // edge, just collapse it.
     BasicBlock *SinglePred = BB->getSinglePredecessor();
 
     // Don't merge if BB's address is taken.
     if (!SinglePred || SinglePred == BB || BB->hasAddressTaken())
       continue;
 
+    // Make an effort to skip unreachable blocks.
+    if (DT && !DT->isReachableFromEntry(BB))
+      continue;
+
     BranchInst *Term = dyn_cast<BranchInst>(SinglePred->getTerminator());
     if (Term && !Term->isConditional()) {
       Changed = true;
       LLVM_DEBUG(dbgs() << "To merge:\n" << *BB << "\n\n\n");
 
       // Merge BB into SinglePred and delete it.
-      MergeBlockIntoPredecessor(BB);
+      MergeBlockIntoPredecessor(BB, /* DTU */ nullptr, LI, /* MSSAU */ nullptr,
+                                /* MemDep */ nullptr,
+                                /* PredecessorWithTwoSuccessors */ false, DT);
       Preds.insert(SinglePred);
 
       if (IsHugeFunc) {
         // Update FreshBBs to optimize the merged BB.
         FreshBBs.insert(SinglePred);
         FreshBBs.erase(BB);
       }
     }

 ///   cond.false:
 ///     %z = call i64 @llvm.cttz.i64(i64 %A, i1 true)
 ///     br label %cond.end
 ///   cond.end:
 ///     %ctz = phi i64 [ 64, %entry ], [ %z, %cond.false ]
 ///
 /// If the transform is performed, return true and set ModifiedDT to true.
 static bool despeculateCountZeros(IntrinsicInst *CountZeros,
+                                  LoopInfo &LI,
                                   const TargetLowering *TLI,
                                   const DataLayout *DL, ModifyDT &ModifiedDT,
                                   SmallSet<BasicBlock *, 32> &FreshBBs,
                                   bool IsHugeFunc) {
   // If a zero input is undefined, it doesn't make sense to despeculate that.
   if (match(CountZeros->getOperand(1), m_One()))
     return false;
 
static bool despeculateCountZeros(IntrinsicInst *CountZeros,
   // Create another block after the count zero intrinsic. A PHI will be added
   // in this block to select the result of the intrinsic or the bit-width
   // constant if the input to the intrinsic is zero.
   BasicBlock::iterator SplitPt = ++(BasicBlock::iterator(CountZeros));
   BasicBlock *EndBlock = CallBlock->splitBasicBlock(SplitPt, "cond.end");
   if (IsHugeFunc)
     FreshBBs.insert(EndBlock);
 
+  // Update the LoopInfo. The new blocks are in the same loop as the start
+  // block.
+  if (Loop *L = LI.getLoopFor(StartBlock)) {
+    L->addBasicBlockToLoop(CallBlock, LI);
+    L->addBasicBlockToLoop(EndBlock, LI);
+  }
+
   // Set up a builder to create a compare, conditional branch, and PHI.
   IRBuilder<> Builder(CountZeros->getContext());
   Builder.SetInsertPoint(StartBlock->getTerminator());
   Builder.SetCurrentDebugLocation(CountZeros->getDebugLoc());
 
   // Replace the unconditional branch that was created by the first split with
   // a compare against zero and a conditional branch.
   Value *Zero = Constant::getNullValue(Ty);
case Intrinsic::strip_invariant_group: {
 
       replaceAllUsesWith(II, ArgVal, FreshBBs, IsHugeFunc);
       II->eraseFromParent();
       return true;
     }
     case Intrinsic::cttz:
     case Intrinsic::ctlz:
       // If counting zeros is expensive, try to avoid it.
-      return despeculateCountZeros(II, TLI, DL, ModifiedDT, FreshBBs,
+      return despeculateCountZeros(II, *LI, TLI, DL, ModifiedDT, FreshBBs,
                                    IsHugeFunc);
     case Intrinsic::fshl:
     case Intrinsic::fshr:
       return optimizeFunnelShift(II);
     case Intrinsic::dbg_assign:
     case Intrinsic::dbg_value:
       return fixupDbgValue(II);
     case Intrinsic::masked_gather:
bool CodeGenPrepare::dupRetToEnableTailCallOpts(BasicBlock *BB,
                                                 ModifyDT &ModifiedDT) {
   if (!BB->getTerminator())
     return false;
 
   ReturnInst *RetI = dyn_cast<ReturnInst>(BB->getTerminator());
   if (!RetI)
     return false;
 
+  assert(LI->getLoopFor(BB) == nullptr && "A return block cannot be in a loop");
+
   PHINode *PN = nullptr;
   ExtractValueInst *EVI = nullptr;
   BitCastInst *BCI = nullptr;
   Value *V = RetI->getReturnValue();
   if (V) {
     BCI = dyn_cast<BitCastInst>(V);
     if (BCI)
       V = BCI->getOperand(0);
while (LargeOffsetGEP != LargeOffsetGEPs.end()) {
         if (auto *BaseI = dyn_cast<Instruction>(OldBase)) {
           // If the base of the struct is an instruction, the new base will be
           // inserted close to it.
           NewBaseInsertBB = BaseI->getParent();
           if (isa<PHINode>(BaseI))
             NewBaseInsertPt = NewBaseInsertBB->getFirstInsertionPt();
           else if (InvokeInst *Invoke = dyn_cast<InvokeInst>(BaseI)) {
             NewBaseInsertBB =
-                SplitEdge(NewBaseInsertBB, Invoke->getNormalDest());
+                SplitEdge(NewBaseInsertBB, Invoke->getNormalDest(), DT.get(), LI);
             NewBaseInsertPt = NewBaseInsertBB->getFirstInsertionPt();
           } else
             NewBaseInsertPt = std::next(BaseI->getIterator());
         } else {
           // If the current base is an argument or global value, the new base
           // will be inserted to the entry block.
           NewBaseInsertBB = &BaseGEP->getFunction()->getEntryBlock();
           NewBaseInsertPt = NewBaseInsertBB->getFirstInsertionPt();
bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
   // If the true or false blocks do not contain a sunken instruction, that
   // block and its branch may be optimized away. In that case, one side of the
   // first branch will point directly to select.end, and the corresponding PHI
   // predecessor block will be the start block.
 
   // First, we split the block containing the select into 2 blocks.
   BasicBlock *StartBlock = SI->getParent();
   BasicBlock::iterator SplitPt = ++(BasicBlock::iterator(LastSI));
   BasicBlock *EndBlock = StartBlock->splitBasicBlock(SplitPt, "select.end");
   if (IsHugeFunc)
     FreshBBs.insert(EndBlock);
+  Loop *L = LI->getLoopFor(StartBlock);
+  if (L)
+    L->addBasicBlockToLoop(EndBlock, *LI);
   BFI->setBlockFreq(EndBlock, BFI->getBlockFreq(StartBlock).getFrequency());
-
   // Delete the unconditional branch that was just created by the split.
   StartBlock->getTerminator()->eraseFromParent();
 
   // These are the new basic blocks for the conditional branch.
   // At least one will become an actual new basic block.
   BasicBlock *TrueBlock = nullptr;
   BasicBlock *FalseBlock = nullptr;
   BranchInst *TrueBranch = nullptr;
   BranchInst *FalseBranch = nullptr;
 
   // Sink expensive instructions into the conditional blocks to avoid executing
   // them speculatively.
   for (SelectInst *SI : ASI) {
     if (sinkSelectOperand(TTI, SI->getTrueValue())) {
       if (TrueBlock == nullptr) {
         TrueBlock = BasicBlock::Create(SI->getContext(), "select.true.sink",
                                        EndBlock->getParent(), EndBlock);
         TrueBranch = BranchInst::Create(EndBlock, TrueBlock);
         if (IsHugeFunc)
           FreshBBs.insert(TrueBlock);
+        if (L)
+          L->addBasicBlockToLoop(TrueBlock, *LI);
         TrueBranch->setDebugLoc(SI->getDebugLoc());
       }
       auto *TrueInst = cast<Instruction>(SI->getTrueValue());
       TrueInst->moveBefore(TrueBranch);
     }
     if (sinkSelectOperand(TTI, SI->getFalseValue())) {
       if (FalseBlock == nullptr) {
         FalseBlock = BasicBlock::Create(SI->getContext(), "select.false.sink",
                                         EndBlock->getParent(), EndBlock);
         if (IsHugeFunc)
           FreshBBs.insert(FalseBlock);
+        if (L)
+          L->addBasicBlockToLoop(FalseBlock, *LI);
         FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
         FalseBranch->setDebugLoc(SI->getDebugLoc());
       }
       auto *FalseInst = cast<Instruction>(SI->getFalseValue());
       FalseInst->moveBefore(FalseBranch);
     }
   }
 
   // If there was nothing to sink, then arbitrarily choose the 'false' side
   // for a new input value to the PHI.
   if (TrueBlock == FalseBlock) {
     assert(TrueBlock == nullptr &&
            "Unexpected basic block transform while optimizing select");
 
     FalseBlock = BasicBlock::Create(SI->getContext(), "select.false",
                                     EndBlock->getParent(), EndBlock);
     if (IsHugeFunc)
       FreshBBs.insert(FalseBlock);
+    if (L)
+      L->addBasicBlockToLoop(FalseBlock, *LI);
     auto *FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
     FalseBranch->setDebugLoc(SI->getDebugLoc());
   }
 
   // Insert the real conditional branch based on the original condition.
   // If we did not create a new block for one of the 'true' or 'false' paths
   // of the condition, it means that side of the branch goes to the end block
   // directly and the path originates from the start block from the point of