[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

2182

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

590

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

650

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

650

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

632

Weird indentation

6946

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;
 
bool CodeGenPrepare::runOnFunction(Function &F) {
   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);
 
       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:
while (MadeChange) {
     if (!LargeOffsetGEPMap.empty())
       MadeChange |= splitLargeGEPOffsets();
     MadeChange |= optimizePhiTypes(F);
 
     if (MadeChange)
       eliminateFallThrough(F);
 
     // 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)
if (InValPhi && InValPhi->getParent() == BB) {
           PN.addIncoming(InVal, Pred);
       }
     }
   }
 
   // The PHIs are now updated, change everything that refers to BB to use
   // DestBB and remove BB.
   BB->replaceAllUsesWith(DestBB);
+  LI->removeBlock(BB);
   BB->eraseFromParent();
   ++NumBlocksElim;
 
   LLVM_DEBUG(dbgs() << "AFTER:\n" << *DestBB << "\n\n\n");
 }
 
 // Computes a map of base pointer relocation instructions to corresponding
 // derived pointer relocation instructions given a vector of all relocate calls

 ///   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:
BFI->setBlockFreq(
         BB,
         (BFI->getBlockFreq(BB) - BFI->getBlockFreq(TailCallBB)).getFrequency());
     ModifiedDT = ModifyDT::ModifyBBDT;
     Changed = true;
     ++NumRetsDup;
   }
 
   // If we eliminated all predecessors of the block, delete the block now.
-  if (Changed && !BB->hasAddressTaken() && pred_empty(BB))
+  if (Changed && !BB->hasAddressTaken() && pred_empty(BB)) {
+    LI->removeBlock(BB);
     BB->eraseFromParent();
+  }
 
   return Changed;
 }
 
 //===----------------------------------------------------------------------===//
 // Memory Optimization
 //===----------------------------------------------------------------------===//
 
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
for (auto &BB : F) {
     LLVM_DEBUG(dbgs() << "Before branch condition splitting\n"; BB.dump());
 
     // Create a new BB.
     auto *TmpBB =
         BasicBlock::Create(BB.getContext(), BB.getName() + ".cond.split",
                            BB.getParent(), BB.getNextNode());
     if (IsHugeFunc)
       FreshBBs.insert(TmpBB);
+    if (Loop *L = LI->getLoopFor(&BB))
+      L->addBasicBlockToLoop(TmpBB, *LI);
 
     // Update original basic block by using the first condition directly by the
     // branch instruction and removing the no longer needed and/or instruction.
     Br1->setCondition(Cond1);
     LogicOp->eraseFromParent();
 
     // Depending on the condition we have to either replace the true or the
     // false successor of the original branch instruction.