[DAG] Extend visitSCALAR_TO_VECTOR optimization to truncated vector.

[niravd]

Summary

Allow SCALAR_TO_VECTOR of EXTRACT_VECTOR_ELT to reduce to
EXTRACT_SUBVECTOR of vector shuffle when output is smaller. Marginally
improves vector shuffle computations.

[efriedma]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15751–15753

The only way for VT.getScalarType() == InVal.getValueType() to be false is if legalization promoted it to a larger integer type... and I can't why that should block the transform (the extra bits get discarded anyway).

15757

Do you need to check that the element type matches? And in that case, isn't the "VT.getVectorNumElements() == InVecT.getVectorNumElements()" check redundant (since it implies VT == InVecT)?

[niravd]

Update to weaken type restrictions

[efriedma]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15767

I think you need to verify that VT and InVecT have the same element type here?

[niravd]

Add type check in truncate case.

[efriedma]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15767

"VT.getVectorNumElements() != InVecT.getVectorNumElements()" check is redundant.

[niravd]

Remove redundant check

[niravd]

Unfold unrelated commit.

[efriedma]

LGTM

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15766

I would probably rearrange this to move the element type check into the outer if statement, for the sake of clarity... but either way is fine, I guess.

[niravd]

Looks like this is raising an assertion. Reopening until I can look at it again.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

while (StoreNodes.size() > 1) {
     if (IsConstantSrc) {
       LSBaseSDNode *FirstInChain = StoreNodes[0].MemNode;
       unsigned FirstStoreAS = FirstInChain->getAddressSpace();
       unsigned FirstStoreAlign = FirstInChain->getAlignment();
       unsigned LastLegalType = 1;
       unsigned LastLegalVectorType = 1;
       bool LastIntegerTrunc = false;
       bool NonZero = false;
+      unsigned FirstZeroAfterNonZero = NumConsecutiveStores;
       for (unsigned i = 0; i < NumConsecutiveStores; ++i) {
         StoreSDNode *ST = cast<StoreSDNode>(StoreNodes[i].MemNode);
         SDValue StoredVal = ST->getValue();
-
-        if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(StoredVal)) {
-          NonZero |= !C->isNullValue();
-        } else if (ConstantFPSDNode *C =
-                       dyn_cast<ConstantFPSDNode>(StoredVal)) {
-          NonZero |= !C->getConstantFPValue()->isNullValue();
+        bool IsElementZero = false;
+        if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(StoredVal))
+          IsElementZero = C->isNullValue();
+        else if (ConstantFPSDNode *C = dyn_cast<ConstantFPSDNode>(StoredVal))
+          IsElementZero = C->getConstantFPValue()->isNullValue();
+        if (IsElementZero) {
+          if (NonZero && FirstZeroAfterNonZero == NumConsecutiveStores)
+            FirstZeroAfterNonZero = i;
         } else {
-          // Non-constant.
-          break;
+          if (NonZero) FirstZeroAfterNonZero = i;
+          NonZero = true;
         }
 
         // Find a legal type for the constant store.
         unsigned SizeInBits = (i + 1) * ElementSizeBytes * 8;
         EVT StoreTy = EVT::getIntegerVT(Context, SizeInBits);
         bool IsFast = false;
         if (TLI.isTypeLegal(StoreTy) &&
             TLI.canMergeStoresTo(FirstStoreAS, StoreTy, DAG) &&
if (IsConstantSrc) {
                                      FirstStoreAlign, &IsFast) &&
               IsFast)
             LastLegalVectorType = i + 1;
         }
       }
 
       // Check if we found a legal integer type that creates a meaningful merge.
       if (LastLegalType < 2 && LastLegalVectorType < 2) {
-        StoreNodes.erase(StoreNodes.begin(), StoreNodes.begin() + 1);
+        // We have already checked that candidate stores are in order
+        // and of correct shape. While there is no mergable sequence
+        // from the beggining one may start later in the sequence. The
+        // only reason a merge of size N could have failed where
+        // another of the same size would not have does not is if the
+        // alignment has improved or we've dropped a non-zero
+        // value. Drop as many candidates as we can here.
+        unsigned NumSkip = 1;
+        while ((NumSkip < NumConsecutiveStores) && (NumSkip < FirstZeroAfterNonZero) &&
+               (StoreNodes[NumSkip].MemNode->getAlignment() <= FirstStoreAlign)) {
+          NumSkip++;
+        }
+        StoreNodes.erase(StoreNodes.begin(), StoreNodes.begin() + NumSkip);
         continue;
       }
 
       bool UseVector = (LastLegalVectorType > LastLegalType) && !NoVectors;
       unsigned NumElem = (UseVector) ? LastLegalVectorType : LastLegalType;
 
       bool Merged = MergeStoresOfConstantsOrVecElts(
           StoreNodes, MemVT, NumElem, true, UseVector, LastIntegerTrunc);
