[VPlan] Use predicate in VPInstruction::computeCost for selects. (llvm#170278)

In some cases, the lowering a select depends on the predicate. If the
condition of a select is a compare instruction, thread the predicate
through to the TTI hook.

PR: llvm#170278

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -1007,9 +1007,8 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
 
   switch (getOpcode()) {
   case Instruction::Select: {
-    // TODO: It may be possible to improve this by analyzing where the
-    // condition operand comes from.
-    CmpInst::Predicate Pred = CmpInst::BAD_ICMP_PREDICATE;
+    llvm::CmpPredicate Pred = CmpInst::BAD_ICMP_PREDICATE;
+    match(getOperand(0), m_Cmp(Pred, m_VPValue(), m_VPValue()));
     auto *CondTy = Ctx.Types.inferScalarType(getOperand(0));
     auto *VecTy = Ctx.Types.inferScalarType(getOperand(1));
     if (!vputils::onlyFirstLaneUsed(this)) {

llvm/test/Transforms/LoopVectorize/AArch64/select-costs.ll

@@ -73,3 +73,26 @@ exit:
   %1 = select i1 %all.off, i32 1, i32 %0
   ret i32 %1
 }
+
+define i32 @select_vpinst_for_tail_folding(i8 %n) {
+; CHECK: LV: Checking a loop in 'select_vpinst_for_tail_folding'
+; CHECK: Cost of 1 for VF 2: EMIT vp<{{.+}}> = select vp<{{.+}}>, ir<%red.next>, ir<%red>
+; CHECK: Cost of 1 for VF 4: EMIT vp<{{.+}}> = select vp<{{.+}}>, ir<%red.next>, ir<%red>
+; CHECK: LV: Selecting VF: 4
+
+entry:
+  %c = icmp ne i8 %n, 0
+  %ext = zext i1 %c to i32
+  br label %loop
+
+loop:
+  %iv = phi i32 [ %ext, %entry ], [ %iv.next, %loop ]
+  %red = phi i32 [ 0, %entry ], [ %red.next, %loop ]
+  %iv.next = add i32 %iv, 1
+  %red.next = mul i32 %red, %iv
+  %ec = icmp eq i32 %iv, 12
+  br i1 %ec, label %exit, label %loop
+
+exit:
+  ret i32 %red.next
+}