[InstCombine] Use known bits to simplify mask in foldSelectICmpAnd

[juliannagele]

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/MM8wz2

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Julian Nagele (juliannagele)

Changes

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 --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
+}

[dtcxzyw]

Can you please provide a generalized alive2 proof?

[juliannagele]

Thanks, I incorporated the code comments. For the alive proof did you have a specific generalization in mind? I'm not quite sure how to use/represent known bits and isPowerOf2 there.

[dtcxzyw]

Thanks, I incorporated the code comments. For the alive proof did you have a specific generalization in mind? I'm not quite sure how to use/represent known bits and isPowerOf2 there.

You can use @llvm.assume:

%ctpop = call i32 @llvm.ctpop.i32(i32 %pow2)
%is_pow2 = icmp eq i32 %ctpop, 1
call void @llvm.assume(i1 %is_pow2)

[juliannagele]

Thanks! Here's a generalized version assume-ing that the rhs and bitmask are powers of 2 https://alive2.llvm.org/ce/z/RszbRS

[dtcxzyw]

Can you add some negative tests (e.g., AndMask & MaxValueOfX is not a power of 2)?

[juliannagele]

Sure, added two negative tests for the not-a-power-of-2 cases.

[dtcxzyw]

LGTM. Thank you!