[llvm] [DAG] isKnownNeverZero - add ISD::VECTOR_SHUFFLE handling (PR #186335)

[via llvm-commits]

https://github.com/absarkar created https://github.com/llvm/llvm-project/pull/186335

Closed #183351 

>From 6e70e8c5bbe447255d6644f62daf91b1b24f9f13 Mon Sep 17 00:00:00 2001
From: Abdullah Sarkar <fa2818 at nyu.edu>
Date: Fri, 13 Mar 2026 04:25:14 -0400
Subject: [PATCH] [DAG] isKnownNeverZero - add ISD::VECTOR_SHUFFLE handling

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 25 +++++++++++++++++++
 llvm/test/CodeGen/X86/known-never-zero.ll     | 17 +++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 1c5b2d00fe83c..ef042e5ca9e54 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6295,6 +6295,31 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, const APInt &DemandedElts,
     break;
   }
 
+  case ISD::VECTOR_SHUFFLE: {
+    assert(!Op.getValueType().isScalableVector());
+    unsigned NumElts = DemandedElts.getBitWidth();
+    // Collect the known bits that are shared by every vector element referenced
+    // by the shuffle.
+    APInt DemandedLHS, DemandedRHS;
+    const ShuffleVectorSDNode *SVN = cast<ShuffleVectorSDNode>(Op);
+    assert(NumElts == SVN->getMask().size() && "Unexpected vector size");
+    if (!getShuffleDemandedElts(NumElts, SVN->getMask(), DemandedElts,
+                                DemandedLHS, DemandedRHS))
+      return false;
+
+    for (int i = 0; i < (int)NumElts; ++i)
+	    if (DemandedElts[i] && SVN->getMaskElt(i) < 0)
+		    return false;
+
+    if (!!DemandedLHS && !isKnownNeverZero(Op.getOperand(0), DemandedLHS, Depth + 1))
+	    return false;
+
+    if (!!DemandedRHS && !isKnownNeverZero(Op.getOperand(1), DemandedRHS, Depth + 1))
+	    return false;
+    
+    return true;
+  }
+
   case ISD::UADDSAT:
   case ISD::UMAX:
     return isKnownNeverZero(Op.getOperand(1), DemandedElts, Depth + 1) ||
diff --git a/llvm/test/CodeGen/X86/known-never-zero.ll b/llvm/test/CodeGen/X86/known-never-zero.ll
index 8327a90bdeda5..15210d8278a40 100644
--- a/llvm/test/CodeGen/X86/known-never-zero.ll
+++ b/llvm/test/CodeGen/X86/known-never-zero.ll
@@ -2821,3 +2821,20 @@ define i32 @test_sext_demanded_elts(<4 x i32> %a0, ptr %p) {
   ret i32 %res
 }
 
+
+define <4 x i32> @known_nonzero_shuffle(<4 x i32> %a, <4 x i32> %b) {
+; X86-LABEL: known_nonzero_shuffle:
+; X86:       # %bb.0:
+; X86-NEXT:    pshufd {{.*#+}} xmm1 = xmm1[1,3,2,3]
+; X86-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[0,2,2,3]
+; X86-NEXT:    punpckldq {{.*#+}} xmm0 = xmm0[0],xmm1[0],xmm0[1],xmm1[1]
+; X86-NEXT:    retl
+;
+; X64-LABEL: known_nonzero_shuffle:
+; X64:       # %bb.0:
+; X64-NEXT:    vblendps {{.*#+}} xmm0 = xmm0[0],xmm1[1],xmm0[2],xmm1[3]
+; X64-NEXT:    retq
+	%shuf = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+	ret <4 x i32> %shuf
+}
+