[CodeGenPrepare] Refactor optimizeSelectInst (NFC)

[chill]

Summary

Refactor to use BasicBlockUtils functions and make life easier for
a subsequent patch for updating the dominator tree.

[chill]

Ping?

[dmgreen]

This looks good as far as I can tell. The formatting needs updating a little in places.

llvm/lib/CodeGen/CodeGenPrepare.cpp

7004

Formatting.

7030

Formatting.

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

 /// Updates DTU and LI if given.
 Instruction *SplitBlockAndInsertIfThen(Value *Cond, Instruction *SplitBefore,
                                        bool Unreachable,
                                        MDNode *BranchWeights = nullptr,
                                        DomTreeUpdater *DTU = nullptr,
                                        LoopInfo *LI = nullptr,
                                        BasicBlock *ThenBlock = nullptr);
 
+/// Similar to SplitBlockAndInsertIfThen, but the inserted block is on the false
+/// path of the branch.
+Instruction *SplitBlockAndInsertIfElse(Value *Cond, Instruction *SplitBefore,
+                                       bool Unreachable,
+                                       MDNode *BranchWeights = nullptr,
+                                       DomTreeUpdater *DTU = nullptr,
+                                       LoopInfo *LI = nullptr,
+                                       BasicBlock *ElseBlock = nullptr);
+
 /// SplitBlockAndInsertIfThenElse is similar to SplitBlockAndInsertIfThen,
 /// but also creates the ElseBlock.
 /// Before:
 ///   Head
 ///   SplitBefore
 ///   Tail
 /// After:
 ///   Head

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
   // %cmp should be frozen, otherwise it may introduce undefined behavior.
   // In addition, we may sink instructions that produce %c or %d from
   // the entry block into the destination(s) of the new branch.
   // 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.
+  // Collect values that go on the true side and the values that go on the false
+  // side.
+  SmallVector<Instruction *> TrueInstrs, FalseInstrs;
+  for (SelectInst *SI : ASI) {
+    if (Value *V = SI->getTrueValue(); sinkSelectOperand(TTI, V))
+      TrueInstrs.push_back(cast<Instruction>(V));
+    if (Value *V = SI->getFalseValue(); sinkSelectOperand(TTI, V))
+      FalseInstrs.push_back(cast<Instruction>(V));
+  }
+
+  // Split the select block, according to how many (if any) values go on each
+  // side.
   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.
+  IRBuilder<> IB(SI);
+  auto *CondFr = IB.CreateFreeze(SI->getCondition(), SI->getName() + ".frozen");
+
   BasicBlock *TrueBlock = nullptr;
   BasicBlock *FalseBlock = nullptr;
+  BasicBlock *EndBlock = nullptr;
   BranchInst *TrueBranch = nullptr;
   BranchInst *FalseBranch = nullptr;
+  if (TrueInstrs.size() == 0) {
+    FalseBranch = cast<BranchInst>(SplitBlockAndInsertIfElse(
+        CondFr, &*SplitPt, false, nullptr, nullptr, LI));
+    FalseBlock = FalseBranch->getParent();
+    EndBlock = cast<BasicBlock>(FalseBranch->getOperand(0));
+  } else if (FalseInstrs.size() == 0) {
+    TrueBranch = cast<BranchInst>(SplitBlockAndInsertIfThen(
+        CondFr, &*SplitPt, false, nullptr, nullptr, LI));
+    TrueBlock = TrueBranch->getParent();
+    EndBlock = cast<BasicBlock>(TrueBranch->getOperand(0));
+  } else {
+    Instruction *ThenTerm = nullptr;
+    Instruction *ElseTerm = nullptr;
+    SplitBlockAndInsertIfThenElse(CondFr, &*SplitPt, &ThenTerm, &ElseTerm,
+                                  nullptr, nullptr, LI);
+    TrueBranch = cast<BranchInst>(ThenTerm);
+    FalseBranch = cast<BranchInst>(ElseTerm);
+    TrueBlock = TrueBranch->getParent();
+    FalseBlock = FalseBranch->getParent();
+    EndBlock = cast<BasicBlock>(TrueBranch->getOperand(0));
+  }
+
+  EndBlock->setName("select.end");
+  if (TrueBlock)
+    TrueBlock->setName("select.true.sink");
+  if (FalseBlock)
+    FalseBlock->setName(FalseInstrs.size() == 0 ? "select.false"
+                                                : "select.false.sink");
 
-  // 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 (IsHugeFunc) {
+    if (TrueBlock)
+    if (FalseBlock)
+    FreshBBs.insert(EndBlock);
   }
 
-  // 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");
+  BFI->setBlockFreq(EndBlock, BFI->getBlockFreq(StartBlock).getFrequency());
 
