[llvm] Port isNonZeroShift to SelectionDAG (PR #146125)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/146125

>From 4123e5a896776f16a920177467d7cd1d350d19d7 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Fri, 27 Jun 2025 13:58:29 -0400
Subject: [PATCH] Port isNonZeroShift to SelectionDAG

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 59 ++++++++++++++++---
 1 file changed, 52 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index fe9a6ea3e77e6..ff256dd428d2a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5871,6 +5871,55 @@ bool SelectionDAG::isKnownNeverZeroFloat(SDValue Op) const {
       Op, [](ConstantFPSDNode *C) { return !C->isZero(); });
 }
 
+static bool isNonZeroShift(const SDValue I, const KnownBits &KnownVal,
+                           unsigned Depth) {
+  auto ShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
+    switch (I.getOpcode()) {
+    case Instruction::Shl:
+      return Lhs.shl(Rhs);
+    case Instruction::LShr:
+      return Lhs.lshr(Rhs);
+    case Instruction::AShr:
+      return Lhs.ashr(Rhs);
+    default:
+      llvm_unreachable("Unknown Shift Opcode");
+    }
+  };
+
+  auto InvShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
+    switch (I.getOpcode()) {
+    case ISD::SHL:
+      return Lhs.lshr(Rhs);
+    case ISD::SRA:
+    case ISD::SRL:
+      return Lhs.shl(Rhs);
+    default:
+      llvm_unreachable("Unknown Shift Opcode");
+    }
+  };
+
+  if (KnownVal.isUnknown())
+    return false;
+
+  KnownBits KnownCnt = computeKnownBits(I.getOperand(1), Depth + 1);
+  APInt MaxShift = KnownCnt.getMaxValue();
+  unsigned NumBits = KnownVal.getBitWidth();
+  if (MaxShift.uge(NumBits))
+    return false;
+
+  if (!ShiftOp(KnownVal.One, MaxShift).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 (InvShiftOp(KnownVal.Zero, NumBits - MaxShift)
+          .eq(InvShiftOp(APInt::getAllOnes(NumBits), NumBits - MaxShift)) &&
+      isKnownNeverZero(I.getOperand(0), Depth + 1))
+    return true;
+
+  return false;
+}
+
 bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
   if (Depth >= MaxRecursionDepth)
     return false; // Limit search depth.
@@ -5906,9 +5955,7 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     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())
+    if (isNonZeroShift(Op, ValKnown, Depth))
       return true;
     break;
   }
@@ -5968,10 +6015,8 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     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())
+
+    if (isNonZeroShift(Op, ValKnown, Depth))
       return true;
     break;
   }