[DAG] Relax type restriction for store merge

[niravd]

Summary

Allow stores of bitcastable types to be merged by peeking through BITCAST nodes and recasting stored values constant and vector extract nodes as necessary.

[niravd]

Add missing check for typesize equality and explicitly avoid truncating stores.

[niravd]

Resolve outstanding comments and rebase over latned extracted NFCI patch.

[niravd]

Ignore this last patch. It should have gone to another diff. Will overwrite with correct version presently.

[RKSimon]

This looks like the D34472 patch?

[niravd]

overwrite accidental patch.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12324

Val.getOpcode()

12748

StVal.getOpcode()

[niravd]

Simpify, add comments, and address nits discussed.

[niravd]

Addressed missed nits.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12448

It'd probably be better if you avoid the ternary inside the if().

if (MemVT.isInteger() && !MemVT.bitsEq(Other->getMemoryVT()))
   return false;
if (!MemVT.isInteger() && Other->getMemoryVT() != MemVT)
   return false;

or

if ((MemVT.isInteger() && !MemVT.bitsEq(Other->getMemoryVT())) ||
    (!MemVT.isInteger() && Other->getMemoryVT() != MemVT))
   return false;

That or pull it out as a bool variable in the line above.

12458

Another ternary inside an if()

[niravd]

extract ternary operators from conditionals

[niravd]

Cleanup the mergestore helper function casting and add missing extract_subvector casting logic.

[niravd]

Rework to regularize and simplify bitcasting logic.

[niravd]

Rebase past landed NFC commits.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

483

A lot of these comment changes look like NFCs that can go in separately.

12316

Please can you create a static helper that performs this? We have a lot of peek through bitcasts in this patch alone....

12318

Val.getValueSizeInBits()

12346

Merge these ifs() to reduce indentation?

