!fixup introduce SafeUDivMode to SCEVExpander.

Author: fhahn

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

@@ -124,6 +124,11 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
   /// "expanded" form.
   bool LSRMode;
 
+  /// When true, rewrite any divisors of UDiv expressions that may be 0 to
+  /// umax(Divisor, 1) to avoid introducing UB. If the divisor may be poison,
+  /// freeze it first.
+  bool SafeUDivMode = false;
+
   typedef IRBuilder<InstSimplifyFolder, IRBuilderCallbackInserter> BuilderType;
   BuilderType Builder;
 
@@ -300,6 +305,9 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
   /// location and their operands are defined at this location.
   bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint) const;
 
+  static bool isSafeToExpand(const SCEV *S, bool CanonicalMode,
+                             ScalarEvolution &SE);
+
   /// Insert code to directly compute the specified SCEV expression into the
   /// program.  The code is inserted into the specified block.
   Value *expandCodeFor(const SCEV *SH, Type *Ty, BasicBlock::iterator I);
@@ -418,13 +426,11 @@ 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);
 
+  void setSafeUDivMode() { SafeUDivMode = true; }
+
 private:
   LLVMContext &getContext() const { return SE.getContext(); }
 

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

@@ -665,6 +665,12 @@ Value *SCEVExpander::visitUDivExpr(const SCEVUDivExpr *S) {
   }
 
   Value *RHS = expand(S->getRHS());
+  if (SafeUDivMode && !SE.isKnownNonZero(S->getRHS())) {
+    if (!isGuaranteedNotToBePoison(RHS))
+      RHS = Builder.CreateFreeze(RHS);
+    RHS = Builder.CreateIntrinsic(RHS->getType(), Intrinsic::umax,
+                                  {RHS, ConstantInt::get(RHS->getType(), 1)});
+  }
   return InsertBinop(Instruction::UDiv, LHS, RHS, SCEV::FlagAnyWrap,
                      /*IsSafeToHoist*/ SE.isKnownNonZero(S->getRHS()));
 }
@@ -2341,15 +2347,15 @@ struct SCEVFindUnsafe {
 };
 } // namespace
 
-static bool isSafeToExpand(const SCEV *S, bool CanonicalMode,
-                           ScalarEvolution &SE) {
+bool SCEVExpander::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);
+  return isSafeToExpand(S, CanonicalMode, SE);
 }
 
 bool SCEVExpander::isSafeToExpandAt(const SCEV *S,
@@ -2374,47 +2380,6 @@ 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

@@ -916,9 +916,6 @@ 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);
 }
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -866,8 +866,31 @@ VPlanPtr VPlan::createInitialVPlan(const SCEV *TripCount, ScalarEvolution &SE,
   VPIRBasicBlock *Entry = new VPIRBasicBlock(TheLoop->getLoopPreheader());
   VPBasicBlock *VecPreheader = new VPBasicBlock("vector.ph");
   auto Plan = std::make_unique<VPlan>(Entry, VecPreheader);
-  Plan->TripCount =
-      vputils::getOrCreateVPValueForSCEVExpr(*Plan, TripCount, SE);
+
+  bool NeedsSafeUDivMode = false;
+  {
+    SmallVector<BasicBlock *> Exiting;
+    TheLoop->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.
+    NeedsSafeUDivMode =
+        Exiting.size() != 1 && any_of(Exiting, [TheLoop, &SE](BasicBlock *E) {
+          if (TheLoop->getHeader() == E)
+            return false;
+          const SCEV *EC = SE.getExitCount(TheLoop, E);
+          if (isa<SCEVCouldNotCompute>(EC))
+            return false;
+          return !SCEVExpander::isSafeToExpand(EC, true, SE);
+        });
+  }
+
+  Plan->TripCount = vputils::getOrCreateVPValueForSCEVExpr(*Plan, TripCount, SE,
+                                                           NeedsSafeUDivMode);
   // Create VPRegionBlock, with empty header and latch blocks, to be filled
   // during processing later.
   VPBasicBlock *HeaderVPBB = new VPBasicBlock("vector.body");

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -2771,9 +2771,15 @@ class VPExpandSCEVRecipe : public VPSingleDefRecipe {
   const SCEV *Expr;
   ScalarEvolution &SE;
 
+  /// When set to true, set SafeUDivMode when expanding the SCEV to avoid
+  /// introducing UB.
+  bool SafeUDivMode;
+
 public:
-  VPExpandSCEVRecipe(const SCEV *Expr, ScalarEvolution &SE)
-      : VPSingleDefRecipe(VPDef::VPExpandSCEVSC, {}), Expr(Expr), SE(SE) {}
+  VPExpandSCEVRecipe(const SCEV *Expr, ScalarEvolution &SE,
+                     bool SafeUDivMode = false)
+      : VPSingleDefRecipe(VPDef::VPExpandSCEVSC, {}), Expr(Expr), SE(SE),
+        SafeUDivMode(SafeUDivMode) {}
 
   ~VPExpandSCEVRecipe() override = default;
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -2932,6 +2932,8 @@ void VPExpandSCEVRecipe::execute(VPTransformState &State) {
   assert(!State.Instance && "cannot be used in per-lane");
   const DataLayout &DL = State.CFG.PrevBB->getDataLayout();
   SCEVExpander Exp(SE, DL, "induction");
+  if (SafeUDivMode)
+    Exp.setSafeUDivMode();
 
   Value *Res = Exp.expandCodeFor(Expr, Expr->getType(),
                                  &*State.Builder.GetInsertPoint());

llvm/lib/Transforms/Vectorize/VPlanUtils.cpp

@@ -23,7 +23,8 @@ bool vputils::onlyFirstPartUsed(const VPValue *Def) {
 }
 
 VPValue *vputils::getOrCreateVPValueForSCEVExpr(VPlan &Plan, const SCEV *Expr,
-                                                ScalarEvolution &SE) {
+                                                ScalarEvolution &SE,
+                                                bool SafeUDivMode) {
   if (auto *Expanded = Plan.getSCEVExpansion(Expr))
     return Expanded;
   VPValue *Expanded = nullptr;
@@ -32,7 +33,7 @@ VPValue *vputils::getOrCreateVPValueForSCEVExpr(VPlan &Plan, const SCEV *Expr,
   else if (auto *E = dyn_cast<SCEVUnknown>(Expr))
     Expanded = Plan.getOrAddLiveIn(E->getValue());
   else {
-    Expanded = new VPExpandSCEVRecipe(Expr, SE);
+    Expanded = new VPExpandSCEVRecipe(Expr, SE, SafeUDivMode);
     Plan.getPreheader()->appendRecipe(Expanded->getDefiningRecipe());
   }
   Plan.addSCEVExpansion(Expr, Expanded);

llvm/lib/Transforms/Vectorize/VPlanUtils.h

@@ -24,7 +24,8 @@ bool onlyFirstPartUsed(const VPValue *Def);
 /// pre-header already contains a recipe expanding \p Expr, return it. If not,
 /// create a new one.
 VPValue *getOrCreateVPValueForSCEVExpr(VPlan &Plan, const SCEV *Expr,
-                                       ScalarEvolution &SE);
+                                       ScalarEvolution &SE,
+                                       bool SaveUDivMode = false);
 
 /// Returns true if \p VPV is uniform after vectorization.
 inline bool isUniformAfterVectorization(const VPValue *VPV) {