[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

@@ -2341,12 +2341,17 @@ struct SCEVFindUnsafe {
 };
 } // namespace
 
-bool SCEVExpander::isSafeToExpand(const SCEV *S) const {
+static bool isSafeToExpand(const SCEV *S, bool CanonicalMode,
+                           ScalarEvolution &SE) {
   SCEVFindUnsafe Search(SE, CanonicalMode);
   visitAll(S, Search);
   return !Search.IsUnsafe;
 }
 
+bool SCEVExpander::isSafeToExpand(const SCEV *S) const {
+  return ::isSafeToExpand(S, CanonicalMode, SE);
+}
+
 bool SCEVExpander::isSafeToExpandAt(const SCEV *S,
                                     const Instruction *InsertionPoint) const {
   if (!isSafeToExpand(S))
@@ -2369,6 +2374,47 @@ bool SCEVExpander::isSafeToExpandAt(const SCEV *S,
   return false;
 }
 
+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 || all_of(Exiting, [L, &SE](BasicBlock *E) {
+        if (L->getHeader() == E)
+          return true;
+        const SCEV *EC = SE.getExitCount(L, E);
+        if (isa<SCEVCouldNotCompute>(EC))
+          return true;
+        return ::isSafeToExpand(EC, true, SE);
+      }))
+    return Expr;
+
+  return SCEVUDivRewriter::rewrite(Expr, SE);
+}
+
 void SCEVExpanderCleaner::cleanup() {
   // Result is used, nothing to remove.
   if (ResultUsed)

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);
 }