if (MemVT != Val.getValueType() &&
    (Val.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
     Val.getOpcode() == ISD::EXTRACT_SUBVECTOR) {

[niravd]

Pull out some NFC changes and use peekThroughBitcast function

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12321

Why isn't this CFP->getValueAPF().bitcastToAPInt().zextOrTrunc(8 * ElementSizeBytes) ?

12445

Pull the ternary operator out of the if()

12473

NFC change

[niravd]

Excise NFC change and address comments.

[RKSimon]

LGTM - thanks

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
 
     /// This is a helper function for MergeConsecutiveStores.
     /// Stores that may be merged are placed in StoreNodes.
     void getStoreMergeCandidates(StoreSDNode *St,
                                  SmallVectorImpl<MemOpLink> &StoreNodes);
 
     /// Helper function for MergeConsecutiveStores. Checks if
     /// Candidate stores have indirect dependency through their
     /// operands. \return True if safe to merge
     bool checkMergeStoreCandidatesForDependencies(
         SmallVectorImpl<MemOpLink> &StoreNodes, unsigned NumStores);
 
     /// Merge consecutive store operations into a wide store.
     /// This optimization uses wide integers or vectors when possible.
     /// \return number of stores that were merged into a merged store (the
     /// affected nodes are stored as a prefix in \p StoreNodes).
     bool MergeConsecutiveStores(StoreSDNode *N);
if (IsConstantSrc) {
         if (MemVT.getScalarType().isInteger())
           if (auto *CFP = dyn_cast<ConstantFPSDNode>(St->getValue()))
             Val = DAG.getConstant(
                 (uint32_t)CFP->getValueAPF().bitcastToAPInt().getZExtValue(),
                 SDLoc(CFP), MemVT);
         BuildVector.push_back(Val);
       }
       StoredVal = DAG.getBuildVector(Ty, DL, BuildVector);
     } else {
       SmallVector<SDValue, 8> Ops;
       for (unsigned i = 0; i < NumStores; ++i) {
         StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
         SDValue Val = St->getValue();
         // All operands of BUILD_VECTOR / CONCAT_VECTOR must have the same type.
+        while (Val.getValueType() != MemVT && Val->getOpcode() == ISD::BITCAST)
+          Val = Val.getOperand(0);
+        // If we see an extract of a Bitcast, and we could remove both
+        // bitcasts, do so.
+        if (Val.getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
+            Val.getOperand(0).getOpcode() == ISD::BITCAST) {
+          SDValue Vec = Val.getOperand(0).getOperand(0);
+          if (Vec.getValueType().getVectorElementType() == MemVT)
+            Val = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SDLoc(Val), MemVT, Vec,
+                              Val.getOperand(1));
+        }
+        // TODO: add this check for EXTRACT_SUBVECTOR as well.
         if (Val.getValueType() != MemVT)
           return false;
+
         Ops.push_back(Val);
       }
 
       // Build the extracted vector elements back into a vector.
       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 SizeInBits = NumStores * ElementSizeBytes * 8;
     APInt StoreInt(SizeInBits, 0);
 
     // Construct a single integer constant which is made of the smaller
     // constant inputs.
     bool IsLE = DAG.getDataLayout().isLittleEndian();
     for (unsigned i = 0; i < NumStores; ++i) {
bool DAGCombiner::MergeStoresOfConstantsOrVecElts(
   return true;
 }
 
 void DAGCombiner::getStoreMergeCandidates(
     StoreSDNode *St, SmallVectorImpl<MemOpLink> &StoreNodes) {
   // This holds the base pointer, index, and the offset in bytes from the base
   // pointer.
   BaseIndexOffset BasePtr = BaseIndexOffset::match(St->getBasePtr(), DAG);
+
   EVT MemVT = St->getMemoryVT();
 
+  SDValue Val = St->getValue();
+  while (Val.getValueType() != MemVT && Val->getOpcode() == ISD::BITCAST)
+    Val = Val.getOperand(0);
+
   // We must have a base and an offset.
   if (!BasePtr.getBase().getNode())
     return;
 
   // Do not handle stores to undef base pointers.
   if (BasePtr.getBase().isUndef())
     return;
 
   bool IsConstantSrc = isa<ConstantSDNode>(St->getValue()) ||
                        isa<ConstantFPSDNode>(St->getValue());
-  bool IsExtractVecSrc =
-      (St->getValue().getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
-       St->getValue().getOpcode() == ISD::EXTRACT_SUBVECTOR);
+
+  bool IsExtractVecSrc = (Val.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
+                          Val.getOpcode() == ISD::EXTRACT_SUBVECTOR);
   bool IsLoadSrc = isa<LoadSDNode>(St->getValue());
   BaseIndexOffset LBasePtr;
   // Match on loadbaseptr if relevant.
   if (IsLoadSrc)
     LBasePtr = BaseIndexOffset::match(
         cast<LoadSDNode>(St->getValue())->getBasePtr(), DAG);
 
   auto CandidateMatch = [&](StoreSDNode *Other, BaseIndexOffset &Ptr,
                             int64_t &Offset) -> bool {
+    bool EquivType = (Other->getMemoryVT() == MemVT) ||
+                     (MemVT.isInteger() && MemVT.bitsEq(Other->getMemoryVT()));
     if (Other->isVolatile() || Other->isIndexed())
       return false;
-    // We can merge constant floats to equivalent integers
-    if (Other->getMemoryVT() != MemVT)
-      if (!(MemVT.isInteger() && MemVT.bitsEq(Other->getMemoryVT()) &&
-            isa<ConstantFPSDNode>(Other->getValue())))
-        return false;
     if (IsLoadSrc) {
       // The Load's Base Ptr must also match
+      if (!EquivType)
+        return false;
       if (LoadSDNode *OtherLd = dyn_cast<LoadSDNode>(Other->getValue())) {
         auto LPtr = BaseIndexOffset::match(OtherLd->getBasePtr(), DAG);
         if (!(LBasePtr.equalBaseIndex(LPtr, DAG)))
           return false;
       } else
         return false;
     }
-    if (IsConstantSrc)
-      if (!(isa<ConstantSDNode>(Other->getValue()) ||
+    if (IsConstantSrc) {
+      // May truncate things converted to integers
+      if (!EquivType || !(isa<ConstantSDNode>(Other->getValue()) ||
                           isa<ConstantFPSDNode>(Other->getValue())))
         return false;
-    if (IsExtractVecSrc)
-      if (!(Other->getValue().getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
-            Other->getValue().getOpcode() == ISD::EXTRACT_SUBVECTOR))
+    }
+    if (IsExtractVecSrc) {
+      SDValue Val = Other->getValue();
+      // Peel off bitcasts.
+      bool MatchVT = false;
+      while (Val.getOpcode() == ISD::BITCAST) {
+        if (Val.getValueType() == MemVT)
+          MatchVT = true;
+        Val = Val.getOperand(0);
+      }
+      if (!(Val.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
+            Val.getOpcode() == ISD::EXTRACT_SUBVECTOR))
         return false;
+    }
+
     Ptr = BaseIndexOffset::match(Other->getBasePtr(), DAG);
     return (BasePtr.equalBaseIndex(Ptr, DAG, Offset));
   };
   // We looking for a root node which is an ancestor to all mergable
   // stores. We search up through a load, to our root and then down
   // through all children. For instance we will find Store{1,2,3} if
   // St is Store1, Store2. or Store3 where the root is not a load
   // which always true for nonvolatile ops. TODO: Expand
while (StoreNodes.size() > 1) {
     if (IsExtractVecSrc) {
       LSBaseSDNode *FirstInChain = StoreNodes[0].MemNode;
       unsigned FirstStoreAS = FirstInChain->getAddressSpace();
       unsigned FirstStoreAlign = FirstInChain->getAlignment();
       unsigned NumStoresToMerge = 1;
       bool IsVec = MemVT.isVector();
       for (unsigned i = 0; i < NumConsecutiveStores; ++i) {
         StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
-        unsigned StoreValOpcode = St->getValue().getOpcode();
+        SDValue StVal = St->getValue();
+        // peek through bitcasts
+        while (StVal->getOpcode() == ISD::BITCAST)
+          StVal = StVal.getOperand(0);
+
+        unsigned StoreValOpcode = StVal.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 (StoreValOpcode != ISD::EXTRACT_VECTOR_ELT &&
             StoreValOpcode != ISD::EXTRACT_SUBVECTOR)
           return RV;

test/CodeGen/X86/MergeConsecutiveStores.ll

define void @merge_bitcast(<4 x i32> %v, float* %ptr) {
   %idx3 = getelementptr inbounds float, float* %ptr, i64 3
   store float %f0, float* %idx0, align 4
   store float %f1, float* %idx1, align 4
   store float %f2, float* %idx2, align 4
   store float %f3, float* %idx3, align 4
   ret void
 
 ; CHECK-LABEL: merge_bitcast
-; CHECK:      vmovd %xmm0, (%rdi)
-; CHECK-NEXT: vpextrd $1, %xmm0, 4(%rdi)
-; CHECK-NEXT: vpextrd $2, %xmm0, 8(%rdi)
-; CHECK-NEXT: vpextrd $3, %xmm0, 12(%rdi)
+; CHECK:      vmovups %xmm0, (%rdi)
 ; CHECK-NEXT: retq
 }