Add additional known bits info for self multiply

Author: Macsen Casaus

llvm/lib/Support/KnownBits.cpp

@@ -888,11 +888,30 @@ 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);
+  // Self multiplying
+  if (NoUndefSelfMultiply) {
+    // bit[1] is guaranteed to be zero.
+    if (BitWidth > 1) {
+      assert(Res.One[1] == 0 &&
+             "Self-multiplication failed Quadratic Reciprocity!");
+      Res.Zero.setBit(1);
+    }
+
+    // If input bit[0] is set, then output bit[2] is guaranteed to be zero.
+    if (BitWidth > 2 && LHS.One[0])
+      Res.Zero.setBit(2);
+
+    // If input bit[0] is clear, then output bit[3] is guaranteed to be zero.
+    if (BitWidth > 3 && LHS.Zero[0])
+      Res.Zero.setBit(3);
+
+    // If input % 4 == 2, then output bit[3] and bit[4] are guarantted to be
+    // zero.
+    if (BitWidth > 3 && LHS.Zero[0] && LHS.One[1]) {
+      Res.Zero.setBit(3);
+      if (BitWidth > 4)
+        Res.Zero.setBit(4);
+    }
   }
 
   return Res;

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

@@ -1246,10 +1246,7 @@ define i1 @extract_value_smul_fail(i8 %xx, i8 %yy) {
 
 define i8 @known_self_mul_bit_0_set(i8 noundef %x) {
 ; CHECK-LABEL: @known_self_mul_bit_0_set(
-; CHECK-NEXT:    [[BIT_0_SET:%.*]] = or i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[SELF_MUL:%.*]] = mul i8 [[BIT_0_SET]], [[BIT_0_SET]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[SELF_MUL]], 4
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    ret i8 0
 ;
   %bit_0_set = or i8 %x, 1
   %self_mul = mul i8 %bit_0_set, %bit_0_set
@@ -1259,10 +1256,7 @@ define i8 @known_self_mul_bit_0_set(i8 noundef %x) {
 
 define i8 @known_self_mul_bit_0_unset(i8 noundef %x) {
 ; CHECK-LABEL: @known_self_mul_bit_0_unset(
-; CHECK-NEXT:    [[BIT_0_UNSET:%.*]] = and i8 [[X:%.*]], -2
-; CHECK-NEXT:    [[SELF_MUL:%.*]] = mul i8 [[BIT_0_UNSET]], [[BIT_0_UNSET]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[SELF_MUL]], 8
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    ret i8 0
 ;
   %bit_0_unset = and i8 %x, -2
   %self_mul = mul i8 %bit_0_unset, %bit_0_unset
@@ -1272,11 +1266,7 @@ define i8 @known_self_mul_bit_0_unset(i8 noundef %x) {
 
 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:    [[LOWER_2_UNSET:%.*]] = and i8 [[X:%.*]], -4
-; CHECK-NEXT:    [[BIT_1_SET_BIT_0_UNSET:%.*]] = or disjoint i8 [[LOWER_2_UNSET]], 2
-; CHECK-NEXT:    [[SELF_MUL:%.*]] = mul i8 [[BIT_1_SET_BIT_0_UNSET]], [[BIT_1_SET_BIT_0_UNSET]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[SELF_MUL]], 24
-; CHECK-NEXT:    ret i8 [[R]]
+; 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
@@ -1287,11 +1277,7 @@ define i8 @known_self_mul_bit_1_set_bit_0_unset(i8 noundef %x) {
 
 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:    [[LOWER_2_UNSET:%.*]] = and i4 [[X:%.*]], -4
-; CHECK-NEXT:    [[BIT_1_SET_BIT_0_UNSET:%.*]] = or disjoint i4 [[LOWER_2_UNSET]], 2
-; CHECK-NEXT:    [[SELF_MUL:%.*]] = mul i4 [[BIT_1_SET_BIT_0_UNSET]], [[BIT_1_SET_BIT_0_UNSET]]
-; CHECK-NEXT:    [[R:%.*]] = and i4 [[SELF_MUL]], -8
-; CHECK-NEXT:    ret i4 [[R]]
+; 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