[llvm] [ConstantRange] Improve `shlWithNoWrap` (PR #101800)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/101800

Closes https://github.com/dtcxzyw/llvm-tools/issues/22.


>From c10b3d4d4ff0ed15766f5872dc6fc740dca7a42d Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 3 Aug 2024 15:54:57 +0800
Subject: [PATCH 1/3] [ConstantRange] Add pre-commit tests. NFC.

---
 .../Transforms/CorrelatedValuePropagation/shl.ll  | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll b/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll
index 8b4dbc98425bf..b5b943a20bff2 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll
@@ -474,3 +474,18 @@ define i1 @shl_nuw_nsw_test4(i32 %x, i32 range(i32 0, 32) %k) {
   %cmp = icmp eq i64 %shl, -9223372036854775808
   ret i1 %cmp
 }
+
+define i1 @shl_nuw_nsw_test5(i32 %x) {
+; CHECK-LABEL: @shl_nuw_nsw_test5(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SHL:%.*]] = shl nuw nsw i32 768, [[X:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw i32 [[SHL]], 1846
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[ADD]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+entry:
+  %shl = shl nuw nsw i32 768, %x
+  %add = add nuw i32 %shl, 1846
+  %cmp = icmp sgt i32 %add, 0
+  ret i1 %cmp
+}

>From 763bea5cd1f1b5a7005fb2ac9aea0da6227c5d79 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 3 Aug 2024 16:04:46 +0800
Subject: [PATCH 2/3] [ConstantRange] Improve shlWithNoWarp

---
 llvm/lib/IR/ConstantRange.cpp                 | 28 +++++++++++++++----
 .../CorrelatedValuePropagation/shl.ll         |  7 ++---
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index 50b211a302e8f..24c86641f76df 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -1624,12 +1624,30 @@ ConstantRange ConstantRange::shlWithNoWrap(const ConstantRange &Other,
     return getEmpty();
 
   ConstantRange Result = shl(Other);
+  KnownBits Known = toKnownBits();
+
+  if (NoWrapKind & OverflowingBinaryOperator::NoSignedWrap) {
+    ConstantRange ShAmtRange = Other;
+    if (isAllNonNegative())
+      ShAmtRange = ShAmtRange.intersectWith(
+          ConstantRange(APInt::getZero(getBitWidth()),
+                        APInt(getBitWidth(), Known.countMaxLeadingZeros())),
+          Unsigned);
+    else if (isAllNegative())
+      ShAmtRange = ShAmtRange.intersectWith(
+          ConstantRange(APInt::getZero(getBitWidth()),
+                        APInt(getBitWidth(), Known.countMaxLeadingOnes())),
+          Unsigned);
+    Result = Result.intersectWith(sshl_sat(ShAmtRange), RangeType);
+  }
 
-  if (NoWrapKind & OverflowingBinaryOperator::NoSignedWrap)
-    Result = Result.intersectWith(sshl_sat(Other), RangeType);
-
-  if (NoWrapKind & OverflowingBinaryOperator::NoUnsignedWrap)
-    Result = Result.intersectWith(ushl_sat(Other), RangeType);
+  if (NoWrapKind & OverflowingBinaryOperator::NoUnsignedWrap) {
+    ConstantRange ShAmtRange =
+        getNonEmpty(APInt::getZero(getBitWidth()),
+                    APInt(getBitWidth(), Known.countMaxLeadingZeros() + 1));
+    Result = Result.intersectWith(
+        ushl_sat(Other.intersectWith(ShAmtRange, Unsigned)), RangeType);
+  }
 
   return Result;
 }
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll b/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll
index b5b943a20bff2..a55081e1604e6 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/shl.ll
@@ -105,7 +105,7 @@ exit:
 define i8 @test5(i8 %b) {
 ; CHECK-LABEL: @test5(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw nsw i8 0, [[B:%.*]]
-; CHECK-NEXT:    ret i8 [[SHL]]
+; CHECK-NEXT:    ret i8 0
 ;
   %shl = shl i8 0, %b
   ret i8 %shl
@@ -479,9 +479,8 @@ define i1 @shl_nuw_nsw_test5(i32 %x) {
 ; CHECK-LABEL: @shl_nuw_nsw_test5(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw nsw i32 768, [[X:%.*]]
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw i32 [[SHL]], 1846
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[ADD]], 0
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[SHL]], 1846
+; CHECK-NEXT:    ret i1 true
 ;
 entry:
   %shl = shl nuw nsw i32 768, %x

>From 9f7da466a27b14629831b9e3d8a3d8ea3b51df1b Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 3 Aug 2024 16:09:16 +0800
Subject: [PATCH 3/3] [ConstantRange] Early exit

---
 llvm/lib/IR/ConstantRange.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index 24c86641f76df..062e0e7a3b251 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -1624,6 +1624,9 @@ ConstantRange ConstantRange::shlWithNoWrap(const ConstantRange &Other,
     return getEmpty();
 
   ConstantRange Result = shl(Other);
+  if (!NoWrapKind)
+    return Result;
+
   KnownBits Known = toKnownBits();
 
   if (NoWrapKind & OverflowingBinaryOperator::NoSignedWrap) {