[llvm] [SCEV] Introduce UDiv::mayTriggerUB (NFC) (PR #217064)

[Ramkumar Ramachandra via llvm-commits]

https://github.com/artagnon updated https://github.com/llvm/llvm-project/pull/217064

>From 862d66987f31cb872d25087553685074a484ebf0 Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <artagnon at tenstorrent.com>
Date: Tue, 18 Aug 2026 16:43:58 +0100
Subject: [PATCH] [SCEV] Introduce UDiv::mayTriggerUB (NFC)

Use it to clean up the logic in SCEVExpander.
---
 .../Analysis/ScalarEvolutionExpressions.h     |  6 ++++
 .../Utils/ScalarEvolutionExpander.cpp         | 31 ++++++-------------
 2 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h b/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h
index ebee63963c701..0bf5a4ebf343b 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h
@@ -319,6 +319,12 @@ class SCEVUDivExpr : public SCEV {
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const SCEV *S) { return S->getSCEVType() == scUDivExpr; }
+
+  /// Returns true if the expression may trigger undefined-behavior.
+  bool mayTriggerUB(ScalarEvolution &SE) const {
+    return !SE.isKnownNonZero(getRHS()) ||
+           !ScalarEvolution::isGuaranteedNotToBePoison(getRHS());
+  }
 };
 
 /// This node represents a polynomial recurrence on the trip count
diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index a7ca2104e59c3..da9e05c7f4a7f 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -717,20 +717,18 @@ Value *SCEVExpander::visitUDivExpr(SCEVUseT<const SCEVUDivExpr *> S) {
   const SCEV *RHSExpr = S->getRHS();
   Value *RHS = expand(RHSExpr);
   if (SafeUDivMode) {
-    bool GuaranteedNotPoison =
-        ScalarEvolution::isGuaranteedNotToBePoison(RHSExpr);
-    if (!GuaranteedNotPoison)
+    if (!ScalarEvolution::isGuaranteedNotToBePoison(RHSExpr))
       RHS = Builder.CreateFreeze(RHS);
 
     // We need an umax if either RHSExpr is not known to be zero, or if it is
     // not guaranteed to be non-poison. In the later case, the frozen poison may
     // be 0.
-    if (!SE.isKnownNonZero(RHSExpr) || !GuaranteedNotPoison)
+    if (S->mayTriggerUB(SE))
       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()));
+                     /*IsSafeToHoist=*/!S->mayTriggerUB(SE));
 }
 
 /// Determine if this is a well-behaved chain of instructions leading back to
@@ -1667,20 +1665,12 @@ Value *SCEVExpander::expand(SCEVUse S) {
 
   // We can move insertion point only if there is no div or rem operations
   // otherwise we are risky to move it over the check for zero denominator.
-  auto SafeToHoist = [](const SCEV *S) {
-    return !SCEVExprContains(S, [](const SCEV *S) {
-              if (const auto *D = dyn_cast<SCEVUDivExpr>(S)) {
-                if (const auto *SC = dyn_cast<SCEVConstant>(D->getRHS()))
-                  // Division by non-zero constants can be hoisted.
-                  return SC->getValue()->isZero();
-                // All other divisions should not be moved as they may be
-                // divisions by zero and should be kept within the
-                // conditions of the surrounding loops that guard their
-                // execution (see PR35406).
-                return true;
-              }
-              return false;
-            });
+  auto SafeToHoist = [this](const SCEV *S) {
+    return !SCEVExprContains(S, [this](const SCEV *S) {
+      const auto *D = dyn_cast<SCEVUDivExpr>(S);
+      // TODO: Why is this RHS-constant check necessary?
+      return D && (!isa<SCEVConstant>(D->getRHS()) || D->mayTriggerUB(SE));
+    });
   };
   if (SafeToHoist(S)) {
     for (Loop *L = SE.LI.getLoopFor(Builder.GetInsertBlock());;
@@ -2507,8 +2497,7 @@ struct SCEVFindUnsafe {
 
   bool follow(const SCEV *S) {
     if (const SCEVUDivExpr *D = dyn_cast<SCEVUDivExpr>(S)) {
-      if (!SE.isKnownNonZero(D->getRHS()) ||
-          !SE.isGuaranteedNotToBePoison(D->getRHS())) {
+      if (D->mayTriggerUB(SE)) {
         IsUnsafe = true;
         return false;
       }