WIP: Make MergeConsecutiveStores look at other stores on same chain

[arsenm]

Summary

When combiner AA is enabled, look at stores on the same chain.
Non-aliasing stores are moved to the same chain so the existing
code fails because it expects to find an adajcent store on a consecutive
chain.

There are some ordering issues by adding this so the optimal store merge
is not found, and I'm not sure how best to solve them.

Because of how DAGCombiner tries these store combines,
MergeConsecutiveStores doesn't see the correct set of stores on the chain
when it visits the other stores. Each store individually has its chain
fixed before trying to merge consecutive stores, and then tries to merge
stores from that point before the other stores have been processed to
have their chains fixed.

Suppose you have 4 32-bit stores that should be merged into 1 vector
store. One store is visited first, fixing the chain. What happens is
because not all of the store chains have yet been fixed, 2 of the stores
are merged. The other 2 stores later have their chains fixed,
but because the other stores were already merged, they have different
memory types and merging the two different sized stores is not
supported and would be more difficult to handle.

I'm able to get what I want to happen by only performing
MergeConsecutiveStores on the legalized DAG, at which point
all of the nonaliasing stores have had their chains fixed. However,
it is also desirable to have MergeConsecutivStores process stores of illegal
types like it does now, so I'm not sure how to best ensure all stores
have their chains fixed before store merging happens. Should
FindBetterChain search through consecutive chains and try to handle
other consecutive stores at the same time? Should there be some DAG
preprocess step to fix all store chains or something else?

[hfinkel]

Looking for stores sharing the same chain for merging seems like the right thing to do. This is what you should get from memcpy legalization, etc. anyway.

Shouldn't, however, this be part of the visitor loop below that walks the chain instead of only searching the top (or bottom, depending on how you look at it) store?

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10787

Why? Expanding memcpy, etc. also creates these kinds of parallel stores.

[arsenm]

In D11166#212218, @hfinkel wrote:

Looking for stores sharing the same chain for merging seems like the right thing to do. This is what you should get from memcpy legalization, etc. anyway.

Shouldn't, however, this be part of the visitor loop below that walks the chain instead of only searching the top (or bottom, depending on how you look at it) store?

The idea was if we can rely on FindBetterChain canonicalizing the chain layout, we wouldn't have to worry about stores on other chains and could delete the existing loop here. I'll post another rough version in a minute which more or less copies the loop in MergeConsecutiveStores and changes it so that visitSTORE tries to FindBetterChain on possible store merging candidates, which should make MergeConsecutiveStores only need to consider other users of the same chain. Since this only happens if combiner-alias-analysis is enabled, my questions are why isn't it on by default? Would it be OK to say that if you want consecutive store merging combiner AA should be enabled, or should I try to keep these supporting both?

I would prefer to canonicalize how non-aliasing load/store chains are arranged, rather than having MergeConsecutiveStores do it because I would like to write target DAG combines for more exotic load/store merging. For example I would like to merge certain non-adjacent loads and certain cases of load one address space and store to another.

[arsenm]

Try to call FindBetterChain on possible store merging candidates during visitSTORE

[arsenm]

ping

[hfinkel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10787

As I said before, this comment is misleading. That situation also commonly occurs when expanding memcpy intrinsics, expanding vector stores of illegal type, etc.

10804

Based on your review comments, you're relying here on the fact that, when AA is being used, FindBetterChain canonicalizes non-aliasing stores to be on the same chain. This needs at least a comment.

11400

Could we be consistent and do this even when UseAA is false? This code ends up calling isAlias, but isAlias does things when UseAA is false (it just does more things when UseAA is true). If we did that, could be remove the special case for UseAA in getStoreMergeAndAliasCandidates (because we'd again have a single canonical form for the store chains).

14222

We don't seem to use this for anything other than this. Remove it?

[arsenm]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11400

That just exposes more instances of the same fundamental problem, that different combines currently expect different chain layouts. Ideally we could probably do this, however that will require more work. Doing this without AA will move more chains, but most of the code is not ready for this (which is why I assume CombinerAA is not on by default now).

If you remove the UseAA check here (and in current code), you many of the same test failures as you do if you enable combiner-aa by default. Eventually when more combines are fixed to look at the same chain, we could probably do this.

[arsenm]

Fix comment. Add test updates

[hfinkel]

LGTM.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11400

Doing this without AA will move more chains, but most of the code is not ready for this (which is why I assume CombinerAA is not on by default now).

It is on by default for some targets (whatever targets return true from useAA() in their TargetSubtargetInfo subclass).

If you remove the UseAA check here (and in current code), you many of the same test failures as you do if you enable combiner-aa by default. Eventually when more combines are fixed to look at the same chain, we could probably do this.

Okay, please add a FIXME explaining this here.

[arsenm]

r246307

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     SDValue visitFCEIL(SDNode *N);
     SDValue visitFTRUNC(SDNode *N);
     SDValue visitFFLOOR(SDNode *N);
     SDValue visitFMINNUM(SDNode *N);
     SDValue visitFMAXNUM(SDNode *N);
     SDValue visitBRCOND(SDNode *N);
     SDValue visitBR_CC(SDNode *N);
     SDValue visitLOAD(SDNode *N);
+
+    SDValue replaceStoreChain(StoreSDNode *ST, SDValue BetterChain);
+
     SDValue visitSTORE(SDNode *N);
     SDValue visitINSERT_VECTOR_ELT(SDNode *N);
     SDValue visitEXTRACT_VECTOR_ELT(SDNode *N);
     SDValue visitBUILD_VECTOR(SDNode *N);
     SDValue visitCONCAT_VECTORS(SDNode *N);
     SDValue visitEXTRACT_SUBVECTOR(SDNode *N);
     SDValue visitVECTOR_SHUFFLE(SDNode *N);
     SDValue visitSCALAR_TO_VECTOR(SDNode *N);
private:
 
     /// Return true if there is any possibility that the two addresses overlap.
     bool isAlias(LSBaseSDNode *Op0, LSBaseSDNode *Op1) const;
 
     /// Walk up chain skipping non-aliasing memory nodes, looking for a better
     /// chain (aliasing node.)
     SDValue FindBetterChain(SDNode *N, SDValue Chain);
 
+    /// Do FindBetterChain for a store and any possibly adjacent stores on
+    /// consecutive chains.
+    bool findBetterNeighborChains(StoreSDNode *St);
+
     /// Holds a pointer to an LSBaseSDNode as well as information on where it
     /// is located in a sequence of memory operations connected by a chain.
     struct MemOpLink {
       MemOpLink (LSBaseSDNode *N, int64_t Offset, unsigned Seq):
       MemNode(N), OffsetFromBase(Offset), SequenceNum(Seq) { }
       // Ptr to the mem node.
       LSBaseSDNode *MemNode;
       // Offset from the base ptr.
if (BasePtr.Base.getOpcode() == ISD::UNDEF)
     return;
 
   // Walk up the chain and look for nodes with offsets from the same
   // base pointer. Stop when reaching an instruction with a different kind
   // or instruction which has a different base pointer.
   EVT MemVT = St->getMemoryVT();
   unsigned Seq = 0;
   StoreSDNode *Index = St;
+
+
+  bool UseAA = CombinerAA.getNumOccurrences() > 0 ? CombinerAA
+                                                  : DAG.getSubtarget().useAA();
+
+  if (UseAA) {
+    // Look at other users of the same chain. Stores on the same chain do not
+    // alias.
+    // TODO: This should be the common (only?) case if combiner-aa is
+    // enabled. Should we skip or remove the code looking at consecutive chains?
+    SDValue Chain = St->getChain();
+    for (auto I = Chain->use_begin(), E = Chain->use_end(); I != E; ++I) {
+      if (StoreSDNode *OtherST = dyn_cast<StoreSDNode>(*I)) {
+
+        if (OtherST->isVolatile() || OtherST->isIndexed())
+          continue;
+
+        if (OtherST->getMemoryVT() != MemVT)
+          continue;
+
+        BaseIndexOffset Ptr = BaseIndexOffset::match(OtherST->getBasePtr());
+        StoreNodes.push_back(MemOpLink(OtherST, Ptr.Offset, Seq++));
+      }
+    }
+
+    return;
+  }
+
   while (Index) {
     // If the chain has more than one use, then we can't reorder the mem ops.
     if (Index != St && !SDValue(Index, 0)->hasOneUse())
       break;
 
     // Find the base pointer and offset for this memory node.
     BaseIndexOffset Ptr = BaseIndexOffset::match(Index->getBasePtr());
 
for (unsigned i = 0; i < NumElem ; ++i) {
     StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
     DAG.ReplaceAllUsesOfValueWith(SDValue(St, 0), St->getChain());
     deleteAndRecombine(St);
   }
 
   return true;
 }
 
+SDValue DAGCombiner::replaceStoreChain(StoreSDNode *ST, SDValue BetterChain) {
+  SDLoc SL(ST);
+  SDValue ReplStore;
+
+  // Replace the chain to avoid dependency.
+  if (ST->isTruncatingStore()) {
+    ReplStore = DAG.getTruncStore(BetterChain, SL, ST->getValue(),
+                                  ST->getBasePtr(), ST->getMemoryVT(),
+                                  ST->getMemOperand());
+  } else {
+    ReplStore = DAG.getStore(BetterChain, SL, ST->getValue(), ST->getBasePtr(),
+                             ST->getMemOperand());
+  }
+
+  // Create token to keep both nodes around.
+  SDValue Token = DAG.getNode(ISD::TokenFactor, SL,
+                              MVT::Other, ST->getChain(), ReplStore);
+
+  // Make sure the new and old chains are cleaned up.
+  AddToWorklist(Token.getNode());
+
+  // Don't add users to work list.
+  return CombineTo(ST, Token, false);
+}
+
 SDValue DAGCombiner::visitSTORE(SDNode *N) {
   StoreSDNode *ST  = cast<StoreSDNode>(N);
   SDValue Chain = ST->getChain();
   SDValue Value = ST->getValue();
   SDValue Ptr   = ST->getBasePtr();
 
   // If this is a store of a bit convert, store the input value if the
   // resultant store does not need a higher alignment than the original.
SDValue DAGCombiner::visitSTORE(SDNode *N) {
   bool UseAA = CombinerAA.getNumOccurrences() > 0 ? CombinerAA
                                                   : DAG.getSubtarget().useAA();
 #ifndef NDEBUG
   if (CombinerAAOnlyFunc.getNumOccurrences() &&
       CombinerAAOnlyFunc != DAG.getMachineFunction().getName())
     UseAA = false;
 #endif
   if (UseAA && ST->isUnindexed()) {
-    DEBUG(
-      dbgs() << "Look for better chain: ";
-      ST->dump(&DAG);
-    );
-    // Walk up chain skipping non-aliasing memory nodes.
-    SDValue BetterChain = FindBetterChain(N, Chain);
-
-    // If there is a better chain.
-    if (Chain != BetterChain) {
-      DEBUG(
-        dbgs() << "Found better chain: ";
-        BetterChain->dump(&DAG);
-      );
-
-      SDValue ReplStore;
-
-      // Replace the chain to avoid dependency.
-      if (ST->isTruncatingStore()) {
-        ReplStore = DAG.getTruncStore(BetterChain, SDLoc(N), Value, Ptr,
-                                      ST->getMemoryVT(), ST->getMemOperand());
-      } else {
-        ReplStore = DAG.getStore(BetterChain, SDLoc(N), Value, Ptr,
-                                 ST->getMemOperand());
-      }
-
-      // Create token to keep both nodes around.
-      SDValue Token = DAG.getNode(ISD::TokenFactor, SDLoc(N),
-                                  MVT::Other, Chain, ReplStore);
-
-      // Make sure the new and old chains are cleaned up.
-      AddToWorklist(Token.getNode());
-
-      // Don't add users to work list.
-      return CombineTo(N, Token, false);
-    } else {
-      DEBUG(dbgs() << "Did not find better chain\n");
+    // Walk up chain skipping non-aliasing memory nodes, on this store and any
+    // adjacent stores.
+    if (findBetterNeighborChains(ST)) {
+      // replaceStoreChain uses CombineTo, which handled all of the worklist
+      // manipulation. Return the original node to not do anything else.
+      return SDValue(ST, 0);
     }
   }
 
   // Try transforming N to an indexed store.
   if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N))
     return SDValue(N, 0);
 
   // FIXME: is there such a thing as a truncating indexed store?
if (Aliases.size() == 1)
     return Aliases[0];
 
   DEBUG(dbgs() << "Found better chain with " << Aliases.size() << " aliases\n");
 
   // Construct a custom tailored token factor.
   return DAG.getNode(ISD::TokenFactor, SDLoc(N), MVT::Other, Aliases);
 }
 
+bool DAGCombiner::findBetterNeighborChains(StoreSDNode* St) {
+  // 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())
+    return false;
+
+  // Do not handle stores to undef base pointers.
+  if (BasePtr.Base.getOpcode() == ISD::UNDEF)
+    return false;
+
+  SmallVector<StoreSDNode *, 8> ChainedStores;
+  SmallVector<LoadSDNode *, 8> ChainedLoads;
+
+  ChainedStores.push_back(St);
+
+  // Walk up the chain and look for nodes with offsets from the same
+  // base pointer. Stop when reaching an instruction with a different kind
+  // or instruction which has a different base pointer.
+  EVT MemVT = St->getMemoryVT();
+  unsigned Seq = 0;
+  StoreSDNode *Index = St;
+  while (Index) {
+    // If the chain has more than one use, then we can't reorder the mem ops.
+    if (Index != St && !SDValue(Index, 0)->hasOneUse())
+      break;
+
+    // Find the base pointer and offset for this memory node.
+    BaseIndexOffset Ptr = BaseIndexOffset::match(Index->getBasePtr());
+
+    // Check that the base pointer is the same as the original one.
+    if (!Ptr.equalBaseIndex(BasePtr))
+      break;
+
+    if (Index->isVolatile() || Index->isIndexed())
+      break;
+
+    // Find the next memory operand in the chain. If the next operand in the
+    // chain is a store then move up and continue the scan with the next
+    // memory operand. If the next operand is a load save it and use alias
+    // information to check if it interferes with anything.
+    SDNode *NextInChain = Index->getChain().getNode();
+    while (true) {
+      if (StoreSDNode *STn = dyn_cast<StoreSDNode>(NextInChain)) {
+        // We found a store node. Use it for the next iteration.
+        ChainedStores.push_back(STn);
+        Index = STn;
+        break;
+      } else if (LoadSDNode *Ldn = dyn_cast<LoadSDNode>(NextInChain)) {
+        ChainedLoads.push_back(Ldn);
+        NextInChain = Ldn->getChain().getNode();
+        continue;
+      } else {
+        Index = nullptr;
+        break;
+      }
+    }
+  }
+
+  bool MadeChange = false;
+  SmallVector<std::pair<StoreSDNode *, SDValue>, 8> BetterChains;
+
+  for (StoreSDNode *ChainedStore : ChainedStores) {
+    SDValue Chain = ChainedStore->getChain();
+    SDValue BetterChain = FindBetterChain(ChainedStore, Chain);
+
+    if (Chain != BetterChain) {
+      MadeChange = true;
+      BetterChains.push_back(std::make_pair(ChainedStore, BetterChain));
+    }
+  }
+
+  // Do all replacements after finding the replacements to make to avoid making
+  // the chains more complicated by introducing new TokenFactors.
+  for (auto Replacement : BetterChains)
+    replaceStoreChain(Replacement.first, Replacement.second);
+
+  return MadeChange;
+}
+
 /// This is the entry point for the file.
 void SelectionDAG::Combine(CombineLevel Level, AliasAnalysis &AA,
                            CodeGenOpt::Level OptLevel) {
   /// This is the main entry point to this class.
   DAGCombiner(*this, AA, OptLevel).Run(Level);
 }