[WIP]Forming branch from select aggressively in loop

[junbuml]

Summary

When a SelectInst is in a critical path of a loop, we aggressively turn the SelectInst into a branch
so that we rely more on the branch predictor.

This change simply check if the SelectInst is in a single latch of a loop in isFormingBranchFromSelectProfitable().
Most of the changes in this patch is to update LoopInfo when there is any change in CFG.

With this change, I observed 5% performance gain in spec2000/vpr in kryo without any regression in other benchmarks in spec2000/spec2006.

I am sending this out early to get early feedback about the basic idea of this change as I do the tuning/testing in parallel. This is also more test cases.

[jmolloy]

Hi Jun,

This feels a bit fishy to me. In particular, you don't ever check if the select is actually on the critical path, unlike your comment. For example:

%1 = load %a
%2 = select %1, %3
store %2

This would be turned into a branch based on your patch, but isn't on the critical path.

[junbuml]

This is a WIP patch and I posted this to get any early feedback about the overall direction and any detail that need to be considered. Basic idea about this was to use an explicit branch when a SelectInst is in frequently executed blocks (e.g., I simply see if SelectInst is in a single latch block) with the hope of HW branch predictor mostly take the right branch during the execution so that we could avoid executing speculated code for most of the loop iterations. I didn't mean to detect any dependence. Sorry for the confusion.

[junbuml]

This was an experimental change to see the impact of aggressive use of branch from SelectInst. Basically, checking if a SelectInst is in a loop latch would be not enough information to determine the profitability of turning a SelectInst into a branch. However, I doubt if the current heuristic used in isFormingBranchFromSelectProfitable(), which simply check if one of select's operands is expensive, is reasonable enough to determine profitability of this transformation.

As a part of improving this transformation, I was trying to see other good place for this in MI level. For me, EarlyIfConv seems to be feasible for this using MachineTraceMetrics. Current EarlyIfConv performs the exactly opposite way of transformation (from branch to select) based on a heuristic which check if extra cycles added in the critical path when using select is more than half of a mispredict penalty. Based on this existing framework in EarlyIfConv, I'm trying to do the select to branch transformation using similar, but opposite way of this heuristic. Can anyone give me any feedback about this approach?

lib/CodeGen/CodeGenPrepare.cpp

class TypePromotionTransaction;
 
   class CodeGenPrepare : public FunctionPass {
     const TargetMachine *TM;
     const TargetSubtargetInfo *SubtargetInfo;
     const TargetLowering *TLI;
     const TargetRegisterInfo *TRI;
     const TargetTransformInfo *TTI;
     const TargetLibraryInfo *TLInfo;
-    const LoopInfo *LI;
+    LoopInfo *LI;
+
     std::unique_ptr<BlockFrequencyInfo> BFI;
     std::unique_ptr<BranchProbabilityInfo> BPI;
 
     /// As we scan instructions optimizing them, this is the next instruction
     /// to optimize. Transforms that can invalidate this should update it.
     BasicBlock::iterator CurInstIterator;
 
     /// Keeps track of non-local addresses that have been sunk into a block.
for (Function::iterator I = std::next(F.begin()), E = F.end(); I != E;) {
     if (Term && !Term->isConditional()) {
       Changed = true;
       DEBUG(dbgs() << "To merge:\n"<< *SinglePred << "\n\n\n");
       // Remember if SinglePred was the entry block of the function.
       // If so, we will need to move BB back to the entry position.
       bool isEntry = SinglePred == &SinglePred->getParent()->getEntryBlock();
       MergeBasicBlockIntoOnlyPred(BB, nullptr);
 
+      if (LI->getLoopFor(SinglePred))
+        LI->removeBlock(SinglePred);
+
       if (isEntry && BB != &BB->getParent()->getEntryBlock())
         BB->moveBefore(&BB->getParent()->getEntryBlock());
 
       // We have erased a block. Update the iterator.
       I = BB->getIterator();
     }
   }
   return Changed;
while (Indirect != End) {
           PHINode::Create(IndPHI->getType(), 2, "merge", &*MergeInsert);
       MergePHI->addIncoming(NewIndPHI, Target);
       MergePHI->addIncoming(DirPHI, DirectSucc);
 
       IndPHI->replaceAllUsesWith(MergePHI);
       IndPHI->eraseFromParent();
     }
 
+    if (Loop * L = LI->getLoopFor(Target)) {
+      L->addBasicBlockToLoop(BodyBlock, *LI);
+      L->addBasicBlockToLoop(DirectSucc, *LI);
+    }
     Changed = true;
   }
 
   return Changed;
 }
 
 /// Eliminate blocks that contain only PHI nodes, debug info directives, and an
 /// unconditional branch. Passes before isel (e.g. LSR/loopsimplify) often split
