[DAG] Ensure more Legal BUILD_VECTOR elements types in shuffle->And combine

[dmgreen]

Summary

This is a followup to D131350, which caused another problem for i64 types being split into i32 on i32 targets. This patch tries to make sure that either Illegal types are OK, or that the element types of a buildvector are legal and bigger than or equal to the size of the original elements.

[RKSimon]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (!IsInLaneMask &&
         (!DemandedLHS.isNullValue() || !DemandedRHS.isNullValue()) &&
         (DemandedLHS.isNullValue() ||
          DAG.MaskedVectorIsZero(N0, DemandedLHS)) &&
         (DemandedRHS.isNullValue() ||
          DAG.MaskedVectorIsZero(N1, DemandedRHS))) {
       SDLoc DL(N);
       EVT IntVT = VT.changeVectorElementTypeToInteger();
       EVT IntSVT = VT.getVectorElementType().changeTypeToInteger();
+      // Transform the type to a legal type so that the buildvector constant
+      // elements are not illegal. Make sure that the result is larger than the
+      // original type, incase the value is split into two (eg i64->i32).
+      if (!TLI.isTypeLegal(IntSVT) && LegalTypes)
         IntSVT = TLI.getTypeToTransformTo(*DAG.getContext(), IntSVT);
+      if (IntSVT.getSizeInBits() >= IntVT.getScalarSizeInBits()) {
         SDValue ZeroElt = DAG.getConstant(0, DL, IntSVT);
         SDValue AllOnesElt = DAG.getAllOnesConstant(DL, IntSVT);
         SmallVector<SDValue, 16> AndMask(NumElts, DAG.getUNDEF(IntSVT));
         for (int I = 0; I != (int)NumElts; ++I)
           if (0 <= Mask[I])
             AndMask[I] = Mask[I] == I ? AllOnesElt : ZeroElt;
 
         // See if a clear mask is legal instead of going via
         // XformToShuffleWithZero which loses UNDEF mask elements.
         if (TLI.isVectorClearMaskLegal(ClearMask, IntVT))
           return DAG.getBitcast(
               VT, DAG.getVectorShuffle(IntVT, DL, DAG.getBitcast(IntVT, N0),
                                       DAG.getConstant(0, DL, IntVT), ClearMask));
 
         if (TLI.isOperationLegalOrCustom(ISD::AND, IntVT))
           return DAG.getBitcast(
               VT, DAG.getNode(ISD::AND, DL, IntVT, DAG.getBitcast(IntVT, N0),
                               DAG.getBuildVector(IntVT, DL, AndMask)));
       }
     }
+  }
 
   // Attempt to combine a shuffle of 2 inputs of 'scalar sources' -
   // BUILD_VECTOR or SCALAR_TO_VECTOR into a single BUILD_VECTOR.
   if (Level < AfterLegalizeDAG && TLI.isTypeLegal(VT))
     if (SDValue Res = combineShuffleOfScalars(SVN, DAG, TLI))
       return Res;
 
   // If this shuffle only has a single input that is a bitcasted shuffle,

llvm/test/CodeGen/ARM/vector-store.ll

 ; CHECK-NEXT:    strb r2, [r0, #2]
 ; CHECK-NEXT:    vmov.32 r1, d16[0]
 ; CHECK-NEXT:    strh r1, [r0]
 ; CHECK-NEXT:    add sp, #4
 ; CHECK-NEXT:    bx lr
   store <3 x i8> zeroinitializer, <3 x i8> *%p, align 4
   ret void
 }
+
+define void @v3i64shuffle(<3 x i64> *%p, <3 x i64> %a) {
+; CHECK-LABEL: v3i64shuffle:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vmov.i32 q8, #0x0
+; CHECK-NEXT:    ldrd r12, r1, [sp, #8]
+; CHECK-NEXT:    vmov d18, r2, r3
+; CHECK-NEXT:    vorr d19, d16, d16
+; CHECK-NEXT:    str r1, [r0, #20]
+; CHECK-NEXT:    vst1.32 {d18, d19}, [r0]!
+; CHECK-NEXT:    str.w r12, [r0]
+; CHECK-NEXT:    bx lr
+  %b = shufflevector <3 x i64> %a, <3 x i64> zeroinitializer, <3 x i32> <i32 0, i32 3, i32 2>
+  store <3 x i64> %b, <3 x i64> *%p, align 4
+  ret void
+}
+