[llvm] [LLVM][LV] Improve UF calculation for vscale based scalar loops. (PR #146102)

[Paul Walker via llvm-commits]

================
@@ -423,7 +423,26 @@ static bool hasIrregularType(Type *Ty, const DataLayout &DL) {
 /// ElementCount to include loops whose trip count is a function of vscale.
 static ElementCount getSmallConstantTripCount(ScalarEvolution *SE,
                                               const Loop *L) {
-  return ElementCount::getFixed(SE->getSmallConstantTripCount(L));
+  if (unsigned ExpectedTC = SE->getSmallConstantTripCount(L))
+    return ElementCount::getFixed(ExpectedTC);
+
+  const SCEV *BTC = SE->getBackedgeTakenCount(L);
+
+  if (isa<SCEVCouldNotCompute>(BTC))
+    return ElementCount::getFixed(0);
+
+  const SCEV *ExitCount = SE->getTripCountFromExitCount(BTC, BTC->getType(), L);
+
+  if (isa<SCEVVScale>(ExitCount))
+    return ElementCount::getScalable(1);
+
+  if (auto *Mul = dyn_cast<SCEVMulExpr>(ExitCount))
+    if (Mul->getNumOperands() == 2 && isa<SCEVConstant>(Mul->getOperand(0)) &&
----------------
paulwalker-arm wrote:

Fixed both in terms of requiring NUW flags and ensuring the constant fits within the range of ElementCount.

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