Lower consecutive select instructions correctly.

[danielcdh]

Summary

If consecutive select instructions are lowered separately in CGP, it will introduce redundant condition check and branches that cannot be removed by later optimization phases. This patch lowers all consecutive select instructions at the same to to avoid inefficent code as demonstrated in https://llvm.org/bugs/show_bug.cgi?id=29095

[vsk]

lib/CodeGen/CodeGenPrepare.cpp

4545

Assert SI1->getParent() == SI2->getParent()?

4549

Use 'auto It', a const_iterator, and clang-format?

4553

Is SI2 not dependent on SI1 if, e.g SI2->getTrueValue() == (add SI 1)? I'm fuzzy on what 'dependent' means in this context, sorry!

4556

This would be better off as an llvm_unreachable; that should let you remove the return.

4571

I'm worried that this is quadratic in the number of SelectInsts in a BB. Is there a way to perform the isDependent check by looking at the users of a SelectInst instead of traversing the BB?

test/CodeGen/X86/codegen-prepare-select.ll

8 ↗

(On Diff #70041)

Please move these checks closer to the IR which produces the expected assembly.

Similarly, please move the NOT tests closer to the IR which used to produce the corresponding assembly.

40 ↗

(On Diff #70041)

It would help me follow along if you could reduce this test a bit more. E.g, besides the branch_weights, I'm not sure if any of the metadata here is really needed. Similarly, it's not clear if/why some of the operations in the loop are really needed (mul, ashr, sext). If they really are needed, it would help to have a small explanatory comment.

[danielcdh]

update

lib/CodeGen/CodeGenPrepare.cpp

4553

function removed

4556

removed

test/CodeGen/X86/codegen-prepare-select.ll

8 ↗

(On Diff #70041)

remove the entire test and added to some legacy tests.

[vsk]

Thanks! I found this much easier to follow.

It might be worth adding negative tests (for >1 jae), where the conditions of adjacent SelectInst's differ, or where there is overlap between the true/false values.

[danielcdh]

add more test

[danielcdh]

correct comment

[davidxl]

lib/CodeGen/CodeGenPrepare.cpp

4549

Always push SI first and start from the next instruction?

4558

Early return true if ASI size is 1?

4559

Probably use a different variable name from the input parameter.

4560

I1 --> TI
I2 --> FI

4564

Add a comment here for explanation.

4613

Instead of skipping the optimization, It is better to keep the current behavior if those selects can not be grouped -- optimize them one by one. Skipping optimization completely need more discussion and can be done as follow up.

[danielcdh]

update

lib/CodeGen/CodeGenPrepare.cpp

4558

code removed

4564

code removed

4613

Had a better fix to make sure that the PHI can be lowered even there is dependency.

[davidxl]

lib/CodeGen/CodeGenPrepare.cpp

4619

Perhaps use a local var LastSI to be ASI.back() ?

4734

Add a comment about using reverse iterator.

4741

Add a comment here describing the situation.

[danielcdh]

update

[davidxl]

lib/CodeGen/CodeGenPrepare.cpp

4738

There is probably no need to do reverse iteration. I suggest create a helper function

Value *getTrueOrFalseValue(SelectInst *SI, bool isTrue, const SmallPtrSet<...>& Selects) {

Value *V; 
do {
   V = (isTrue? SI->getTrueValue() : SI->getFalseValue());
   VI = dyn_cast<SelectInst>(V);
   if (!VI || !Selects.count(VI))
      break;
   SI = VI;
 } while (true);
 return V;

}

Then the code below can be simplified a lot and easier to read:

PN->addIncoming(getTrueOrFalseValue(SI, true, INS), TrueBlock);
PN->addIncoming(getTrueOfFalseValue(SI, false, INS), FalseBlock);

[danielcdh]

refactor the helper function

[davidxl]

lib/CodeGen/CodeGenPrepare.cpp

4586

is defined

4589

make it static.

4594

Before the loop save SI to SI0. Here an assert can be added to make SI and SI0 have same cond.

[danielcdh]

update

[davidxl]

ltgm

lib/CodeGen/CodeGenPrepare.cpp

static bool sinkSelectOperand(const TargetTransformInfo *TTI, Value *V) {
   auto *I = dyn_cast<Instruction>(V);
   // If it's safe to speculatively execute, then it should not have side
   // 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) {
   // 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,
   // then we want to replace the select with a branch.
   uint64_t TrueWeight, FalseWeight;
   if (SI->extractProfMetadata(TrueWeight, FalseWeight)) {
     uint64_t Max = std::max(TrueWeight, FalseWeight);
     uint64_t Sum = TrueWeight + FalseWeight;
     if (Sum != 0) {
       auto Probability = BranchProbability::getBranchProbability(Max, Sum);
       if (Probability > TLI->getPredictableBranchThreshold())
         return true;
     }
   }
 
   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 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;
 }
 
+/// If \p isTrue is true, return the true value of \p SI, otherwise return
+/// false value of \p SI. If the true/false value of \p SI is defined by any
+/// select instructions in \p Selects, look through the defining select
+/// instruction until the true/false value is not defined in \p Selects.
+static Value *getTrueOrFalseValue(
+    SelectInst *SI, bool isTrue,
+    const SmallPtrSet<const Instruction *, 2> &Selects) {
+  Value *V;
+
+  for (SelectInst *DefSI = SI; DefSI != nullptr && Selects.count(DefSI);
+       DefSI = dyn_cast<SelectInst>(V)) {
+    assert(DefSI.getCondition() == SI->getCondition() &&
+           "The condition of DefSI does not match with SI");
+    V = (isTrue ? DefSI->getTrueValue() : DefSI->getFalseValue());
+  }
+  return V;
+}
 
 /// If we have a SelectInst that will likely profit from branch prediction,
 /// turn it into a branch.
 bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
+  // Find all consecutive select instructions that share the same condition.
+  SmallVector<SelectInst *, 2> ASI;
+  ASI.push_back(SI);
+  for (BasicBlock::iterator It = ++BasicBlock::iterator(SI);
+       It != SI->getParent()->end(); ++It) {
+    SelectInst *I = dyn_cast<SelectInst>(&*It);
+    if (I && SI->getCondition() == I->getCondition()) {
+      ASI.push_back(I);
+    } else {
+      break;
+    }
+  }
+
+  SelectInst *LastSI = ASI.back();
+  // Increment the current iterator to skip all the rest of select instructions
+  // because they will be either "not lowered" or "all lowered" to branch.
+  CurInstIterator = std::next(LastSI->getIterator());
+
   bool VectorCond = !SI->getCondition()->getType()->isIntegerTy(1);
 
   // Can we convert the 'select' to CF ?
   if (DisableSelectToBranch || OptSize || !TLI || VectorCond ||
       SI->getMetadata(LLVMContext::MD_unpredictable))
     return false;
 
   TargetLowering::SelectSupportKind SelectKind;
bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
   // 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.
   BasicBlock *StartBlock = SI->getParent();
-  BasicBlock::iterator SplitPt = ++(BasicBlock::iterator(SI));
+  BasicBlock::iterator SplitPt = ++(BasicBlock::iterator(LastSI));
   BasicBlock *EndBlock = StartBlock->splitBasicBlock(SplitPt, "select.end");
 
   // 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);
-    auto *TrueBranch = BranchInst::Create(EndBlock, TrueBlock);
+        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);
-    auto *FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
+        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",
if (TrueBlock == nullptr) {
     TrueBlock = StartBlock;
   } else if (FalseBlock == nullptr) {
     BranchInst::Create(TrueBlock, EndBlock, SI->getCondition(), SI);
     FalseBlock = StartBlock;
   } else {
     BranchInst::Create(TrueBlock, FalseBlock, SI->getCondition(), 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 (auto It = ASI.rbegin(); It != ASI.rend(); ++It) {
+    SelectInst *SI = *It;
     // The select itself is replaced with a PHI Node.
     PHINode *PN = PHINode::Create(SI->getType(), 2, "", &EndBlock->front());
     PN->takeName(SI);
-  PN->addIncoming(SI->getTrueValue(), TrueBlock);
-  PN->addIncoming(SI->getFalseValue(), FalseBlock);
+    PN->addIncoming(getTrueOrFalseValue(SI, true, INS), TrueBlock);
+    PN->addIncoming(getTrueOrFalseValue(SI, false, INS), FalseBlock);
 
     SI->replaceAllUsesWith(PN);
     SI->eraseFromParent();
+    INS.erase(SI);
+    ++NumSelectsExpanded;
+  }
 
   // Instruct OptimizeBlock to skip to the next block.
   CurInstIterator = StartBlock->end();
-  ++NumSelectsExpanded;
   return true;
 }
 
 static bool isBroadcastShuffle(ShuffleVectorInst *SVI) {
   SmallVector<int, 16> Mask(SVI->getShuffleMask());
   int SplatElem = -1;
   for (unsigned i = 0; i < Mask.size(); ++i) {
     if (SplatElem != -1 && Mask[i] != -1 && Mask[i] != SplatElem)

test/CodeGen/X86/pseudo_cmov_lower2.ll

entry:
   %c1 = icmp slt i32 %p1, 0
   %d2 = select i1 %c1, double %p2, double %p3
   %c2 = icmp sge i32 %p1, 0
   %d3 = select i1 %c2, double %p4, double %p3
   %d4 = select i1 %c2, double %d3, double %d2
   %d5 = fdiv double %d4, %d3
   ret double %d5
 }
+
+; This test checks that only a single jae gets generated in the final code
+; for lowering the CMOV pseudos that get created for this IR.  The tricky part
+; of this test is that it tests the special code in CodeGenPrepare.
+;
+; CHECK-LABEL: foo5:
+; CHECK: jb
+; CHECK-NOT: jb
+define double @foo5(float %p1, double %p2, double %p3) nounwind {
+entry:
+  %c1 = fcmp oge float %p1, 0.000000e+00
+  %d0 = fadd double %p2, 1.25e0
+  %d1 = fadd double %p3, 1.25e0
+  %d2 = select i1 %c1, double %d0, double %d1, !prof !0
+  %d3 = select i1 %c1, double %d2, double %p2, !prof !0
+  %d4 = select i1 %c1, double %d3, double %p3, !prof !0
+  %d5 = fsub double %d2, %d3
+  %d6 = fadd double %d5, %d4
+  ret double %d6
+}
+
+; We should expand select instructions into 3 conditional branches as their
+; condtions are different.
+;
+; CHECK-LABEL: foo6:
+; CHECK: jb
+; CHECK: jae
+; CHECK: jae
+define double @foo6(float %p1, double %p2, double %p3) nounwind {
+entry:
+  %c1 = fcmp oge float %p1, 0.000000e+00
+  %c2 = fcmp oge float %p1, 1.000000e+00
+  %c3 = fcmp oge float %p1, 2.000000e+00
+  %d0 = fadd double %p2, 1.25e0
+  %d1 = fadd double %p3, 1.25e0
+  %d2 = select i1 %c1, double %d0, double %d1, !prof !0
+  %d3 = select i1 %c2, double %d2, double %p2, !prof !0
+  %d4 = select i1 %c3, double %d3, double %p3, !prof !0
+  %d5 = fsub double %d2, %d3
+  %d6 = fadd double %d5, %d4
+  ret double %d6
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2000}