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

[Florian Hahn via llvm-commits]

================
@@ -423,7 +423,33 @@ 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->hasNoUnsignedWrap())
+      return ElementCount::getFixed(0);
+
+    if (Mul->getNumOperands() == 2 && isa<SCEVConstant>(Mul->getOperand(0)) &&
+        isa<SCEVVScale>(Mul->getOperand(1))) {
+      ConstantInt *Scale = cast<SCEVConstant>(Mul->getOperand(0))->getValue();
----------------
fhahn wrote:

Could we use SCEV pattern matching here to make this a bit easier to read?

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