[X86] LowerAndToBT - fold ICMP_ZERO(AND(X,AND(Y,SHL(1,Z)))) -> BT(AND(X,Y),Z) patterns

[RKSimon]

Use m_ReassociatableAnd matcher to handle any permutation of a 3 op AND chain that involves a bit test

Fix 1 of 2 for #147216

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: Simon Pilgrim (RKSimon)

Changes

Use m_ReassociatableAnd matcher to handle any permutation of a 3 op AND chain that involves a bit test

Fix 1 of 2 for #147216

---

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

2 Files Affected:

• (modified) llvm/lib/Target/X86/X86ISelLowering.cpp (+8-1)
• (modified) llvm/test/CodeGen/X86/known-pow2.ll (+16-31)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 5dce0980c3fed..7418cff763d2d 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -24370,7 +24370,7 @@ static SDValue LowerAndToBT(SDValue And, ISD::CondCode CC, const SDLoc &dl,
   assert(And.getValueType().isScalarInteger() && "Scalar type expected");
 
   APInt AndRHSVal;
-  SDValue Shl, Src, BitNo;
+  SDValue Shl, Src, Mask, BitNo;
   if (sd_match(And,
                m_And(m_TruncOrSelf(m_Value(Src)),
                      m_TruncOrSelf(m_AllOf(m_Value(Shl),
@@ -24384,6 +24384,10 @@ static SDValue LowerAndToBT(SDValue And, ISD::CondCode CC, const SDLoc &dl,
       if (Known.countMinLeadingZeros() < (BitWidth - AndBitWidth))
         return SDValue();
     }
+  } else if (sd_match(And,
+                      m_ReassociatableAnd(m_Value(Src), m_Value(Mask),
+                                          m_Shl(m_One(), m_Value(BitNo))))) {
+    // (Src & Mask & (1 << BitNo)) ==/!= 0
   } else if (sd_match(And,
                       m_And(m_TruncOrSelf(m_Srl(m_Value(Src), m_Value(BitNo))),
                             m_One()))) {
@@ -24402,6 +24406,9 @@ static SDValue LowerAndToBT(SDValue And, ISD::CondCode CC, const SDLoc &dl,
     return SDValue();
   }
 
+  if (Mask)
+    Src = DAG.getNode(ISD::AND, dl, Src.getValueType(), Src, Mask);
+
   // Remove any bit flip.
   if (isBitwiseNot(Src)) {
     Src = Src.getOperand(0);
diff --git a/llvm/test/CodeGen/X86/known-pow2.ll b/llvm/test/CodeGen/X86/known-pow2.ll
index 92b176b7a4bbb..b4dd00125aab5 100644
--- a/llvm/test/CodeGen/X86/known-pow2.ll
+++ b/llvm/test/CodeGen/X86/known-pow2.ll
@@ -797,17 +797,13 @@ define i1 @pow2_and_fail0(i32 %x, i32 %y) {
   ret i1 %r
 }
 
-define i1 @pow2_and_fail1(i32 %x, i32 %y) {
-; CHECK-LABEL: pow2_and_fail1:
+define i1 @pow2_andnot_3op(i32 %x, i32 %y) {
+; CHECK-LABEL: pow2_andnot_3op:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl %esi, %ecx
-; CHECK-NEXT:    movl $1, %eax
-; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
-; CHECK-NEXT:    shll %cl, %eax
 ; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    andl %eax, %edi
-; CHECK-NEXT:    testl $-2, %edi
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    andl $-2, %edi
+; CHECK-NEXT:    btl %esi, %edi
+; CHECK-NEXT:    setae %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %nyy = sub i32 1, %yy
@@ -817,17 +813,13 @@ define i1 @pow2_and_fail1(i32 %x, i32 %y) {
   ret i1 %r
 }
 
-define i1 @pow2_and_fail2(i32 %x, i32 %y, i32 %z) {
-; CHECK-LABEL: pow2_and_fail2:
+define i1 @pow2_and_3op(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: pow2_and_3op:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl %esi, %ecx
-; CHECK-NEXT:    movl $1, %eax
-; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
-; CHECK-NEXT:    shll %cl, %eax
-; CHECK-NEXT:    andl %edx, %eax
 ; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %eax
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    andl %edx, %edi
+; CHECK-NEXT:    btl %esi, %edi
+; CHECK-NEXT:    setae %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %d = and i32 %yy, %z
@@ -856,13 +848,9 @@ define i1 @pow2_though_zext(i32 %x, i16 %y) {
 define i1 @pow2_and_i20(i20 %num, i20 %shift) {
 ; CHECK-LABEL: pow2_and_i20:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl %esi, %ecx
-; CHECK-NEXT:    movl $1, %eax
-; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
-; CHECK-NEXT:    shll %cl, %eax
-; CHECK-NEXT:    andl %edi, %eax
-; CHECK-NEXT:    testl $1048575, %eax # imm = 0xFFFFF
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    andl $1048575, %edi # imm = 0xFFFFF
+; CHECK-NEXT:    btl %esi, %edi
+; CHECK-NEXT:    setae %al
 ; CHECK-NEXT:    retq
   %mask = shl nuw i20 1, %shift
   %bit = and i20 %mask, %num
@@ -873,13 +861,10 @@ define i1 @pow2_and_i20(i20 %num, i20 %shift) {
 define i1 @pow2_and_i50(i50 %num, i50 %shift) {
 ; CHECK-LABEL: pow2_and_i50:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq %rsi, %rcx
-; CHECK-NEXT:    movl $1, %eax
-; CHECK-NEXT:    # kill: def $cl killed $cl killed $rcx
-; CHECK-NEXT:    shlq %cl, %rax
+; CHECK-NEXT:    movabsq $1125899906842623, %rax # imm = 0x3FFFFFFFFFFFF
 ; CHECK-NEXT:    andq %rdi, %rax
-; CHECK-NEXT:    shlq $14, %rax
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    btq %rsi, %rax
+; CHECK-NEXT:    setae %al
 ; CHECK-NEXT:    retq
   %mask = shl nuw i50 1, %shift
   %bit = and i50 %mask, %num

[phoebewang]

LGTM.

[phoebewang]

Question: why don't we need to handle (Src & (1 << BitNo)) ==/!= 0 here?

[RKSimon]

The m_Shl variant above (with m_TruncOrSelf wrappers) already handles it