Index: llvm/lib/CodeGen/CodeGenPrepare.cpp
===================================================================
--- llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -304,7 +304,7 @@
   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;
@@ -583,6 +583,8 @@
   while (MadeChange) {
     MadeChange = false;
     DT.reset();
+    LI->releaseMemory();
+    LI->analyze(getDT(F));
 
     for (BasicBlock &BB : llvm::make_early_inc_range(F)) {
       if (FuncIterated && !FreshBBs.contains(&BB))
@@ -643,6 +645,12 @@
   NewGEPBases.clear();
   SunkAddrs.clear();
 
+#ifndef NDEBUG
+// From this point on, the LoopInfo might not be valid.
+  if (VerifyLoopInfo)
+    LI->verify(getDT(F));
+#endif
+
   if (!DisableBranchOpts) {
     MadeChange = false;
     // Use a set vector to get deterministic iteration order. The order the
@@ -2135,6 +2143,7 @@
 ///
 /// 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,
@@ -2174,6 +2183,13 @@
   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());
@@ -2348,7 +2364,7 @@
     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:
@@ -2431,6 +2447,8 @@
   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;
@@ -6939,8 +6957,10 @@
   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();
 
@@ -6961,6 +6981,8 @@
         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());
@@ -6972,6 +6994,8 @@
                                         EndBlock->getParent(), EndBlock);
         if (IsHugeFunc)
           FreshBBs.insert(FalseBlock);
+        if (L)
+          L->addBasicBlockToLoop(FalseBlock, *LI);
         FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
         FalseBranch->setDebugLoc(SI->getDebugLoc());
       }
@@ -6990,6 +7014,8 @@
                                     EndBlock->getParent(), EndBlock);
     if (IsHugeFunc)
       FreshBBs.insert(FalseBlock);
+    if (L)
+      L->addBasicBlockToLoop(FalseBlock, *LI);
     auto *FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
     FalseBranch->setDebugLoc(SI->getDebugLoc());
   }