void CodeGenPrepare::eliminateMostlyEmptyBlock(BasicBlock *BB) {
   // just collapse it.
   if (BasicBlock *SinglePred = DestBB->getSinglePredecessor()) {
     if (SinglePred != DestBB) {
       // Remember if SinglePred was the entry block of the function.  If so, we
       // will need to move BB back to the entry position.
       bool isEntry = SinglePred == &SinglePred->getParent()->getEntryBlock();
       MergeBasicBlockIntoOnlyPred(DestBB, nullptr);
 
+      if (LI->getLoopFor(SinglePred))
+        LI->removeBlock(SinglePred);
+
       if (isEntry && BB != &BB->getParent()->getEntryBlock())
         BB->moveBefore(&BB->getParent()->getEntryBlock());
 
       DEBUG(dbgs() << "AFTER:\n" << *DestBB << "\n\n\n");
       return;
     }
   }
 
if (InValPhi && InValPhi->getParent() == BB) {
           PN->addIncoming(InVal, *PI);
       }
     }
   }
 
   // The PHIs are now updated, change everything that refers to BB to use
   // DestBB and remove BB.
   BB->replaceAllUsesWith(DestBB);
+  if (LI->getLoopFor(BB))
+    LI->removeBlock(BB);
+
   BB->eraseFromParent();
   ++NumBlocksElim;
 
   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

 //  br label %else2
 //
 //else2:                                            ; preds = %else, %cond.load1
 //  %res.phi.else3 = phi <16 x i32> [ %11, %cond.load1 ], [ %res.phi.else, %else ]
 //  %12 = extractelement <16 x i1> %mask, i32 2
 //  %13 = icmp eq i1 %12, true
 //  br i1 %13, label %cond.load4, label %else5
 //
