From c2d18d700500991eff15ac674dacdeb3e1d8c21d Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Mon, 8 Mar 2021 11:44:15 +0000
Subject: [PATCH] [KnownBits] Add min/max shift amount handling to
 shl/lshr/ashr KnownBits helpers

Pulled out of the original D90479 patch - also includes the "impossible shift amount" filtering from computeKnownBitsFromShiftOperator.

Differential Revision: https://reviews.llvm.org/D90479
---
 llvm/lib/Support/KnownBits.cpp                | 75 +++++++++++++++++++
 .../Transforms/InstSimplify/icmp-constant.ll  |  2 +-
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
index d7265d03d27d..6c7aaad968f5 100644
--- a/llvm/lib/Support/KnownBits.cpp
+++ b/llvm/lib/Support/KnownBits.cpp
@@ -187,6 +187,31 @@ KnownBits KnownBits::shl(const KnownBits &LHS, const KnownBits &RHS) {
     MinTrailingZeros = std::min(MinTrailingZeros, BitWidth);
   }
 
+  // If the maximum shift is in range, then find the common bits from all
+  // possible shifts.
+  APInt MaxShiftAmount = RHS.getMaxValue();
+  if (MaxShiftAmount.ult(BitWidth) && !LHS.isUnknown()) {
+    uint64_t ShiftAmtZeroMask = (~RHS.Zero).getZExtValue();
+    uint64_t ShiftAmtOneMask = RHS.One.getZExtValue();
+    assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
+    Known.Zero.setAllBits();
+    Known.One.setAllBits();
+    for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
+                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+         ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
+      // Skip if the shift amount is impossible.
+      if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
+          (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
+        continue;
+      KnownBits SpecificShift;
+      SpecificShift.Zero = LHS.Zero << ShiftAmt;
+      SpecificShift.One = LHS.One << ShiftAmt;
+      Known = KnownBits::commonBits(Known, SpecificShift);
+      if (Known.isUnknown())
+        break;
+    }
+  }
+
   Known.Zero.setLowBits(MinTrailingZeros);
   return Known;
 }
@@ -215,6 +240,31 @@ KnownBits KnownBits::lshr(const KnownBits &LHS, const KnownBits &RHS) {
     MinLeadingZeros = std::min(MinLeadingZeros, BitWidth);
   }
 
+  // If the maximum shift is in range, then find the common bits from all
+  // possible shifts.
+  APInt MaxShiftAmount = RHS.getMaxValue();
+  if (MaxShiftAmount.ult(BitWidth) && !LHS.isUnknown()) {
+    uint64_t ShiftAmtZeroMask = (~RHS.Zero).getZExtValue();
+    uint64_t ShiftAmtOneMask = RHS.One.getZExtValue();
+    assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
+    Known.Zero.setAllBits();
+    Known.One.setAllBits();
+    for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
+                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+         ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
+      // Skip if the shift amount is impossible.
+      if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
+          (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
+        continue;
+      KnownBits SpecificShift = LHS;
+      SpecificShift.Zero.lshrInPlace(ShiftAmt);
+      SpecificShift.One.lshrInPlace(ShiftAmt);
+      Known = KnownBits::commonBits(Known, SpecificShift);
+      if (Known.isUnknown())
+        break;
+    }
+  }
+
   Known.Zero.setHighBits(MinLeadingZeros);
   return Known;
 }
@@ -248,6 +298,31 @@ KnownBits KnownBits::ashr(const KnownBits &LHS, const KnownBits &RHS) {
     }
   }
 
+  // If the maximum shift is in range, then find the common bits from all
+  // possible shifts.
+  APInt MaxShiftAmount = RHS.getMaxValue();
+  if (MaxShiftAmount.ult(BitWidth) && !LHS.isUnknown()) {
+    uint64_t ShiftAmtZeroMask = (~RHS.Zero).getZExtValue();
+    uint64_t ShiftAmtOneMask = RHS.One.getZExtValue();
+    assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
+    Known.Zero.setAllBits();
+    Known.One.setAllBits();
+    for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
+                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+         ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
+      // Skip if the shift amount is impossible.
+      if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
+          (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
+        continue;
+      KnownBits SpecificShift = LHS;
+      SpecificShift.Zero.ashrInPlace(ShiftAmt);
+      SpecificShift.One.ashrInPlace(ShiftAmt);
+      Known = KnownBits::commonBits(Known, SpecificShift);
+      if (Known.isUnknown())
+        break;
+    }
+  }
+
   Known.Zero.setHighBits(MinLeadingZeros);
   Known.One.setHighBits(MinLeadingOnes);
   return Known;
diff --git a/llvm/test/Transforms/InstSimplify/icmp-constant.ll b/llvm/test/Transforms/InstSimplify/icmp-constant.ll
index f8d5bdc89ceb..9db63e64b914 100644
--- a/llvm/test/Transforms/InstSimplify/icmp-constant.ll
+++ b/llvm/test/Transforms/InstSimplify/icmp-constant.ll
@@ -799,7 +799,7 @@ define i1 @eq_shl_by_constant_produces_poison(i8 %x) {
 
 define i1 @eq_shl_by_variable_produces_poison(i8 %x) {
 ; CHECK-LABEL: @eq_shl_by_variable_produces_poison(
-; CHECK-NEXT:    ret i1 false
+; CHECK-NEXT:    ret i1 poison
 ;
   %clear_high_bit = and i8 %x, 127                 ; 0x7f
   %set_next_high_bits = or i8 %clear_high_bit, 112 ; 0x70