[SCEV] Preserve information in applyLoopGuards

[alonkom]

Summary

Preserve both the information on divisibility of an expression (Expr%M == 0) and minimum value of the expression (Expr > N) regardless of their order.
Fixing #55645

[alonkom]

Ping. Can you please address this 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 Max SCEV expression with a constant
+    // operand. If so, return in \p SCTy the SCEV type and in \p LHS the
+    // non-constant operand and in \p RHS the constant operand
+    auto IsMaxSCEVWithConstant = [&](const SCEV *Expr, SCEVTypes &SCTy,
+                                     const SCEV *&LHS, const SCEV *&RHS) {
+      auto Max = dyn_cast<SCEVMinMaxExpr>(Expr);
+      if (!Max)
+        return false;
+      SCTy = cast<SCEVMinMaxExpr>(Expr)->getSCEVType();
+      if (SCTy != scUMaxExpr && SCTy != scSMaxExpr)
+        return false;
+      LHS = Max->getOperand(0);
+      RHS = Max->getOperand(1);
+      if (isa<SCEVConstant>(LHS))
+        std::swap(LHS, RHS);
+      return isa<SCEVConstant>(RHS);
+    };
+
+    // Return a new SCEV that aligns up \p Expr to the multiple of by \p
+    // Divisor.
+    auto AlignUpToMultipleOfDivisor = [&](const SCEV *Expr,
+                                          const SCEV *Divisor) {
+      if (!isKnownNonNegative(Expr) || !isKnownPositive(Divisor))
+        return Expr;
+      auto Rem = getURemExpr(Expr, Divisor);
+      if (!Rem->isZero()) {
+        // return the SCEV: Expr + Divisor - Expr % Divisor
+        return getAddExpr(getMinusSCEV(Divisor, Rem), Expr);
+      }
+      return Expr;
+    };
+
     // 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)) {
-          auto Multiple = getMulExpr(getUDivExpr(URemLHS, URemRHS), URemRHS);
+          // Check whether LHSUnknown has already been rewritten.
+          // In that case we want to maintiain both the divisibility property,
+          // and the maximum property.
+          auto I = RewriteMap.find(LHSUnknown);
+          const SCEV *RewrittenLHS =
+              I != RewriteMap.end() ? I->second : LHSUnknown;
+          const SCEV *MaxLHS, *MaxRHS;
+          SCEVTypes SCTy;
+          if (IsMaxSCEVWithConstant(RewrittenLHS, SCTy, MaxLHS, MaxRHS)) {
+            auto AlignedUp = AlignUpToMultipleOfDivisor(MaxRHS, URemRHS);
+            SmallVector<const SCEV *> Ops = {MaxLHS, AlignedUp};
+            RewrittenLHS = getMinMaxExpr(SCTy, Ops);
+          }
+          auto Multiple =
+              getMulExpr(getUDivExpr(RewrittenLHS, URemRHS), URemRHS);
           RewriteMap[LHSUnknown] = Multiple;
+          if (RewrittenLHS == LHSUnknown)
             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);
     }
 
     // Limit to expressions that can be rewritten.
     if (!isa<SCEVUnknown>(LHS) && !isa<SCEVZeroExtendExpr>(LHS))
       return;
 
     // 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;
 