-static void scalarizeMaskedLoad(CallInst *CI) {
+static void scalarizeMaskedLoad(CallInst *CI, LoopInfo &LI) {
   Value *Ptr  = CI->getArgOperand(0);
   Value *Alignment = CI->getArgOperand(1);
   Value *Mask = CI->getArgOperand(2);
   Value *Src0 = CI->getArgOperand(3);
 
   unsigned AlignVal = cast<ConstantInt>(Alignment)->getZExtValue();
   VectorType *VecType = dyn_cast<VectorType>(CI->getType());
   assert(VecType && "Unexpected return type of masked load intrinsic");
if (isa<ConstantVector>(Mask)) {
     CI->replaceAllUsesWith(NewI);
     CI->eraseFromParent();
     return;
   }
 
   PHINode *Phi = nullptr;
   Value *PrevPhi = UndefVal;
 
+  Loop *L = LI.getLoopFor(CI->getParent());
+
   for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
 
     // Fill the "else" block, created in the previous iteration
     //
     //  %res.phi.else3 = phi <16 x i32> [ %11, %cond.load1 ], [ %res.phi.else, %else ]
     //  %mask_1 = extractelement <16 x i1> %mask, i32 Idx
     //  %to_load = icmp eq i1 %mask_1, true
     //  br i1 %to_load, label %cond.load, label %else
for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
     VResult = Builder.CreateInsertElement(VResult, Load, Builder.getInt32(Idx));
 
     // Create "else" block, fill it in the next iteration
     BasicBlock *NewIfBlock =
         CondBlock->splitBasicBlock(InsertPt->getIterator(), "else");
     Builder.SetInsertPoint(InsertPt);
     Instruction *OldBr = IfBlock->getTerminator();
     BranchInst::Create(CondBlock, NewIfBlock, Cmp, OldBr);
+
+    if (L) {
+      L->addBasicBlockToLoop(CondBlock, LI);
+      L->addBasicBlockToLoop(NewIfBlock, LI);
+    }
+
     OldBr->eraseFromParent();
     PrevIfBlock = IfBlock;
     IfBlock = NewIfBlock;
   }
 
   Phi = Builder.CreatePHI(VecType, 2, "res.phi.select");
   Phi->addIncoming(VResult, CondBlock);
   Phi->addIncoming(PrevPhi, PrevIfBlock);

 //   br i1 %7, label %cond.store1, label %else2
 //
 // cond.store1:                                      ; preds = %else
 //   %8 = extractelement <16 x i32> %val, i32 1
 //   %9 = getelementptr i32* %1, i32 1
 //   store i32 %8, i32* %9
 //   br label %else2
 //   . . .
-static void scalarizeMaskedStore(CallInst *CI) {
+static void scalarizeMaskedStore(CallInst *CI, LoopInfo &LI) {
   Value *Src = CI->getArgOperand(0);
   Value *Ptr  = CI->getArgOperand(1);
   Value *Alignment = CI->getArgOperand(2);
   Value *Mask = CI->getArgOperand(3);
 
   unsigned AlignVal = cast<ConstantInt>(Alignment)->getZExtValue();
   VectorType *VecType = dyn_cast<VectorType>(Src->getType());
   assert(VecType && "Unexpected data type in masked store intrinsic");
for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
       Value *Gep =
           Builder.CreateInBoundsGEP(EltTy, FirstEltPtr, Builder.getInt32(Idx));
       Builder.CreateAlignedStore(OneElt, Gep, AlignVal);
     }
     CI->eraseFromParent();
     return;
   }
 
+  Loop *L = LI.getLoopFor(CI->getParent());
   for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
 
     // Fill the "else" block, created in the previous iteration
     //
     //  %mask_1 = extractelement <16 x i1> %mask, i32 Idx
     //  %to_store = icmp eq i1 %mask_1, true
     //  br i1 %to_store, label %cond.store, label %else
     //
for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
 
     // Create "else" block, fill it in the next iteration
     BasicBlock *NewIfBlock =
         CondBlock->splitBasicBlock(InsertPt->getIterator(), "else");
     Builder.SetInsertPoint(InsertPt);
     Instruction *OldBr = IfBlock->getTerminator();
     BranchInst::Create(CondBlock, NewIfBlock, Cmp, OldBr);
     OldBr->eraseFromParent();
+
+    if (L) {
+      L->addBasicBlockToLoop(CondBlock, LI);
+      L->addBasicBlockToLoop(NewIfBlock, LI);
+    }
+
     IfBlock = NewIfBlock;
   }
   CI->eraseFromParent();
 }
 
 // Translate a masked gather intrinsic like
 // <16 x i32 > @llvm.masked.gather.v16i32( <16 x i32*> %Ptrs, i32 4,
 //                               <16 x i1> %Mask, <16 x i32> %Src)

 // cond.load1:
 // % Ptr1 = extractelement <16 x i32*> %Ptrs, i32 1
 // % Load1 = load i32, i32* % Ptr1, align 4
 // % Res1 = insertelement <16 x i32> %res.phi.else, i32 % Load1, i32 1
 // br label %else2
 // . . .
 // % Result = select <16 x i1> %Mask, <16 x i32> %res.phi.select, <16 x i32> %Src
 // ret <16 x i32> %Result
