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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Julian Nagele (juliannagele)

<details>
<summary>Changes</summary>

We currently do not make use of known bits when trying to decompose a select/icmp bittest and folding it into an and. This means we lose some folds when additional information, for instance via range attribute or metadata, would allow us to conclude that the resulting mask is in fact a power of two. See also https://alive2.llvm.org/ce/z/neLGhH



---
Full diff: https://github.com/llvm/llvm-project/pull/128741.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+8-5) 
- (modified) llvm/test/Transforms/InstCombine/select-icmp-and.ll (+10) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index dca969e160bb7..dc317da97e13c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/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
diff --git a/llvm/test/Transforms/InstCombine/select-icmp-and.ll b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
index 16fb3f34047ee..c84e925658046 100644
--- a/llvm/test/Transforms/InstCombine/select-icmp-and.ll
+++ b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
@@ -912,3 +912,13 @@ 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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/128741