-    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());
-  }
+  static const unsigned MD[] = {
+      LLVMContext::MD_prof, LLVMContext::MD_unpredictable,
+      LLVMContext::MD_make_implicit, LLVMContext::MD_dbg};
+  StartBlock->getTerminator()->copyMetadata(*SI, MD);
+
+  // Sink expensive instructions into the conditional blocks to avoid executing
+  // them speculatively.
+  for (Instruction *I : TrueInstrs)
+    I->moveBefore(TrueBranch);
+  for (Instruction *I : FalseInstrs)
+    I->moveBefore(FalseBranch);
 
-  // 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.
-  BasicBlock *TT, *FT;
-  if (TrueBlock == nullptr) {
-    TT = EndBlock;
-    FT = FalseBlock;
+  if (TrueBlock == nullptr)
     TrueBlock = StartBlock;
-  } else if (FalseBlock == nullptr) {
-    TT = TrueBlock;
-    FT = EndBlock;
+  else if (FalseBlock == nullptr)
     FalseBlock = StartBlock;
-  } else {
-    TT = TrueBlock;
-    FT = FalseBlock;
-  }
-  IRBuilder<> IB(SI);
-  auto *CondFr = IB.CreateFreeze(SI->getCondition(), SI->getName() + ".frozen");
-  IB.CreateCondBr(CondFr, TT, FT, SI);
 
   SmallPtrSet<const Instruction *, 2> INS;
   INS.insert(ASI.begin(), ASI.end());
   // Use reverse iterator because later select may use the value of the
   // earlier select, and we need to propagate value through earlier select
   // to get the PHI operand.
   for (SelectInst *SI : llvm::reverse(ASI)) {
     // The select itself is replaced with a PHI Node.

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

Instruction *llvm::SplitBlockAndInsertIfThen(Value *Cond,
                                              BasicBlock *ThenBlock) {
   SplitBlockAndInsertIfThenElse(
       Cond, SplitBefore, &ThenBlock, /* ElseBlock */ nullptr,
       /* UnreachableThen */ Unreachable,
       /* UnreachableElse */ false, BranchWeights, DTU, LI);
   return ThenBlock->getTerminator();
 }
 
+Instruction *llvm::SplitBlockAndInsertIfElse(Value *Cond,
+                                             Instruction *SplitBefore,
+                                             bool Unreachable,
+                                             MDNode *BranchWeights,
+                                             DomTreeUpdater *DTU, LoopInfo *LI,
+                                             BasicBlock *ElseBlock) {
+  SplitBlockAndInsertIfThenElse(
+      Cond, SplitBefore, /* ThenBlock */ nullptr, &ElseBlock,
+      /* UnreachableThen */ false,
+      /* UnreachableElse */ Unreachable, BranchWeights, DTU, LI);
+  return ElseBlock->getTerminator();
+}
+
 void llvm::SplitBlockAndInsertIfThenElse(Value *Cond, Instruction *SplitBefore,
                                          Instruction **ThenTerm,
                                          Instruction **ElseTerm,
                                          MDNode *BranchWeights,
                                          DomTreeUpdater *DTU, LoopInfo *LI) {
   BasicBlock *ThenBlock = nullptr;
   BasicBlock *ElseBlock = nullptr;
   SplitBlockAndInsertIfThenElse(

llvm/test/CodeGen/ARM/aes-erratum-fix.ll

 ; CHECK-FIX-NEXT:    beq .LBB13_2
 ; CHECK-FIX-NEXT:  @ %bb.1:
 ; CHECK-FIX-NEXT:    vmov.8 d16[0], r1
 ; CHECK-FIX-NEXT:  .LBB13_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB13_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.8 d0[0], r1
-; CHECK-FIX-NEXT:  .LBB13_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB13_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aese.8 q8, q0
 ; CHECK-FIX-NEXT:    aesmc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r2]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = load <16 x i8>, <16 x i8>* %3, align 8
   %6 = insertelement <16 x i8> %5, i8 %1, i64 0
   %7 = select i1 %0, <16 x i8> %6, <16 x i8> %5
   %8 = insertelement <16 x i8> %2, i8 %1, i64 0

 ; CHECK-FIX-NEXT:    beq .LBB19_2
 ; CHECK-FIX-NEXT:  @ %bb.1:
 ; CHECK-FIX-NEXT:    vmov.16 d16[0], r1
 ; CHECK-FIX-NEXT:  .LBB19_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB19_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.16 d0[0], r1
-; CHECK-FIX-NEXT:  .LBB19_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB19_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aese.8 q8, q0
 ; CHECK-FIX-NEXT:    aesmc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r2]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = bitcast <16 x i8>* %3 to <8 x i16>*
   %6 = load <8 x i16>, <8 x i16>* %5, align 8
   %7 = insertelement <8 x i16> %6, i16 %1, i64 0
   %8 = select i1 %0, <8 x i16> %7, <8 x i16> %6

 ; CHECK-FIX-NEXT:    beq .LBB25_2
 ; CHECK-FIX-NEXT:  @ %bb.1:
 ; CHECK-FIX-NEXT:    vmov.32 d16[0], r1
 ; CHECK-FIX-NEXT:  .LBB25_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB25_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.32 d0[0], r1
-; CHECK-FIX-NEXT:  .LBB25_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB25_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aese.8 q8, q0
 ; CHECK-FIX-NEXT:    aesmc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r2]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = bitcast <16 x i8>* %3 to <4 x i32>*
   %6 = load <4 x i32>, <4 x i32>* %5, align 8
   %7 = insertelement <4 x i32> %6, i32 %1, i64 0
   %8 = select i1 %0, <4 x i32> %7, <4 x i32> %6

 ; CHECK-FIX-NEXT:    vmov.32 d16[0], r2
 ; CHECK-FIX-NEXT:    vmov.32 d16[1], r3
 ; CHECK-FIX-NEXT:  .LBB31_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB31_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.32 d0[0], r2
 ; CHECK-FIX-NEXT:    vmov.32 d0[1], r3
