[ConstantRange] Improve shlWithNoWrap

[dtcxzyw]

Closes dtcxzyw/llvm-tools#22.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-ir

Author: Yingwei Zheng (dtcxzyw)

Changes

Closes dtcxzyw/llvm-tools#22.

---

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

2 Files Affected:

• (modified) llvm/lib/IR/ConstantRange.cpp (+25-4)
• (modified) llvm/test/Transforms/CorrelatedValuePropagation/shl.ll (+15-1)

diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index 50b211a302e8f..062e0e7a3b251 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -1624,12 +1624,33 @@ ConstantRange ConstantRange::shlWithNoWrap(const ConstantRange &Other,
     return getEmpty();
 
   ConstantRange Result = shl(Other);
+  if (!NoWrapKind)
+    return Result;
 
-  if (NoWrapKind & OverflowingBinaryOperator::NoSignedWrap)
-    Result = Result.intersectWith(sshl_sat(Other), RangeType);
+  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::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 8b4dbc98425bf..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
@@ -474,3 +474,17 @@ 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 nsw i32 [[SHL]], 1846
+; CHECK-NEXT:    ret i1 true
+;
+entry:
+  %shl = shl nuw nsw i32 768, %x
+  %add = add nuw i32 %shl, 1846
+  %cmp = icmp sgt i32 %add, 0
+  ret i1 %cmp
+}

[antoniofrighetto]

LGTM, thanks!

[nikic]

Does this make the results optimal in some cases? If so, please adjust the exhaustive test. If not, please add some unit tests to cover some cases that are expected to work.

[dtcxzyw]

Do you guys have a plan to treat empty ranges as poison?

[AZero13]

Is this a miscompile or?

[dtcxzyw]

Compile-time impact: http://llvm-compile-time-tracker.com/compare.php?from=23609a383c3a0a389339ec6f9784d130bf2f0994&to=a9efe755565400ee488aa4fc294d6a9a5de9fe1a&stat=instructions%3Au

[goldsteinn]

What is this logic doing? How is LHSMax << RHSMaxWithoutOV not the max value if we have nuw?

[dtcxzyw]

Consider the following case:

%x = [3, 5)
%y = Full
%shl = shl nuw i8 %x, %y

We have

4 << ctlz(4) = 128
3 << ctlz(3) = 192

[goldsteinn]

Ahh makes sense. If the range spans a power of 2, won't the MaxShl always be (Pow2Floor(LHSMax) - 1) << AsHighAsItCanGet? In which case you can drop the APIntOps::umax

[dtcxzyw]

Compile-time impact: http://llvm-compile-time-tracker.com/compare.php?from=23609a383c3a0a389339ec6f9784d130bf2f0994&to=677d8eea077a3fdd9f0dde377663efdcb71d4449&stat=instructions%3Au

[nikic]

LGTM