DAGCombiner: fix use-after-free when merging consecutive stores

[nhaehnle]

Summary

Have MergeConsecutiveStores explicitly return information about the stores
that were merged, so that we can safely determine whether the starting
node has been freed.

[nhaehnle]

Ping.

[chandlerc]

Nirav, can you take a look at this?

[niravd]

Do you have a test case for this?

The moving of StoreNodes to the local context seems like it could be a real fix to a nasty problem. We never seem to clear it and that may cause incorrect merges (though i have no example).

Otherwise this seems like a small worthwhile cleanup. As I understand it, when we merge stores we look for candidate nodes that match St but do not require that St is one of the merged nodes which makes it possible for a merge to leave the original store without marking it to be checked again. This may leave some stores that were merge candidates but were not chosen untouched.
I think a comment explaining this would be helpful.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11703

Please kill the comment on the preceding line.

12327

Use StoreNodes.size() instead of NumChanged (or better yet an iterator). This lets you keep the return type above in MergeConsecutiveStores.

[niravd]

The moving of StoreNodes to the local context seems like it could be a real fix to a nasty problem. We never seem to clear it and that may cause incorrect merges (though i have no example).

Ignore this.

[niravd]

Looking at this further, it seems like this is just stopping earlier than otherwise and potentially masking an issue. Nicolai, if you have a testcase you can share I'd be very interested at looking at it as I think it may be related to a bug I've run into separately.

[nhaehnle]

Sorry for the slow response, had some other bugs to deal with. I'll address the comments now.

There are lots of lit tests that hit this when you apply D25313, unfortunately I didn't take a note of which ones precisely because there were so many... I don't think they manifest as actual miscompiles because the deleted nodes land in the Recycler, and nothing allocates a new node in the meantime. But it's certainly fragile...

[nhaehnle]

Keep returning a boolean.

Since not all nodes are necessarily merged, the success-path now needs to
truncate the StoreNodes vector to just those nodes that were actually
removed.

[nhaehnle]

To follow up, these are the tests with asan failures (when D25313 is applied) that are fixed by applying this patch (AMDGPU and X86 backends only):

LLVM :: CodeGen/AMDGPU/amdgpu.private-memory.ll
LLVM :: CodeGen/AMDGPU/invariant-load-no-alias-store.ll
LLVM :: CodeGen/AMDGPU/merge-stores.ll
LLVM :: CodeGen/AMDGPU/vector-alloca.ll
LLVM :: CodeGen/X86/2012-11-28-merge-store-alias.ll
LLVM :: CodeGen/X86/2012-12-1-merge-multiple.ll
LLVM :: CodeGen/X86/MergeConsecutiveStores.ll
LLVM :: CodeGen/X86/avx-vextractf128.ll
LLVM :: CodeGen/X86/merge_store.ll
LLVM :: CodeGen/X86/pr18023.ll
LLVM :: CodeGen/X86/stores-merging.ll
LLVM :: CodeGen/X86/vector-merge-store-fp-constants.ll

[nhaehnle]

Simplify using any_of.

[niravd]

I'm still not convinced this fixes the problem. It reduces the number of merges that we do but they should should have been safe as is. . Let's take some time to look at the failing test cases and see if we can find where the problem is coming from.

[niravd]

I'm still not convinced this fixes the problem. It reduces the number of merges that we do but they should should have been safe as is. . Let's take some time to look at the failing test cases and see if we can find where the problem is coming from.

[nhaehnle]

This doesn't reduce the number of merges. The patch has 0 changes in any of the test/CodeGen tests across all backends, and I don't have any out-of-tree tests with changes either. There is no real functional change.

Even before this patch, the code satisfied: The store nodes that are deleted for the merge are exactly a prefix of the StoreNodes vector. The patch (a) truncates the StoreNodes vector to that prefix and (b) passes the vector out so that we can safely decide whether ST was deleted without implicitly relying on the behaviour of the SDNode recycler.

Look, I'll do you one better and do a llvm-lit test/CodeGen run with the following change in a debug build:

if (any_of(StoreNodes,
           [ST](const MemOpLink &Link) { return Link.MemNode == ST; })) {
  assert(ST->getOpcode() == ISD::DELETED_NODE);
  break;
}

assert(ST->getOpcode() != ISD::DELETED_NODE);

Of course the assertions re-introduce the problematic reliance on the recycler behaviour, but the point is that neither assertion ever triggers on any of the in-tree CodeGen tests.

[niravd]

LGTM modulo previous comment.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     /// need to be narrowed in order to match.
     bool isAndLoadExtLoad(ConstantSDNode *AndC, LoadSDNode *LoadN,
                           EVT LoadResultTy, EVT &ExtVT, EVT &LoadedVT,
                           bool &NarrowLoad);
 
     /// 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);
