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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-analysis

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

<details>
<summary>Changes</summary>

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 https://github.com/llvm/llvm-project/issues/155941

---
Full diff: https://github.com/llvm/llvm-project/pull/156013.diff


4 Files Affected:

- (modified) llvm/lib/Analysis/ScalarEvolution.cpp (+10) 
- (modified) llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info-apply-to-adds.ll (+3-3) 
- (modified) llvm/test/Analysis/ScalarEvolution/max-backedge-taken-count-guard-info-apply-to-adds.ll (+1-1) 
- (modified) llvm/test/Transforms/LoopVectorize/miniters.ll (+2-4) 


``````````diff
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index aa2bcf7917537..27e6e7e8806c3 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/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()) {
diff --git a/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info-apply-to-adds.ll b/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info-apply-to-adds.ll
index ba859f2e3eec9..75014f3a58eb6 100644
--- a/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info-apply-to-adds.ll
+++ b/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:
diff --git a/llvm/test/Analysis/ScalarEvolution/max-backedge-taken-count-guard-info-apply-to-adds.ll b/llvm/test/Analysis/ScalarEvolution/max-backedge-taken-count-guard-info-apply-to-adds.ll
index 635126c9262cf..951b07272dd4b 100644
--- a/llvm/test/Analysis/ScalarEvolution/max-backedge-taken-count-guard-info-apply-to-adds.ll
+++ b/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
 ;
diff --git a/llvm/test/Transforms/LoopVectorize/miniters.ll b/llvm/test/Transforms/LoopVectorize/miniters.ll
index a0fd48d510f24..6d06a03d0d018 100644
--- a/llvm/test/Transforms/LoopVectorize/miniters.ll
+++ b/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:

``````````

</details>


https://github.com/llvm/llvm-project/pull/156013