[LV] Don't vectorize if trip count expansion may introduce UB.

Introduce a utility to check if a SCEV expansion may introduce UB
(couldn't find a similar utility after a quick glance) and use to the
avoid vectorizing when expanding the trip count introduces UB.

Fixes #89958.

Author: fhahn

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

@@ -418,6 +418,13 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
   BasicBlock::iterator findInsertPointAfter(Instruction *I,
                                             Instruction *MustDominate) const;
 
+  /// If \p L contains exits which may execute conditionally and contain UDiv
+  /// expressions with divisors that can be 0, expanding \p BTC may introduce
+  /// new UB. In that case, rewrite UDiv(A, B) to  UDiv(A, UMAX(1, B)). If B is
+  /// 0, that exit cannot be taken.
+  static const SCEV *rewriteExpressionToRemoveUB(const SCEV *BTC, Loop *L,
+                                                 ScalarEvolution &SE);
+
 private:
   LLVMContext &getContext() const { return SE.getContext(); }
 

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

@@ -161,6 +161,52 @@ SCEVExpander::findInsertPointAfter(Instruction *I,
   return IP;
 }
 
+namespace {
+struct SCEVUDivRewriter : public SCEVRewriteVisitor<SCEVUDivRewriter> {
+  SCEVUDivRewriter(ScalarEvolution &SE) : SCEVRewriteVisitor(SE) {}
+
+  static const SCEV *rewrite(const SCEV *Scev, ScalarEvolution &SE) {
+    SCEVUDivRewriter Rewriter(SE);
+    return Rewriter.visit(Scev);
+  }
+
+  const SCEV *visitUDivExpr(const SCEVUDivExpr *Expr) {
+    auto *LHS = visit(Expr->getLHS());
+    auto *RHS = visit(Expr->getRHS());
+    return SE.getUDivExpr(LHS, SE.getUMaxExpr(SE.getOne(RHS->getType()), RHS));
+  }
+};
+} // namespace
+
+const SCEV *SCEVExpander::rewriteExpressionToRemoveUB(const SCEV *Expr, Loop *L,
+                                                      ScalarEvolution &SE) {
+  SmallVector<BasicBlock *> Exiting;
+  L->getExitingBlocks(Exiting);
+  // Check if exit count for any exit that may execute unconditionally may in
+  // introduce UB. Note that we can skip checks in the header or if there's a
+  // single exit, as in those cases we know that the exit count will be
+  // evaluated in each loop iteration. There are other cases where the exiting
+  // block executes on each loop iteration, but we don't have a cheap way to
+  // check at the moment.
+
+  if (Exiting.size() > 1 && any_of(Exiting, [L, &SE](BasicBlock *E) {
+        if (L->getHeader() == E)
+          return false;
+        const SCEV *EC = SE.getExitCount(L, E);
+        return !isa<SCEVCouldNotCompute>(EC) &&
+               SCEVExprContains(EC, [&SE](const SCEV *Op) {
+                 auto *UDiv = dyn_cast<SCEVUDivExpr>(Op);
+                 if (!UDiv)
+                   return false;
+
+                 return !SE.isKnownNonZero(UDiv->getOperand(1));
+               });
+      })) {
+    return SCEVUDivRewriter::rewrite(Expr, SE);
+  }
+  return Expr;
+}
+
 BasicBlock::iterator
 SCEVExpander::GetOptimalInsertionPointForCastOf(Value *V) const {
   // Cast the argument at the beginning of the entry block, after

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -984,6 +984,9 @@ const SCEV *createTripCountSCEV(Type *IdxTy, PredicatedScalarEvolution &PSE,
   assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCount) && "Invalid loop count");
 
   ScalarEvolution &SE = *PSE.getSE();
+  if (OrigLoop)
+    BackedgeTakenCount = SCEVExpander::rewriteExpressionToRemoveUB(
+        BackedgeTakenCount, OrigLoop, SE);
   return SE.getTripCountFromExitCount(BackedgeTakenCount, IdxTy, OrigLoop);
 }