[SCEV] Make getUDivExactExpr handle non-nuw multiplies correctly

[efriedma]

Summary

To avoid regressions, make ScalarEvolution::createSCEV a bit more clever.

Also get rid of some useless code in ScalarEvolution::howFarToZero which was hiding this bug.

No new testcase because it's impossible to actually expose this bug: we don't have any in-tree users of getUDivExactExpr besides the two functions I just mentioned, and they both dodged the problem.

[efriedma]

Uploaded wrong version of the patch.

[sanjoy]

Can I ask you to please to (in a later commit) change createSCEV to use getUDivExactExpr for udiv exact instructions, and add some tests around what you're changing here?

lib/Analysis/ScalarEvolution.cpp

7248

I think we can basically drop this whole comment.

lib/Analysis/ScalarEvolution.cpp

 /// it's not exact because the udiv may be clearing bits.
 const SCEV *ScalarEvolution::getUDivExactExpr(const SCEV *LHS,
                                               const SCEV *RHS) {
   // TODO: we could try to find factors in all sorts of things, but for now we
   // just deal with u/exact (multiply, constant). See SCEVDivision towards the
   // end of this file for inspiration.
 
   const SCEVMulExpr *Mul = dyn_cast<SCEVMulExpr>(LHS);
-  if (!Mul)
+  if (!Mul || !Mul->hasNoUnsignedWrap())
     return getUDivExpr(LHS, RHS);
 
   if (const SCEVConstant *RHSCst = dyn_cast<SCEVConstant>(RHS)) {
     // If the mulexpr multiplies by a constant, then that constant must be the
     // first element of the mulexpr.
     if (const auto *LHSCst = dyn_cast<SCEVConstant>(Mul->getOperand(0))) {
       if (LHSCst == RHSCst) {
         SmallVector<const SCEV *, 2> Operands;
case Instruction::And:
         unsigned BitWidth = A.getBitWidth();
         APInt KnownZero(BitWidth, 0), KnownOne(BitWidth, 0);
         computeKnownBits(BO->LHS, KnownZero, KnownOne, getDataLayout(),
                          0, &AC, nullptr, &DT);
 
         APInt EffectiveMask =
             APInt::getLowBitsSet(BitWidth, BitWidth - LZ - TZ).shl(TZ);
         if ((LZ != 0 || TZ != 0) && !((~A & ~KnownZero) & EffectiveMask)) {
-          const SCEV *MulCount = getConstant(ConstantInt::get(
-              getContext(), APInt::getOneBitSet(BitWidth, TZ)));
+          const SCEV *MulCount = getConstant(APInt::getOneBitSet(BitWidth, TZ));
+          const SCEV *LHS = getSCEV(BO->LHS);
+          const SCEV *ShiftedLHS = nullptr;
+          if (auto *LHSMul = dyn_cast<SCEVMulExpr>(LHS)) {
+            if (auto *OpC = dyn_cast<SCEVConstant>(LHSMul->getOperand(0))) {
+              // For an expression like (x * 8) & 8, simplify the multiply.
+              unsigned MulZeros = OpC->getAPInt().countTrailingZeros();
+              unsigned GCD = std::min(MulZeros, TZ);
+              APInt DivAmt = APInt::getOneBitSet(BitWidth, TZ - GCD);
+              SmallVector<const SCEV*, 4> MulOps;
+              MulOps.push_back(getConstant(OpC->getAPInt().lshr(GCD)));
+              MulOps.append(LHSMul->op_begin() + 1, LHSMul->op_end());
+              auto *NewMul = getMulExpr(MulOps, LHSMul->getNoWrapFlags());
+              ShiftedLHS = getUDivExpr(NewMul, getConstant(DivAmt));
+            }
+          }
+          if (!ShiftedLHS)
+            ShiftedLHS = getUDivExpr(LHS, MulCount);
           return getMulExpr(
               getZeroExtendExpr(
-                  getTruncateExpr(
-                      getUDivExactExpr(getSCEV(BO->LHS), MulCount),
+                  getTruncateExpr(ShiftedLHS,
                       IntegerType::get(getContext(), BitWidth - LZ - TZ)),
                   BO->LHS->getType()),
               MulCount);
         }
       }
       break;
 
     case Instruction::Or:
ScalarEvolution::howFarToZero(const SCEV *V, const Loop *L, bool ControlsExit,
   // Distance.
   if (!CountDown) {
     const APInt &StepV = StepC->getAPInt();
     // StepV.isPowerOf2() returns true if StepV is an positive power of two.  It
     // also returns true if StepV is maximally negative (eg, INT_MIN), but that
     // case is not handled as this code is guarded by !CountDown.
     if (StepV.isPowerOf2() &&
         GetMinTrailingZeros(Distance) >= StepV.countTrailingZeros()) {
       // Here we've constrained the equation to be of the form
       //
       //   2^(N + k) * Distance' = (StepV == 2^N) * X (mod 2^W)  ... (0)
       //
       // where we're operating on a W bit wide integer domain and k is
       // non-negative.  The smallest unsigned solution for X is the trip count.
       //
       // (0) is equivalent to:
       //
if (StepV.isPowerOf2() &&
       //
       // then from (2), we get X = Val URem i8 128 (k = 0 in this case).
       //
       // Note: It is tempting to solve (3) by setting X = Val, but Val is not
       // necessarily the smallest unsigned value of X that satisfies (3).
       // E.g. if Val is i8 -127 then the smallest value of X that satisfies (3)
       // is i8 1, not i8 -127
 
-      const auto *ModuloResult = getUDivExactExpr(Distance, Step);
-
-      // Since SCEV does not have a URem node, we construct one using a truncate
-      // and a zero extend.
-
-      unsigned NarrowWidth = StepV.getBitWidth() - StepV.countTrailingZeros();
-      auto *NarrowTy = IntegerType::get(getContext(), NarrowWidth);
-      auto *WideTy = Distance->getType();
-
-      const SCEV *Limit =
-          getZeroExtendExpr(getTruncateExpr(ModuloResult, NarrowTy), WideTy);
+      const auto *Limit = getUDivExactExpr(Distance, Step);
       return ExitLimit(Limit, Limit, false, Predicates);
     }
   }
 
   // If the condition controls loop exit (the loop exits only if the expression
   // is true) and the addition is no-wrap we can use unsigned divide to
   // compute the backedge count.  In this case, the step may not divide the
   // distance, but we don't care because if the condition is "missed" the loop

test/Analysis/ScalarEvolution/pr24757.ll

 ; RUN: opt -analyze -scalar-evolution < %s | FileCheck %s
 
-; CHECK: Loop %bb1: backedge-taken count is (zext i7 (trunc i8 %a.promoted to i7) to i8)
+; CHECK: Loop %bb1: backedge-taken count is ((2 * %a.promoted) /u 2)
 
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-apple-macosx10.10.0"
 
 @a = global i8 -127, align 1
 @b = common global i32 0, align 4
 
 declare void @use(i32)