[llvm] [SelectionDAG] Use getAllOnes to evaluate isKnownNeverZero (PR #92923)

[via llvm-commits]

https://github.com/AtariDreams updated https://github.com/llvm/llvm-project/pull/92923

>From 39b49498543f5feaee218524fbe8d8837c88dfe8 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Tue, 21 May 2024 10:03:46 -0400
Subject: [PATCH] [SelectionDAG] Use getAllOnes to complete SHL
 isKnownNeverZero check

This check was ported from ValueTracking to SelectionDAG.
---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 777bbf071732e..c25d756b5511d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5483,10 +5483,23 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     // 1 << X is never zero.
     if (ValKnown.One[0])
       return true;
+
+    if (ValKnown.isUnknown())
+      break;
+
     // 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())
+    unsigned NumBits = ValKnown.getBitWidth();
+    if (MaxCnt.uge(NumBits))
+      break;
+    if (!ValKnown.One.shl(MaxCnt).isZero())
+      return true;
+
+    // If all of the bits shifted out are known to be zero, and Val is known
+    // non-zero then at least one non-zero bit must remain.
+    if (ValKnown.Zero.lshr(NumBits - MaxCnt)
+            .eq(APInt::getAllOnes(NumBits).lshr(NumBits - MaxCnt)) &&
+        isKnownNeverZero(Op.getOperand(0), Depth + 1))
       return true;
     break;
   }