[DAGCombine] Fold Splat(bitcast(buildvector(x,..))) to splat(x)

[dmgreen]

Summary

This adds a fold which fixes a couple of regressions found in D138872 (one for MVE, another for AArch64 NEON, both from intrinsic code). The backend could be caught not knowing how to simplify either splat(bitcast(buildvector(x,..))) or splat(bitcast(scalar_to_vector(x))). This only handles lane 0 splats, which are only valid under LE, and needs to be a little careful with the types it creates for the new buildvector.

[RKSimon]

Could the ARM/AArch64 tests be extended to ensure we have BE coverage?

[dmgreen]

Rebased onto BE tests added to mve-vdup.ll.

[RKSimon]

LGTM - cheers

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if ((!LegalOperations || TLI.isOperationLegal(ISD::BUILD_VECTOR, VT)) &&
         N0.hasOneUse()) {
       if (N0.getOpcode() == ISD::SCALAR_TO_VECTOR && SplatIndex == 0)
         return DAG.getSplatBuildVector(VT, SDLoc(N), N0.getOperand(0));
 
       if (N0.getOpcode() == ISD::INSERT_VECTOR_ELT)
         if (auto *Idx = dyn_cast<ConstantSDNode>(N0.getOperand(2)))
           if (Idx->getAPIntValue() == SplatIndex)
             return DAG.getSplatBuildVector(VT, SDLoc(N), N0.getOperand(1));
+
+      // Look through a bitcast if LE and splatting lane 0, through to a
+      // scalar_to_vector or a build_vector.
+      if (N0.getOpcode() == ISD::BITCAST && N0.getOperand(0).hasOneUse() &&
+          SplatIndex == 0 && DAG.getDataLayout().isLittleEndian() &&
+          (N0.getOperand(0).getOpcode() == ISD::SCALAR_TO_VECTOR ||
+           N0.getOperand(0).getOpcode() == ISD::BUILD_VECTOR)) {
+        EVT N00VT = N0.getOperand(0).getValueType();
+        if (VT.getScalarSizeInBits() <= N00VT.getScalarSizeInBits() &&
+            VT.isInteger() && N00VT.isInteger()) {
+          EVT InVT =
+              TLI.getTypeToTransformTo(*DAG.getContext(), VT.getScalarType());
+          SDValue Op = DAG.getZExtOrTrunc(N0.getOperand(0).getOperand(0),
+                                          SDLoc(N), InVT);
+          return DAG.getSplatBuildVector(VT, SDLoc(N), Op);
+        }
+      }
     }
 
     // If this is a bit convert that changes the element type of the vector but
     // not the number of vector elements, look through it.  Be careful not to
     // look though conversions that change things like v4f32 to v2f64.
     SDNode *V = N0.getNode();
     if (V->getOpcode() == ISD::BITCAST) {
       SDValue ConvInput = V->getOperand(0);

llvm/test/CodeGen/AArch64/arm64-dup.ll

; CHECK-NEXT:    ret
   %tmp3 = add <4 x i32> %C, <i32 8421376, i32 8421376, i32 8421376, i32 8421376>
   %tmp5 = xor <4 x i32> %tmp3, %tmp4
   ret <4 x i32> %tmp5
 }
 
 define <8 x i16> @bitcast_i64_v8i16(i64 %a) {
 ; CHECK-LABEL: bitcast_i64_v8i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d0, x0
-; CHECK-NEXT:    dup.8h v0, v0[0]
+; CHECK-NEXT:    dup.8h v0, w0
 ; CHECK-NEXT:    ret
   %b = bitcast i64 %a to <4 x i16>
   %r = shufflevector <4 x i16> %b, <4 x i16> poison, <8 x i32> zeroinitializer
   ret <8 x i16> %r
 }
 
 define <8 x i16> @bitcast_i64_v8i16_lane1(i64 %a) {
 ; CHECK-LABEL: bitcast_i64_v8i16_lane1:

llvm/test/CodeGen/AArch64/arm64-neon-3vdiff.ll

entry:
   %vmull3.i.i = bitcast <16 x i8> %vmull2.i.i to i128
   ret i128 %vmull3.i.i
 }
 
 define <8 x i16> @cmplx_mul_combined_re_im(<8 x i16> noundef %a, i64 %scale.coerce) {
 ; CHECK-LABEL: cmplx_mul_combined_re_im:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    lsr x8, x0, #16
+; CHECK-NEXT:    adrp x9, .LCPI196_0
 ; CHECK-NEXT:    fmov d4, x0
 ; CHECK-NEXT:    rev32 v5.8h, v0.8h
-; CHECK-NEXT:    fmov d1, x8
-; CHECK-NEXT:    adrp x8, .LCPI196_0
-; CHECK-NEXT:    dup v1.8h, v1.h[0]
-; CHECK-NEXT:    ldr q3, [x8, :lo12:.LCPI196_0]
+; CHECK-NEXT:    dup v1.8h, w8
+; CHECK-NEXT:    ldr q3, [x9, :lo12:.LCPI196_0]
 ; CHECK-NEXT:    sqneg v2.8h, v1.8h
 ; CHECK-NEXT:    tbl v1.16b, { v1.16b, v2.16b }, v3.16b
 ; CHECK-NEXT:    sqdmull v2.4s, v0.4h, v4.h[0]
 ; CHECK-NEXT:    sqdmull2 v0.4s, v0.8h, v4.h[0]
 ; CHECK-NEXT:    sqdmlal v2.4s, v5.4h, v1.4h
 ; CHECK-NEXT:    sqdmlal2 v0.4s, v5.8h, v1.8h
 ; CHECK-NEXT:    uzp2 v0.8h, v2.8h, v0.8h
 ; CHECK-NEXT:    ret

llvm/test/CodeGen/Thumb2/mve-vdup.ll

; CHECK-BE-NEXT:    bx lr
   %b = bitcast i64 %a to <4 x i16>
   %r = shufflevector <4 x i16> %b, <4 x i16> poison, <8 x i32> zeroinitializer
   ret <8 x i16> %r
 }
 
 define arm_aapcs_vfpcc <8 x i16> @bitcast_i128_v8i16(i128 %a) {
 ; CHECK-LE-LABEL: bitcast_i128_v8i16:
 ; CHECK-LE:       @ %bb.0:
-; CHECK-LE-NEXT:    vmov.32 q0[0], r0
-; CHECK-LE-NEXT:    vmov.u16 r0, q0[0]
 ; CHECK-LE-NEXT:    vdup.16 q0, r0
 ; CHECK-LE-NEXT:    bx lr
 ;
 ; CHECK-BE-LABEL: bitcast_i128_v8i16:
 ; CHECK-BE:       @ %bb.0:
 ; CHECK-BE-NEXT:    vmov.32 q0[0], r0
 ; CHECK-BE-NEXT:    vrev32.16 q0, q0
 ; CHECK-BE-NEXT:    vmov.u16 r0, q0[0]
; CHECK-BE-NEXT:    bx lr
   %b = bitcast <2 x i64> %a to <8 x i16>
   %r = shufflevector <8 x i16> %b, <8 x i16> poison, <8 x i32> zeroinitializer
   ret <8 x i16> %r
 }
 
 define arm_aapcs_vfpcc <8 x i16> @other_max_case(i32 %blockSize) {
 ; CHECK-LE-LABEL: other_max_case:
 ; CHECK-LE:       @ %bb.0:
-; CHECK-LE-NEXT:    vmov.32 q0[0], r0
-; CHECK-LE-NEXT:    vmov.u16 r0, q0[0]
 ; CHECK-LE-NEXT:    vdup.16 q0, r0
 ; CHECK-LE-NEXT:    bx lr
 ;
 ; CHECK-BE-LABEL: other_max_case:
 ; CHECK-BE:       @ %bb.0:
 ; CHECK-BE-NEXT:    vmov.32 q0[0], r0
 ; CHECK-BE-NEXT:    vrev32.16 q0, q0
 ; CHECK-BE-NEXT:    vmov.u16 r0, q0[0]
 ; CHECK-BE-NEXT:    vdup.16 q1, r0
 ; CHECK-BE-NEXT:    vrev64.16 q0, q1
 ; CHECK-BE-NEXT:    bx lr
   %vec.blockSize = bitcast i32 %blockSize to <2 x i16>
   %.splat2 = shufflevector <2 x i16> %vec.blockSize, <2 x i16> poison, <8 x i32> zeroinitializer
   ret <8 x i16> %.splat2
 }

llvm/test/CodeGen/WebAssembly/simd-shuffle-bitcast.ll

 ; RUN: llc < %s -asm-verbose=false -verify-machineinstrs -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -mattr=+simd128 | FileCheck %s
 
 ; Test that a splat shuffle of an fp-to-int bitcasted vector correctly
-; optimizes and lowers to a single splat instruction. Without a custom
-; DAG combine, this ends up doing both a splat and a shuffle.
+; optimizes and lowers to a single splat instruction.
 
 target triple = "wasm32-unknown-unknown"
 
 ; CHECK-LABEL: f32x4_splat:
 ; CHECK-NEXT: .functype f32x4_splat (f32) -> (v128){{$}}
 ; CHECK-NEXT: f32x4.splat $push[[R:[0-9]+]]=, $0{{$}}
 ; CHECK-NEXT: return $pop[[R]]{{$}}
 define <4 x i32> @f32x4_splat(float %x) {
   %vecinit = insertelement <4 x float> undef, float %x, i32 0
   %a = bitcast <4 x float> %vecinit to <4 x i32>
   %b = shufflevector <4 x i32> %a, <4 x i32> undef, <4 x i32> zeroinitializer
   ret <4 x i32> %b
 }
 
 ; CHECK-LABEL: not_a_vec:
 ; CHECK-NEXT: .functype not_a_vec (i64, i64) -> (v128){{$}}
