[DAGCombiner] Require same type of splat & element for build_vector (#88284)

Only allow to change build_vector to concat_vector when the splat type
and vector element type is same. It's to fix assertion of failing to
bitcast types of different sizes.

Author: fzou1

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -23428,7 +23428,10 @@ SDValue DAGCombiner::visitBUILD_VECTOR(SDNode *N) {
   // If X is a build_vector itself, the concat can become a larger build_vector.
   // TODO: Maybe this is useful for non-splat too?
   if (!LegalOperations) {
-    if (SDValue Splat = cast<BuildVectorSDNode>(N)->getSplatValue()) {
+    SDValue Splat = cast<BuildVectorSDNode>(N)->getSplatValue();
+    // Only change build_vector to a concat_vector if the splat value type is
+    // same as the vector element type.
+    if (Splat && Splat.getValueType() == VT.getVectorElementType()) {
       Splat = peekThroughBitcasts(Splat);
       EVT SrcVT = Splat.getValueType();
       if (SrcVT.isVector()) {
@@ -23437,8 +23440,8 @@ SDValue DAGCombiner::visitBUILD_VECTOR(SDNode *N) {
                                      SrcVT.getVectorElementType(), NumElts);
         if (!LegalTypes || TLI.isTypeLegal(NewVT)) {
           SmallVector<SDValue, 8> Ops(N->getNumOperands(), Splat);
-          SDValue Concat = DAG.getNode(ISD::CONCAT_VECTORS, SDLoc(N),
-                                       NewVT, Ops);
+          SDValue Concat =
+              DAG.getNode(ISD::CONCAT_VECTORS, SDLoc(N), NewVT, Ops);
           return DAG.getBitcast(VT, Concat);
         }
       }

llvm/test/CodeGen/X86/buildvec-bitcast.ll

@@ -0,0 +1,25 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=x86_64 -mattr=avx512bw | FileCheck %s
+
+; Verify that the DAGCombiner doesn't change build_vector to concat_vectors if
+; the vector element type is different than splat type. The example here:
+;   v8i1 = build_vector (i8 (bitcast (v8i1 X))), ..., (i8 (bitcast (v8i1 X))))
+
+define <8 x i1> @foo(<8 x i1> %mask.i1) {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vxorps %xmm0, %xmm0, %xmm0
+; CHECK-NEXT:    retq
+entry:
+  %0 = and <8 x i1> %mask.i1, <i1 true, i1 false, i1 false, i1 false, i1 false, i1 false, i1 false, i1 false>
+  %1 = bitcast <8 x i1> %0 to i8
+  %2 = icmp ne i8 %1, 0
+  %insert54 = insertelement <8 x i1> zeroinitializer, i1 %2, i64 0
+  %splat55 = shufflevector <8 x i1> %insert54, <8 x i1> zeroinitializer, <8 x i32> zeroinitializer
+  %3 = and <8 x i1> %0, %splat55
+  br label %end
+
+end:                           ; preds = %entry
+  %4 = select <8 x i1> %3, <8 x i1> zeroinitializer, <8 x i1> zeroinitializer
+  ret <8 x i1> %4
+}