[llvm] 2ba7822 - [SCEV] Try smaller ZExts when using loop guard info.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2023-06-09T20:05:50+01:00
New Revision: 2ba78229e42e1f22fc6356389ab1660ad77a3f59

URL: https://github.com/llvm/llvm-project/commit/2ba78229e42e1f22fc6356389ab1660ad77a3f59
DIFF: https://github.com/llvm/llvm-project/commit/2ba78229e42e1f22fc6356389ab1660ad77a3f59.diff

LOG: [SCEV] Try smaller ZExts when using loop guard info.

If we didn't find the extact ZExt expr in the rewrite map, check if
there's an entry for a smaller ZExt we can use instead.

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D149786

Added: 
    

Modified: 
    llvm/lib/Analysis/ScalarEvolution.cpp
    llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 7272fdee26d24..901ce76f80b33 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -14958,9 +14958,25 @@ class SCEVLoopGuardRewriter : public SCEVRewriteVisitor<SCEVLoopGuardRewriter> {
 
   const SCEV *visitZeroExtendExpr(const SCEVZeroExtendExpr *Expr) {
     auto I = Map.find(Expr);
-    if (I == Map.end())
+    if (I == Map.end()) {
+      // If we didn't find the extact ZExt expr in the map, check if there's an
+      // entry for a smaller ZExt we can use instead.
+      Type *Ty = Expr->getType();
+      const SCEV *Op = Expr->getOperand(0);
+      unsigned Bitwidth = Ty->getScalarSizeInBits() / 2;
+      while (Bitwidth % 8 == 0 && Bitwidth >= 8 &&
+             Bitwidth > Op->getType()->getScalarSizeInBits()) {
+        Type *NarrowTy = IntegerType::get(SE.getContext(), Bitwidth);
+        auto *NarrowExt = SE.getZeroExtendExpr(Op, NarrowTy);
+        auto I = Map.find(NarrowExt);
+        if (I != Map.end())
+          return SE.getZeroExtendExpr(I->second, Ty);
+        Bitwidth = Bitwidth / 2;
+      }
+
       return SCEVRewriteVisitor<SCEVLoopGuardRewriter>::visitZeroExtendExpr(
           Expr);
+    }
     return I->second;
   }
 

diff  --git a/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info.ll b/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info.ll
index c0704d7423ae6..8e0b29cc62d14 100644
--- a/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info.ll
+++ b/llvm/test/Analysis/ScalarEvolution/backedge-taken-count-guard-info.ll
@@ -9,14 +9,14 @@ define void @loop_guard_improves_exact_backedge_taken_count_1(i32 %conv) {
 ; CHECK-NEXT:    %conv8 = zext i32 %and to i64
 ; CHECK-NEXT:    --> (zext i1 (trunc i32 %conv to i1) to i64) U: [0,2) S: [0,2)
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
-; CHECK-NEXT:    --> {0,+,1}<nuw><nsw><%loop> U: [0,2) S: [0,2) Exits: (zext i1 (trunc i32 %conv to i1) to i64) LoopDispositions: { %loop: Computable }
+; CHECK-NEXT:    --> {0,+,1}<nuw><nsw><%loop> U: [0,1) S: [0,1) Exits: 0 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
-; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%loop> U: [1,3) S: [1,3) Exits: (1 + (zext i1 (trunc i32 %conv to i1) to i64))<nuw><nsw> LoopDispositions: { %loop: Computable }
+; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%loop> U: [1,2) S: [1,2) Exits: 1 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:  Determining loop execution counts for: @loop_guard_improves_exact_backedge_taken_count_1
-; CHECK-NEXT:  Loop %loop: backedge-taken count is (zext i1 (trunc i32 %conv to i1) to i64)
-; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is 1
-; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is (zext i1 (trunc i32 %conv to i1) to i64)
-; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is (zext i1 (trunc i32 %conv to i1) to i64)
+; CHECK-NEXT:  Loop %loop: backedge-taken count is 0
+; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is 0
+; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is 0
+; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 0
 ; CHECK-NEXT:   Predicates:
 ; CHECK:       Loop %loop: Trip multiple is 1
 ;
@@ -54,7 +54,7 @@ define void @loop_guard_improves_exact_backedge_taken_count_2(i32 %conv) {
 ; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is (zext i1 (trunc i32 %conv to i1) to i64)
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is (zext i1 (trunc i32 %conv to i1) to i64)
 ; CHECK-NEXT:   Predicates:
-; CHECK:       Loop %loop: Trip multiple is 1
+; CHECK:       Loop %loop: Trip multiple is 2
 ;
 entry:
   %and = and i32 %conv, 1