[SDAG] Expand MergedConsecutiveStores to better handle Giving up in Chain Analysis

[niravd]

Summary

For large blocks of code, FindBetterChain will give up leaving memory operations in their initial configuration. A common idiom for this results in a long chain of stores which are generally mergable. Augment the search to consider these stores as merge candidates.

[niravd]

Add testcase and minor cleanup.

[niravd]

Ping

[niravd]

Ping.

[niravd]

Ping.

[aditya_nandakumar]

Hi Nirav
I tried reproducing it with the patch and the issue is still not fixed. I'll try to get back with a repro case (for open source targets).

[aditya_nandakumar]

Nirav - I posted the reason for the regression and a quick fix for the issue I saw in https://reviews.llvm.org/D31914.

[niravd]

I'm a little concerned about the selection candidates having a chain loop dependence after mething. Can you get a llvm testcase I can look at myself. Alternatively I can probably cobble one myself if I can get a snapshot of the DAG right before the relevant merge

In D31068#723192, @aditya_nandakumar wrote:

Nirav - I posted the reason for the regression and a quick fix for the issue I saw in https://reviews.llvm.org/D31914.

[niravd]

Rebasing.

[niravd]

Ping.

[niravd]

Ping.

[courbet]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12217

Please rebase, you'll have to go through CandidateMatch's Offset parameter now.

test/CodeGen/X86/stores-merging2.ll

16 ↗

(On Diff #97480)

This test fails to fail before this change when rebased on master. I guess this is because the value of GatherAllAliasesMaxDepth has been changed from 6 to 18.

The test quite brittle against that kind of config changes BTW. You should at least leave a comment for the future explaining how to fix the test if changing the value.

[niravd]

Rebase to head. Some changes to alias make this require a much longer chain of stores. The need is clearly less than when I wrote this patch a year ago, but still has a minor positive effect.

[courbet]

In D31068#1267984, @niravd wrote:

Rebase to head. Some changes to alias make this require a much longer chain of stores. The need is clearly less than when I wrote this patch a year ago, but still has a minor positive effect.

Do you think likely to happen in practice ?

[niravd]

It's certainly less likely than before, but it does happen in the wild, but D53552 should always parallelize these store chains, so this this should be droppable.

In D31068#1272588, @courbet wrote:

In D31068#1267984, @niravd wrote:

Rebase to head. Some changes to alias make this require a much longer chain of stores. The need is clearly less than when I wrote this patch a year ago, but still has a minor positive effect.

Do you think likely to happen in practice ?

In D31068#1272588, @courbet wrote:

In D31068#1267984, @niravd wrote:

Rebase to head. Some changes to alias make this require a much longer chain of stores. The need is clearly less than when I wrote this patch a year ago, but still has a minor positive effect.

Do you think likely to happen in practice ?

[niravd]

This is no longer necessary now that r346432 has landed.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     /// true if the (and (load x) c) pattern matches an extload.  ExtVT returns
     /// the type of the loaded value to be extended.  LoadedVT returns the type
     /// of the original loaded value.  NarrowLoad returns whether the load would
     /// need to be narrowed in order to match.
     bool isAndLoadExtLoad(ConstantSDNode *AndC, LoadSDNode *LoadN,
                           EVT LoadResultTy, EVT &ExtVT, EVT &LoadedVT,
                           bool &NarrowLoad);
 
+    //// Helper function for MergeConsecutiveStores which merges the
+    //// component store chains.
+    SDValue getMergeStoreChains(SmallVectorImpl<MemOpLink> &StoreNodes,
+                                unsigned NumStores);
+
     /// 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 NumStores,
                                          bool IsConstantSrc, bool UseVector);
 
if (Use->getOpcode() == ISD::MUL) { // We have another multiply use.
         return true;
     }
   }
 
   // Didn't find a case where this would be profitable.
   return false;
 }
 
+SDValue DAGCombiner::getMergeStoreChains(SmallVectorImpl<MemOpLink> &StoreNodes,
+                                         unsigned NumStores) {
+  SmallVector<SDValue, 8> Chains;
+  SmallPtrSet<const SDNode *, 8> Visited;
+  SDLoc StoreDL(StoreNodes[0].MemNode);
+
+  for (unsigned i = 0; i < NumStores; ++i) {
+    Visited.insert(StoreNodes[i].MemNode);
+  }
+
+  // don't include nodes that are children
+  for (unsigned i = 0; i < NumStores; ++i) {
+    if (Visited.count(StoreNodes[i].MemNode->getChain().getNode()) == 0)
+      Chains.push_back(StoreNodes[i].MemNode->getChain());
+  }
+
+  assert(Chains.size() > 0 && "Chain should have generated a chain");
+  return DAG.getNode(ISD::TokenFactor, StoreDL, MVT::Other, Chains);
+}
+
 bool DAGCombiner::MergeStoresOfConstantsOrVecElts(
                   SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT,
                   unsigned NumStores, bool IsConstantSrc, bool UseVector) {
   // Make sure we have something to merge.
   if (NumStores < 2)
     return false;
 
   int64_t ElementSizeBytes = MemVT.getSizeInBits() / 8;
for (unsigned i = 0; i < NumStores; ++i) {
       }
     }
 
     // Create the new Load and Store operations.
     EVT StoreTy = EVT::getIntegerVT(*DAG.getContext(), SizeInBits);
     StoredVal = DAG.getConstant(StoreInt, DL, StoreTy);
   }
 
-  SmallVector<SDValue, 8> Chains;
-
-  // Gather all Chains we're inheriting. As generally all chains are
-  // equal, do minor check to remove obvious redundancies.
-  Chains.push_back(StoreNodes[0].MemNode->getChain());
-  for (unsigned i = 1; i < NumStores; ++i)
-    if (StoreNodes[0].MemNode->getChain() != StoreNodes[i].MemNode->getChain())
-      Chains.push_back(StoreNodes[i].MemNode->getChain());
-
   LSBaseSDNode *FirstInChain = StoreNodes[0].MemNode;
-  SDValue NewChain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, Chains);
+  SDValue NewChain = getMergeStoreChains(StoreNodes, NumStores);
   SDValue NewStore = DAG.getStore(NewChain, DL, StoredVal,
                                   FirstInChain->getBasePtr(),
                                   FirstInChain->getPointerInfo(),
                                   FirstInChain->getAlignment());
 
   // Replace all merged stores with the new store.
   for (unsigned i = 0; i < NumStores; ++i)
     CombineTo(StoreNodes[i].MemNode, NewStore);
void DAGCombiner::getStoreMergeCandidates(
   // We must have a base and an offset.
   if (!BasePtr.Base.getNode())
     return;
 
   // Do not handle stores to undef base pointers.
   if (BasePtr.Base.isUndef())
     return;
 
+  bool IsLoadSrc = isa<LoadSDNode>(St->getValue());
+  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);
+  auto CandidateMatch = [&](StoreSDNode *Other, BaseIndexOffset &Ptr) -> bool {
+    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;
+    Ptr = BaseIndexOffset::match(Other->getBasePtr(), DAG);
+    if (!Ptr.equalBaseIndex(BasePtr))
+      return false;
+    if (IsLoadSrc)
+      return isa<LoadSDNode>(Other->getValue());
+    if (IsConstantSrc)
+      return (isa<ConstantSDNode>(Other->getValue()) ||
+              isa<ConstantFPSDNode>(Other->getValue()));
+    if (IsExtractVecSrc)
+      return (Other->getValue().getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
+              Other->getValue().getOpcode() == ISD::EXTRACT_SUBVECTOR);
+    return false;
+  };
+
+  // In the case FindBetterChain gives up, we capture direct chains of Stores
+
+  StoreSDNode *STChain = St;
+
+  int64_t LastOffset = BasePtr.Offset;
+  while (true) {
+    BaseIndexOffset Ptr;
+    if (StoreSDNode *Chain = dyn_cast<StoreSDNode>(STChain->getChain())) {
+      if (CandidateMatch(Chain, Ptr) && Ptr.Offset <= LastOffset) {
+        StoreNodes.push_back(MemOpLink(STChain, LastOffset));
+        LastOffset = Ptr.Offset;
+        STChain = Chain;
+        continue;
+      }
+    }
+    break;
+  }
+
+  SDNode *RootNode = (STChain->getChain()).getNode();
+
   // 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
   // the search to find all valid candidates through multiple layers of loads.
   //
   // Root
   // |-------|-------|
   // Load    Load    Store3
   // |       |
-  // Store1   Store2
+  // Store1  Store2
   //
   // FIXME: We should be able to climb and
   // descend TokenFactors to find candidates as well.
 
-  SDNode *RootNode = (St->getChain()).getNode();
-
   // Set of Parents of Candidates
   std::set<SDNode *> CandidateParents;
 
   if (LoadSDNode *Ldn = dyn_cast<LoadSDNode>(RootNode)) {
     RootNode = Ldn->getChain().getNode();
     for (auto I = RootNode->use_begin(), E = RootNode->use_end(); I != E; ++I)
       if (I.getOperandNo() == 0 && isa<LoadSDNode>(*I)) // walk down chain
         CandidateParents.insert(*I);
   } else
     CandidateParents.insert(RootNode);
 
-  bool IsLoadSrc = isa<LoadSDNode>(St->getValue());
-  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);
-  auto CorrectValueKind = [&](StoreSDNode *Other) -> bool {
-    if (IsLoadSrc)
-      return isa<LoadSDNode>(Other->getValue());
-    if (IsConstantSrc)
-      return (isa<ConstantSDNode>(Other->getValue()) ||
-              isa<ConstantFPSDNode>(Other->getValue()));
-    if (IsExtractVecSrc)
-      return (Other->getValue().getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
-              Other->getValue().getOpcode() == ISD::EXTRACT_SUBVECTOR);
-    return false;
-  };
-
   // check all parents of mergable children
   for (auto P = CandidateParents.begin(); P != CandidateParents.end(); ++P)
     for (auto I = (*P)->use_begin(), E = (*P)->use_end(); I != E; ++I)
       if (I.getOperandNo() == 0)
         if (StoreSDNode *OtherST = dyn_cast<StoreSDNode>(*I)) {
-          if (OtherST->isVolatile() || OtherST->isIndexed())
-            continue;
-          // We can merge constant floats to equivalent integers
-          if (OtherST->getMemoryVT() != MemVT)
-            if (!(MemVT.isInteger() && MemVT.bitsEq(OtherST->getMemoryVT()) &&
-                  isa<ConstantFPSDNode>(OtherST->getValue())))
-              continue;
-          BaseIndexOffset Ptr =
-              BaseIndexOffset::match(OtherST->getBasePtr(), DAG);
-          if (Ptr.equalBaseIndex(BasePtr) && CorrectValueKind(OtherST))
+          BaseIndexOffset Ptr;
+          if (CandidateMatch(OtherST, Ptr))
             StoreNodes.push_back(MemOpLink(OtherST, Ptr.Offset));
         }
 }
 
 // We need to check that merging these stores does not cause a loop
 // in the DAG. Any store candidate may depend on another candidate
 // indirectly through its operand (we already consider dependencies
 // through the chain). Check in parallel by searching up from
bool DAGCombiner::MergeConsecutiveStores(StoreSDNode *St) {
   // We add +1 here because the LastXXX variables refer to location while
   // the NumElem refers to array/index size.
   unsigned NumElem = std::min(NumConsecutiveStores, LastConsecutiveLoad + 1);
   NumElem = std::min(LastLegalType, NumElem);
 
   if (NumElem < 2)
     return false;
 
-  // Collect the chains from all merged stores. Because the common case
-  // all chains are the same, check if we match the first Chain.
-  SmallVector<SDValue, 8> MergeStoreChains;
-  MergeStoreChains.push_back(StoreNodes[0].MemNode->getChain());
-  for (unsigned i = 1; i < NumElem; ++i)
-    if (StoreNodes[0].MemNode->getChain() != StoreNodes[i].MemNode->getChain())
-      MergeStoreChains.push_back(StoreNodes[i].MemNode->getChain());
-
   // 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);
   } else {
     unsigned SizeInBits = NumElem * ElementSizeBytes * 8;
     JointMemOpVT = EVT::getIntegerVT(Context, SizeInBits);
   }
 
   SDLoc LoadDL(LoadNodes[0].MemNode);
   SDLoc StoreDL(StoreNodes[0].MemNode);
 
   // The merged loads are required to have the same incoming chain, so
   // using the first's chain is acceptable.
   SDValue NewLoad = DAG.getLoad(JointMemOpVT, LoadDL, FirstLoad->getChain(),
                                 FirstLoad->getBasePtr(),
                                 FirstLoad->getPointerInfo(), FirstLoadAlign);
 
-  SDValue NewStoreChain =
-    DAG.getNode(ISD::TokenFactor, StoreDL, MVT::Other, MergeStoreChains);
+  SDValue NewStoreChain = getMergeStoreChains(StoreNodes, NumElem);
 
   AddToWorklist(NewStoreChain.getNode());
 
   SDValue NewStore =
       DAG.getStore(NewStoreChain, StoreDL, NewLoad, FirstInChain->getBasePtr(),
                    FirstInChain->getPointerInfo(), FirstStoreAlign);
 
   // Transfer chain users from old loads to the new load.