[InstCombine] Add additional known bits info for self multiply (#151202)

Closes #151043 
https://alive2.llvm.org/ce/z/JyMSk8

Author: macsencasaus

llvm/lib/Support/KnownBits.cpp

@@ -888,11 +888,19 @@ KnownBits KnownBits::mul(const KnownBits &LHS, const KnownBits &RHS,
   Res.Zero |= (~BottomKnown).getLoBits(ResultBitsKnown);
   Res.One = BottomKnown.getLoBits(ResultBitsKnown);
 
-  // If we're self-multiplying then bit[1] is guaranteed to be zero.
-  if (NoUndefSelfMultiply && BitWidth > 1) {
-    assert(Res.One[1] == 0 &&
-           "Self-multiplication failed Quadratic Reciprocity!");
-    Res.Zero.setBit(1);
+  if (NoUndefSelfMultiply) {
+    // If X has at least TZ trailing zeroes, then bit (2 * TZ + 1) must be zero.
+    unsigned TwoTZP1 = 2 * TrailZero0 + 1;
+    if (TwoTZP1 < BitWidth)
+      Res.Zero.setBit(TwoTZP1);
+
+    // If X has exactly TZ trailing zeros, then bit (2 * TZ + 2) must also be
+    // zero.
+    if (TrailZero0 < BitWidth && LHS.One[TrailZero0]) {
+      unsigned TwoTZP2 = TwoTZP1 + 1;
+      if (TwoTZP2 < BitWidth)
+        Res.Zero.setBit(TwoTZP2);
+    }
   }
 
   return Res;

llvm/test/Transforms/InstCombine/known-bits.ll

@@ -1244,6 +1244,48 @@ define i1 @extract_value_smul_fail(i8 %xx, i8 %yy) {
   ret i1 %r
 }
 
+define i8 @known_self_mul_bit_0_set(i8 noundef %x) {
+; CHECK-LABEL: @known_self_mul_bit_0_set(
+; CHECK-NEXT:    ret i8 0
+;
+  %bit_0_set = or i8 %x, 1
+  %self_mul = mul i8 %bit_0_set, %bit_0_set
+  %r = and i8 %self_mul, 4
+  ret i8 %r
+}
+
+define i8 @known_self_mul_bit_0_unset(i8 noundef %x) {
+; CHECK-LABEL: @known_self_mul_bit_0_unset(
+; CHECK-NEXT:    ret i8 0
+;
+  %bit_0_unset = and i8 %x, -2
+  %self_mul = mul i8 %bit_0_unset, %bit_0_unset
+  %r = and i8 %self_mul, 8
+  ret i8 %r
+}
+
+define i8 @known_self_mul_bit_1_set_bit_0_unset(i8 noundef %x) {
+; CHECK-LABEL: @known_self_mul_bit_1_set_bit_0_unset(
+; CHECK-NEXT:    ret i8 0
+;
+  %lower_2_unset = and i8 %x, -4
+  %bit_1_set_bit_0_unset = or disjoint i8 %lower_2_unset, 2
+  %self_mul = mul i8 %bit_1_set_bit_0_unset, %bit_1_set_bit_0_unset
+  %r = and i8 %self_mul, 24
+  ret i8 %r
+}
+
+define i4 @known_self_mul_bit_1_set_bit_0_unset_i4(i4 noundef %x) {
+; CHECK-LABEL: @known_self_mul_bit_1_set_bit_0_unset_i4(
+; CHECK-NEXT:    ret i4 0
+;
+  %lower_2_unset = and i4 %x, -4
+  %bit_1_set_bit_0_unset = or disjoint i4 %lower_2_unset, 2
+  %self_mul = mul i4 %bit_1_set_bit_0_unset, %bit_1_set_bit_0_unset
+  %r = and i4 %self_mul, 24
+  ret i4 %r
+}
+
 define i8 @known_reduce_or(<2 x i8> %xx) {
 ; CHECK-LABEL: @known_reduce_or(
 ; CHECK-NEXT:    ret i8 1