+    /// \return number of stores that were merged into a merged store (always
+    /// a prefix of \p StoreNode).
+    bool MergeStoresOfConstantsOrVecElts(
+        SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT, unsigned NumStores,
 
     /// This is a helper function for MergeConsecutiveStores.
     /// Stores that may be merged are placed in StoreNodes.
     /// Loads that may alias with those stores are placed in AliasLoadNodes.
     void getStoreMergeAndAliasCandidates(
         StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes,
         SmallVectorImpl<LSBaseSDNode*> &AliasLoadNodes);
 
     /// 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);
 
     /// Merge consecutive store operations into a wide store.
     /// This optimization uses wide integers or vectors when possible.
-    /// \return True if some memory operations were changed.
-    bool MergeConsecutiveStores(StoreSDNode *N);
+    /// \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,
+                                SmallVectorImpl<MemOpLink> &StoreNodes);
 
     /// \brief Try to transform a truncation where C is a constant:
     ///     (trunc (and X, C)) -> (and (trunc X), (trunc C))
     ///
     /// \p N needs to be a truncation and its first operand an AND. Other
     /// requirements are checked by the function (e.g. that trunc is
     /// single-use) and if missed an empty SDValue is returned.
     SDValue distributeTruncateThroughAnd(SDNode *N);
for (unsigned i = 0; i < NumStores; ++i) {
       // get CSEed and the net result is that X is now a use of St.
       // Since we know that St is redundant, just iterate.
       while (!St->use_empty())
         DAG.ReplaceAllUsesWith(SDValue(St, 0), St->getChain());
       deleteAndRecombine(St);
     }
   }
 
+  StoreNodes.erase(StoreNodes.begin() + NumStores, StoreNodes.end());
   return true;
 }
 
 void DAGCombiner::getStoreMergeAndAliasCandidates(
     StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes,
     SmallVectorImpl<LSBaseSDNode*> &AliasLoadNodes) {
   // This holds the base pointer, index, and the offset in bytes from the base
   // pointer.
bool DAGCombiner::checkMergeStoreCandidatesForDependencies(
   // search through DAG. We can stop early if we find a storenode
   for (unsigned i = 0; i < StoreNodes.size(); ++i) {
     if (SDNode::hasPredecessorHelper(StoreNodes[i].MemNode, Visited, Worklist))
       return false;
   }
   return true;
 }
 
-bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
+bool DAGCombiner::MergeConsecutiveStores(
+    StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes) {
   if (OptLevel == CodeGenOpt::None)
     return false;
 
   EVT MemVT = St->getMemoryVT();
   int64_t ElementSizeBytes = MemVT.getSizeInBits() / 8;
   bool NoVectors = DAG.getMachineFunction().getFunction()->hasFnAttribute(
       Attribute::NoImplicitFloat);
 
bool DAGCombiner::MergeConsecutiveStores(
   if (Chain->hasOneUse() && Chain->use_begin()->getOpcode() == ISD::STORE)
     return false;
 
   // Save the LoadSDNodes that we find in the chain.
   // We need to make sure that these nodes do not interfere with
   // any of the store nodes.
   SmallVector<LSBaseSDNode*, 8> AliasLoadNodes;
 
-  // Save the StoreSDNodes that we find in the chain.
-  SmallVector<MemOpLink, 8> StoreNodes;
-
   getStoreMergeAndAliasCandidates(St, StoreNodes, AliasLoadNodes);
 
   // Check if there is anything to merge.
   if (StoreNodes.size() < 2)
     return false;
 
   // only do dependence check in AA case
   bool UseAA = CombinerAA.getNumOccurrences() > 0 ? CombinerAA
for (unsigned i = 0; i < NumElem; ++i) {
       if (StoreNodes[i].MemNode == LatestOp)
         continue;
       StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
       DAG.ReplaceAllUsesOfValueWith(SDValue(St, 0), St->getChain());
       deleteAndRecombine(St);
     }
   }
 
+  StoreNodes.erase(StoreNodes.begin() + NumElem, StoreNodes.end());
   return true;
 }
 
 SDValue DAGCombiner::replaceStoreChain(StoreSDNode *ST, SDValue BetterChain) {
   SDLoc SL(ST);
   SDValue ReplStore;
 
   // Replace the chain to avoid dependency.
return DAG.getTruncStore(Chain, SDLoc(N), Value.getOperand(0),
                              Ptr, ST->getMemoryVT(), ST->getMemOperand());
   }
 
   // Only perform this optimization before the types are legal, because we
   // don't want to perform this optimization on every DAGCombine invocation.
   if (!LegalTypes) {
     bool EverChanged = false;
 
-    do {
+    for (;;) {
       // There can be multiple store sequences on the same chain.
       // Keep trying to merge store sequences until we are unable to do so
       // or until we merge the last store on the chain.
-      bool Changed = MergeConsecutiveStores(ST);
+      SmallVector<MemOpLink, 8> StoreNodes;
+      bool Changed = MergeConsecutiveStores(ST, StoreNodes);
       EverChanged |= Changed;
       if (!Changed) break;
-    } while (ST->getOpcode() != ISD::DELETED_NODE);
+
+      bool ChangedST = false;
+      for (const auto &Link : StoreNodes) {
+        if (Link.MemNode == ST) {
+          ChangedST = true;
+          break;
+        }
+      }
+      if (ChangedST)
+        break;
+    }
 
     if (EverChanged)
       return SDValue(N, 0);
   }
 
   // Turn 'store float 1.0, Ptr' -> 'store int 0x12345678, Ptr'
   //
   // Make sure to do this only after attempting to merge stores in order to