-static void scalarizeMaskedGather(CallInst *CI) {
+static void scalarizeMaskedGather(CallInst *CI, LoopInfo &LI) {
   Value *Ptrs = CI->getArgOperand(0);
   Value *Alignment = CI->getArgOperand(1);
   Value *Mask = CI->getArgOperand(2);
   Value *Src0 = CI->getArgOperand(3);
 
   VectorType *VecType = dyn_cast<VectorType>(CI->getType());
 
   assert(VecType && "Unexpected return type of masked load intrinsic");
if (IsConstMask) {
     CI->replaceAllUsesWith(NewI);
     CI->eraseFromParent();
     return;
   }
 
   PHINode *Phi = nullptr;
   Value *PrevPhi = UndefVal;
 
+  Loop *L = LI.getLoopFor(CI->getParent());
   for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
 
     // Fill the "else" block, created in the previous iteration
     //
     //  %Mask1 = extractelement <16 x i1> %Mask, i32 1
     //  %ToLoad1 = icmp eq i1 %Mask1, true
     //  br i1 %ToLoad1, label %cond.load, label %else
     //
VResult = Builder.CreateInsertElement(VResult, Load, Builder.getInt32(Idx),
                                           "Res" + Twine(Idx));
 
     // Create "else" block, fill it in the next iteration
     BasicBlock *NewIfBlock = CondBlock->splitBasicBlock(InsertPt, "else");
     Builder.SetInsertPoint(InsertPt);
     Instruction *OldBr = IfBlock->getTerminator();
     BranchInst::Create(CondBlock, NewIfBlock, Cmp, OldBr);
     OldBr->eraseFromParent();
+
+    if (L) {
+      L->addBasicBlockToLoop(CondBlock, LI);
+      L->addBasicBlockToLoop(NewIfBlock, LI);
+    }
+
     PrevIfBlock = IfBlock;
     IfBlock = NewIfBlock;
   }
 
   Phi = Builder.CreatePHI(VecType, 2, "res.phi.select");
   Phi->addIncoming(VResult, CondBlock);
   Phi->addIncoming(PrevPhi, PrevIfBlock);
   Value *NewI = Builder.CreateSelect(Mask, Phi, Src0);

 // br i1 % ToStore1, label %cond.store1, label %else2
 //
 // cond.store1:
 // % Elt1 = extractelement <16 x i32> %Src, i32 1
 // % Ptr1 = extractelement <16 x i32*> %Ptrs, i32 1
 // store i32 % Elt1, i32* % Ptr1, align 4
 // br label %else2
 //   . . .
-static void scalarizeMaskedScatter(CallInst *CI) {
+static void scalarizeMaskedScatter(CallInst *CI, LoopInfo &LI) {
   Value *Src = CI->getArgOperand(0);
   Value *Ptrs = CI->getArgOperand(1);
   Value *Alignment = CI->getArgOperand(2);
   Value *Mask = CI->getArgOperand(3);
 
   assert(isa<VectorType>(Src->getType()) &&
          "Unexpected data type in masked scatter intrinsic");
   assert(isa<VectorType>(Ptrs->getType()) &&
for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
                                                    "Elt" + Twine(Idx));
       Value *Ptr = Builder.CreateExtractElement(Ptrs, Builder.getInt32(Idx),
                                                 "Ptr" + Twine(Idx));
       Builder.CreateAlignedStore(OneElt, Ptr, AlignVal);
     }
     CI->eraseFromParent();
     return;
   }
+
+  Loop *L = LI.getLoopFor(CI->getParent());
   for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
     // Fill the "else" block, created in the previous iteration
     //
     //  % Mask1 = extractelement <16 x i1> % Mask, i32 Idx
     //  % ToStore = icmp eq i1 % Mask1, true
     //  br i1 % ToStore, label %cond.store, label %else
     //
     Value *Predicate = Builder.CreateExtractElement(Mask,
for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
     Builder.CreateAlignedStore(OneElt, Ptr, AlignVal);
 
     // Create "else" block, fill it in the next iteration
     BasicBlock *NewIfBlock = CondBlock->splitBasicBlock(InsertPt, "else");
     Builder.SetInsertPoint(InsertPt);
     Instruction *OldBr = IfBlock->getTerminator();
     BranchInst::Create(CondBlock, NewIfBlock, Cmp, OldBr);
     OldBr->eraseFromParent();
+
+    if (L) {
+      L->addBasicBlockToLoop(CondBlock, LI);
+      L->addBasicBlockToLoop(NewIfBlock, LI);
+    }
+
     IfBlock = NewIfBlock;
   }
   CI->eraseFromParent();
 }
 
 /// If counting leading or trailing zeros is an expensive operation and a zero
 /// input is defined, add a check for zero to avoid calling the intrinsic.
 ///

 ///     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,
                                   const TargetLowering *TLI,
                                   const DataLayout *DL,
-                                  bool &ModifiedDT) {
+                                  bool &ModifiedDT,
+                                  LoopInfo &LI) {
   if (!TLI || !DL)
     return false;
 
   // If a zero input is undefined, it doesn't make sense to despeculate that.
   if (match(CountZeros->getOperand(1), m_One()))
     return false;
 
   // If it's cheap to speculate, there's nothing to do.
static bool despeculateCountZeros(IntrinsicInst *CountZeros,
   BasicBlock *CallBlock = StartBlock->splitBasicBlock(CountZeros, "cond.false");
 
   // 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");
 
+  Loop *L = LI.getLoopFor(StartBlock);
+
   // 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);
static bool despeculateCountZeros(IntrinsicInst *CountZeros,
   Value *BitWidth = Builder.getInt(APInt(SizeInBits, SizeInBits));
   PN->addIncoming(BitWidth, StartBlock);
   PN->addIncoming(CountZeros, CallBlock);
 
   // We are explicitly handling the zero case, so we can set the intrinsic's
   // undefined zero argument to 'true'. This will also prevent reprocessing the
   // intrinsic; we only despeculate when a zero input is defined.
   CountZeros->setArgOperand(1, Builder.getTrue());
+
+  if (L) {
+    L->addBasicBlockToLoop(CallBlock, LI);
+    L->addBasicBlockToLoop(EndBlock, LI);
+  }
+
   ModifiedDT = true;
   return true;
 }
 
 bool CodeGenPrepare::optimizeCallInst(CallInst *CI, bool& ModifiedDT) {
   BasicBlock *BB = CI->getParent();
 
   // Lower inline assembly if we can.
case Intrinsic::objectsize: {
         CurInstIterator = BB->begin();
         SunkAddrs.clear();
       }
       return true;
     }
     case Intrinsic::masked_load: {
       // Scalarize unsupported vector masked load
       if (!TTI->isLegalMaskedLoad(CI->getType())) {
-        scalarizeMaskedLoad(CI);
+        scalarizeMaskedLoad(CI, *LI);
         ModifiedDT = true;
         return true;
       }
       return false;
     }
     case Intrinsic::masked_store: {
       if (!TTI->isLegalMaskedStore(CI->getArgOperand(0)->getType())) {
-        scalarizeMaskedStore(CI);
+        scalarizeMaskedStore(CI, *LI);
         ModifiedDT = true;
         return true;
       }
       return false;
     }
     case Intrinsic::masked_gather: {
       if (!TTI->isLegalMaskedGather(CI->getType())) {
-        scalarizeMaskedGather(CI);
+        scalarizeMaskedGather(CI, *LI);
         ModifiedDT = true;
         return true;
       }
       return false;
     }
     case Intrinsic::masked_scatter: {
       if (!TTI->isLegalMaskedScatter(CI->getArgOperand(0)->getType())) {
-        scalarizeMaskedScatter(CI);
+        scalarizeMaskedScatter(CI, *LI);
         ModifiedDT = true;
         return true;
       }
       return false;
     }
     case Intrinsic::aarch64_stlxr:
     case Intrinsic::aarch64_stxr: {
       ZExtInst *ExtVal = dyn_cast<ZExtInst>(CI->getArgOperand(0));
if (II) {
     case Intrinsic::invariant_group_barrier:
       II->replaceAllUsesWith(II->getArgOperand(0));
       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);
+      return despeculateCountZeros(II, TLI, DL, ModifiedDT, *LI);
     }
 
     if (TLI) {
       SmallVector<Value*, 2> PtrOps;
       Type *AccessTy;
       if (TLI->getAddrModeArguments(II, PtrOps, AccessTy))
         while (!PtrOps.empty()) {
           Value *PtrVal = PtrOps.pop_back_val();
for (unsigned i = 0, e = TailCalls.size(); i != e; ++i) {
 
     // Duplicate the return into CallBB.
     (void)FoldReturnIntoUncondBranch(RetI, BB, CallBB);
     ModifiedDT = Changed = true;
     ++NumRetsDup;
   }
 
   // If we eliminated all predecessors of the block, delete the block now.
-  if (Changed && !BB->hasAddressTaken() && pred_begin(BB) == pred_end(BB))
+  if (Changed && !BB->hasAddressTaken() && pred_begin(BB) == pred_end(BB)) {
+    if (LI->getLoopFor(BB))
+      LI->removeBlock(BB);
     BB->eraseFromParent();
+  }
 
   return Changed;
 }
 
 //===----------------------------------------------------------------------===//
 // Memory Optimization
 //===----------------------------------------------------------------------===//
 
static bool sinkSelectOperand(const TargetTransformInfo *TTI, Value *V) {
   // effects; therefore, it's safe to sink and possibly *not* execute.
   return I && I->hasOneUse() && isSafeToSpeculativelyExecute(I) &&
          TTI->getUserCost(I) >= TargetTransformInfo::TCC_Expensive;
 }
 
 /// Returns true if a SelectInst should be turned into an explicit branch.
 static bool isFormingBranchFromSelectProfitable(const TargetTransformInfo *TTI,
                                                 const TargetLowering *TLI,
-                                                SelectInst *SI) {
+                                                SelectInst *SI,
+                                                const LoopInfo *LI) {
   // If even a predictable select is cheap, then a branch can't be cheaper.
   if (!TLI->isPredictableSelectExpensive())
     return false;
 
   // FIXME: This should use the same heuristics as IfConversion to determine
   // whether a select is better represented as a branch.
 
   // If metadata tells us that the select condition is obviously predictable,
static bool isFormingBranchFromSelectProfitable(const TargetTransformInfo *TTI,
   CmpInst *Cmp = dyn_cast<CmpInst>(SI->getCondition());
 
   // If a branch is predictable, an out-of-order CPU can avoid blocking on its
   // comparison condition. If the compare has more than one use, there's
   // probably another cmov or setcc around, so it's not worth emitting a branch.
   if (!Cmp || !Cmp->hasOneUse())
     return false;
 
+  // If the select is in a critical path of a loop, we aggressively turn it into
+  // a branch so that we rely more on the branch predictor.
+  if (Loop *L = LI->getLoopFor(SI->getParent())) {
+    if (L->getLoopLatch() == SI->getParent())
+      return true;
+  }
+
   // If either operand of the select is expensive and only needed on one side
   // of the select, we should form a branch.
   if (sinkSelectOperand(TTI, SI->getTrueValue()) ||
       sinkSelectOperand(TTI, SI->getFalseValue()))
     return true;
 
   return false;
 }
bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
   if (VectorCond)
     SelectKind = TargetLowering::VectorMaskSelect;
   else if (SI->getType()->isVectorTy())
     SelectKind = TargetLowering::ScalarCondVectorVal;
   else
     SelectKind = TargetLowering::ScalarValSelect;
 
   if (TLI->isSelectSupported(SelectKind) &&
-      !isFormingBranchFromSelectProfitable(TTI, TLI, SI))
+      !isFormingBranchFromSelectProfitable(TTI, TLI, SI, LI))
     return false;
 
   ModifiedDT = true;
 
   // Transform a sequence like this:
   //    start:
   //       %cmp = cmp uge i32 %a, %b
   //       %sel = select i1 %cmp, i32 %c, i32 %d
bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
   // 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");
+  Loop *L = LI->getLoopFor(StartBlock);
+
+  if (L)
+    L->addBasicBlockToLoop(EndBlock, *LI);
 
   // 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);
+        if (L)
+          L->addBasicBlockToLoop(TrueBlock, *LI);
+
         TrueBranch = BranchInst::Create(EndBlock, TrueBlock);
       }
       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 (L)
+          L->addBasicBlockToLoop(FalseBlock, *LI);
+
         FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
       }
       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 (L)
+      L->addBasicBlockToLoop(FalseBlock, *LI);
+
     BranchInst::Create(EndBlock, FalseBlock);
   }
 
   // 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
   // view of the new PHI.
