[llvm] [DAG] isKnownNeverZero - add ISD::SHL DemandedElts handling and tests (PR #183772)

[via llvm-commits]

https://github.com/Joeljm1 updated https://github.com/llvm/llvm-project/pull/183772

>From 9b3c3d5100cf124af25b571d0cdb4fc712a7dfc5 Mon Sep 17 00:00:00 2001
From: Joel Joseph Mathews <joeljosephcl10 at gmail.com>
Date: Fri, 27 Feb 2026 22:28:38 +0530
Subject: [PATCH 1/2] [DAG] isKnownNeverZero - add ISD::SHL DemandedElts
 handling and tests

added DemandedElts to ISD::SHL case in isKnowNeverZero and required
tests
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  4 +-
 .../AArch64/AArch64SelectionDAGTest.cpp       | 50 +++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 4a2bd811b5214..b1e2b07f0c1fb 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6209,13 +6209,13 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, const APInt &DemandedElts,
 
   case ISD::SHL: {
     if (Op->getFlags().hasNoSignedWrap() || Op->getFlags().hasNoUnsignedWrap())
-      return isKnownNeverZero(Op.getOperand(0), Depth + 1);
+      return isKnownNeverZero(Op.getOperand(0),DemandedElts, Depth + 1);
     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();
+    APInt MaxCnt = computeKnownBits(Op.getOperand(1),DemandedElts ,Depth + 1).getMaxValue();
     if (MaxCnt.ult(ValKnown.getBitWidth()) &&
         !ValKnown.One.shl(MaxCnt).isZero())
       return true;
diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index b0c48e8c97995..90a765dc16433 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -1574,4 +1574,54 @@ TEST_F(AArch64SelectionDAGTest, KnownNeverZero_Select) {
   EXPECT_FALSE(DAG->isKnownNeverZero(VSelect444Big, DemandAll));
   EXPECT_TRUE(DAG->isKnownNeverZero(VSelect4444, DemandAll));
 }
+
+TEST_F(AArch64SelectionDAGTest, KnownNeverZero_SHL_DemandedElts) {
+  SDLoc Loc;
+  EVT VT = EVT::getVectorVT(Context, MVT::i32, 4);
+
+  // Vector: < -2147483648, 1, 0 -8 >
+  SDValue V1 = DAG->getConstant(-2147483648, Loc, MVT::i32);
+  SDValue V2 = DAG->getConstant(1, Loc, MVT::i32);
+  SDValue V3 = DAG->getConstant(0, Loc, MVT::i32);
+  SDValue V4 = DAG->getConstant(-8, Loc, MVT::i32);
+
+  SDValue Vec = DAG->getBuildVector(VT, Loc, {V1, V2, V3, V4});
+
+  SDValue Shift1 = DAG->getConstant(1, Loc, VT);
+  SDValue Shift31 = DAG->getConstant(31, Loc, VT);
+  SDValue ShiftOverflow = DAG->getConstant(32, Loc, VT);
+  SDValue ShiftUnkown = DAG->getExternalSymbol("unknown", VT);
+
+  SDValue Op1 = DAG->getNode(ISD::SHL, Loc, VT, Vec, Shift1);
+  SDValue Op31 = DAG->getNode(ISD::SHL, Loc, VT, Vec, Shift31);
+  SDValue Op32 = DAG->getNode(ISD::SHL, Loc, VT, Vec, ShiftOverflow);
+  SDValue OpUnkown = DAG->getNode(ISD::SHL, Loc, VT, Vec, ShiftUnkown);
+  // Lane 0: -2147483648 << 1 = 0
+  APInt Lane0(4, 1);
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op1, Lane0));
+  EXPECT_FALSE(DAG->isKnownNeverZero(OpUnkown, Lane0));
+  // Lane 1: 1 << 31 = -2147483648
+  APInt Lane1(4, 2);
+  EXPECT_TRUE(DAG->isKnownNeverZero(Op31, Lane1));
+  // Lane 2: 0 << 1 = 0
+  APInt Lane2(4, 4);
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op1, Lane2));
+  EXPECT_FALSE(DAG->isKnownNeverZero(OpUnkown, Lane2));
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op32, Lane2));
+  // Lane 3: -8 << 1 = -16
+  APInt Lane3(4, 8);
+  EXPECT_TRUE(DAG->isKnownNeverZero(Op1, Lane3));
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op32, Lane3));
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op31, Lane3));
+  // lane 2 and 4
+  APInt Lane1AND4(4, -6);
+  EXPECT_TRUE(DAG->isKnownNeverZero(Op1, Lane1AND4));
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op31, Lane1AND4));
+  EXPECT_FALSE(DAG->isKnownNeverZero(OpUnkown, Lane1AND4));
+  // all lanes
+  APInt LaneAll(4, -1);
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op1, LaneAll));
+  EXPECT_FALSE(DAG->isKnownNeverZero(OpUnkown, LaneAll));
+  EXPECT_FALSE(DAG->isKnownNeverZero(Op32, LaneAll));
+}
 } // end namespace llvm

>From cedf95e471e32fd9c00cafefc3de1796bd2b9824 Mon Sep 17 00:00:00 2001
From: Joeljm1 <143259392+Joeljm1 at users.noreply.github.com>
Date: Fri, 27 Feb 2026 20:54:52 +0300
Subject: [PATCH 2/2] Update llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

Add a missing DemandedElts in isKnownNeverZero ISD::SHL case

Co-authored-by: Simon Pilgrim <git at redking.me.uk>
---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index b1e2b07f0c1fb..34a8e3c6d1cc5 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6210,7 +6210,7 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, const APInt &DemandedElts,
   case ISD::SHL: {
     if (Op->getFlags().hasNoSignedWrap() || Op->getFlags().hasNoUnsignedWrap())
       return isKnownNeverZero(Op.getOperand(0),DemandedElts, Depth + 1);
-    KnownBits ValKnown = computeKnownBits(Op.getOperand(0), Depth + 1);
+    KnownBits ValKnown = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
     // 1 << X is never zero.
     if (ValKnown.One[0])
       return true;