[SCEV] Preserve divisibility and min/max information in applyLoopGuards

[alonkom]

Summary

applyLoopGuards doesn't always preserve information when there are multiple assumes.
This patch tries to deal with multiple assumes regarding a SCEV's divisibility and min/max values, and rewrite it into a SCEV that still preserves all of the information.
For example, let the trip count of the loop be TC. Consider the 3 following assumes:

1. __builtin_assume(TC % 8 == 0);
2. __builtin_assume(TC > 0);
3. __builtin_assume(TC < 100);

Before this patch, depending on the assume processing order applyLoopGuards could create the following SCEV:
max(min((8 * (TC / 8)) , 99), 1)

Looking at this SCEV, it doesn't preserve the divisibility by 8 information.

After this patch, depending on the assume processing order applyLoopGuards could create the following SCEV:
max(min((8 * (TC / 8)) , 96), 8)

By aligning up 1 to 8, and aligning down 99 to 96, the new SCEV still preserves all of the original assumes.

[caojoshua]

What is the relation to https://reviews.llvm.org/D128701? Seems maybe you should close the other one.

Can you make the commit title more specific, and add more info in the commit message? In particular, I would like to see something like:

after applyLoopGuards, turns <A> into <B>

where A and B are SCEV expressions.

Other information like, what the problem was before, and how you address it, is also helpful. I would suggest looking at the git log of ScalarEvolution.cpp for examples of descriptive commit messages. This will save the reviewers time, and will be helpful when people look at the git log/blame.

[alonkom]

In D141850#4057996, @caojoshua wrote:

What is the relation to https://reviews.llvm.org/D128701? Seems maybe you should close the other one.

Can you make the commit title more specific, and add more info in the commit message? In particular, I would like to see something like:

after applyLoopGuards, turns <A> into <B>

where A and B are SCEV expressions.

Other information like, what the problem was before, and how you address it, is also helpful. I would suggest looking at the git log of ScalarEvolution.cpp for examples of descriptive commit messages. This will save the reviewers time, and will be helpful when people look at the git log/blame.

Thanks, updated the commit message.

[caojoshua]

The improved trip multiples from the test results look good. Ordering in applyLoopGuards is an issue. However, I think we can simplify this down a bit. What if we always applied min/max first, before we apply divisibility guards? For example, given:

__builtin_assume(TC % 8 == 0);
__builtin_assume(TC > 0);
__builtin_assume(TC < 100);

1. apply max: umax(1, TC)
2. apply min: umin(100, umax(1, TC))
3. apply divisibility info: 8 * (umin(100, umax(1, TC))) / 8

This makes divisibility info obvious. And traversing the SCEV, we can still see TC > 0 and TC < 100. I believe that if we always apply max/min first, we will never lose this info. This approach seems much simpler and easier to understand.

FYI. I haven't been around that long, and this change is non-trivial enough where I prefer a longstanding developer to give the final approval. I will still be around to give my thoughts.

llvm/lib/Analysis/ScalarEvolution.cpp

15042

SequentialMinMax SCEVs can be applied here as well.

15043

unused var

15047

Constant should always be on the left side. Lets not change that.

15062

Have not thought about it too deeply yet, but I'm concerned that we may need to take a look at when this is legal given Expr's NoWrapFlags. Same for the equivalent getMinusSCEV below

15171

Shouldn't swap the constant. Assume constant is on left.

[alonkom]

In D141850#4064477, @caojoshua wrote:

The improved trip multiples from the test results look good. Ordering in applyLoopGuards is an issue. However, I think we can simplify this down a bit. What if we always applied min/max first, before we apply divisibility guards? For example, given:

__builtin_assume(TC % 8 == 0);
__builtin_assume(TC > 0);
__builtin_assume(TC < 100);

1. apply max: umax(1, TC)
2. apply min: umin(100, umax(1, TC))
3. apply divisibility info: 8 * (umin(100, umax(1, TC))) / 8

This makes divisibility info obvious. And traversing the SCEV, we can still see TC > 0 and TC < 100. I believe that if we always apply max/min first, we will never lose this info. This approach seems much simpler and easier to understand.

FYI. I haven't been around that long, and this change is non-trivial enough where I prefer a longstanding developer to give the final approval. I will still be around to give my thoughts.

Thanks for the reply.
This works in this example, but what would happen if we had only these assumes:

__builtin_assume(TC % 8 == 0);
__builtin_assume(TC > 0);

In that case, we would have created the following SCEV:

8 * (umax(1, TC) / 8)

since umax(1, TC) may still be between [1,7], when we divide it by 8, and multiply by 8, we get 0. So this doesn't preserve the fact that TC > 0.
In order to do that we must align up 1 to 8, and then umax(8, TC) / 8 is always > 0.

[caojoshua]

In D141850#4064581, @alonkom wrote:

In D141850#4064477, @caojoshua wrote:

The improved trip multiples from the test results look good. Ordering in applyLoopGuards is an issue. However, I think we can simplify this down a bit. What if we always applied min/max first, before we apply divisibility guards? For example, given:

__builtin_assume(TC % 8 == 0);
__builtin_assume(TC > 0);
__builtin_assume(TC < 100);

1. apply max: umax(1, TC)
2. apply min: umin(100, umax(1, TC))
3. apply divisibility info: 8 * (umin(100, umax(1, TC))) / 8

This makes divisibility info obvious. And traversing the SCEV, we can still see TC > 0 and TC < 100. I believe that if we always apply max/min first, we will never lose this info. This approach seems much simpler and easier to understand.

FYI. I haven't been around that long, and this change is non-trivial enough where I prefer a longstanding developer to give the final approval. I will still be around to give my thoughts.

Thanks for the reply.
This works in this example, but what would happen if we had only these assumes:

__builtin_assume(TC % 8 == 0);
__builtin_assume(TC > 0);

In that case, we would have created the following SCEV:

8 * (umax(1, TC) / 8)

since umax(1, TC) may still be between [1,7], when we divide it by 8, and multiply by 8, we get 0. So this doesn't preserve the fact that TC > 0.
In order to do that we must align up 1 to 8, and then umax(8, TC) / 8 is always > 0.

umax(1, TC) implies TC <= 1, which holds in this case.

I think what you meant is in the case that, TC > 0, we get: 8 * (umin(1, TC)) / 8. If TC is in [0, 7], the expression evaluates to 0, and TC > 0 is lost.

Your point comes across that we can't just apply guards in a certain order. I have some thoughts, but I'll think about it and write it down later.

[caojoshua]

After this patch, depending on the assume processing order applyLoopGuards could create the following SCEV:
max(min((8 * (TC / 8)) , 96), 8)

This example looks wrong. I think min/max should be switched. Should be

min(max((8 * (TC / 8)), 96), 8)

Please update description.

---

In terms of overall approach, I'm not sure. Feels a bit hacky to have custom logic to check that an expressions is a min/max of mul/div. I'll let others chime in here.

llvm/lib/Analysis/ScalarEvolution.cpp

15163

typo: wether

I'm not sure what the divisor B in \p DividesBy means.

I think this paragraph needs to be more clear. What does it mean to be composed on Min/Max SCEVs?

[mkazantsev]

Some nits, and potentially a bug in formula.

llvm/lib/Analysis/ScalarEvolution.cpp

15042

I'd prefer it to be a separate patch, if it's legal at all. Need to check carefully how poison flows in these formulae.

15045

MinMax can have more than 2 operands. Check that it is exactly 2 of them?

15054

greater or equal

15057

If I'm reading this correctly, this is a guarantee of no-overflow for computations you are doing. Maybe add this comment explicitly?

15062

This computation is inconsistent. If Rem is constant 0, you'll return Expr. But if Rem is effectively zero, but not a constant (e.g. some complex expression which is always zero), you return Expr + Divisor. Bug?

15068

less or equal

15076

I think you can safely drop check for Rem->isZero() here as it will be trivially simplified away in getMinusSCEV

15084

auto

15165

auto

15179

if (auto *MinMax = dyn_cast<SCEVMinMaxExpr>(Expr))

15188

auto

[alonkom]

In D141850#4072675, @caojoshua wrote:

After this patch, depending on the assume processing order applyLoopGuards could create the following SCEV:
max(min((8 * (TC / 8)) , 96), 8)

This example looks wrong. I think min/max should be switched. Should be

min(max((8 * (TC / 8)), 96), 8)

Please update description.

---

In terms of overall approach, I'm not sure. Feels a bit hacky to have custom logic to check that an expressions is a min/max of mul/div. I'll let others chime in here.