for (auto &BB : F) {
     }
 
     // Note: No point in getting fancy here, since the DT info is never
     // available to CodeGenPrepare.
     ModifiedDT = true;
 
     MadeChange = true;
 
+    if (Loop *L = LI->getLoopFor(TmpBB))
+      L->addBasicBlockToLoop(TmpBB, *LI);
+
     DEBUG(dbgs() << "After branch condition splitting\n"; BB.dump();
           TmpBB->dump());
   }
   return MadeChange;
 }

test/CodeGen/AArch64/aarch64-aggressive-select-to-branch.ll

+; RUN: opt  -codegenprepare -S -mcpu=kryo < %s  | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-linaro-linux-gnueabi"
+
+; This test if a SelectInst in a loop latch is turned into an explicit branch
+; in codegenprepare pass. This test case was an input IR of codegenprepare
+; when compiling the C code below in -O3.
+;
+;struct s1 {int idx; int cost;};
+;void foo(struct s1 **G, int n) {
+; int j, i;
+; while (j < n) {
+;   if (G[j+1]->cost < G[j]->cost)
+;     j++;
+;   if (G[j]->cost > G[i]->cost)
+;     break;
+;  }
+;}
+;
+
+%struct.s1 = type { i32, i32 }
+;
+; CHECK-LABEL: @foo
+; CHECK-LABEL: while.body:
+; CHECK-NOT: select i1
+; CHECK: [[CMP:%.*]] = icmp slt i32 %l2, %l4
+; CHECK: br i1 [[CMP]], label %select.end, label %select.false
+;
+define void @foo(i32 %n, %struct.s1** nocapture %G) local_unnamed_addr #0 {
+entry:
+  ;%0 = load i32, i32* @n, align 4
+  ;%1 = load %struct.s1**, %struct.s1*** @G, align 8
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.body, %entry
+  %j.0 = phi i32 [ undef, %entry ], [ %add.j.0, %while.body ]
+  %cmp = icmp slt i32 %j.0, %n
+  br i1 %cmp, label %while.body, label %while.end
+
+while.body:                                       ; preds = %while.cond
+  %add = add nsw i32 %j.0, 1
+  %idxprom = sext i32 %add to i64
+  %arrayidx = getelementptr inbounds %struct.s1*, %struct.s1** %G, i64 %idxprom
+  %l = load %struct.s1*, %struct.s1** %arrayidx, align 8
+  %cost = getelementptr inbounds %struct.s1, %struct.s1* %l, i64 0, i32 0
+  %l2 = load i32, i32* %cost, align 4
+  %idxprom1 = sext i32 %j.0 to i64
+  %arrayidx2 = getelementptr inbounds %struct.s1*, %struct.s1** %G, i64 %idxprom1
+  %l3 = load %struct.s1*, %struct.s1** %arrayidx2, align 8
+  %cost3 = getelementptr inbounds %struct.s1, %struct.s1* %l3, i64 0, i32 0
+  %l4 = load i32, i32* %cost3, align 4
+  %cmp4 = icmp slt i32 %l2, %l4
+  %add.j.0 = select i1 %cmp4, i32 %add, i32 %j.0
+  %idxprom5 = sext i32 %add.j.0 to i64
+  %arrayidx6 = getelementptr inbounds %struct.s1*, %struct.s1** %G, i64 %idxprom5
+  %l5 = load %struct.s1*, %struct.s1** %arrayidx6, align 8
+  %cost7 = getelementptr inbounds %struct.s1, %struct.s1* %l5, i64 0, i32 0
+  %l6 = load i32, i32* %cost7, align 4
+  %l7 = load %struct.s1*, %struct.s1** %G, align 8
+  %cost10 = getelementptr inbounds %struct.s1, %struct.s1* %l7, i64 0, i32 0
+  %l8 = load i32, i32* %cost10, align 4
+  %cmp11 = icmp sgt i32 %l6, %l8
+  br i1 %cmp11, label %while.end, label %while.cond
+
+while.end:                                        ; preds = %while.body, %while.cond
+  ret void
+}
+