[llvm] [DAGCombiner] Don't ignore N2's undef elements in `foldVSelectOfConstants` (PR #129272)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/129272

Since N2 will be reused in the fold, we cannot skip N2's undef elements if the corresponding element in N1 is well-defined.
For example:
```
t2: v4i32 = BUILD_VECTOR Constant:i32<0>, Constant:i32<0>, Constant:i32<0>, Constant:i32<0>
t24: v4i32 = BUILD_VECTOR undef:i32, undef:i32, Constant:i32<1>, undef:i32
t11: v4i32 = vselect t8, t2, t10
```
Before this patch, we fold t11 into:
```
t26: v4i32 = sign_extend t8
t27: v4i32 = add t26, t24
```
The last element of t27 is incorrect.

Closes https://github.com/llvm/llvm-project/issues/129181.


>From e4f4fad2ce3bed3213be0d19fd39436709adfcbb Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 1 Mar 2025 00:33:27 +0800
Subject: [PATCH 1/2] [X86] Add pre-commit tests. NFC.

---
 llvm/test/CodeGen/X86/vselect-constants.ll | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/llvm/test/CodeGen/X86/vselect-constants.ll b/llvm/test/CodeGen/X86/vselect-constants.ll
index 901f7e4a00eb5..f5924605ef3ff 100644
--- a/llvm/test/CodeGen/X86/vselect-constants.ll
+++ b/llvm/test/CodeGen/X86/vselect-constants.ll
@@ -302,3 +302,19 @@ define i32 @wrong_min_signbits(<2 x i16> %x) {
   %t1 = bitcast <2 x i16> %sel to i32
   ret i32 %t1
 }
+
+define i32 @pr129181() {
+; SSE-LABEL: pr129181:
+; SSE:       # %bb.0: # %entry
+; SSE-NEXT:    retq
+;
+; AVX-LABEL: pr129181:
+; AVX:       # %bb.0: # %entry
+; AVX-NEXT:    retq
+entry:
+  %x = insertelement <4 x i32> zeroinitializer, i32 0, i32 0
+  %cmp = icmp ult <4 x i32> %x, splat (i32 1)
+  %sel = select <4 x i1> %cmp, <4 x i32> zeroinitializer, <4 x i32> <i32 0, i32 0, i32 1, i32 poison>
+  %reduce = tail call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %sel)
+  ret i32 %reduce
+}

>From 1bc745841ee18b901b6e47edebefe4fd864c5e40 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 1 Mar 2025 01:07:39 +0800
Subject: [PATCH 2/2] [DAGCombiner] Don't ignore N2's undef elements in
 `foldVSelectOfConstants`

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 5 +++--
 llvm/test/CodeGen/X86/vselect-constants.ll    | 2 ++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index f197ae61550a9..343a93785766c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -12535,9 +12535,10 @@ SDValue DAGCombiner::foldVSelectOfConstants(SDNode *N) {
   for (unsigned i = 0; i != Elts; ++i) {
     SDValue N1Elt = N1.getOperand(i);
     SDValue N2Elt = N2.getOperand(i);
-    if (N1Elt.isUndef() || N2Elt.isUndef())
+    if (N1Elt.isUndef())
       continue;
-    if (N1Elt.getValueType() != N2Elt.getValueType()) {
+    // N2 should not contain undef values since it will be reused in the fold.
+    if (N2Elt.isUndef() || N1Elt.getValueType() != N2Elt.getValueType()) {
       AllAddOne = false;
       AllSubOne = false;
       break;
diff --git a/llvm/test/CodeGen/X86/vselect-constants.ll b/llvm/test/CodeGen/X86/vselect-constants.ll
index f5924605ef3ff..34bda718db8f6 100644
--- a/llvm/test/CodeGen/X86/vselect-constants.ll
+++ b/llvm/test/CodeGen/X86/vselect-constants.ll
@@ -306,10 +306,12 @@ define i32 @wrong_min_signbits(<2 x i16> %x) {
 define i32 @pr129181() {
 ; SSE-LABEL: pr129181:
 ; SSE:       # %bb.0: # %entry
+; SSE-NEXT:    xorl %eax, %eax
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: pr129181:
 ; AVX:       # %bb.0: # %entry
+; AVX-NEXT:    xorl %eax, %eax
 ; AVX-NEXT:    retq
 entry:
   %x = insertelement <4 x i32> zeroinitializer, i32 0, i32 0