merge vector stores into wider vector stores (PR21711)

[spatel]

Summary

This is, hopefully, the final patch needed to resolve PR21711 ( http://llvm.org/bugs/show_bug.cgi?id=21711 ).

The 'f3' test case in that report presents a situation where we have two 128-bit stores extracted from a 256-bit source vector. Instead of producing this:

vmovaps %xmm0, (%rdi)
vextractf128    $1, %ymm0, 16(%rdi)

This patch merges the 128-bit stores into a single 256-bit store:

vmovups %ymm0, (%rdi)

To minimize changes, this patch is limited to handling the single pattern of extract_subvectors feeding into stores, but I've included test cases for the other store merge patterns that could be handled similarly: build_vectors of constants and vector loads feeding into vector stores that could be widened.

[hfinkel]

LGTM.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9860

Indenting looks odd here.

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9860

Thanks, Hal! Fixed indent and checked in at r227242.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

struct MemOpLink {
       unsigned SequenceNum;
     };
 
     /// This is a helper function for MergeConsecutiveStores. When the source
     /// elements of the consecutive stores are all constants or all extracted
     /// vector elements, try to merge them into one larger store.
     /// \return True if a merged store was created.
     bool MergeStoresOfConstantsOrVecElts(SmallVectorImpl<MemOpLink> &StoreNodes,
-                                         EVT MemVT, unsigned NumElem,
+                                         EVT MemVT, unsigned StoresToMerge,
                                          bool IsConstantSrc, bool UseVector);
     
     /// Merge consecutive store operations into a wide store.
     /// This optimization uses wide integers or vectors when possible.
     /// \return True if some memory operations were changed.
     bool MergeConsecutiveStores(StoreSDNode *N);
 
     /// \brief Try to transform a truncation where C is a constant:
static BaseIndexOffset match(SDValue Ptr) {
 
     int64_t Off = cast<ConstantSDNode>(Offset)->getSExtValue();
     return BaseIndexOffset(Base, Index, Off, IsIndexSignExt);
   }
 };
 
 bool DAGCombiner::MergeStoresOfConstantsOrVecElts(
                   SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT,
-                  unsigned NumElem, bool IsConstantSrc, bool UseVector) {
+                  unsigned StoresToMerge, bool IsConstantSrc, bool UseVector) {
   // Make sure we have something to merge.
-  if (NumElem < 2)
+  if (StoresToMerge < 2)
     return false;
   
   int64_t ElementSizeBytes = MemVT.getSizeInBits() / 8;
   LSBaseSDNode *FirstInChain = StoreNodes[0].MemNode;
   unsigned EarliestNodeUsed = 0;
   
-  for (unsigned i=0; i < NumElem; ++i) {
+  for (unsigned i=0; i < StoresToMerge; ++i) {
     // Find a chain for the new wide-store operand. Notice that some
     // of the store nodes that we found may not be selected for inclusion
     // in the wide store. The chain we use needs to be the chain of the
     // earliest store node which is *used* and replaced by the wide store.
     if (StoreNodes[i].SequenceNum > StoreNodes[EarliestNodeUsed].SequenceNum)
       EarliestNodeUsed = i;
   }
   
   // The earliest Node in the DAG.
   LSBaseSDNode *EarliestOp = StoreNodes[EarliestNodeUsed].MemNode;
   SDLoc DL(StoreNodes[0].MemNode);
   
   SDValue StoredVal;
   if (UseVector) {
-    // Find a legal type for the vector store.
-    EVT Ty = EVT::getVectorVT(*DAG.getContext(), MemVT, NumElem);
+    bool IsVec = MemVT.isVector();
+    unsigned Elts = StoresToMerge;
+    if (IsVec) {
+      // When merging vector stores, get the total number of elements.
+      Elts *= MemVT.getVectorNumElements();
+    }
+    // Get the type for the merged vector store.
+    EVT Ty = EVT::getVectorVT(*DAG.getContext(), MemVT.getScalarType(), Elts);
     assert(TLI.isTypeLegal(Ty) && "Illegal vector store");
+
     if (IsConstantSrc) {
       // A vector store with a constant source implies that the constant is
       // zero; we only handle merging stores of constant zeros because the zero
       // can be materialized without a load.
       // It may be beneficial to loosen this restriction to allow non-zero
       // store merging.
       StoredVal = DAG.getConstant(0, Ty);
     } else {
       SmallVector<SDValue, 8> Ops;
-      for (unsigned i = 0; i < NumElem ; ++i) {
+      for (unsigned i = 0; i < StoresToMerge ; ++i) {
         StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
         SDValue Val = St->getValue();
-        // All of the operands of a BUILD_VECTOR must have the same type.
+        // All operands of BUILD_VECTOR / CONCAT_VECTOR must have the same type.
         if (Val.getValueType() != MemVT)
           return false;
         Ops.push_back(Val);
       }
-      
       // Build the extracted vector elements back into a vector.
-      StoredVal = DAG.getNode(ISD::BUILD_VECTOR, DL, Ty, Ops);
+      StoredVal = DAG.getNode(IsVec ? ISD::CONCAT_VECTORS : ISD::BUILD_VECTOR,
+                              DL, Ty, Ops);
     }
   } else {
     // We should always use a vector store when merging extracted vector
     // elements, so this path implies a store of constants.
     assert(IsConstantSrc && "Merged vector elements should use vector store");
 
-    unsigned StoreBW = NumElem * ElementSizeBytes * 8;
+    unsigned StoreBW = StoresToMerge * ElementSizeBytes * 8;
     APInt StoreInt(StoreBW, 0);
     
     // Construct a single integer constant which is made of the smaller
     // constant inputs.
     bool IsLE = TLI.isLittleEndian();
-    for (unsigned i = 0; i < NumElem ; ++i) {
-      unsigned Idx = IsLE ? (NumElem - 1 - i) : i;
+    for (unsigned i = 0; i < StoresToMerge ; ++i) {
+      unsigned Idx = IsLE ? (StoresToMerge - 1 - i) : i;
       StoreSDNode *St  = cast<StoreSDNode>(StoreNodes[Idx].MemNode);
       SDValue Val = St->getValue();
       StoreInt <<= ElementSizeBytes*8;
       if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(Val)) {
         StoreInt |= C->getAPIntValue().zext(StoreBW);
       } else if (ConstantFPSDNode *C = dyn_cast<ConstantFPSDNode>(Val)) {
         StoreInt |= C->getValueAPF().bitcastToAPInt().zext(StoreBW);
       } else {
SDValue NewStore = DAG.getStore(EarliestOp->getChain(), DL, StoredVal,
                                   FirstInChain->getBasePtr(),
                                   FirstInChain->getPointerInfo(),
                                   false, false,
                                   FirstInChain->getAlignment());
   
   // Replace the first store with the new store
   CombineTo(EarliestOp, NewStore);
   // Erase all other stores.
-  for (unsigned i = 0; i < NumElem ; ++i) {
+  for (unsigned i = 0; i < StoresToMerge ; ++i) {
     if (StoreNodes[i].MemNode == EarliestOp)
       continue;
     StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
     // ReplaceAllUsesWith will replace all uses that existed when it was
     // called, but graph optimizations may cause new ones to appear. For
     // example, the case in pr14333 looks like
     //
     //  St's chain -> St -> another store -> X
     //
     // And the only difference from St to the other store is the chain.
     // When we change it's chain to be St's chain they become identical,
     // get CSEed and the net result is that X is now a use of St.
     // Since we know that St is redundant, just iterate.
     while (!St->use_empty())
       DAG.ReplaceAllUsesWith(SDValue(St, 0), St->getChain());
     deleteAndRecombine(St);
   }
   
   return true;
 }
 
 bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
-  EVT MemVT = St->getMemoryVT();
-  int64_t ElementSizeBytes = MemVT.getSizeInBits()/8;
   bool NoVectors = DAG.getMachineFunction().getFunction()->getAttributes().
     hasAttribute(AttributeSet::FunctionIndex, Attribute::NoImplicitFloat);
 
-  // Don't merge vectors into wider inputs.
-  if (MemVT.isVector() || !MemVT.isSimple())
-    return false;
-
   // Perform an early exit check. Do not bother looking at stored values that
   // are not constants, loads, or extracted vector elements.
   SDValue StoredVal = St->getValue();
   bool IsLoadSrc = isa<LoadSDNode>(StoredVal);
   bool IsConstantSrc = isa<ConstantSDNode>(StoredVal) ||
                        isa<ConstantFPSDNode>(StoredVal);
-  bool IsExtractVecEltSrc = (StoredVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT);
+  bool IsExtractVecSrc = (StoredVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
+                             StoredVal.getOpcode() == ISD::EXTRACT_SUBVECTOR);
    
-  if (!IsConstantSrc && !IsLoadSrc && !IsExtractVecEltSrc)
+  if (!IsConstantSrc && !IsLoadSrc && !IsExtractVecSrc)
     return false;
 
+  EVT MemVT = St->getMemoryVT();
+  int64_t ElementSizeBytes = MemVT.getSizeInBits()/8;
+
+  // 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.
+  // There's no such thing as a vector constant node; that merge case should be
+  // handled by looking through a BUILD_VECTOR source with all constant inputs.
+  if (MemVT.isVector() && IsLoadSrc)
+    return false;
+  
+  if (!MemVT.isSimple())
+    return false;
+  
+
   // Only look at ends of store sequences.
   SDValue Chain = SDValue(St, 0);
   if (Chain->hasOneUse() && Chain->use_begin()->getOpcode() == ISD::STORE)
     return false;
 
   // This holds the base pointer, index, and the offset in bytes from the base
   // pointer.
   BaseIndexOffset BasePtr = BaseIndexOffset::match(St->getBasePtr());
if (IsConstantSrc) {
     unsigned NumElem = UseVector ? LastLegalVectorType : LastLegalType;
 
     return MergeStoresOfConstantsOrVecElts(StoreNodes, MemVT, NumElem,
                                            true, UseVector);
   }
 
   // When extracting multiple vector elements, try to store them
   // in one vector store rather than a sequence of scalar stores.
-  if (IsExtractVecEltSrc) {
-    unsigned NumElem = 0;
+  if (IsExtractVecSrc) {
+    unsigned StoresToMerge = 0;
+    bool IsVec = MemVT.isVector();
     for (unsigned i = 0; i < LastConsecutiveStore + 1; ++i) {
       StoreSDNode *St  = cast<StoreSDNode>(StoreNodes[i].MemNode);
-      SDValue StoredVal = St->getValue();
+      unsigned StoreValOpcode = St->getValue().getOpcode();
       // This restriction could be loosened.
       // Bail out if any stored values are not elements extracted from a vector.
       // It should be possible to handle mixed sources, but load sources need
       // more careful handling (see the block of code below that handles
       // consecutive loads).
-      if (StoredVal.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
+      if (StoreValOpcode != ISD::EXTRACT_VECTOR_ELT &&
+          StoreValOpcode != ISD::EXTRACT_SUBVECTOR)
         return false;
       
       // Find a legal type for the vector store.
-      EVT Ty = EVT::getVectorVT(*DAG.getContext(), MemVT, i+1);
-      if (TLI.isTypeLegal(Ty))
-        NumElem = i + 1;
+      unsigned Elts = i + 1;
+      if (IsVec) {
+        // When merging vector stores, get the total number of elements.
+        Elts *= MemVT.getVectorNumElements();
+      }
+      if (TLI.isTypeLegal(EVT::getVectorVT(*DAG.getContext(),
+                                           MemVT.getScalarType(), Elts)))
+        StoresToMerge = i + 1;
     }
 
-    return MergeStoresOfConstantsOrVecElts(StoreNodes, MemVT, NumElem,
+    return MergeStoresOfConstantsOrVecElts(StoreNodes, MemVT, StoresToMerge,
                                            false, true);
   }
 
   // Below we handle the case of multiple consecutive stores that
   // come from multiple consecutive loads. We merge them into a single
   // wide load and a single wide store.
 
   // Look for load nodes which are used by the stored values.

test/CodeGen/X86/MergeConsecutiveStores.ll

define void @merge_vec_element_store(<8 x float> %v, float* %ptr) {
   ret void
 
 ; CHECK-LABEL: merge_vec_element_store
 ; CHECK: vmovups
 ; CHECK-NEXT: vzeroupper
 ; CHECK-NEXT: retq
 }
 
+; PR21711 - Merge vector stores into wider vector stores.
+define void @merge_vec_extract_stores(<8 x float> %v1, <8 x float> %v2, <4 x float>* %ptr) {
+  %idx0 = getelementptr inbounds <4 x float>* %ptr, i64 3
+  %idx1 = getelementptr inbounds <4 x float>* %ptr, i64 4
+  %idx2 = getelementptr inbounds <4 x float>* %ptr, i64 5
+  %idx3 = getelementptr inbounds <4 x float>* %ptr, i64 6
+  %shuffle0 = shufflevector <8 x float> %v1, <8 x float> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %shuffle1 = shufflevector <8 x float> %v1, <8 x float> undef, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
+  %shuffle2 = shufflevector <8 x float> %v2, <8 x float> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %shuffle3 = shufflevector <8 x float> %v2, <8 x float> undef, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
+  store <4 x float> %shuffle0, <4 x float>* %idx0, align 16
+  store <4 x float> %shuffle1, <4 x float>* %idx1, align 16
+  store <4 x float> %shuffle2, <4 x float>* %idx2, align 16
+  store <4 x float> %shuffle3, <4 x float>* %idx3, align 16
+  ret void
+
+; CHECK-LABEL: merge_vec_extract_stores
+; CHECK:      vmovups %ymm0, 48(%rdi)
+; CHECK-NEXT: vmovups %ymm1, 80(%rdi)
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+}
+
+; Merging vector stores when sourced from vector loads is not currently handled.
+define void @merge_vec_stores_from_loads(<4 x float>* %v, <4 x float>* %ptr) {
+  %load_idx0 = getelementptr inbounds <4 x float>* %v, i64 0
+  %load_idx1 = getelementptr inbounds <4 x float>* %v, i64 1
+  %v0 = load <4 x float>* %load_idx0
+  %v1 = load <4 x float>* %load_idx1
+  %store_idx0 = getelementptr inbounds <4 x float>* %ptr, i64 0
+  %store_idx1 = getelementptr inbounds <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-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>* %ptr, i64 3
+  %idx1 = getelementptr inbounds <4 x i32>* %ptr, i64 4
+  store <4 x i32> <i32 0, i32 0, i32 0, i32 0>, <4 x i32>* %idx0, align 16
+  store <4 x i32> <i32 0, i32 0, i32 0, i32 0>, <4 x i32>* %idx1, align 16
+  ret void
+
+; CHECK-LABEL: merge_vec_stores_of_constants
+; CHECK:      vxorps
+; CHECK-NEXT: vmovaps
+; CHECK-NEXT: vmovaps
+; CHECK-NEXT: retq
+}
+
 ; This is a minimized test based on real code that was failing.
 ; We could merge stores (and loads) like this...
 
 define void @merge_vec_element_and_scalar_load([6 x i64]* %array) {
   %idx0 = getelementptr inbounds [6 x i64]* %array, i64 0, i64 0
   %idx1 = getelementptr inbounds [6 x i64]* %array, i64 0, i64 1
   %idx4 = getelementptr inbounds [6 x i64]* %array, i64 0, i64 4
   %idx5 = getelementptr inbounds [6 x i64]* %array, i64 0, i64 5