merge consecutive stores of extracted vector elements (PR21711)

[spatel]

Summary

This is a 2nd try at the same optimization as http://reviews.llvm.org/D6698. That patch was checked in at r224611, but reverted at r225031 because it caused a failure.

The cause of the crash was not recognizing consecutive stores that have mixed source values (loads and vector element extracts), so I've just added a check to bail out if any store value is not coming from a vector element extract.

[hfinkel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9939

How does this cause a crash if you remove this check?

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9939

For code like the second test case, we merge the stores and later during legalization get an error like this for the vector load:
Operand not processed?
0x5c4c770: v2i64,ch = load 0x5c4ac60, 0x5c4c660, 0x5c45850<LD16[%sunkaddr115](tbaa=<0x5b6c2a8>)> [ORD=61] [ID=1]

I haven't figured out exactly what happens to cause that.

[hfinkel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9939

Is that load supposed to be dead? What did you do with users of the load's chain result (which I assume might include at least one of the stores you're merging).

9998

Maybe if the load appeared in the chain in between these stores somewhere, you'll cause a cycle?

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9939

I didn't do anything with the load chain because I wasn't expecting any load-based sources in this code path. :)

Based on the later code (starting around line 9857) that handles sequences of loads, we need to replace the existing loads:

// Replace one of the loads with the new load.
LoadSDNode *Ld = cast<LoadSDNode>(LoadNodes[0].MemNode);
DAG.ReplaceAllUsesOfValueWith(SDValue(Ld, 1),
                              SDValue(NewLoad.getNode(), 1));

// Remove the rest of the load chains.

There are also a lot of other clauses down there that need to be met when replacing loads, so I was trying to handle only vector element extracts with this patch (this is similar to handling constant sources for the stores as in the clause above this one).

If it's ok with you, I can make this clearer in the comment here and then take a shot at the mixed source case later. I'm guessing the mixed store source case (some vector element extracts and some loads) isn't too common.

[spatel]

Ping.

Updated comment to hopefully better explain the mixed source case that is not currently handled.

[spatel]

Ping * 2.

[qcolombet]

Hi Sanjay,

This looks almost good. Indeed, I believe we can refactor the code to avoid some code duplication. See my inline comment.

Thanks,
-Quentin

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9950

From this point, AFAICT, the code is exactly very similar to the constant version.
Please refractor the code with a helper method.

[spatel]

Thanks, Quentin. Yes, there's a big block of copy/paste here. I've moved the duplicated logic into a helper function in this version of the patch.

[qcolombet]

Hi Sanjay,

LGTM with some nitpicks.

No need for another round of review.

Thanks,
-Quentin

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

384

Use SmallVectorImpl<MemOpLink> instead.

9662

I would do one more refactoring here.
I would just check for UseVector and having an assert that if IsConstantSrc is true then UseVector must be true as well.

That way instead of having:
if (!IsConstantSrc) {

// look for type.
// loop.
// set res val

} else if (UseVector) {

// look for type.
// set res val.

} else {
}

We would have:
if (UseVector) {

// look for type.
if (!IsConstantSrc) {
  // loop.
  // set val.
} else {
 // set val.
}

} else {
}

The reasons why I suggest this (I am fine if you do not want to do that), are:

1. The type creation will not be duplicated as well as the assert.
2. At first glanced I found confusing to have !IsConstantSrc and UseVector ending in different blocks. Indeed, per the comment, if it is not a constant, this is a vector so having a vector not using vector is confusing.

The alternative would be to add a comment on the actual usage of UseVector in the function prototype.

[spatel]

In D6850#111943, @qcolombet wrote:

LGTM with some nitpicks.

Thanks, Quentin! I implemented your 2nd refactoring suggestion; I think that makes it easier to understand.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     /// single-use) and if missed an empty SDValue is returned.
     SDValue distributeTruncateThroughAnd(SDNode *N);
 
   public:
     DAGCombiner(SelectionDAG &D, AliasAnalysis &A, CodeGenOpt::Level OL)
         : DAG(D), TLI(D.getTargetLoweringInfo()), Level(BeforeLegalizeTypes),
           OptLevel(OL), LegalOperations(false), LegalTypes(false), AA(A) {
       AttributeSet FnAttrs =
           DAG.getMachineFunction().getFunction()->getAttributes();
       ForCodeSize =
           FnAttrs.hasAttribute(AttributeSet::FunctionIndex,
                                Attribute::OptimizeForSize) ||
           FnAttrs.hasAttribute(AttributeSet::FunctionIndex, Attribute::MinSize);
     }
 
     /// Runs the dag combiner on all nodes in the work list
     void Run(CombineLevel AtLevel);
bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
   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 or loads.
+  // are not constants, loads, or extracted vector elements.
   SDValue StoredVal = St->getValue();
   bool IsLoadSrc = isa<LoadSDNode>(StoredVal);
-  if (!isa<ConstantSDNode>(StoredVal) && !isa<ConstantFPSDNode>(StoredVal) &&
-      !IsLoadSrc)
+  bool IsConstantSrc = isa<ConstantSDNode>(StoredVal) ||
+                       isa<ConstantFPSDNode>(StoredVal);
+  bool IsExtractVecEltSrc = (StoredVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT);
+   
+  if (!IsConstantSrc && !IsLoadSrc && !IsExtractVecEltSrc)
     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());
 
   // We must have a base and an offset.
   if (!BasePtr.Base.getNode())