+    // Return true if Expr is (A /u B) * B while B is a constant, and return B
+    // in DividesBy.
+    auto IsKnownToDivideBy = [&](const SCEV *Expr, const SCEV *&DividesBy) {
+      if (auto Mul = dyn_cast<SCEVMulExpr>(Expr)) {
+        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;
+          }
+        }
+      }
+      return false;
+    };
+
+    const SCEV *DividesBy = nullptr;
+    IsKnownToDivideBy(RewrittenLHS, DividesBy);
+
     const SCEV *RewrittenRHS = nullptr;
     switch (Predicate) {
     case CmpInst::ICMP_ULT:
       RewrittenRHS =
           getUMinExpr(RewrittenLHS, getMinusSCEV(RHS, getOne(RHS->getType())));
       break;
     case CmpInst::ICMP_SLT:
       RewrittenRHS =
           getSMinExpr(RewrittenLHS, getMinusSCEV(RHS, getOne(RHS->getType())));
       break;
     case CmpInst::ICMP_ULE:
       RewrittenRHS = getUMinExpr(RewrittenLHS, RHS);
       break;
     case CmpInst::ICMP_SLE:
       RewrittenRHS = getSMinExpr(RewrittenLHS, RHS);
       break;
-    case CmpInst::ICMP_UGT:
-      RewrittenRHS =
-          getUMaxExpr(RewrittenLHS, getAddExpr(RHS, getOne(RHS->getType())));
-      break;
-    case CmpInst::ICMP_SGT:
-      RewrittenRHS =
-          getSMaxExpr(RewrittenLHS, getAddExpr(RHS, getOne(RHS->getType())));
-      break;
-    case CmpInst::ICMP_UGE:
-      RewrittenRHS = getUMaxExpr(RewrittenLHS, RHS);
-      break;
-    case CmpInst::ICMP_SGE:
-      RewrittenRHS = getSMaxExpr(RewrittenLHS, RHS);
+    case CmpInst::ICMP_UGT: {
+      auto ModifiedRHS = getAddExpr(RHS, getOne(RHS->getType()));
+      ModifiedRHS = DividesBy
+                        ? AlignUpToMultipleOfDivisor(ModifiedRHS, DividesBy)
+                        : ModifiedRHS;
+      RewrittenRHS = getUMaxExpr(RewrittenLHS, ModifiedRHS);
+      break;
+    }
+    case CmpInst::ICMP_SGT: {
+      auto ModifiedRHS = getAddExpr(RHS, getOne(RHS->getType()));
+      ModifiedRHS = DividesBy
+                        ? AlignUpToMultipleOfDivisor(ModifiedRHS, DividesBy)
+                        : ModifiedRHS;
+      RewrittenRHS = getSMaxExpr(RewrittenLHS, ModifiedRHS);
+      break;
+    }
+    case CmpInst::ICMP_UGE: {
+      auto ModifiedRHS =
+          DividesBy ? AlignUpToMultipleOfDivisor(RHS, DividesBy) : RHS;
+      RewrittenRHS = getUMaxExpr(RewrittenLHS, ModifiedRHS);
+      break;
+    }
+    case CmpInst::ICMP_SGE: {
+      auto ModifiedRHS =
+          DividesBy ? AlignUpToMultipleOfDivisor(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()));

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

for.body:
   %cmp2 = icmp ult i32 %inc, %num
   br i1 %cmp2, label %for.body, label %exit
 
 exit:
   ret void
 }
 
 define void @test_trip_multiple_4_ugt_5_order_swapped(i32 %num) {
-; TODO: Trip multiple can be 4, it is missed due to the processing order of the assumes.
 ; CHECK-LABEL: @test_trip_multiple_4_ugt_5_order_swapped
 ; CHECK:       Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
-; 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
for.body:
   %cmp2 = icmp slt i32 %inc, %num
   br i1 %cmp2, label %for.body, label %exit
 
 exit:
   ret void
 }
 
 define void @test_trip_multiple_4_sgt_5_order_swapped(i32 %num) {
-; TODO: Trip multiple can be 4, it is missed due to the processing order of the assumes.
 ; CHECK-LABEL: @test_trip_multiple_4_sgt_5_order_swapped
 ; CHECK:       Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: max backedge-taken count is 2147483646
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
-; 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
for.body:
   %cmp2 = icmp ult i32 %inc, %num
   br i1 %cmp2, label %for.body, label %exit
 
 exit:
   ret void
 }
 
 define void @test_trip_multiple_4_uge_5_order_swapped(i32 %num) {
-; TODO: Trip multiple can be 4, it is missed due to the processing order of the assumes.
 ; CHECK-LABEL: @test_trip_multiple_4_uge_5_order_swapped
 ; CHECK:       Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
-; 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
 
 for.body:
   %i.010 = phi i32 [ 0, %entry ], [ %inc, %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_sge_5(i32 %num) {
-; TODO: Trip multiple can be 4, it is missed due to the processing order of the assumes.
 ; CHECK-LABEL: @test_trip_multiple_4_sge_5
 ; CHECK:       Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: max backedge-taken count is 2147483646
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
 ; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %u = urem i32 %num, 4
exit:
   ret void
 }
 
 define void @test_trip_multiple_4_sge_5_order_swapped(i32 %num) {
 ; CHECK-LABEL: @test_trip_multiple_4_sge_5_order_swapped
 ; CHECK:       Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: max backedge-taken count is 2147483646
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
-; 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