[SelectionDAG] Improve `isKnownToBeAPowerOfTwo`

[goldstein.w.n]

Summary

Add additional cases for:
select, vselect, {u,s}{min,max}, and, casts, rotl, rotr

And improve handling of constants and shifts.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

4069

Looking through ANY_EXTEND is probably broken

4121–4122

try RHS first

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

4124

Add comments describing the and(x,neg(x)) pattern

4130

Do you need to use matchUnaryPredicate or could you just call isNullOrNullSplat ?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

4078

We can't safely use ANY_EXTEND, afaict you just need this:

while (Val.getOpcode() == ISD::ZERO_EXTEND)
  Val = Val.getOperand(0);

[goldstein.w.n]

Remove any-extend, comment and case.

[RKSimon]

llvm/test/CodeGen/X86/known-pow2.ll

626

Any luck in triaging this?

[goldstein.w.n]

Rebase

[goldstein.w.n]

llvm/test/CodeGen/X86/known-pow2.ll

626

I think for vecs (on X86) doing X & Y != 0 is just not preferable to X & Y == Y.
I fixed this up in D157014 and changed the rest of the vselect tests to use ne so we can properly see the change.

[RKSimon]

LGTM

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

 bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val, unsigned Depth) const {
   if (Depth >= MaxRecursionDepth)
     return false; // Limit search depth.
 
   EVT OpVT = Val.getValueType();
   unsigned BitWidth = OpVT.getScalarSizeInBits();
 
   // Is the constant a known power of 2?
-  if (ConstantSDNode *Const = dyn_cast<ConstantSDNode>(Val))
-    return Const->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2();
+  if (ISD::matchUnaryPredicate(Val, [BitWidth](ConstantSDNode *C) {
+        return C->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2();
+      }))
+    return true;
 
   // A left-shift of a constant one will have exactly one bit set because
   // shifting the bit off the end is undefined.
   if (Val.getOpcode() == ISD::SHL) {
     auto *C = isConstOrConstSplat(Val.getOperand(0));
     if (C && C->getAPIntValue() == 1)
       return true;
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1) &&
+           isKnownNeverZero(Val, Depth);
   }
 
   // Similarly, a logical right-shift of a constant sign-bit will have exactly
   // one bit set.
   if (Val.getOpcode() == ISD::SRL) {
     auto *C = isConstOrConstSplat(Val.getOperand(0));
     if (C && C->getAPIntValue().isSignMask())
       return true;
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1) &&
+           isKnownNeverZero(Val, Depth);
   }
 
+  if (Val.getOpcode() == ISD::ROTL || Val.getOpcode() == ISD::ROTR)
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1);
+
   // Are all operands of a build vector constant powers of two?
   if (Val.getOpcode() == ISD::BUILD_VECTOR)
     if (llvm::all_of(Val->ops(), [BitWidth](SDValue E) {
           if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(E))
             return C->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2();
           return false;
         }))
       return true;
 
   // Is the operand of a splat vector a constant power of two?
   if (Val.getOpcode() == ISD::SPLAT_VECTOR)
     if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(Val->getOperand(0)))
       if (C->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2())
         return true;
 
   // vscale(power-of-two) is a power-of-two for some targets
   if (Val.getOpcode() == ISD::VSCALE &&
       getTargetLoweringInfo().isVScaleKnownToBeAPowerOfTwo() &&
       isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1))
     return true;
 