I think the example is correct.
Even before this patch:
TC > 0 is translated to max (TC, 1)
TC < 99 is translated to min (TC, 98)

llvm/lib/Analysis/ScalarEvolution.cpp

15042

They are not created in this function, so I prefer only handling MinMax at this point.

15042

can you explain what isn't legal here?
just checking if this is a min/max of 2 operands, when the 2nd one is constant.

15043

it is used as an output of the function

15057

I will only handle constants for now.

15062

I will only support constants for now

15062

I will only support constants for now

15084

can't use auto since there's a recursive call here

15163

rephrased it a bit.

15165

can't use auto since there's a recursive call here

15188

can't use auto, due to recursive call.

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

15042

I mean "I would prefer adding sequential min/max in a separate patch". Meaning, so far so good here.

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

15165

auto HasDivisibiltyInfo = [&](const SCEV *Expr, const SCEV *&DividesBy) -> bool {...
?

[caojoshua]

I think the example is correct.
Even before this patch:
TC > 0 is translated to max (TC, 1)
TC < 99 is translated to min (TC, 98)

My mistake. You're correct.

[alonkom]

llvm/lib/Analysis/ScalarEvolution.cpp

15165

I'm still getting this build error:
use of ‘HasDivisibiltyInfo’ before deduction of ‘auto’

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

15224

How do you account for RHS being SINT_MIN and similar cases here?

[mkazantsev]

llvm/test/Analysis/ScalarEvolution/trip-multiple-guard-info.ll

397

Please precommit tests as is, I want to see what exactly this patch changes.

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

15216

auto *

[alonkom]

llvm/test/Analysis/ScalarEvolution/trip-multiple-guard-info.ll

397

https://reviews.llvm.org/D143337

[mkazantsev]

I think supporting smin/smax is a bug, this code is only written for unsigned values.

llvm/lib/Analysis/ScalarEvolution.cpp

15042

nit: auto *

15064

Do you check bit width anywhere? What if it doesn't fit into unsigned?

Better use APInt unless you have a reason not to.

15100

Init with nullptr, makes it easier to debug.

15105

Do you really support smin here? Your code is only correct for unsigned values. It interprets all values as non-negative.

[alonkom]

llvm/lib/Analysis/ScalarEvolution.cpp

15224

I wonder how this worked before my patch.
if we have assume (N < SINT_MIN) then the following SCEV was generated:
smin(N, SINT_MIN - 1) which would overflow.

[alonkom]

llvm/lib/Analysis/ScalarEvolution.cpp

15105

I check if the operands are non-negative inside GetPreviousSCEV/GetNextSCEVD

[mkazantsev]

Mostly looks good, but I think I found a bug in IsMinMaxSCEVWithConstant . Also please rebase on top of your tests.

llvm/lib/Analysis/ScalarEvolution.cpp

15048

1. If I remember correctly, SCEV operands are always sorted by kind. It means that constants always go first. RHS can't be a constant because then LHS is also a constant and the whole thing should be folded away. So it should be trivially false, and I guess this code part isn't covered by tests if it wasn't caught.
2. Usually matcher methods aren't supposed to change operands if they return false. In your case it doesn't matter, but generally better to follow the standard practice and do all checks before modifying the operands.

15105

Ah ok, then it makes sense. Can we assert on that here?

15196

if (auto *MinMax = dyn_cast<SCEVMinMaxExpr>(Expr)) { ...

15224

Well, if there is a check somewhere that RHS is non-negative, then it makes sense. Otherwise it's possibly broken. :)

llvm/test/Analysis/ScalarEvolution/trip-multiple-guard-info.ll

397

Rebase patch on top of it?

[alonkom]

@mkazantsev Let me know if you have any other comments

llvm/lib/Analysis/ScalarEvolution.cpp

15048

1. Nice catch! It didn't hurt the divisibility info, but some min/max info was lost. I Fixed this and added a unit-test since the lit only checks the divisibility.
2. Done.

15224

Anyway, the change I've made only applies to non-negative values

[mkazantsev]

Ok, thanks, let's give it a try.

llvm/lib/Analysis/ScalarEvolution.cpp

15069

nit: {} not needed

15081–15088

This code is copy-paste-ish (see lambda above), maybe factor out?

15095

If you see any problems with compile time from this patch, this is a potential place where things can be limited. Not sure if it's worth it in practice.

15108

nit: auto*

15135

Is it possible that the map contains both LHSUnknown -> RewrittenLHS and RewrittenLHS -> SomeOtherLHS? Should this be a loop?

I'm ok if it's a separate patch.

[alonkom]

llvm/lib/Analysis/ScalarEvolution.cpp

15095

I don't expect more than 3 assumes per value in the real-world, so I don't believe it'll affect compilation time drastically, but let's see.

15135

The entire function assumes 1 level of nesting

[caojoshua]

Why was this reverted?

[alonkom]

In D141850#4152680, @caojoshua wrote:

Why was this reverted?

Failed on an assertion. Uploading a fix now.
(Also some buildbots failed on time-outs, but I'm not sure this is due to my change)

llvm/lib/Analysis/ScalarEvolution.cpp

auto MatchRangeCheckIdiom = [this, Predicate, LHS, RHS, &RewriteMap,
           getConstant(ExactRegion.getUnsignedMin()),
           getUMinExpr(RewrittenLHS, getConstant(ExactRegion.getUnsignedMax())));
       ExprsToRewrite.push_back(LHSUnknown);
       return true;
     };
     if (MatchRangeCheckIdiom())
       return;
 
+    // Return true if \p Expr is a MinMax SCEV expression with a constant
+    // operand. If so, return in \p SCTy the SCEV type and in \p RHS the
+    // non-constant operand and in \p LHS the constant operand.
+    auto IsMinMaxSCEVWithConstant = [&](const SCEV *Expr, SCEVTypes &SCTy,
+                                        const SCEV *&LHS, const SCEV *&RHS) {
+      if (auto *MinMax = dyn_cast<SCEVMinMaxExpr>(Expr)) {
+        if (MinMax->getNumOperands() != 2)
+          return false;
+        SCTy = MinMax->getSCEVType();
+        if (!isa<SCEVConstant>(MinMax->getOperand(0)))
+          return false;
+        LHS = MinMax->getOperand(0);
+        RHS = MinMax->getOperand(1);
+        return true;
+      }
+      return false;
+    };
+
+    // Checks whether Expr is a non-negative constant, and Divisor is a positive
+    // constant, and returns their APInt in ExprVal and in DivisorVal.
+    auto GetNonNegExprAndPosDivisor = [&](const SCEV *Expr, const SCEV *Divisor,
+                                          APInt &ExprVal, APInt &DivisorVal) {
+      if (!isKnownNonNegative(Expr) || !isKnownPositive(Divisor))
+        return false;
+      auto *ConstExpr = dyn_cast<SCEVConstant>(Expr);
+      auto *ConstDivisor = dyn_cast<SCEVConstant>(Divisor);
+      if (!ConstExpr || !ConstDivisor)
+        return false;
+      ExprVal = ConstExpr->getAPInt();
+      DivisorVal = ConstDivisor->getAPInt();
+      return true;
+    };
+
+    // Return a new SCEV that modifies \p Expr to the closest number divides by
+    // \p Divisor and greater or equal than Expr.
+    // For now, only handle constant Expr and Divisor.
+    auto GetNextSCEVDividesByDivisor = [&](const SCEV *Expr,
+                                           const SCEV *Divisor) {
+      APInt ExprVal;
+      APInt DivisorVal;
+      if (!GetNonNegExprAndPosDivisor(Expr, Divisor, ExprVal, DivisorVal))
+        return Expr;
+      APInt Rem = ExprVal.urem(DivisorVal);
+      if (!Rem.isZero())
+        // return the SCEV: Expr + Divisor - Expr % Divisor
+        return getConstant(ExprVal + DivisorVal - Rem);
+      return Expr;
+    };
+
+    // Return a new SCEV that modifies \p Expr to the closest number divides by
+    // \p Divisor and less or equal than Expr.
+    // For now, only handle constant Expr and Divisor.
+    auto GetPreviousSCEVDividesByDivisor = [&](const SCEV *Expr,
+                                               const SCEV *Divisor) {
+      APInt ExprVal;
+      APInt DivisorVal;
+      if (!GetNonNegExprAndPosDivisor(Expr, Divisor, ExprVal, DivisorVal))
+        return Expr;
+      APInt Rem = ExprVal.urem(DivisorVal);
+      // return the SCEV: Expr - Expr % Divisor
+      return getConstant(ExprVal - Rem);
+    };
+
+    // Apply divisibilty by \p Divisor on MinMaxExpr with constant values,
+    // recursively. This is done by aligning up/down the constant value to the
+    // Divisor.
+    std::function<const SCEV *(const SCEV *, const SCEV *)>
+        ApplyDivisibiltyOnMinMaxExpr = [&](const SCEV *MinMaxExpr,
+                                           const SCEV *Divisor) {
+          const SCEV *MinMaxLHS = nullptr, *MinMaxRHS = nullptr;
+          SCEVTypes SCTy;
+          if (!IsMinMaxSCEVWithConstant(MinMaxExpr, SCTy, MinMaxLHS, MinMaxRHS))
+            return MinMaxExpr;
+          auto IsMin =
+              isa<SCEVSMinExpr>(MinMaxExpr) || isa<SCEVUMinExpr>(MinMaxExpr);
+          assert(isKnownNonNegative(MinMaxLHS) &&
+                 "Expected non-negative operand!");
+          auto *DivisibleExpr =
+              IsMin ? GetPreviousSCEVDividesByDivisor(MinMaxLHS, Divisor)
+                    : GetNextSCEVDividesByDivisor(MinMaxLHS, Divisor);
+          SmallVector<const SCEV *> Ops = {
+              ApplyDivisibiltyOnMinMaxExpr(MinMaxRHS, Divisor), DivisibleExpr};
+          return getMinMaxExpr(SCTy, Ops);
+        };
+
     // If we have LHS == 0, check if LHS is computing a property of some unknown
     // SCEV %v which we can rewrite %v to express explicitly.
     const SCEVConstant *RHSC = dyn_cast<SCEVConstant>(RHS);
     if (Predicate == CmpInst::ICMP_EQ && RHSC &&
         RHSC->getValue()->isNullValue()) {
       // If LHS is A % B, i.e. A % B == 0, rewrite A to (A /u B) * B to
       // explicitly express that.
       const SCEV *URemLHS = nullptr;
       const SCEV *URemRHS = nullptr;
       if (matchURem(LHS, URemLHS, URemRHS)) {
         if (const SCEVUnknown *LHSUnknown = dyn_cast<SCEVUnknown>(URemLHS)) {
-          const auto *Multiple = getMulExpr(getUDivExpr(URemLHS, URemRHS), URemRHS);
+          auto I = RewriteMap.find(LHSUnknown);
+          const SCEV *RewrittenLHS =
+              I != RewriteMap.end() ? I->second : LHSUnknown;
+          RewrittenLHS = ApplyDivisibiltyOnMinMaxExpr(RewrittenLHS, URemRHS);
+          const auto *Multiple =
+              getMulExpr(getUDivExpr(RewrittenLHS, URemRHS), URemRHS);
           RewriteMap[LHSUnknown] = Multiple;
           ExprsToRewrite.push_back(LHSUnknown);
           return;
         }
       }
     }
 
     // Do not apply information for constants or if RHS contains an AddRec.
     if (isa<SCEVConstant>(LHS) || containsAddRecurrence(RHS))
       return;
 
     // If RHS is SCEVUnknown, make sure the information is applied to it.
     if (!isa<SCEVUnknown>(LHS) && isa<SCEVUnknown>(RHS)) {
       std::swap(LHS, RHS);
       Predicate = CmpInst::getSwappedPredicate(Predicate);
     }
 
     // Check whether LHS has already been rewritten. In that case we want to
     // chain further rewrites onto the already rewritten value.
     auto I = RewriteMap.find(LHS);
     const SCEV *RewrittenLHS = I != RewriteMap.end() ? I->second : LHS;
 
+    // Check for the SCEV expression (A /u B) * B while B is a constant, inside
+    // \p Expr. The check is done recuresively on \p Expr, which is assumed to
+    // be a composition of Min/Max SCEVs. Return whether the SCEV expression (A
+    // /u B) * B was found, and return the divisor B in \p DividesBy. For
+    // example, if Expr = umin (umax ((A /u 8) * 8, 16), 64), return true since
+    // (A /u 8) * 8 matched the pattern, and return the constant SCEV 8 in \p
+    // DividesBy.
+    std::function<bool(const SCEV *, const SCEV *&)> HasDivisibiltyInfo =
+        [&](const SCEV *Expr, const SCEV *&DividesBy) {
+          if (auto *Mul = dyn_cast<SCEVMulExpr>(Expr)) {
+            if (Mul->getNumOperands() != 2)
+              return false;
+            auto *MulLHS = Mul->getOperand(0);
+            auto *MulRHS = Mul->getOperand(1);
+            if (isa<SCEVConstant>(MulLHS))
+              std::swap(MulLHS, MulRHS);
+            if (auto *Div = dyn_cast<SCEVUDivExpr>(MulLHS)) {
+              if (Div->getOperand(1) == MulRHS) {
+                DividesBy = MulRHS;
+                return true;
+              }
+            }
+          }
+          if (auto *MinMax = dyn_cast<SCEVMinMaxExpr>(Expr)) {
+            return HasDivisibiltyInfo(MinMax->getOperand(0), DividesBy) ||
+                   HasDivisibiltyInfo(MinMax->getOperand(1), DividesBy);
+          }
+          return false;
+        };
+
+    // Return true if Expr known to divide by \p DividesBy.
+    std::function<bool(const SCEV *, const SCEV *&)> IsKnownToDivideBy =
+        [&](const SCEV *Expr, const SCEV *DividesBy) {
+          if (getURemExpr(Expr, DividesBy)->isZero())
+            return true;
+          if (auto *MinMax = dyn_cast<SCEVMinMaxExpr>(Expr)) {
+            return IsKnownToDivideBy(MinMax->getOperand(0), DividesBy) &&
+                   IsKnownToDivideBy(MinMax->getOperand(1), DividesBy);
+          }
+          return false;
+        };
+
+    const SCEV *DividesBy = nullptr;
+    if (HasDivisibiltyInfo(RewrittenLHS, DividesBy))
+      // Check that the whole expression is divided by DividesBy
+      DividesBy =
+          IsKnownToDivideBy(RewrittenLHS, DividesBy) ? DividesBy : nullptr;
+
     const SCEV *RewrittenRHS = nullptr;
     switch (Predicate) {
     case CmpInst::ICMP_ULT: {
       if (RHS->getType()->isPointerTy())
         break;
       const SCEV *One = getOne(RHS->getType());
-      RewrittenRHS =
-          getUMinExpr(RewrittenLHS, getMinusSCEV(getUMaxExpr(RHS, One), One));
+      auto *ModifiedRHS = getMinusSCEV(getUMaxExpr(RHS, One), One);
+      ModifiedRHS =
+          DividesBy ? GetPreviousSCEVDividesByDivisor(ModifiedRHS, DividesBy)
+                    : ModifiedRHS;
+      RewrittenRHS = getUMinExpr(RewrittenLHS, ModifiedRHS);
       break;
     }
-    case CmpInst::ICMP_SLT:
-      RewrittenRHS =
-          getSMinExpr(RewrittenLHS, getMinusSCEV(RHS, getOne(RHS->getType())));
+    case CmpInst::ICMP_SLT: {
+      auto *ModifiedRHS = getMinusSCEV(RHS, getOne(RHS->getType()));
+      ModifiedRHS =
+          DividesBy ? GetPreviousSCEVDividesByDivisor(ModifiedRHS, DividesBy)
+                    : ModifiedRHS;
+      RewrittenRHS = getSMinExpr(RewrittenLHS, ModifiedRHS);
       break;
-    case CmpInst::ICMP_ULE:
-      RewrittenRHS = getUMinExpr(RewrittenLHS, RHS);
+    }
+    case CmpInst::ICMP_ULE: {
+      auto *ModifiedRHS =
+          DividesBy ? GetPreviousSCEVDividesByDivisor(RHS, DividesBy) : RHS;
+      RewrittenRHS = getUMinExpr(RewrittenLHS, ModifiedRHS);
       break;
-    case CmpInst::ICMP_SLE:
-      RewrittenRHS = getSMinExpr(RewrittenLHS, RHS);
+    }
+    case CmpInst::ICMP_SLE: {
+      auto *ModifiedRHS =
+          DividesBy ? GetPreviousSCEVDividesByDivisor(RHS, DividesBy) : RHS;
+      RewrittenRHS = getSMinExpr(RewrittenLHS, ModifiedRHS);
       break;
-    case CmpInst::ICMP_UGT:
-      RewrittenRHS =
-          getUMaxExpr(RewrittenLHS, getAddExpr(RHS, getOne(RHS->getType())));
+    }
+    case CmpInst::ICMP_UGT: {
+      auto *ModifiedRHS = getAddExpr(RHS, getOne(RHS->getType()));
+      ModifiedRHS = DividesBy
+                        ? GetNextSCEVDividesByDivisor(ModifiedRHS, DividesBy)
+                        : ModifiedRHS;
+      RewrittenRHS = getUMaxExpr(RewrittenLHS, ModifiedRHS);
       break;
-    case CmpInst::ICMP_SGT:
-      RewrittenRHS =
-          getSMaxExpr(RewrittenLHS, getAddExpr(RHS, getOne(RHS->getType())));
+    }
+    case CmpInst::ICMP_SGT: {
+      auto *ModifiedRHS = getAddExpr(RHS, getOne(RHS->getType()));
+      ModifiedRHS = DividesBy
+                        ? GetNextSCEVDividesByDivisor(ModifiedRHS, DividesBy)
+                        : ModifiedRHS;
+      RewrittenRHS = getSMaxExpr(RewrittenLHS, ModifiedRHS);
       break;
-    case CmpInst::ICMP_UGE:
-      RewrittenRHS = getUMaxExpr(RewrittenLHS, RHS);
+    }
+    case CmpInst::ICMP_UGE: {
+      auto *ModifiedRHS =
+          DividesBy ? GetNextSCEVDividesByDivisor(RHS, DividesBy) : RHS;
+      RewrittenRHS = getUMaxExpr(RewrittenLHS, ModifiedRHS);
       break;
-    case CmpInst::ICMP_SGE:
-      RewrittenRHS = getSMaxExpr(RewrittenLHS, RHS);
+    }
+    case CmpInst::ICMP_SGE: {
+      auto *ModifiedRHS =
+          DividesBy ? GetNextSCEVDividesByDivisor(RHS, DividesBy) : RHS;
+      RewrittenRHS = getSMaxExpr(RewrittenLHS, ModifiedRHS);
       break;
+    }
     case CmpInst::ICMP_EQ:
       if (isa<SCEVConstant>(RHS))
         RewrittenRHS = RHS;
       break;
     case CmpInst::ICMP_NE:
       if (isa<SCEVConstant>(RHS) &&
-          cast<SCEVConstant>(RHS)->getValue()->isNullValue())
-        RewrittenRHS = getUMaxExpr(RewrittenLHS, getOne(RHS->getType()));
+          cast<SCEVConstant>(RHS)->getValue()->isNullValue()) {
+        auto *ModifiedRHS = getOne(RHS->getType());
+        ModifiedRHS = DividesBy
+                          ? GetNextSCEVDividesByDivisor(ModifiedRHS, DividesBy)
+                          : ModifiedRHS;
+        RewrittenRHS = getUMaxExpr(RewrittenLHS, ModifiedRHS);
+      }
       break;
     default:
       break;
     }
 
     if (RewrittenRHS) {
       RewriteMap[LHS] = RewrittenRHS;
       if (LHS == RewrittenLHS)

llvm/test/Analysis/ScalarEvolution/trip-multiple-guard-info.ll

 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_ugt_5_order_swapped
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 2
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %u = urem i32 %num, 4
   %cmp.1 = icmp ugt i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   %cmp = icmp eq i32 %u, 0
   tail call void @llvm.assume(i1 %cmp)
   br label %for.body

 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_sgt_5_order_swapped
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is 2147483646
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 2
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %u = urem i32 %num, 4
   %cmp.1 = icmp sgt i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   %cmp = icmp eq i32 %u, 0
   tail call void @llvm.assume(i1 %cmp)
   br label %for.body

 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_uge_5_order_swapped
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 1
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %u = urem i32 %num, 4
   %cmp = icmp eq i32 %u, 0
   %cmp.1 = icmp uge i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   tail call void @llvm.assume(i1 %cmp)
   br label %for.body

 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_sge_5_order_swapped
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is 2147483646
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 1
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %u = urem i32 %num, 4
   %cmp = icmp eq i32 %u, 0
   %cmp.1 = icmp sge i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   tail call void @llvm.assume(i1 %cmp)
   br label %for.body
for.body:
   %inc = add nuw nsw i32 %i.010, 1
   %cmp2 = icmp ult i32 %inc, %num
   br i1 %cmp2, label %for.body, label %exit
 
 exit:
   ret void
 }
 
 define void @test_trip_multiple_4_upper_lower_bounds(i32 %num) {
 ; CHECK-LABEL: 'test_trip_multiple_4_upper_lower_bounds'
 ; CHECK-NEXT:  Classifying expressions for: @test_trip_multiple_4_upper_lower_bounds
 ; CHECK-NEXT:    %u = urem i32 %num, 4
 ; CHECK-NEXT:    --> (zext i2 (trunc i32 %num to i2) to i32) U: [0,4) S: [0,4)
 ; CHECK-NEXT:    %i.010 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
 ; CHECK-NEXT:    --> {0,+,1}<nuw><nsw><%for.body> U: [0,-2147483648) S: [0,-2147483648) Exits: (-1 + %num) LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_upper_lower_bounds
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 1
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %cmp.1 = icmp uge i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   %u = urem i32 %num, 4
   %cmp = icmp eq i32 %u, 0
   tail call void @llvm.assume(i1 %cmp)
   %cmp.2 = icmp ult i32 %num, 59000

 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_upper_lower_bounds_swapped1
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 1
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %cmp.1 = icmp uge i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   %u = urem i32 %num, 4
   %cmp = icmp eq i32 %u, 0
   tail call void @llvm.assume(i1 %cmp)
   %cmp.2 = icmp ult i32 %num, 59000

 ; CHECK-NEXT:    %inc = add nuw nsw i32 %i.010, 1
 ; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%for.body> U: [1,-2147483648) S: [1,-2147483648) Exits: %num LoopDispositions: { %for.body: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_trip_multiple_4_upper_lower_bounds_swapped2
 ; CHECK-NEXT:  Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: constant max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: symbolic max backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %for.body: Trip multiple is 1
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %cmp.1 = icmp uge i32 %num, 5
   tail call void @llvm.assume(i1 %cmp.1)
   %cmp.2 = icmp ult i32 %num, 59000
   tail call void @llvm.assume(i1 %cmp.2)
   %u = urem i32 %num, 4
   %cmp = icmp eq i32 %u, 0

llvm/unittests/Analysis/ScalarEvolutionTest.cpp

runWithSE(*M, "foo", [](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
     auto *ScevIV = SE.getSCEV(getInstructionByName(F, "iv"));
     const Loop *L = cast<SCEVAddRecExpr>(ScevIV)->getLoop();
 
     const SCEV *ITC = SE.getConstantMaxTripCountFromArray(L);
     EXPECT_TRUE(isa<SCEVCouldNotCompute>(ITC));
   });
 }
 
+TEST_F(ScalarEvolutionsTest, ApplyLoopGuards) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(
+      "declare void @llvm.assume(i1)\n"
+      "define void @test(i32 %num) {\n"
+      "entry:\n"
+      "  %u = urem i32 %num, 4\n"
+      "  %cmp = icmp eq i32 %u, 0\n"
+      "  tail call void @llvm.assume(i1 %cmp)\n"
+      "  %cmp.1 = icmp ugt i32 %num, 0\n"
+      "  tail call void @llvm.assume(i1 %cmp.1)\n"
+      "  br label %for.body\n"
+      "for.body:\n"
+      "  %i.010 = phi i32 [ 0, %entry ], [ %inc, %for.body ]\n"
+      "  %inc = add nuw nsw i32 %i.010, 1\n"
+      "  %cmp2 = icmp ult i32 %inc, %num\n"
+      "  br i1 %cmp2, label %for.body, label %exit\n"
+      "exit:\n"
+      "  ret void\n"
+      "}\n",
+      Err, C);
+
+  ASSERT_TRUE(M && "Could not parse module?");
+  ASSERT_TRUE(!verifyModule(*M) && "Must have been well formed!");
+
+  runWithSE(*M, "test", [](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+    auto *TCScev = SE.getSCEV(getArgByName(F, "num"));
+    auto *ApplyLoopGuardsTC = SE.applyLoopGuards(TCScev, *LI.begin());
+    // Assert that the new TC is (4 * ((4 umax %num) /u 4))
+    APInt Four(32, 4);
+    auto *Constant4 = SE.getConstant(Four);
+    auto *Max = SE.getUMaxExpr(TCScev, Constant4);
+    auto *Mul = SE.getMulExpr(SE.getUDivExpr(Max, Constant4), Constant4);
+    ASSERT_TRUE(Mul == ApplyLoopGuardsTC);
+  });
+}
+
 }  // end namespace llvm