[PATCH] D25601: DAGCombiner: fix use-after-free when merging consecutive stores

[Justin Bogner via llvm-commits]

Nicolai Hähnle <nhaehnle at gmail.com> writes:
>>>> -    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);
>>>
>>> What do you think of returning the SmallVector instead of the bool, then
>>> just doing `EverChanged |= !StoreNodes.empty()`? It seems a bit simpler
>>> to me, and move semantics should mean it's just as efficient.
>>
>> A SmallVector move cannot be as efficient because it cannot avoid the
>> copy in the small case. I don't feel strongly about it though.

That's a good point about the SmallVector move, but I'm pretty sure this
case would fall under NRVO and copy elision anyway, so it probably
doesn't even come up.

> Another point to consider is that encoding the size of the SmallVector
> as part of a function signature feels a bit odd, especially in this
> case where it'd affect two functions whose return type must then agree
> on the size of the SmallVector. I'd really prefer not to make that
> change.

I don't really see why, but I don't feel too strongly about it.

> nhaehnle updated this revision to Diff 76759.
> nhaehnle added a comment.
>
> Simplify using any_of.

This version LGTM.

> https://reviews.llvm.org/D25601
>
> Files:
>   lib/CodeGen/SelectionDAG/DAGCombiner.cpp
>
>
> Index: lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> ===================================================================
> --- lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> +++ lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> @@ -450,10 +450,11 @@
>      /// 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,
> +        bool IsConstantSrc, bool UseVector);
>  
>      /// This is a helper function for MergeConsecutiveStores.
>      /// Stores that may be merged are placed in StoreNodes.
> @@ -470,8 +471,10 @@
>  
>      /// 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))
> @@ -11513,6 +11516,7 @@
>      }
>    }
>  
> +  StoreNodes.erase(StoreNodes.begin() + NumStores, StoreNodes.end());
>    return true;
>  }
>  
> @@ -11655,7 +11659,8 @@
>    return true;
>  }
>  
> -bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
> +bool DAGCombiner::MergeConsecutiveStores(
> +    StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes) {
>    if (OptLevel == CodeGenOpt::None)
>      return false;
>  
> @@ -11699,9 +11704,6 @@
>    // 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.
> @@ -12060,6 +12062,7 @@
>      }
>    }
>  
> +  StoreNodes.erase(StoreNodes.begin() + NumElem, StoreNodes.end());
>    return true;
>  }
>  
> @@ -12305,14 +12308,19 @@
>    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);
> +
> +      if (any_of(StoreNodes,
> +                 [ST](const MemOpLink &Link) { return Link.MemNode == ST; }))
> +        break;
> +    }
>  
>      if (EverChanged)
>        return SDValue(N, 0);