[llvm] [SCEV] Infer NUW on AR in getUDivExpr (PR #217133)

[Ramkumar Ramachandra via llvm-commits]

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

>From d9cb18718aa01cbc5a3d3c207caa600e82f64656 Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <artagnon at tenstorrent.com>
Date: Wed, 19 Aug 2026 08:53:15 +0100
Subject: [PATCH] [SCEV] Infer NUW on AR in getUDivExpr

---
 llvm/lib/Analysis/ScalarEvolution.cpp | 26 +++++---------------------
 1 file changed, 5 insertions(+), 21 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 52a03d2cad62d..a2c0d31362bee 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -3520,27 +3520,17 @@ const SCEV *ScalarEvolution::getUDivExpr(SCEVUse LHS, SCEVUse RHS) {
     if (!RHSC->getValue()->isZero()) {
       // Determine if the division can be folded into the operands of
       // its operands.
-      // TODO: Generalize this to non-constants by using known-bits information.
-      Type *Ty = LHS->getType();
-      unsigned LZ = RHSC->getAPInt().countl_zero();
-      unsigned MaxShiftAmt = getTypeSizeInBits(Ty) - LZ - 1;
-      // For non-power-of-two values, effectively round the value up to the
-      // nearest power of two.
-      if (!RHSC->getAPInt().isPowerOf2())
-        ++MaxShiftAmt;
-      IntegerType *ExtTy =
-        IntegerType::get(getContext(), getTypeSizeInBits(Ty) + MaxShiftAmt);
       if (const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(LHS))
         if (const SCEVConstant *Step =
             dyn_cast<SCEVConstant>(AR->getStepRecurrence(*this))) {
           // {X,+,N}/C --> {X/C,+,N/C} if safe and N/C can be folded.
           const APInt &StepInt = Step->getAPInt();
           const APInt &DivInt = RHSC->getAPInt();
-          if (!StepInt.urem(DivInt) &&
-              getZeroExtendExpr(AR, ExtTy) ==
-              getAddRecExpr(getZeroExtendExpr(AR->getStart(), ExtTy),
-                            getZeroExtendExpr(Step, ExtTy),
-                            AR->getLoop(), SCEV::FlagAnyWrap)) {
+          bool NoWrap = hasFlags(proveNoUnsignedWrapViaInduction(AR),
+                                 SCEVNoWrapFlags::FlagNUW) ||
+                        proveNoWrapByVaryingStart<SCEVZeroExtendExpr>(
+                            AR->getStart(), Step, AR->getLoop());
+          if (!StepInt.urem(DivInt) && NoWrap) {
             SmallVector<SCEVUse, 4> Operands;
             for (const SCEV *Op : AR->operands())
               Operands.push_back(getUDivExpr(Op, RHS));
@@ -3551,12 +3541,6 @@ const SCEV *ScalarEvolution::getUDivExpr(SCEVUse LHS, SCEVUse RHS) {
           const APInt *StartRem;
           if (!DivInt.urem(StepInt) && match(getURemExpr(AR->getStart(), Step),
                                              m_scev_APInt(StartRem))) {
-            bool NoWrap =
-                getZeroExtendExpr(AR, ExtTy) ==
-                getAddRecExpr(getZeroExtendExpr(AR->getStart(), ExtTy),
-                              getZeroExtendExpr(Step, ExtTy), AR->getLoop(),
-                              SCEV::FlagAnyWrap);
-
             // With N <= C and both N, C as powers-of-2, the transformation
             // {X,+,N}/C => {(X - X%N),+,N}/C preserves division results even
             // if wrapping occurs, as the division results remain equivalent for