-; CHECK-FIX-NEXT:  .LBB31_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB31_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aese.8 q8, q0
 ; CHECK-FIX-NEXT:    aesmc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r1]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = bitcast <16 x i8>* %3 to <2 x i64>*
   %6 = load <2 x i64>, <2 x i64>* %5, align 8
   %7 = insertelement <2 x i64> %6, i64 %1, i64 0
   %8 = select i1 %0, <2 x i64> %7, <2 x i64> %6

 ; CHECK-FIX-NEXT:    beq .LBB59_2
 ; CHECK-FIX-NEXT:  @ %bb.1:
 ; CHECK-FIX-NEXT:    vmov.8 d16[0], r1
 ; CHECK-FIX-NEXT:  .LBB59_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB59_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.8 d0[0], r1
-; CHECK-FIX-NEXT:  .LBB59_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB59_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aesd.8 q8, q0
 ; CHECK-FIX-NEXT:    aesimc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r2]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = load <16 x i8>, <16 x i8>* %3, align 8
   %6 = insertelement <16 x i8> %5, i8 %1, i64 0
   %7 = select i1 %0, <16 x i8> %6, <16 x i8> %5
   %8 = insertelement <16 x i8> %2, i8 %1, i64 0

 ; CHECK-FIX-NEXT:    beq .LBB65_2
 ; CHECK-FIX-NEXT:  @ %bb.1:
 ; CHECK-FIX-NEXT:    vmov.16 d16[0], r1
 ; CHECK-FIX-NEXT:  .LBB65_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB65_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.16 d0[0], r1
-; CHECK-FIX-NEXT:  .LBB65_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB65_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aesd.8 q8, q0
 ; CHECK-FIX-NEXT:    aesimc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r2]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = bitcast <16 x i8>* %3 to <8 x i16>*
   %6 = load <8 x i16>, <8 x i16>* %5, align 8
   %7 = insertelement <8 x i16> %6, i16 %1, i64 0
   %8 = select i1 %0, <8 x i16> %7, <8 x i16> %6

 ; CHECK-FIX-NEXT:    beq .LBB71_2
 ; CHECK-FIX-NEXT:  @ %bb.1:
 ; CHECK-FIX-NEXT:    vmov.32 d16[0], r1
 ; CHECK-FIX-NEXT:  .LBB71_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB71_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.32 d0[0], r1
-; CHECK-FIX-NEXT:  .LBB71_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB71_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aesd.8 q8, q0
 ; CHECK-FIX-NEXT:    aesimc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r2]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = bitcast <16 x i8>* %3 to <4 x i32>*
   %6 = load <4 x i32>, <4 x i32>* %5, align 8
   %7 = insertelement <4 x i32> %6, i32 %1, i64 0
   %8 = select i1 %0, <4 x i32> %7, <4 x i32> %6

 ; CHECK-FIX-NEXT:    vmov.32 d16[0], r2
 ; CHECK-FIX-NEXT:    vmov.32 d16[1], r3
 ; CHECK-FIX-NEXT:  .LBB77_2: @ %select.end
 ; CHECK-FIX-NEXT:    cmp r0, #0
 ; CHECK-FIX-NEXT:    beq .LBB77_4
 ; CHECK-FIX-NEXT:  @ %bb.3:
 ; CHECK-FIX-NEXT:    vmov.32 d0[0], r2
 ; CHECK-FIX-NEXT:    vmov.32 d0[1], r3
-; CHECK-FIX-NEXT:  .LBB77_4: @ %select.end1
+; CHECK-FIX-NEXT:  .LBB77_4: @ %select.end2
 ; CHECK-FIX-NEXT:    aesd.8 q8, q0
 ; CHECK-FIX-NEXT:    aesimc.8 q8, q8
 ; CHECK-FIX-NEXT:    vst1.64 {d16, d17}, [r1]
 ; CHECK-FIX-NEXT:    bx lr
   %5 = bitcast <16 x i8>* %3 to <2 x i64>*
   %6 = load <2 x i64>, <2 x i64>* %5, align 8
   %7 = insertelement <2 x i64> %6, i64 %1, i64 0
   %8 = select i1 %0, <2 x i64> %7, <2 x i64> %6