-; CHECK-NEXT: i64x2.splat $push[[L1:[0-9]+]]=, $0{{$}}
-; CHECK-NEXT: i8x16.shuffle $push[[R:[0-9]+]]=, $pop[[L1]], $2, 0, 1, 2, 3
+; CHECK-NEXT: i32.wrap_i64    $push[[L:[0-9]+]]=, $0
+; CHECK-NEXT: i32x4.splat     $push[[R:[0-9]+]]=, $pop[[L]]
 ; CHECK-NEXT: return $pop[[R]]
 define <4 x i32> @not_a_vec(i128 %x) {
   %a = bitcast i128 %x to <4 x i32>
   %b = shufflevector <4 x i32> %a, <4 x i32> undef, <4 x i32> zeroinitializer
   ret <4 x i32> %b
 }

llvm/test/CodeGen/X86/vector-shuffle-mmx.ll

 ; X32-NEXT:    movsd {{.*#+}} xmm0 = mem[0],zero
 ; X32-NEXT:    unpcklps {{.*#+}} xmm0 = xmm0[0,0,1,1]
 ; X32-NEXT:    movlps %xmm0, (%eax)
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: test2:
 ; X64:       ## %bb.0: ## %entry
 ; X64-NEXT:    movq _tmp_V2i@GOTPCREL(%rip), %rax
-; X64-NEXT:    movq {{.*#+}} xmm0 = mem[0],zero
-; X64-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[0,0,1,1]
+; X64-NEXT:    movd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; X64-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[0,0,0,0]
 ; X64-NEXT:    movq %xmm0, (%rax)
 ; X64-NEXT:    retq
 entry:
   %0 = load <2 x i32>, ptr @tmp_V2i, align 8
   %1 = shufflevector <2 x i32> %0, <2 x i32> undef, <2 x i32> zeroinitializer
   store <2 x i32> %1, ptr @tmp_V2i, align 8
   ret void
 }