[InstCombine] Use known bits to simplify mask in foldSelectICmpAnd

Author: juliannagele

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -120,7 +120,8 @@ static Instruction *foldSelectBinOpIdentity(SelectInst &Sel,
 /// With some variations depending if FC is larger than TC, or the shift
 /// isn't needed, or the bit widths don't match.
 static Value *foldSelectICmpAnd(SelectInst &Sel, ICmpInst *Cmp,
-                                InstCombiner::BuilderTy &Builder) {
+                                InstCombiner::BuilderTy &Builder,
+                                SimplifyQuery &SQ) {
   const APInt *SelTC, *SelFC;
   if (!match(Sel.getTrueValue(), m_APInt(SelTC)) ||
       !match(Sel.getFalseValue(), m_APInt(SelFC)))
@@ -148,11 +149,13 @@ static Value *foldSelectICmpAnd(SelectInst &Sel, ICmpInst *Cmp,
   } else if (auto Res = decomposeBitTestICmp(Cmp->getOperand(0),
                                              Cmp->getOperand(1), Pred)) {
     assert(ICmpInst::isEquality(Res->Pred) && "Not equality test?");
-    if (!Res->Mask.isPowerOf2())
+    AndMask = Res->Mask;
+    V = Res->X;
+    KnownBits Known = computeKnownBits(V, 0, SQ.getWithInstruction(Cmp));
+    AndMask &= Known.getMaxValue();
+    if (!AndMask.isPowerOf2())
       return nullptr;
 
-    V = Res->X;
-    AndMask = Res->Mask;
     Pred = Res->Pred;
     CreateAnd = true;
   } else {
@@ -1957,7 +1960,7 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
           tryToReuseConstantFromSelectInComparison(SI, *ICI, *this))
     return NewSel;
 
-  if (Value *V = foldSelectICmpAnd(SI, ICI, Builder))
+  if (Value *V = foldSelectICmpAnd(SI, ICI, Builder, SQ))
     return replaceInstUsesWith(SI, V);
 
   // NOTE: if we wanted to, this is where to detect integer MIN/MAX

llvm/test/Transforms/InstCombine/select-icmp-and.ll

@@ -915,8 +915,7 @@ define i16 @select_trunc_nuw_bittest_or(i8 %x) {
 
 define i16 @select_icmp_bittest_range(i16 range (i16 0, 512) %a) {
 ; CHECK-LABEL: @select_icmp_bittest_range(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp samesign ult i16 [[A:%.*]], 256
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i16 0, i16 256
+; CHECK-NEXT:    [[RES:%.*]] = and i16 [[A:%.*]], 256
 ; CHECK-NEXT:    ret i16 [[RES]]
 ;
   %cmp = icmp ult i16 %a, 256