[llvm] [SelectionDAG] Fold (icmp eq/ne (shift X, C), 0) -> (icmp eq/ne X, 0) (PR #88801)

[Björn Pettersson via llvm-commits]

================
@@ -4516,6 +4516,35 @@ SDValue TargetLowering::SimplifySetCC(EVT VT, SDValue N0, SDValue N1,
         }
       }
     }
+
+    // Optimize
+    //    (setcc (shift N00, N01C), 0, eq/ne) -> (setcc N00, 0, eq/ne)
+    // If all shifted out bits are known to be zero, then the zero'd ness
+    // doesn't change and we can omit the shift.
+    // If all shifted out bits are equal to at least one bit that isn't
+    // shifted out, then the zero'd ness doesn't change and we can omit the
+    // shift.
+    if ((Cond == ISD::SETEQ || Cond == ISD::SETNE) && C1.isZero() &&
+        N0.hasOneUse() &&
+        (N0.getOpcode() == ISD::SHL || N0.getOpcode() == ISD::SRL ||
+         N0.getOpcode() == ISD::SRA)) {
+      bool IsRightShift = N0.getOpcode() != ISD::SHL;
+      SDValue N00 = N0.getOperand(0);
+      // Quick checks based on exact/nuw/nsw flags.
+      if (IsRightShift ? N0->getFlags().hasExact()
+                       : (N0->getFlags().hasNoUnsignedWrap() ||
+                          N0->getFlags().hasNoSignedWrap()))
+        return DAG.getSetCC(dl, VT, N00, N1, Cond);
+      // More expensive checks based on known bits.
+      if (const APInt *ShAmt = DAG.getValidMaximumShiftAmountConstant(N0)) {
+        KnownBits Known = DAG.computeKnownBits(N00);
+        if (IsRightShift)
+          Known = Known.reverseBits();
+        if (ShAmt->ule(Known.countMinLeadingZeros()) ||
+            ShAmt->ult(Known.countMinSignBits()))
----------------
bjope wrote:

I probably made some bad assumptions here when using computeKnownBits.

Using `DAG.computeKnownBits(N00).countMinSignBits()` isn't always as good as `DAG.ComputeNumSignBits(N00)`. The latter might for example know that the number of sign bits after an arithmetic right shift ,while computKnownBits only can tell number of sign bits if they are known to be zero or one. So I really want to check ComputeNumSignBits, but then as you say we do not have anything for computing number of trailing bits that are equal to each other.
For the left shift it might be better to use ComputeNumSignBits (unless we can trust flags). 
But as you say, for the right shift it seems like computeKnownBits won't help here.

Btw, if adding flags in `visit{SHL,SAR,SHR}` we will depend on iteration order if those will be present or not when visiting the SETCC. And iteration order in DAGCombiner is a bit of a mystery.

https://github.com/llvm/llvm-project/pull/88801