[ConstantRange] makeGuaranteedNoWrapRegion(): `shl` support

Summary:
If all the shifts amount are already poison-producing,
then we can add more poison-producing flags ontop:
https://rise4fun.com/Alive/Ocwi

Otherwise, we should only consider the possible range of shift amts that don't result in poison.

For unsigned range not not overflow, we must not shift out any set bits,
and the actual limit for `x` can be computed by backtransforming
the maximal value we could ever get out of the `shl` - `-1` through
`lshr`. If the `x` is any larger than that then it will overflow.

Likewise for signed range, but just in signed domain..

This is based on the general idea outlined by @nikic in https://reviews.llvm.org/D68672#1714990

Reviewers: nikic, sanjoy

Reviewed By: nikic

Subscribers: hiraditya, llvm-commits, nikic

Tags: #llvm

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@375370 91177308-0d34-0410-b5e6-96231b3b80d8

lib/IR/ConstantRange.cpp

@@ -269,6 +269,27 @@ ConstantRange::makeGuaranteedNoWrapRegion(Instruction::BinaryOps BinOp,
 
     return makeExactMulNSWRegion(Other.getSignedMin())
         .intersectWith(makeExactMulNSWRegion(Other.getSignedMax()));
+
+  case Instruction::Shl: {
+    // For given range of shift amounts, if we ignore all illegal shift amounts
+    // (that always produce poison), what shift amount range is left?
+    ConstantRange ShAmt = Other.intersectWith(
+        ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, (BitWidth - 1) + 1)));
+    if (ShAmt.isEmptySet()) {
+      // If the entire range of shift amounts is already poison-producing,
+      // then we can freely add more poison-producing flags ontop of that.
+      return getFull(BitWidth);
+    }
+    // There are some legal shift amounts, we can compute conservatively-correct
+    // range of no-wrap inputs. Note that by now we have clamped the ShAmtUMax
+    // to be at most bitwidth-1, which results in most conservative range.
+    APInt ShAmtUMax = ShAmt.getUnsignedMax();
+    if (Unsigned)
+      return getNonEmpty(APInt::getNullValue(BitWidth),
+                         APInt::getMaxValue(BitWidth).lshr(ShAmtUMax) + 1);
+    return getNonEmpty(APInt::getSignedMinValue(BitWidth).ashr(ShAmtUMax),
+                       APInt::getSignedMaxValue(BitWidth).ashr(ShAmtUMax) + 1);
+  }
   }
 }
 

unittests/IR/ConstantRangeTest.cpp

@@ -1532,6 +1532,68 @@ TEST(ConstantRange, MakeGuaranteedNoWrapRegion) {
   EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
                 Instruction::Sub, One, OBO::NoUnsignedWrap),
             ConstantRange(APInt::getMinValue(32) + 1, APInt::getMinValue(32)));
+
+  ConstantRange OneLessThanBitWidth(APInt(32, 0), APInt(32, 31) + 1);
+  ConstantRange UpToBitWidth(APInt(32, 0), APInt(32, 32) + 1);
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, UpToBitWidth, OBO::NoUnsignedWrap),
+            ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, OneLessThanBitWidth, OBO::NoUnsignedWrap));
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, UpToBitWidth, OBO::NoSignedWrap),
+            ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, OneLessThanBitWidth, OBO::NoSignedWrap));
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, UpToBitWidth, OBO::NoUnsignedWrap),
+            ConstantRange(APInt(32, 0), APInt(32, 1) + 1));
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, UpToBitWidth, OBO::NoSignedWrap),
+            ConstantRange(APInt(32, -1), APInt(32, 0) + 1));
+
+  EXPECT_EQ(
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, ConstantRange::getFull(32), OBO::NoUnsignedWrap),
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, OneLessThanBitWidth, OBO::NoUnsignedWrap));
+  EXPECT_EQ(
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, ConstantRange::getFull(32), OBO::NoSignedWrap),
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, OneLessThanBitWidth, OBO::NoSignedWrap));
+
+  ConstantRange IllegalShAmt(APInt(32, 32), APInt(32, 0) + 1);
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, IllegalShAmt, OBO::NoUnsignedWrap),
+            ConstantRange::getFull(32));
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl, IllegalShAmt, OBO::NoSignedWrap),
+            ConstantRange::getFull(32));
+
+  EXPECT_EQ(
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, ConstantRange(APInt(32, -32), APInt(32, 16) + 1),
+          OBO::NoUnsignedWrap),
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, ConstantRange(APInt(32, 0), APInt(32, 16) + 1),
+          OBO::NoUnsignedWrap));
+  EXPECT_EQ(
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, ConstantRange(APInt(32, -32), APInt(32, 16) + 1),
+          OBO::NoSignedWrap),
+      ConstantRange::makeGuaranteedNoWrapRegion(
+          Instruction::Shl, ConstantRange(APInt(32, 0), APInt(32, 16) + 1),
+          OBO::NoSignedWrap));
+
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl,
+                ConstantRange(APInt(32, -32), APInt(32, 16) + 1),
+                OBO::NoUnsignedWrap),
+            ConstantRange(APInt(32, 0), APInt(32, 65535) + 1));
+  EXPECT_EQ(ConstantRange::makeGuaranteedNoWrapRegion(
+                Instruction::Shl,
+                ConstantRange(APInt(32, -32), APInt(32, 16) + 1),
+                OBO::NoSignedWrap),
+            ConstantRange(APInt(32, -32768), APInt(32, 32767) + 1));
 }
 
 template<typename Fn>
@@ -1541,6 +1603,8 @@ void TestNoWrapRegionExhaustive(Instruction::BinaryOps BinOp,
   EnumerateConstantRanges(Bits, [&](const ConstantRange &CR) {
     if (CR.isEmptySet())
       return;
+    if (Instruction::isShift(BinOp) && CR.getUnsignedMax().uge(Bits))
+      return;
 
     ConstantRange NoWrap =
         ConstantRange::makeGuaranteedNoWrapRegion(BinOp, CR, NoWrapKind);
@@ -1609,6 +1673,20 @@ TEST(ConstantRange, NoWrapRegionExhaustive) {
         (void) N1.smul_ov(N2, Overflow);
         return Overflow;
       });
+  TestNoWrapRegionExhaustive(Instruction::Shl,
+                             OverflowingBinaryOperator::NoUnsignedWrap,
+                             [](const APInt &N1, const APInt &N2) {
+                               bool Overflow;
+                               (void)N1.ushl_ov(N2, Overflow);
+                               return Overflow;
+                             });
+  TestNoWrapRegionExhaustive(Instruction::Shl,
+                             OverflowingBinaryOperator::NoSignedWrap,
+                             [](const APInt &N1, const APInt &N2) {
+                               bool Overflow;
+                               (void)N1.sshl_ov(N2, Overflow);
+                               return Overflow;
+                             });
 }
 
 TEST(ConstantRange, GetEquivalentICmp) {