[llvm] [SelectionDAG]: Have isKnownNeverZero treat SRL like division if all else fails (PR #89523)

[via llvm-commits]

https://github.com/AtariDreams created https://github.com/llvm/llvm-project/pull/89523

None

>From c24b4bf90e84a4de8ee7249e8c3d8c2d336f5c10 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sat, 20 Apr 2024 21:31:50 -0400
Subject: [PATCH] [SelectionDAG]: Have isKnownNeverZero treat SRL like division
 if all else fails

---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 7dbf83b7adeef0..3b8368a9a73b91 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5443,10 +5443,17 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     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();
+    const KnownBits Shift = computeKnownBits(Op.getOperand(1), Depth + 1);
+    APInt MaxCnt = Shift.getMaxValue();
     if (MaxCnt.ult(ValKnown.getBitWidth()) &&
         !ValKnown.One.lshr(MaxCnt).isZero())
       return true;
+    // We try to see if we can turn it into a division
+    const KnownBits One =
+        KnownBits::makeConstant(APInt(ValKnown.getBitWidth(), 1));
+    if (KnownBits::uge(ValKnown,
+                       KnownBits::lshr(One, Shift, Shift.isNonZero())))
+      return true;
     break;
   }
   case ISD::UDIV:
@@ -5455,6 +5462,10 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     // TODO: For udiv this is also true if Op1 u<= Op0
     if (Op->getFlags().hasExact())
       return isKnownNeverZero(Op.getOperand(0), Depth + 1);
+
+    if (KnownBits::uge(computeKnownBits(Op.getOperand(1), Depth + 1),
+                       computeKnownBits(Op.getOperand(0), Depth + 1)))
+      return true;
     break;
 
   case ISD::ADD: