[SCEV] Rewrite some SCEVAdd sub-expressions using loop guards. (#156013)

Trip count expressions sometimes consist of adding 3 operands, i.e.
(Const + A + B). There may be guard info for A + B, and if so, apply it.

We can probably more generally apply this, but need to be careful w.r.t
compile-time.

Alive2 Proof for changes in miniters.ll:
https://alive2.llvm.org/ce/z/HFfXOx

Fixes #155941.

PR: #156013

Author: fhahn

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -16009,6 +16009,16 @@ const SCEV *ScalarEvolution::LoopGuards::rewrite(const SCEV *Expr) const {
     }
 
     const SCEV *visitAddExpr(const SCEVAddExpr *Expr) {
+      // Trip count expressions sometimes consist of adding 3 operands, i.e.
+      // (Const + A + B). There may be guard info for A + B, and if so, apply
+      // it.
+      // TODO: Could more generally apply guards to Add sub-expressions.
+      if (isa<SCEVConstant>(Expr->getOperand(0)) &&
+          Expr->getNumOperands() == 3) {
+        if (const SCEV *S = Map.lookup(
+                SE.getAddExpr(Expr->getOperand(1), Expr->getOperand(2))))
+          return SE.getAddExpr(Expr->getOperand(0), S);
+      }
       SmallVector<const SCEV *, 2> Operands;
       bool Changed = false;
       for (const auto *Op : Expr->operands()) {

llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info-apply-to-adds.ll

@@ -4,9 +4,9 @@
 define void @ptrtoint_based_trip_count_known_via_guards_applied_to_add_subexpr(ptr %start, ptr %end) {
 ; CHECK-LABEL: 'ptrtoint_based_trip_count_known_via_guards_applied_to_add_subexpr'
 ; CHECK-NEXT:  Determining loop execution counts for: @ptrtoint_based_trip_count_known_via_guards_applied_to_add_subexpr
-; CHECK-NEXT:  Loop %loop: backedge-taken count is ((-4 + (-1 * (ptrtoint ptr %start to i64)) + (ptrtoint ptr %end to i64)) /u 4)
-; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is i64 4611686018427387903
-; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is ((-4 + (-1 * (ptrtoint ptr %start to i64)) + (ptrtoint ptr %end to i64)) /u 4)
+; CHECK-NEXT:  Loop %loop: backedge-taken count is i64 0
+; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is i64 0
+; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is i64 0
 ; CHECK-NEXT:  Loop %loop: Trip multiple is 1
 ;
 entry:

llvm/test/Analysis/ScalarEvolution/max-backedge-taken-count-guard-info-apply-to-adds.ll

@@ -5,7 +5,7 @@ define void @max_btc_improved_by_applying_guards_to_add_subexpr(i32 %low, i32 %h
 ; CHECK-LABEL: 'max_btc_improved_by_applying_guards_to_add_subexpr'
 ; CHECK-NEXT:  Determining loop execution counts for: @max_btc_improved_by_applying_guards_to_add_subexpr
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is (-1 + (zext i32 (1 + (-1 * %low) + %high) to i64))<nsw>
-; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is i64 -1
+; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is i64 7
 ; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is (-1 + (zext i32 (1 + (-1 * %low) + %high) to i64))<nsw>
 ; CHECK-NEXT:  Loop %loop: Trip multiple is 1
 ;

llvm/test/Transforms/LoopVectorize/miniters.ll

@@ -61,8 +61,7 @@ define void @min_iters_known_via_loop_guards_add(i32 %start, i32 %end, ptr %src)
 ; CHECK-NEXT:    [[ADD_1:%.*]] = add i32 [[SUB]], 1
 ; CHECK-NEXT:    [[IV_START:%.*]] = zext i32 [[ADD_1]] to i64
 ; CHECK-NEXT:    [[TMP0:%.*]] = sub i64 101, [[IV_START]]
-; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 [[TMP0]], 4
-; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], [[SCALAR_PH:label %.*]], label %[[VECTOR_PH:.*]]
+; CHECK-NEXT:    br i1 false, [[SCALAR_PH:label %.*]], label %[[VECTOR_PH:.*]]
 ; CHECK:       [[VECTOR_PH]]:
 ;
 ; UNROLL-LABEL: define void @min_iters_known_via_loop_guards_add(
@@ -74,8 +73,7 @@ define void @min_iters_known_via_loop_guards_add(i32 %start, i32 %end, ptr %src)
 ; UNROLL-NEXT:    [[ADD_1:%.*]] = add i32 [[SUB]], 1
 ; UNROLL-NEXT:    [[IV_START:%.*]] = zext i32 [[ADD_1]] to i64
 ; UNROLL-NEXT:    [[TMP0:%.*]] = sub i64 101, [[IV_START]]
-; UNROLL-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 [[TMP0]], 8
-; UNROLL-NEXT:    br i1 [[MIN_ITERS_CHECK]], [[SCALAR_PH:label %.*]], label %[[VECTOR_PH:.*]]
+; UNROLL-NEXT:    br i1 false, [[SCALAR_PH:label %.*]], label %[[VECTOR_PH:.*]]
 ; UNROLL:       [[VECTOR_PH]]:
 ;
 entry: