[DAGCombine]: Fix for regression I brought up in D31068

[aditya_nandakumar]

Summary

This is in continuation of https://reviews.llvm.org/D31068 for addressing a regression that I was seeing in our internal backends.

The DAG looks like this
t0: ch = EntryToken

      t8: i64 = add t6, Constant:i64<2>
    t12: ch = store<ST1[%dst.gep2.i105.2](align=2)> t0, Constant:i8<0>, t8, undef:i64
    t14: i64 = add t6, Constant:i64<3>
  t15: ch = store<ST1[%dst.gep2.i105.3]> t12, Constant:i8<0>, t14, undef:i64
  t17: i64 = add t6, Constant:i64<4>
t18: ch = store<ST1[%dst.gep2.i105.4](align=2)> t15, Constant:i8<0>, t17, undef:i64
  t20: i64 = add t6, Constant:i64<5>
t50: ch = store<ST1[%dst.gep2.i105.5]> t0, Constant:i8<0>, t20, undef:i64

When trying to merge t50, we'd need to walk the chain (Entry t0) and then the uses of the EntryToken and consecutively follow the uses if we keep seeing stores and add that to the worklist.
For reference, before the change in 297695, the DAG looked like

        t8: i64 = add t6, Constant:i64<2>
      t12: ch = store<ST1[%dst.gep2.i105.2](align=2)> t0, Constant:i8<0>, t8, undef:i64
      t14: i64 = add t6, Constant:i64<3>
    t15: ch = store<ST1[%dst.gep2.i105.3]> t12, Constant:i8<0>, t14, undef:i64
    t17: i64 = add t6, Constant:i64<4>
  t18: ch = store<ST1[%dst.gep2.i105.4](align=2)> t15, Constant:i8<0>, t17, undef:i64
  t20: i64 = add t6, Constant:i64<5>
t21: ch = store<ST1[%dst.gep2.i105.5]> t18, Constant:i8<0>, t20, undef:i64

It was straightforward to find candidates for merging stores.

I quickly threw in a hack for the above (extending what was in D31068) and this fixes the regression that I saw.
Looking forward to your comments/feedback.

[niravd]

I suspect the underlying issue is no longer important, but for reference D31068 is a cleaned up version of this.

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 will find chains of
+  // consecutive stores. Do a preprocessing step to walk up the chain
+  // leaving the the head of the chain for the parallel search.
+
+  StoreSDNode *STChain = St;
+  int64_t LastOffset = BasePtr.Offset;
+  while (StoreSDNode *Chain = dyn_cast<StoreSDNode>(STChain->getChain())) {
+    BaseIndexOffset Ptr;
+    //Check Chain.
+    if (!CandidateMatch(Chain, Ptr) || Ptr.Offset > LastOffset)
+      break;
+    StoreNodes.push_back(MemOpLink(STChain, LastOffset));
+    LastOffset = Ptr.Offset;
+    STChain = Chain;
+  }
+
+  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;
-  };
-
+  // If we find a store, also traverse the parent to see if it's a store
+  // we can merge
+  std::vector<SDNode*> WorklistNodes(CandidateParents.begin(), CandidateParents.end());
+  unsigned Idx = 0;
   // 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)
+  do {
+    auto *P = WorklistNodes[Idx];
+    if (StoreSDNode *OSt = dyn_cast<StoreSDNode>(P)) {
+      BaseIndexOffset Ptr;
+      if (CandidateMatch(OSt, Ptr)) {
+        StoreNodes.push_back(MemOpLink(OSt, Ptr.Offset));
+        if (StoreSDNode *OtherSTChainST = dyn_cast<StoreSDNode>(*OSt->use_begin()))
+          WorklistNodes.push_back(OtherSTChainST);
+      }
+    } else {
+      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));
+              if (StoreSDNode *OtherSTChainST = dyn_cast<StoreSDNode>(*OtherST->use_begin()))
+                WorklistNodes.push_back(OtherSTChainST);
+            }
           }
     }
+    ++Idx;
+  } while (Idx < WorklistNodes.size());
+}
 
 // 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
 // non-chain operands of candidates.
 bool DAGCombiner::checkMergeStoreCandidatesForDependencies(
     SmallVectorImpl<MemOpLink> &StoreNodes) {
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.