[SelectionDAG] Add/Improve cases in `isKnownNeverZero`

[goldstein.w.n]

Summary

1. Handle casts a bit more cleanly just with a loop rather than with recursion.

2. Add additional cases for smin/smax

3. For shifts we can also deduce non-zero if the maximum shift amount on the known 1s is non-zero.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5073

Looking through any_extend seems wrong, the high bits may not be 0 as they're undefined

5082

Should this just go through the recursive case like normal?

5132–5135

Should try the RHS operand first as it's canonically simpler

[goldstein.w.n]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5082

Personally think its clearer to have casts simply handled ahead of time, but no strong attachment/opinion so adding back to recursive case.

5132–5135

Hmm? In a SRL we can't canonicalize constants to RHS. Also doing LHS first gives us a fast out (no need to compute RHS if we have negative value, but we can't get same early out with RHS).

[goldstein.w.n]

Remove any_extend, add other casts back to main switch

[goldstein.w.n]

Rebase

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5100

ult? - shift by bitwidth is undefined.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5101

Use isNonZero() instead?

[goldstein.w.n]

Fix shift cnt check

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5101

apint doens't have isNonZero.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5101

I missed that! Would it improve the analysis if we used ValKnown.shl(CntKnown).isNonZero() instead?

[goldstein.w.n]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5101

I don't think so. Known bits needs to actually figure out which is set. I.e if we have:
2 << (x & 7), known bits won't be able to set anything as it doesn't know any ones (value can be 2, 4, 8,...256 but there are no overlapping bits none can be set), whereas this logic still allows us to prove non-zero.

[RKSimon]

LGTM - cheers!

[goldstein.w.n]

Rebase

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
   if (Depth >= MaxRecursionDepth)
     return false; // Limit search depth.
 
   assert(!Op.getValueType().isFloatingPoint() &&
          "Floating point types unsupported - use isKnownNeverZeroFloat");
 
   // If the value is a constant, we can obviously see if it is a zero or not.
   if (ISD::matchUnaryPredicate(Op,
                                [](ConstantSDNode *C) { return !C->isZero(); }))
     return true;
 
   // TODO: Recognize more cases here. Most of the cases are also incomplete to
   // some degree.
   switch (Op.getOpcode()) {
   default:
     break;
 
   case ISD::OR:
     return isKnownNeverZero(Op.getOperand(1), Depth + 1) ||
            isKnownNeverZero(Op.getOperand(0), Depth + 1);
 
   case ISD::VSELECT:
   case ISD::SELECT:
     return isKnownNeverZero(Op.getOperand(1), Depth + 1) &&
            isKnownNeverZero(Op.getOperand(2), Depth + 1);
 
-  case ISD::SHL:
+  case ISD::SHL: {
     if (Op->getFlags().hasNoSignedWrap() || Op->getFlags().hasNoUnsignedWrap())
       return isKnownNeverZero(Op.getOperand(0), Depth + 1);
-
-    // 1 << X is never zero. TODO: This can be expanded if we can bound X.
-    // The expression is really !Known.One[BitWidth-MaxLog2(Known):0].isZero()
-    if (computeKnownBits(Op.getOperand(0), Depth + 1).One[0])
+    KnownBits ValKnown = computeKnownBits(Op.getOperand(0), Depth + 1);
+    // 1 << X is never zero.
+    if (ValKnown.One[0])
+      return true;
+    // If max shift cnt of known ones is non-zero, result is non-zero.
+    APInt MaxCnt = computeKnownBits(Op.getOperand(1), Depth + 1).getMaxValue();
+    if (MaxCnt.ult(ValKnown.getBitWidth()) &&
+        !ValKnown.One.shl(MaxCnt).isZero())
       return true;
     break;
-
+  }
   case ISD::UADDSAT:
   case ISD::UMAX:
     return isKnownNeverZero(Op.getOperand(1), Depth + 1) ||
            isKnownNeverZero(Op.getOperand(0), Depth + 1);
 
+    // TODO for smin/smax: If either operand is known negative/positive
+    // respectively we don't need the other to be known at all.
+  case ISD::SMAX:
+  case ISD::SMIN:
   case ISD::UMIN:
     return isKnownNeverZero(Op.getOperand(1), Depth + 1) &&
            isKnownNeverZero(Op.getOperand(0), Depth + 1);
 
   case ISD::ROTL:
   case ISD::ROTR:
   case ISD::BITREVERSE:
   case ISD::BSWAP:
   case ISD::CTPOP:
   case ISD::ABS:
     return isKnownNeverZero(Op.getOperand(0), Depth + 1);
 
   case ISD::SRA:
-  case ISD::SRL:
+  case ISD::SRL: {
     if (Op->getFlags().hasExact())
       return isKnownNeverZero(Op.getOperand(0), Depth + 1);
-    // Signed >> X is never zero. TODO: This can be expanded if we can bound X.
-    // The expression is really
-    // !Known.One[SignBit:SignBit-(BitWidth-MaxLog2(Known))].isZero()
-    if (computeKnownBits(Op.getOperand(0), Depth + 1).isNegative())
+    KnownBits ValKnown = computeKnownBits(Op.getOperand(0), Depth + 1);
+    if (ValKnown.isNegative())
+      return true;
+    // If max shift cnt of known ones is non-zero, result is non-zero.
+    APInt MaxCnt = computeKnownBits(Op.getOperand(1), Depth + 1).getMaxValue();
+    if (MaxCnt.ult(ValKnown.getBitWidth()) &&
+        !ValKnown.One.lshr(MaxCnt).isZero())
       return true;
     break;
-
+  }
   case ISD::UDIV:
   case ISD::SDIV:
     // div exact can only produce a zero if the dividend is zero.
     // TODO: For udiv this is also true if Op1 u<= Op0
     if (Op->getFlags().hasExact())
       return isKnownNeverZero(Op.getOperand(0), Depth + 1);
     break;
 

llvm/test/CodeGen/X86/known-never-zero.ll

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %edi, %ecx
 ; CHECK-NEXT:    movl $4, %eax
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    shll %cl, %eax
 ; CHECK-NEXT:    addl $4, %esi
 ; CHECK-NEXT:    cmpl %esi, %eax
 ; CHECK-NEXT:    cmovll %eax, %esi
-; CHECK-NEXT:    bsfl %esi, %ecx
-; CHECK-NEXT:    movl $32, %eax
-; CHECK-NEXT:    cmovnel %ecx, %eax
+; CHECK-NEXT:    rep bsfl %esi, %eax
 ; CHECK-NEXT:    retq
   %x = shl nuw i32 4, %xx
   %y = add nuw nsw i32 %yy, 4
   %z = call i32 @llvm.smin.i32(i32 %x, i32 %y)
   %r = call i32 @llvm.cttz.i32(i32 %z, i1 false)
   ret i32 %r
 }
 

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl %edi, %ecx
 ; CHECK-NEXT:    movl $4, %eax
 ; CHECK-NEXT:    # kill: def $cl killed $cl killed $ecx
 ; CHECK-NEXT:    shll %cl, %eax
 ; CHECK-NEXT:    addl $4, %esi
 ; CHECK-NEXT:    cmpl %esi, %eax
 ; CHECK-NEXT:    cmovgl %eax, %esi
-; CHECK-NEXT:    bsfl %esi, %ecx
-; CHECK-NEXT:    movl $32, %eax
-; CHECK-NEXT:    cmovnel %ecx, %eax
+; CHECK-NEXT:    rep bsfl %esi, %eax
 ; CHECK-NEXT:    retq
   %x = shl nuw i32 4, %xx
   %y = add nuw nsw i32 %yy, 4
   %z = call i32 @llvm.smax.i32(i32 %x, i32 %y)
   %r = call i32 @llvm.cttz.i32(i32 %z, i1 false)
   ret i32 %r
 }