[llvm] [DAG] isKnownNeverZero - add ISD::ADD/SUB DemandedElts handling and tests (PR #183958)

[Weiwen He via llvm-commits]

https://github.com/he-weiwen updated https://github.com/llvm/llvm-project/pull/183958

>From 300b6d55f6101246d0b039671cd4baae58db73ce Mon Sep 17 00:00:00 2001
From: he-weiwen <he.weiwen at outlook.com>
Date: Sat, 28 Feb 2026 21:04:14 +0000
Subject: [PATCH] [DAG] isKnownNeverZero - add ISD::ADD/SUB DemandedElts  
 handling and tests

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 12 ++---
 llvm/test/CodeGen/X86/known-never-zero.ll     | 53 +++++++++++++++++++
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 084700c50bd06..8e67df41f5e41 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6337,19 +6337,19 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, const APInt &DemandedElts,
 
   case ISD::ADD:
     if (Op->getFlags().hasNoUnsignedWrap())
-      if (isKnownNeverZero(Op.getOperand(1), Depth + 1) ||
-          isKnownNeverZero(Op.getOperand(0), Depth + 1))
+      if (isKnownNeverZero(Op.getOperand(1), DemandedElts, Depth + 1) ||
+          isKnownNeverZero(Op.getOperand(0), DemandedElts, Depth + 1))
         return true;
     // TODO: There are a lot more cases we can prove for add.
     break;
 
   case ISD::SUB: {
     if (isNullConstant(Op.getOperand(0)))
-      return isKnownNeverZero(Op.getOperand(1), Depth + 1);
+      return isKnownNeverZero(Op.getOperand(1), DemandedElts, Depth + 1);
 
-    std::optional<bool> ne =
-        KnownBits::ne(computeKnownBits(Op.getOperand(0), Depth + 1),
-                      computeKnownBits(Op.getOperand(1), Depth + 1));
+    std::optional<bool> ne = KnownBits::ne(
+        computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1),
+        computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1));
     return ne && *ne;
   }
 
diff --git a/llvm/test/CodeGen/X86/known-never-zero.ll b/llvm/test/CodeGen/X86/known-never-zero.ll
index 12bb486d8dceb..d40a00b0f0767 100644
--- a/llvm/test/CodeGen/X86/known-never-zero.ll
+++ b/llvm/test/CodeGen/X86/known-never-zero.ll
@@ -1140,6 +1140,30 @@ define i32 @add_maybe_zero(i32 %xx, i32 %y) {
   ret i32 %r
 }
 
+define i32 @add_nuw_known_nonzero_vec(<4 x i32> %xx, ptr %p) {
+; X86-LABEL: add_nuw_known_nonzero_vec:
+; X86:       # %bb.0:
+; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; X86-NEXT:    paddd {{\.?LCPI[0-9]+_[0-9]+}}, %xmm0
+; X86-NEXT:    movdqa %xmm0, (%eax)
+; X86-NEXT:    movd %xmm0, %eax
+; X86-NEXT:    rep bsfl %eax, %eax
+; X86-NEXT:    retl
+;
+; X64-LABEL: add_nuw_known_nonzero_vec:
+; X64:       # %bb.0:
+; X64-NEXT:    vpaddd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; X64-NEXT:    vmovdqa %xmm0, (%rdi)
+; X64-NEXT:    vmovd %xmm0, %eax
+; X64-NEXT:    rep bsfl %eax, %eax
+; X64-NEXT:    retq
+  %z = add nuw <4 x i32> %xx, <i32 1, i32 0, i32 0, i32 0>
+  store <4 x i32> %z, ptr %p
+  %e = extractelement <4 x i32> %z, i32 0
+  %r = call i32 @llvm.cttz.i32(i32 %e, i1 false)
+  ret i32 %r
+}
+
 define i32 @sub_known_nonzero_neg_case(i32 %xx) {
 ; X86-LABEL: sub_known_nonzero_neg_case:
 ; X86:       # %bb.0:
@@ -1238,6 +1262,35 @@ define i32 @sub_maybe_zero2(i32 %x) {
   ret i32 %r
 }
 
+define i32 @sub_known_nonzero_ne_vec(<4 x i32> %xx, ptr %p) {
+; X86-LABEL: sub_known_nonzero_ne_vec:
+; X86:       # %bb.0:
+; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; X86-NEXT:    por {{\.?LCPI[0-9]+_[0-9]+}}, %xmm0
+; X86-NEXT:    movd {{.*#+}} xmm1 = [2,0,0,0]
+; X86-NEXT:    psubd %xmm0, %xmm1
+; X86-NEXT:    movdqa %xmm1, (%eax)
+; X86-NEXT:    movd %xmm1, %eax
+; X86-NEXT:    rep bsfl %eax, %eax
+; X86-NEXT:    retl
+;
+; X64-LABEL: sub_known_nonzero_ne_vec:
+; X64:       # %bb.0:
+; X64-NEXT:    vpor {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; X64-NEXT:    vpmovsxbq {{.*#+}} xmm1 = [2,0]
+; X64-NEXT:    vpsubd %xmm0, %xmm1, %xmm0
+; X64-NEXT:    vmovdqa %xmm0, (%rdi)
+; X64-NEXT:    vmovd %xmm0, %eax
+; X64-NEXT:    rep bsfl %eax, %eax
+; X64-NEXT:    retq
+  %u = or <4 x i32> %xx, <i32 1, i32 0, i32 0, i32 0>
+  %z = sub <4 x i32> <i32 2, i32 0, i32 0, i32 0>, %u
+  store <4 x i32> %z, ptr %p
+  %e = extractelement <4 x i32> %z, i32 0
+  %r = call i32 @llvm.cttz.i32(i32 %e, i1 false)
+  ret i32 %r
+}
+
 define i32 @mul_known_nonzero_nsw(i32 %x, i32 %yy) {
 ; X86-LABEL: mul_known_nonzero_nsw:
 ; X86:       # %bb.0: