[DAG] Allow merging of stores of vector loads

[niravd]

Summary

Remove restriction disallowing merging of stores vector loads into larger store of larger vector load.

[niravd]

This is factored from D34569

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13008

Pull out

unsigned NumSrcMemElts = MemVT.isVector() ? MemVT.getVectorNumElements() : 1;
unsigned Elts = (i + 1) * NumSrcMemElts;

as you have a second MemVT.getVectorNumElements() call later on.

13080

unsigned Elts = NumElem * NumSrcMemElts;

[niravd]

Fold away type modification into landed NFC patch

[RKSimon]

LGTM - thanks

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

bool DAGCombiner::MergeConsecutiveStores(StoreSDNode *St) {
   bool IsConstantSrc = isa<ConstantSDNode>(StoredVal) ||
                        isa<ConstantFPSDNode>(StoredVal);
   bool IsExtractVecSrc = (StoredVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
                           StoredVal.getOpcode() == ISD::EXTRACT_SUBVECTOR);
 
   if (!IsConstantSrc && !IsLoadSrc && !IsExtractVecSrc)
     return false;
 
-  // Don't merge vectors into wider vectors if the source data comes from loads.
-  // TODO: This restriction can be lifted by using logic similar to the
-  // ExtractVecSrc case.
-  if (MemVT.isVector() && IsLoadSrc)
-    return false;
-
   SmallVector<MemOpLink, 8> StoreNodes;
   // Find potential store merge candidates by searching through chain sub-DAG
   getStoreMergeCandidates(St, StoreNodes);
 
   // Check if there is anything to merge.
   if (StoreNodes.size() < 2)
     return false;
 
for (unsigned i = 1; i < LoadNodes.size(); ++i) {
       if (CurrAddress - StartAddress != (ElementSizeBytes * i))
         break;
       LastConsecutiveLoad = i;
 
       if (isDereferenceable && !LoadNodes[i].MemNode->isDereferenceable())
         isDereferenceable = false;
 
       // Find a legal type for the vector store.
-      EVT StoreTy = EVT::getVectorVT(Context, MemVT, i + 1);
+      unsigned Elts =
+          (i + 1) * ((MemVT.isVector()) ? MemVT.getVectorNumElements() : 1);
+      EVT StoreTy = EVT::getVectorVT(Context, MemVT.getScalarType(), Elts);
+
       bool IsFastSt, IsFastLd;
       if (TLI.isTypeLegal(StoreTy) &&
           TLI.canMergeStoresTo(FirstStoreAS, StoreTy, DAG) &&
           TLI.allowsMemoryAccess(Context, DL, StoreTy, FirstStoreAS,
                                  FirstStoreAlign, &IsFastSt) &&
           IsFastSt &&
           TLI.allowsMemoryAccess(Context, DL, StoreTy, FirstLoadAS,
                                  FirstLoadAlign, &IsFastLd) &&
if (NumElem < 2) {
       StoreNodes.erase(StoreNodes.begin(), StoreNodes.begin() + 1);
       continue;
     }
 
     // Find if it is better to use vectors or integers to load and store
     // to memory.
     EVT JointMemOpVT;
     if (UseVectorTy) {
-      JointMemOpVT = EVT::getVectorVT(Context, MemVT, NumElem);
+      // Find a legal type for the vector store.
+      unsigned Elts = NumElem;
+      if (MemVT.isVector())
+        Elts *= MemVT.getVectorNumElements();
+      JointMemOpVT = EVT::getVectorVT(Context, MemVT.getScalarType(), Elts);
     } else {
       unsigned SizeInBits = NumElem * ElementSizeBytes * 8;
       JointMemOpVT = EVT::getIntegerVT(Context, SizeInBits);
     }
 
     SDLoc LoadDL(LoadNodes[0].MemNode);
     SDLoc StoreDL(StoreNodes[0].MemNode);
 

test/CodeGen/X86/MergeConsecutiveStores.ll

define void @merge_vec_stores_from_loads(<4 x float>* %v, <4 x float>* %ptr) {
   %v1 = load <4 x float>, <4 x float>* %load_idx1
   %store_idx0 = getelementptr inbounds <4 x float>, <4 x float>* %ptr, i64 0
   %store_idx1 = getelementptr inbounds <4 x float>, <4 x float>* %ptr, i64 1
   store <4 x float> %v0, <4 x float>* %store_idx0, align 16
   store <4 x float> %v1, <4 x float>* %store_idx1, align 16
   ret void
 
 ; CHECK-LABEL: merge_vec_stores_from_loads
-; CHECK:      vmovaps
-; CHECK-NEXT: vmovaps
-; CHECK-NEXT: vmovaps
-; CHECK-NEXT: vmovaps
+; CHECK:      vmovups (%rdi), %ymm0
+; CHECK-NEXT: vmovups %ymm0, (%rsi)
+; CHECK-NEXT: vzeroupper      
 ; CHECK-NEXT: retq
 }
 
 ; Merging vector stores when sourced from a constant vector is not currently handled. 
 define void @merge_vec_stores_of_constants(<4 x i32>* %ptr) {
   %idx0 = getelementptr inbounds <4 x i32>, <4 x i32>* %ptr, i64 3
   %idx1 = getelementptr inbounds <4 x i32>, <4 x i32>* %ptr, i64 4
   store <4 x i32> <i32 0, i32 0, i32 0, i32 0>, <4 x i32>* %idx0, align 16