[InstCombine] Use known bits to simplify mask in foldSelectICmpAnd (#128741)

Make use of known bits when trying to decompose a select/icmp bittest and folding it into an and. This means we can fold when additional information, for instance via a range attribute or metadata, allows us to conclude that the resulting mask is in fact a power of two.

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,
+                                const SimplifyQuery &SQ) {
   const APInt *SelTC, *SelFC;
   if (!match(Sel.getTrueValue(), m_APInt(SelTC)) ||
       !match(Sel.getFalseValue(), m_APInt(SelFC)))
@@ -148,11 +149,14 @@ 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, /*Depth=*/0, SQ.getWithInstruction(&Sel));
+    AndMask &= Known.getMaxValue();
+    if (!AndMask.isPowerOf2())
       return nullptr;
 
-    V = Res->X;
-    AndMask = Res->Mask;
     Pred = Res->Pred;
     CreateAnd = true;
   } else {
@@ -1957,7 +1961,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

@@ -912,3 +912,35 @@ define i16 @select_trunc_nuw_bittest_or(i8 %x) {
   %res = or i16 4, %select
   ret i16 %res
 }
+
+define i16 @select_icmp_bittest_range(i16 range (i16 0, 512) %a) {
+; CHECK-LABEL: @select_icmp_bittest_range(
+; CHECK-NEXT:    [[RES:%.*]] = and i16 [[A:%.*]], 256
+; CHECK-NEXT:    ret i16 [[RES]]
+;
+  %cmp = icmp ult i16 %a, 256
+  %res = select i1 %cmp, i16 0, i16 256
+  ret i16 %res
+}
+
+define i16 @select_icmp_bittest_range_negative_test(i16 range (i16 0, 513) %a) {
+; CHECK-LABEL: @select_icmp_bittest_range_negative_test(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp samesign ult i16 [[A:%.*]], 256
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i16 0, i16 256
+; CHECK-NEXT:    ret i16 [[RES]]
+;
+  %cmp = icmp ult i16 %a, 256
+  %res = select i1 %cmp, i16 0, i16 256
+  ret i16 %res
+}
+
+define i16 @select_icmp_bittest_range_negative_test2(i16 range (i16 0, 512) %a) {
+; CHECK-LABEL: @select_icmp_bittest_range_negative_test2(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp samesign ult i16 [[A:%.*]], 255
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i16 0, i16 255
+; CHECK-NEXT:    ret i16 [[RES]]
+;
+  %cmp = icmp ult i16 %a, 255
+  %res = select i1 %cmp, i16 0, i16 255
+  ret i16 %res
+}