for (unsigned i = 0, e = StoreNodes.size(); i < e; ++i) {
     // Mark this node as useful.
     LastConsecutiveStore = i;
   }
 
   // The node with the lowest store address.
   LSBaseSDNode *FirstInChain = StoreNodes[0].MemNode;
 
   // Store the constants into memory as one consecutive store.
-  if (!IsLoadSrc) {
+  if (IsConstantSrc) {
     unsigned LastLegalType = 0;
     unsigned LastLegalVectorType = 0;
     bool NonZero = false;
     for (unsigned i=0; i<LastConsecutiveStore+1; ++i) {
       StoreSDNode *St  = cast<StoreSDNode>(StoreNodes[i].MemNode);
       SDValue StoredVal = St->getValue();
 
       if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(StoredVal)) {
for (unsigned i = 0; i < NumElem ; ++i) {
       while (!St->use_empty())
         DAG.ReplaceAllUsesWith(SDValue(St, 0), St->getChain());
       deleteAndRecombine(St);
     }
 
     return true;
   }
 
+  // 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;
+    for (unsigned i = 0; i < LastConsecutiveStore + 1; ++i) {
+      StoreSDNode *St  = cast<StoreSDNode>(StoreNodes[i].MemNode);
+      SDValue StoredVal = St->getValue();
+      // 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)
+        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;
+    }
+
+    // Make sure we have a legal type and something to merge.
+    if (NumElem < 2)
+      return false;
+
+    unsigned EarliestNodeUsed = 0;
+    for (unsigned i=0; i < NumElem; ++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;
+
+    // Find a legal type for the vector store.
+    EVT Ty = EVT::getVectorVT(*DAG.getContext(), MemVT, NumElem);
+
+    SmallVector<SDValue, 8> Ops;
+    for (unsigned i = 0; i < NumElem ; ++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.
+      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);
+
+    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) {
+      if (StoreNodes[i].MemNode == EarliestOp)
+        continue;
+      StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
+      while (!St->use_empty())
+        DAG.ReplaceAllUsesWith(SDValue(St, 0), St->getChain());
+      deleteAndRecombine(St);
+    }
+
+    return 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.
   SmallVector<MemOpLink, 8> LoadNodes;
 
   // Find acceptable loads. Loads need to have the same chain (token factor),

test/CodeGen/X86/MergeConsecutiveStores.ll

; <label>:1
   %11 = getelementptr inbounds i8* %.08, i64 2
   %12 = add nsw i32 %.09, -1
   %13 = icmp eq i32 %12, 0
   br i1 %13, label %14, label %1
 
 ; <label>:14
   ret void
 }
+
+; PR21711 ( http://llvm.org/bugs/show_bug.cgi?id=21711 )
+define void @merge_vec_element_store(<8 x float> %v, float* %ptr) {
+  %vecext0 = extractelement <8 x float> %v, i32 0
+  %vecext1 = extractelement <8 x float> %v, i32 1
+  %vecext2 = extractelement <8 x float> %v, i32 2
+  %vecext3 = extractelement <8 x float> %v, i32 3
+  %vecext4 = extractelement <8 x float> %v, i32 4
+  %vecext5 = extractelement <8 x float> %v, i32 5
+  %vecext6 = extractelement <8 x float> %v, i32 6
+  %vecext7 = extractelement <8 x float> %v, i32 7
+  %arrayidx1 = getelementptr inbounds float* %ptr, i64 1
+  %arrayidx2 = getelementptr inbounds float* %ptr, i64 2
+  %arrayidx3 = getelementptr inbounds float* %ptr, i64 3
+  %arrayidx4 = getelementptr inbounds float* %ptr, i64 4
+  %arrayidx5 = getelementptr inbounds float* %ptr, i64 5
+  %arrayidx6 = getelementptr inbounds float* %ptr, i64 6
+  %arrayidx7 = getelementptr inbounds float* %ptr, i64 7
+  store float %vecext0, float* %ptr, align 4
+  store float %vecext1, float* %arrayidx1, align 4
+  store float %vecext2, float* %arrayidx2, align 4
+  store float %vecext3, float* %arrayidx3, align 4
+  store float %vecext4, float* %arrayidx4, align 4
+  store float %vecext5, float* %arrayidx5, align 4
+  store float %vecext6, float* %arrayidx6, align 4
+  store float %vecext7, float* %arrayidx7, align 4
+  ret void
+
+; CHECK-LABEL: merge_vec_element_store
+; CHECK: vmovups
+; CHECK-NEXT: vzeroupper
+; 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
+
+  %a0 = load i64* %idx0, align 8
+  store i64 %a0, i64* %idx4, align 8
+
+  %b = bitcast i64* %idx1 to <2 x i64>*
+  %v = load <2 x i64>* %b, align 8
+  %a1 = extractelement <2 x i64> %v, i32 0
+  store i64 %a1, i64* %idx5, align 8
+  ret void
+
+; CHECK-LABEL: merge_vec_element_and_scalar_load
+; CHECK:      movq  (%rdi), %rax
+; CHECK-NEXT: movq  %rax, 32(%rdi)
+; CHECK-NEXT: movq  8(%rdi), %rax
+; CHECK-NEXT: movq  %rax, 40(%rdi)
+; CHECK-NEXT: retq
+}