[llvm] [VPlan] Don't use the legacy cost model for loop conditions (PR #156864)

[David Sherwood via llvm-commits]

================
@@ -1151,6 +1151,32 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
     return Ctx.TTI.getIndexedVectorInstrCostFromEnd(Instruction::ExtractElement,
                                                     VecTy, Ctx.CostKind, 0);
   }
+  case VPInstruction::BranchOnCount: {
+    // If TC <= VF then this is just a branch.
+    // FIXME: Removing the branch happens in simplifyBranchConditionForVFAndUF
+    // where it checks TC <= VF * UF, but we don't know UF yet. This means in
+    // some cases we get a cost that's too high due to counting a cmp that
+    // later gets removed.
+    // FIXME: The compare could also be removed if TC = M * vscale,
+    // VF = N * vscale, and M <= N. Detecting that would require having the
+    // trip count as a SCEV though.
+    Value *TC = getParent()->getPlan()->getTripCount()->getUnderlyingValue();
+    ConstantInt *TCConst = dyn_cast_if_present<ConstantInt>(TC);
+    if (TCConst && TCConst->getValue().ule(VF.getKnownMinValue()))
+      return 0;
+    // Otherwise BranchOnCount generates ICmpEQ followed by a branch.
+    Type *ValTy = Ctx.Types.inferScalarType(getOperand(0));
+    return Ctx.TTI.getCmpSelInstrCost(Instruction::ICmp, ValTy,
+                                      CmpInst::makeCmpResultType(ValTy),
+                                      CmpInst::ICMP_EQ, Ctx.CostKind);
+  }
+  case Instruction::FCmp:
----------------
david-arm wrote:

The change below is about more than just about the loop conditions. See `VPPredicator::createHeaderMask` for an example where we explicitly introduce a icmp for the current tail-folding mask. In fact, I don't think the cost below is correct because the icmp can have vector inputs.

I think you either need to:
1. Find a way to bail out if the icmp/fcmp isn't used as a branch condition, or
2. Add support for vector types using `ValTY = toVectorTy(ValTy, VF)`, and just make sure the example of `VPPredicator::createHeaderMask` is being tested in this PR.

Sorry about this!

https://github.com/llvm/llvm-project/pull/156864