+  if (Val.getOpcode() == ISD::SMIN || Val.getOpcode() == ISD::SMAX ||
+      Val.getOpcode() == ISD::UMIN || Val.getOpcode() == ISD::UMAX)
+    return isKnownToBeAPowerOfTwo(Val.getOperand(1), Depth + 1) &&
+           isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1);
+
+  if (Val.getOpcode() == ISD::SELECT || Val.getOpcode() == ISD::VSELECT)
+    return isKnownToBeAPowerOfTwo(Val.getOperand(2), Depth + 1) &&
+           isKnownToBeAPowerOfTwo(Val.getOperand(1), Depth + 1);
+
+  if (Val.getOpcode() == ISD::AND) {
+    // Looking for `x & -x` pattern:
+    // If x == 0:
+    //    x & -x -> 0
+    // If x != 0:
+    //    x & -x -> non-zero pow2
+    // so if we find the pattern return whether we know `x` is non-zero.
+    for (unsigned OpIdx = 0; OpIdx < 2; ++OpIdx) {
+      SDValue NegOp = Val.getOperand(OpIdx);
+      if (NegOp.getOpcode() == ISD::SUB &&
+          NegOp.getOperand(1) == Val.getOperand(1 - OpIdx) &&
+          isNullOrNullSplat(NegOp.getOperand(0)))
+        return isKnownNeverZero(Val.getOperand(1 - OpIdx), Depth);
+    }
+  }
+
+  if (Val.getOpcode() == ISD::ZERO_EXTEND)
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1);
+
   // More could be done here, though the above checks are enough
   // to handle some common cases.
   return false;
 }
 
 unsigned SelectionDAG::ComputeNumSignBits(SDValue Op, unsigned Depth) const {
   EVT VT = Op.getValueType();
 

llvm/test/CodeGen/X86/known-pow2.ll

; CHECK-NEXT:    retq
   %r = urem <4 x i32> %x, <i32 9, i32 4, i32 16, i32 64>
   ret <4 x i32> %r
 }
 
 define i1 @pow2_shl(i32 %x, i32 %y) {
 ; CHECK-LABEL: pow2_shl:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %esi, %ecx
-; CHECK-NEXT:    notl %edi
+; CHECK-NEXT:    movl %edi, %eax
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
-; CHECK-NEXT:    shrl %cl, %edi
-; CHECK-NEXT:    testb $4, %dil
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    shrl %cl, %eax
+; CHECK-NEXT:    andl $4, %eax
+; CHECK-NEXT:    shrl $2, %eax
+; CHECK-NEXT:    # kill: def $al killed $al killed $eax
 ; CHECK-NEXT:    retq
   %d = shl nuw nsw i32 4, %y
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 
 define i1 @pow2_shl_fail0(i32 %x, i32 %y) {
; CHECK-NEXT:    retq
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 
 define i1 @pow2_srl(i32 %x, i32 %y) {
 ; CHECK-LABEL: pow2_srl:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %esi, %ecx
+; CHECK-NEXT:    movl %edi, %eax
 ; CHECK-NEXT:    andb $7, %cl
-; CHECK-NEXT:    notl %edi
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
-; CHECK-NEXT:    shll %cl, %edi
-; CHECK-NEXT:    testl $1048576, %edi # imm = 0x100000
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    shll %cl, %eax
+; CHECK-NEXT:    shrl $20, %eax
+; CHECK-NEXT:    andl $1, %eax
+; CHECK-NEXT:    # kill: def $al killed $al killed $eax
 ; CHECK-NEXT:    retq
   %yy = and i32 %y, 7
   %d = lshr i32 1048576, %yy
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 

 
 define i1 @pow2_rotl(i32 %x, i32 %y) {
 ; CHECK-LABEL: pow2_rotl:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %esi, %ecx
 ; CHECK-NEXT:    movl $1048576, %eax # imm = 0x100000
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    roll %cl, %eax
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %eax
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %eax, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %d = call i32 @llvm.fshl.i32(i32 1048576, i32 1048576, i32 %y)
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 
 define i1 @pow2_rotl_fail0(i32 %x, i32 %y) {

 
 define i1 @pow2_rotr(i32 %x, i32 %y) {
 ; CHECK-LABEL: pow2_rotr:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %esi, %ecx
 ; CHECK-NEXT:    movl $1048576, %eax # imm = 0x100000
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    rorl %cl, %eax
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %eax
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %eax, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %d = call i32 @llvm.fshr.i32(i32 1048576, i32 1048576, i32 %y)
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 
 define i1 @pow2_rotr_fail0(i32 %x, i32 %y) {

 ; 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:    cmpl $262144, %eax # imm = 0x40000
 ; CHECK-NEXT:    movl $262144, %ecx # imm = 0x40000
 ; CHECK-NEXT:    cmovbl %eax, %ecx
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %ecx
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %ecx, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %d = call i32 @llvm.umin.i32(i32 %yy, i32 262144)
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 

 ; CHECK-NEXT:    movl $1, %eax
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    shll %cl, %eax
 ; CHECK-NEXT:    movl $-2147483648, %esi # imm = 0x80000000
 ; CHECK-NEXT:    movl %edx, %ecx
 ; CHECK-NEXT:    shrl %cl, %esi
 ; CHECK-NEXT:    cmpl %esi, %eax
 ; CHECK-NEXT:    cmoval %eax, %esi
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %esi
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %esi, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %zz = lshr i32 2147483648, %z
   %d = call i32 @llvm.umax.i32(i32 %yy, i32 %zz)
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }

 ; 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:    cmpl $262144, %eax # imm = 0x40000
 ; CHECK-NEXT:    movl $262144, %ecx # imm = 0x40000
 ; CHECK-NEXT:    cmovll %eax, %ecx
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %ecx
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %ecx, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %d = call i32 @llvm.smin.i32(i32 %yy, i32 262144)
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }
 

 ; CHECK-NEXT:    movl $1, %eax
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    shll %cl, %eax
 ; CHECK-NEXT:    movl $-2147483648, %esi # imm = 0x80000000
 ; CHECK-NEXT:    movl %edx, %ecx
 ; CHECK-NEXT:    shrl %cl, %esi
 ; CHECK-NEXT:    cmpl %esi, %eax
 ; CHECK-NEXT:    cmovgl %eax, %esi
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %esi
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %esi, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %zz = lshr i32 2147483648, %z
   %d = call i32 @llvm.smax.i32(i32 %yy, i32 %zz)
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }

 ; CHECK-NEXT:    movl $1, %edx
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    shll %cl, %edx
 ; CHECK-NEXT:    movl $-2147483648, %r8d # imm = 0x80000000
 ; CHECK-NEXT:    movl %eax, %ecx
 ; CHECK-NEXT:    shrl %cl, %r8d
 ; CHECK-NEXT:    testb $1, %dil
 ; CHECK-NEXT:    cmovnel %edx, %r8d
-; CHECK-NEXT:    notl %esi
-; CHECK-NEXT:    testl %esi, %r8d
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %r8d, %esi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %yy = shl i32 1, %y
   %zz = lshr i32 2147483648, %z
   %d = select i1 %c, i32 %yy, i32 %zz
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }

 ; CHECK-NEXT:    psrld %xmm4, %xmm6
 ; CHECK-NEXT:    pshuflw {{.*#+}} xmm3 = xmm3[0,1,1,1,4,5,6,7]
 ; CHECK-NEXT:    psrld %xmm3, %xmm5
 ; CHECK-NEXT:    punpckhqdq {{.*#+}} xmm5 = xmm5[1],xmm6[1]
 ; CHECK-NEXT:    shufps {{.*#+}} xmm7 = xmm7[0,3],xmm5[0,3]
 ; CHECK-NEXT:    pand %xmm0, %xmm2
 ; CHECK-NEXT:    pandn %xmm7, %xmm0
 ; CHECK-NEXT:    por %xmm2, %xmm0
-; CHECK-NEXT:    pand %xmm0, %xmm1
+; CHECK-NEXT:    pcmpeqd %xmm2, %xmm2
+; CHECK-NEXT:    pand %xmm1, %xmm0
+; CHECK-NEXT:    pxor %xmm1, %xmm1
 ; CHECK-NEXT:    pcmpeqd %xmm1, %xmm0
+; CHECK-NEXT:    pxor %xmm2, %xmm0
 ; CHECK-NEXT:    retq
   %yy = shl <4 x i32> <i32 1, i32 1, i32 1, i32 1>, %y
   %zz = lshr <4 x i32> <i32 2147483648, i32 2147483648, i32 2147483648, i32 2147483648>, %z
   %d = select <4 x i1> %c, <4 x i32> %yy, <4 x i32> %zz
   %and = and <4 x i32> %x, %d
   %r = icmp eq <4 x i32> %and, %d
   ret <4 x i1> %r
 }

 ; CHECK-NEXT:    psrld %xmm4, %xmm6
 ; CHECK-NEXT:    pshuflw {{.*#+}} xmm3 = xmm3[0,1,1,1,4,5,6,7]
 ; CHECK-NEXT:    psrld %xmm3, %xmm5
 ; CHECK-NEXT:    punpckhqdq {{.*#+}} xmm5 = xmm5[1],xmm6[1]
 ; CHECK-NEXT:    shufps {{.*#+}} xmm7 = xmm7[0,3],xmm5[0,3]
 ; CHECK-NEXT:    pand %xmm0, %xmm2
 ; CHECK-NEXT:    pandn %xmm7, %xmm0
 ; CHECK-NEXT:    por %xmm2, %xmm0
-; CHECK-NEXT:    pcmpeqd %xmm2, %xmm2
-; CHECK-NEXT:    pand %xmm0, %xmm1
+; CHECK-NEXT:    pand %xmm1, %xmm0
+; CHECK-NEXT:    pxor %xmm1, %xmm1
 ; CHECK-NEXT:    pcmpeqd %xmm1, %xmm0
-; CHECK-NEXT:    pxor %xmm2, %xmm0
 ; CHECK-NEXT:    retq
   %yy = shl <4 x i32> <i32 1, i32 1, i32 1, i32 1>, %y
   %zz = lshr <4 x i32> <i32 2147483648, i32 2147483648, i32 2147483648, i32 2147483648>, %z
   %d = select <4 x i1> %c, <4 x i32> %yy, <4 x i32> %zz
   %and = and <4 x i32> %x, %d
   %r = icmp ne <4 x i32> %and, %d
   ret <4 x i1> %r
 }

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %esi, %ecx
 ; CHECK-NEXT:    movl $4, %eax
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    shll %cl, %eax
 ; CHECK-NEXT:    movl %eax, %ecx
 ; CHECK-NEXT:    negl %ecx
 ; CHECK-NEXT:    andl %eax, %ecx
-; CHECK-NEXT:    notl %edi
-; CHECK-NEXT:    testl %edi, %ecx
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    testl %ecx, %edi
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %yy = shl nuw nsw i32 4, %y
   %nyy = sub i32 0, %yy
   %d = and i32 %yy, %nyy
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }

 
 define i1 @pow2_though_zext(i32 %x, i16 %y) {
 ; CHECK-LABEL: pow2_though_zext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %esi, %ecx
 ; CHECK-NEXT:    movl $4, %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 $65535, %edi # imm = 0xFFFF
-; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    andl %edi, %eax
+; CHECK-NEXT:    testl $65535, %eax # imm = 0xFFFF
+; CHECK-NEXT:    setne %al
 ; CHECK-NEXT:    retq
   %dd = shl nuw nsw i16 4, %y
   %d = zext i16 %dd to i32
   %and = and i32 %x, %d
   %r = icmp eq i32 %and, %d
   ret i1 %r
 }