[DAGCombiner] When folding (insert_subvector undef, (bitcast (extract_subvector N1, Idx)), Idx) -> (bitcast N1) make sure that N1 has the same total size as the original output

[craig.topper]

Summary

We were only checking the element count, but not the total width. This could cause illegal bitcasts to be created if for example the output was 512-bits, but N1 is 256 bits, and the extraction size was 128-bits.

Fixes PR36199

[niravd]

LGTM.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitINSERT_SUBVECTOR(SDNode *N) {
   // If we are inserting a bitcast value into an undef, with the same
   // number of elements, just use the bitcast input of the extract.
   // i.e. INSERT_SUBVECTOR UNDEF (BITCAST N1) N2 ->
   //        BITCAST (INSERT_SUBVECTOR UNDEF N1 N2)
   if (N0.isUndef() && N1.getOpcode() == ISD::BITCAST &&
       N1.getOperand(0).getOpcode() == ISD::EXTRACT_SUBVECTOR &&
       N1.getOperand(0).getOperand(1) == N2 &&
       N1.getOperand(0).getOperand(0).getValueType().getVectorNumElements() ==
-          VT.getVectorNumElements()) {
+          VT.getVectorNumElements() &&
+      N1.getOperand(0).getOperand(0).getValueType().getSizeInBits() ==
+          VT.getSizeInBits()) {
     return DAG.getBitcast(VT, N1.getOperand(0).getOperand(0));
   }
 
   // If both N1 and N2 are bitcast values on which insert_subvector
   // would makes sense, pull the bitcast through.
   // i.e. INSERT_SUBVECTOR (BITCAST N0) (BITCAST N1) N2 ->
   //        BITCAST (INSERT_SUBVECTOR N0 N1 N2)
   if (N0.getOpcode() == ISD::BITCAST && N1.getOpcode() == ISD::BITCAST) {

llvm/trunk/test/CodeGen/X86/pr36199.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mcpu=skylake-avx512 | FileCheck %s
+
+define void @foo() unnamed_addr #0 {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vaddps %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vinsertf128 $1, %xmm0, %ymm0, %ymm0
+; CHECK-NEXT:    vinsertf64x4 $1, %ymm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovups %zmm0, (%rax)
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+  %1 = fadd <16 x float> undef, undef
+  %bc256 = bitcast <16 x float> %1 to <4 x i128>
+  %2 = extractelement <4 x i128> %bc256, i32 0
+  %3 = bitcast i128 %2 to <4 x float>
+  %4 = shufflevector <4 x float> %3, <4 x float> undef, <16 x i32> <i32 0, i32
+1, i32 2, i32 3, i32 0, i32 1, i32 2, i32 3, i32 0, i32 1, i32 2, i32 3, i32 0,
+i32 1, i32 2, i32 3>
+  store <16 x float> %4, <16 x float>* undef, align 4
+  ret void
+}