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

[Julian Nagele via llvm-commits]

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

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



>From 7f78ed9ad91ae71c78e50ef0c37b1ce44be16d77 Mon Sep 17 00:00:00 2001
From: Julian Nagele <j.nagele at apple.com>
Date: Tue, 25 Feb 2025 15:44:19 +0000
Subject: [PATCH 1/2] [InstCombine] Add test showing bittest which uses range
 attribute

---
 llvm/test/Transforms/InstCombine/select-icmp-and.ll | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/select-icmp-and.ll b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
index 16fb3f34047ee..09a4b0dc64699 100644
--- a/llvm/test/Transforms/InstCombine/select-icmp-and.ll
+++ b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
@@ -912,3 +912,14 @@ 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:    [[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
+}

>From ffd72987fb6c05cb387ac99e17b05decdcfc2926 Mon Sep 17 00:00:00 2001
From: Julian Nagele <j.nagele at apple.com>
Date: Tue, 25 Feb 2025 15:52:21 +0000
Subject: [PATCH 2/2] [InstCombine] Use known bits to simplify mask in
 foldSelectICmpAnd

---
 .../Transforms/InstCombine/InstCombineSelect.cpp    | 13 ++++++++-----
 llvm/test/Transforms/InstCombine/select-icmp-and.ll |  3 +--
 2 files changed, 9 insertions(+), 7 deletions(-)

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 09a4b0dc64699..c84e925658046 100644
--- a/llvm/test/Transforms/InstCombine/select-icmp-and.ll
+++ b/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