if (IsExtractVecSrc) {
         if (TLI.isTypeLegal(Ty) &&
             TLI.canMergeStoresTo(FirstStoreAS, Ty, DAG) &&
             TLI.allowsMemoryAccess(Context, DL, Ty, FirstStoreAS,
                                    FirstStoreAlign, &IsFast) &&
             IsFast)
           NumStoresToMerge = i + 1;
       }
 
+      // Check if we found a legal integer type that creates a meaningful merge.
+      if (NumStoresToMerge < 2) {
+        // We have already checked that candidate stores are in order
+        // and of correct shape. While there is no mergable sequence
+        // from the beggining one may start later in the sequence. The
+        // only reason a merge of size N could have failed where
+        // another of the same size would not have does not is if the
+        // alignment has improved. Drop as many candidates as we can here.
+        unsigned NumSkip = 1;
+        while ((NumSkip < NumConsecutiveStores) &&
+               (StoreNodes[NumSkip].MemNode->getAlignment() <= FirstStoreAlign))
+          NumSkip++;
+
+        StoreNodes.erase(StoreNodes.begin(), StoreNodes.begin() + NumSkip);
+        continue;
+      }
+
       bool Merged = MergeStoresOfConstantsOrVecElts(
           StoreNodes, MemVT, NumStoresToMerge, false, true, false);
       if (!Merged) {
         StoreNodes.erase(StoreNodes.begin(),
                          StoreNodes.begin() + NumStoresToMerge);
         continue;
       }
       // Remove merged stores for next iteration.
