[ValueTracking] Use select condition to help infer bits of arms

If we have something like `(select (icmp ult x, 8), x, y)`, we can use
the `(icmp ult x, 8)` to help compute the knownbits of `x`.

Author: goldsteinn

llvm/lib/Analysis/ValueTracking.cpp

@@ -1023,11 +1023,27 @@ static void computeKnownBitsFromOperator(const Operator *I,
     break;
   }
   case Instruction::Select: {
-    computeKnownBits(I->getOperand(2), Known, Depth + 1, Q);
-    computeKnownBits(I->getOperand(1), Known2, Depth + 1, Q);
-
+    auto ComputeForArm = [&](Value *Arm, bool Invert) {
+      KnownBits Res(Known.getBitWidth());
+      computeKnownBits(Arm, Res, Depth + 1, Q);
+      if (isGuaranteedNotToBeUndef(Arm, Q.AC, Q.CxtI, Q.DT, Depth + 1)) {
+        // See what condition implies about the bits of the two select arms.
+        computeKnownBitsFromCond(Arm, I->getOperand(0), Res, Depth + 1, Q,
+                                 Invert);
+        // We can have conflict if the condition is dead. I.e if we have
+        // (x | 64) < 32 ? (x | 64) : y
+        // we will have conflict at bit 6 from the condition/the `or`.
+        // In that case, we just reset. Its not particularly important
+        // what we do, as this select is going to be simplified soon.
+        if (Res.hasConflict())
+          Res.resetAll();
+      }
+      return Res;
+    };
     // Only known if known in both the LHS and RHS.
-    Known = Known.intersectWith(Known2);
+    Known =
+        ComputeForArm(I->getOperand(1), /*Invert=*/false)
+            .intersectWith(ComputeForArm(I->getOperand(2), /*Invert=*/true));
     break;
   }
   case Instruction::FPTrunc:

llvm/test/Analysis/ValueTracking/knownbits-select-from-cond.ll

@@ -6,7 +6,7 @@ define i8 @select_condition_implies_highbits_op1(i8 %xx, i8 noundef %y) {
 ; CHECK-NEXT:    [[X:%.*]] = and i8 [[XX:%.*]], 15
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ult i8 [[Y:%.*]], 3
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i8 [[Y]], i8 [[X]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[SEL]], 32
+; CHECK-NEXT:    [[R:%.*]] = or disjoint i8 [[SEL]], 32
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x = and i8 %xx, 15
@@ -36,7 +36,7 @@ define i8 @select_condition_implies_highbits_op2(i8 %xx, i8 noundef %y) {
 ; CHECK-NEXT:    [[X:%.*]] = and i8 [[XX:%.*]], 15
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ugt i8 [[Y:%.*]], 3
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i8 [[X]], i8 [[Y]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[SEL]], 32
+; CHECK-NEXT:    [[R:%.*]] = or disjoint i8 [[SEL]], 32
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x = and i8 %xx, 15
@@ -52,7 +52,7 @@ define i8 @select_condition_implies_highbits_op1_and(i8 %xx, i8 noundef %y, i1 %
 ; CHECK-NEXT:    [[COND0:%.*]] = icmp ult i8 [[Y:%.*]], 3
 ; CHECK-NEXT:    [[COND:%.*]] = and i1 [[COND0]], [[OTHER_COND:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i8 [[Y]], i8 [[X]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[SEL]], 32
+; CHECK-NEXT:    [[R:%.*]] = or disjoint i8 [[SEL]], 32
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x = and i8 %xx, 15
@@ -69,7 +69,7 @@ define i8 @select_condition_implies_highbits_op2_or(i8 %xx, i8 noundef %y, i1 %o
 ; CHECK-NEXT:    [[COND0:%.*]] = icmp ugt i8 [[Y:%.*]], 3
 ; CHECK-NEXT:    [[COND:%.*]] = or i1 [[COND0]], [[OTHER_COND:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i8 [[X]], i8 [[Y]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[SEL]], 32
+; CHECK-NEXT:    [[R:%.*]] = or disjoint i8 [[SEL]], 32
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x = and i8 %xx, 15