SDValue DAGCombiner::visitVECTOR_SHUFFLE(SDNode *N) {
   return SDValue();
 }
 
 SDValue DAGCombiner::visitSCALAR_TO_VECTOR(SDNode *N) {
   SDValue InVal = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   // Replace a SCALAR_TO_VECTOR(EXTRACT_VECTOR_ELT(V,C0)) pattern
-  // with a VECTOR_SHUFFLE.
+  // with a VECTOR_SHUFFLE and possible truncate.
   if (InVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT) {
     SDValue InVec = InVal->getOperand(0);
     SDValue EltNo = InVal->getOperand(1);
-
-    // FIXME: We could support implicit truncation if the shuffle can be
-    // scaled to a smaller vector scalar type.
+    auto InVecT = InVec.getValueType();
     ConstantSDNode *C0 = dyn_cast<ConstantSDNode>(EltNo);
-    if (C0 && VT == InVec.getValueType() &&
-        VT.getScalarType() == InVal.getValueType()) {
-      SmallVector<int, 8> NewMask(VT.getVectorNumElements(), -1);
+
+    if (C0) {
+      SmallVector<int, 8> NewMask(InVecT.getVectorNumElements(), -1);
       int Elt = C0->getZExtValue();
       NewMask[0] = Elt;
-
-      if (TLI.isShuffleMaskLegal(NewMask, VT))
-        return DAG.getVectorShuffle(VT, SDLoc(N), InVec, DAG.getUNDEF(VT),
-                                    NewMask);
+      SDValue Val;
+      if (VT.getVectorNumElements() <= InVecT.getVectorNumElements() &&
+          TLI.isShuffleMaskLegal(NewMask, VT)) {
+        Val = DAG.getVectorShuffle(InVecT, SDLoc(N), InVec,
+                                   DAG.getUNDEF(InVecT), NewMask);
+        // If the initial vector is the correct size this shuffle is a
+        // valid result.
+        if (VT == InVecT)
+          return Val;
+        // If not we can truncate the vector.
+        if (VT.getScalarType() == InVecT.getVectorElementType()) {
+          MVT IdxTy = TLI.getVectorIdxTy(DAG.getDataLayout());
+          SDValue ZeroIdx = DAG.getConstant(0, SDLoc(N), IdxTy);
+          EVT SubVT =
+              EVT::getVectorVT(*DAG.getContext(), InVecT.getVectorElementType(),
+                               VT.getVectorNumElements());
+          Val = DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), SubVT, Val,
+                            ZeroIdx);
+          return Val;
+        }
+      }
     }
   }
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitINSERT_SUBVECTOR(SDNode *N) {
   EVT VT = N->getValueType(0);

test/CodeGen/AArch64/arm64-neon-copy.ll

 ; CHECK: ins {{v[0-9]+}}.s[1], {{v[0-9]+}}.s[2]
   %tmp3 = extractelement <4 x float> %tmp1, i32 2
   %tmp4 = insertelement <2 x float> %tmp2, float %tmp3, i32 1
   ret <2 x float> %tmp4
 }
 
 define <1 x double> @ins2f1(<2 x double> %tmp1, <1 x double> %tmp2) {
 ; CHECK-LABEL: ins2f1:
-; CHECK: mov {{d[0-9]+}}, {{v[0-9]+}}.d[1]
+; CHECK: dup {{v[0-9]+}}.2d, {{v[0-9]+}}.d[1]
   %tmp3 = extractelement <2 x double> %tmp1, i32 1
   %tmp4 = insertelement <1 x double> %tmp2, double %tmp3, i32 0
   ret <1 x double> %tmp4
 }
 
 define <8 x i8> @ins8b8(<8 x i8> %tmp1, <8 x i8> %tmp2) {
 ; CHECK-LABEL: ins8b8:
 ; CHECK: ins {{v[0-9]+}}.b[4], {{v[0-9]+}}.b[2]

test/CodeGen/AArch64/neon-scalar-copy.ll

  ; CHECK-NOT: ins {{[vsdh][0-9]+}}
  ; CHECK-NEXT: ret
  %tmp1 = extractelement <8 x half> %v, i32 0
  ret half  %tmp1
 }
 
 define <1 x i8> @test_vector_dup_bv16B(<16 x i8> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_bv16B:
- ; CHECK-NEXT: umov [[W:w[0-9]+]], v0.b[14]
- ; CHECK-NEXT: fmov s0, [[W]]
+ ; CHECK-NEXT: dup v0.16b, v0.b[14]
  ; CHECK-NEXT: ret
  %shuffle.i = shufflevector <16 x i8> %v1, <16 x i8> undef, <1 x i32> <i32 14> 
  ret <1 x i8> %shuffle.i
 }
 
 define <1 x i8> @test_vector_dup_bv8B(<8 x i8> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_bv8B:
  ; CHECK-NEXT: dup v0.8b, v0.b[7]
  ; CHECK-NEXT: ret
  %shuffle.i = shufflevector <8 x i8> %v1, <8 x i8> undef, <1 x i32> <i32 7> 
  ret <1 x i8> %shuffle.i
 }
 
 define <1 x i16> @test_vector_dup_hv8H(<8 x i16> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_hv8H:
- ; CHECK-NEXT: umov [[W:w[0-9]+]], v0.h[7]
- ; CHECK-NEXT: fmov s0, [[W]]
+ ; CHECK-NEXT:  dup v0.8h, v0.h[7]
  ; CHECK-NEXT: ret
  %shuffle.i = shufflevector <8 x i16> %v1, <8 x i16> undef, <1 x i32> <i32 7> 
  ret <1 x i16> %shuffle.i
 }
 
 define <1 x i16> @test_vector_dup_hv4H(<4 x i16> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_hv4H:
  ; CHECK-NEXT: dup v0.4h, v0.h[3]
  ; CHECK-NEXT: ret
  %shuffle.i = shufflevector <4 x i16> %v1, <4 x i16> undef, <1 x i32> <i32 3> 
  ret <1 x i16> %shuffle.i
 }
 
 define <1 x i32> @test_vector_dup_sv4S(<4 x i32> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_sv4S:
- ; CHECK-NEXT: mov  [[W:w[0-9]+]], v0.s[3]
- ; CHECK-NEXT: fmov s0, [[W]]
+ ; CHECK-NEXT: dup v0.4s, v0.s[3]
  ; CHECK-NEXT: ret
  %shuffle = shufflevector <4 x i32> %v1, <4 x i32> undef, <1 x i32> <i32 3> 
  ret <1 x i32> %shuffle
 }
 
 define <1 x i32> @test_vector_dup_sv2S(<2 x i32> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_sv2S:
  ; CHECK-NEXT: dup v0.2s, v0.s[1]
  ; CHECK-NEXT: ret
  %shuffle = shufflevector <2 x i32> %v1, <2 x i32> undef, <1 x i32> <i32 1> 
  ret <1 x i32> %shuffle
 }
 
 define <1 x i64> @test_vector_dup_dv2D(<2 x i64> %v1) #0 {
  ; CHECK-LABEL: test_vector_dup_dv2D:
  ; CHECK-NEXT: ext {{v[0-9]+}}.16b, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b, #8
  ; CHECK-NEXT: ret
  %shuffle.i = shufflevector <2 x i64> %v1, <2 x i64> undef, <1 x i32> <i32 1> 
  ret <1 x i64> %shuffle.i
 }
 
 define <1 x i64> @test_vector_copy_dup_dv2D(<1 x i64> %a, <2 x i64> %c) #0 {
   ; CHECK-LABEL: test_vector_copy_dup_dv2D:
-  ; CHECK-NEXT: {{dup|mov}} {{d[0-9]+}}, {{v[0-9]+}}.d[1]
+  ; CHECK-NEXT: dup v0.2d, v1.d[1]
   ; CHECK-NEXT: ret
   %vget_lane = extractelement <2 x i64> %c, i32 1
   %vset_lane = insertelement <1 x i64> undef, i64 %vget_lane, i32 0
   ret <1 x i64> %vset_lane
 }
 
 ; Undefined behaviour, so we really don't care what actually gets emitted, just
 ; as long as we don't crash (since it could be dynamically unreachable).