[MemorySSA] Extend removeMemoryAccess API to optimize MemoryPhis.

[asbirlea]

Summary

EarlyCSE needs to optimize MemoryPhis after an access is removed and has
special handling for it. This should be handled by MemorySSA instead.
The default remains that MemoryPhis are *not* optimized after an access
is removed.

Once special handling is removed from EarlyCSE, the testcase that uses this
patch is: test/Transforms/EarlyCSE/memoryssa.ll.

[george.burgess.iv]

Thanks for this!

Is there a reason we don't want to have this behavior by default? ISTR we're pretty good about keeping Phis well-optimized, but could be wrong.

Please also include a test-case.

lib/Analysis/MemorySSAUpdater.cpp

1108

It looks like this may result in a use-after-free. Consider a CFG:

|---A
|  / \
| B   C
|/ \ /|
D<--E | 
\     /
 \   /
  \ / 
   F

Control-flow flows down, except for the edge E -> D. (and there's an A->D edge that only looks somewhat awkward. :) )

All BBs are clear, though B has a Def in it. This Def would cause us to generate a Phi in E, D, and F.

So we remove that Def, remembering that the Phis in {D, E} are removal candidates. We check that E only has liveOnEntry first, it does, so we recurse on it. E's removal makes the Phi in D trivial, so we recurse on it, delete it, etc. We then land back at the top level of our stack and check onlySingleValue(D).

If you buy all of this, ISTM we could either fix it by hoisting the onlySingleValue check into the while (!MA->use_empty()) { loop (if the Phi only points to one Def, we're not going to find it otherwise...). I'm willing to bet that the number of cases where the same Phi points to the same value N times is pretty small.

Alternatively, we could takeVector(), completely filter out the Phis that aren't onlySingleValue, from that, and *then* recurse on those Phis.

[asbirlea]

Address comment.

[asbirlea]

Thank you for the comments!

We may want this by default, yes. I think we currently simply didn't come across this, due to little usage. Only EarlyCSE needed it and added the special handling.
I would prefer to do the flag-flip separately though, in case it unearths some problematic usage.

There are already a few tests that cover this scenario: e.g., test/Transforms/EarlyCSE/memoryssa.ll.
(i.e. it has failures without this change and the special handling removed in EarlyCSE)

lib/Analysis/MemorySSAUpdater.cpp

1108

Updated.
Still calling removeMemoryAccess recursively.
AFAICT we can reuse tryRemoveTrivialPhi and not pre-filter for onlySingleValue, but the vector will be TrackingVH since tryRemoveTrivialPhi will remove other phis in the vector being processed.
Both approaches work. Let me know if you prefer to reuse tryRemoveTrivialPhi.

[george.burgess.iv]

I would prefer to do the flag-flip separately though, in case it unearths some problematic usage.

SGTM

There are already a few tests that cover this scenario: e.g., test/Transforms/EarlyCSE/memoryssa.ll.

Please note that in your commit message when this is landed, and I'm OK without a formal test-case here.

lib/Analysis/MemorySSAUpdater.cpp

1108

I'd prefer to centralize this logic as much as reasonable, so I like the tryRemoveTrivialPhi approach. (I forgot that existed -- thanks for catching that :) )

If TrackingVH doesn't work as set keys (I don't think it does, but am happy to be proven wrong), we can probably split the SetVector into SmallVector<TrackingVH, ...> and SmallPtrSet<MemoryPhi *, ...>

[asbirlea]

Address comment: reuse tryRemoveTrivialPhi.

[george.burgess.iv]

lib/Analysis/MemorySSAUpdater.cpp

1107

nit: !PhisToCheck.empty(), please

1108

Why is this a vector of TrackingVH?

I'm under the impression, though I haven't totally proven it to myself, that onlySingleValue(MP) implies that we'll never try to remove/RAUW/etc. MP in tryRemoveTrivialPhi, unless we're specifically calling tryRemoveTrivialPhi on MP. So, we shouldn't need to track anything here.

Happy to have this be an AssertingVH to assert that these nothing in this vector is dropped before we get to it, though.

[asbirlea]

Address comments.

[asbirlea]

lib/Analysis/MemorySSAUpdater.cpp

1108

Oh, I meant to remove the onlySingleValue(MP).

If we're going to use tryRemoveTrivialPhi, which recurses internally, then it makes sense to actually let it remove more than a single MemoryPhi. Otherwise, we could just call removeMemoryAccess.

I meant to call tryRemoveTrivialPhi on the unfiltered list of MemoryPhis, and have that list be TrackingVH, since some may go away (there's a test in that triggers this).

Updated now.

[george.burgess.iv]

LGTM after one nit is addressed.

If I'm just imagining that my nit is possible, I'm happy to have this patch submitted as-is.

Thanks again!

lib/Analysis/MemorySSAUpdater.cpp

1108

...But it can remove more than a single MemoryPhi, no? It'll traverse the users of the MemoryPhis we know we can remove (e.g. onlySingleValue ones); there may be many of these.

In any case, both approaches are functionally equivalent in my mind, and I have no preference either way, so this approach LGTM.

Nit: can we simplify to PhisToOptimize{PhisToCheck.begin(), PhisToCheck.end()}; and skip the loop below?

[asbirlea]

Address nit.
Changed to WeakVH after offline discussion.

Thank you for the review!

include/llvm/Analysis/MemorySSAUpdater.h

MemoryUseOrDef *createMemoryAccessAfter(Instruction *I,
                                           MemoryAccess *InsertPt);
 
   /// Remove a MemoryAccess from MemorySSA, including updating all
   /// definitions and uses.
   /// This should be called when a memory instruction that has a MemoryAccess
   /// associated with it is erased from the program.  For example, if a store or
   /// load is simply erased (not replaced), removeMemoryAccess should be called
   /// on the MemoryAccess for that store/load.
-  void removeMemoryAccess(MemoryAccess *);
+  void removeMemoryAccess(MemoryAccess *, bool OptimizePhis = false);
 
   /// Remove MemoryAccess for a given instruction, if a MemoryAccess exists.
   /// This should be called when an instruction (load/store) is deleted from
   /// the program.
-  void removeMemoryAccess(const Instruction *I) {
+  void removeMemoryAccess(const Instruction *I, bool OptimizePhis = false) {
     if (MemoryAccess *MA = MSSA->getMemoryAccess(I))
-      removeMemoryAccess(MA);
+      removeMemoryAccess(MA, OptimizePhis);
   }
 
   /// Remove all MemoryAcceses in a set of BasicBlocks about to be deleted.
   /// Assumption we make here: all uses of deleted defs and phi must either
   /// occur in blocks about to be deleted (thus will be deleted as well), or
   /// they occur in phis that will simply lose an incoming value.
   /// Deleted blocks still have successor info, but their predecessor edges and
   /// Phi nodes may already be updated. Instructions in DeadBlocks should be

lib/Analysis/MemorySSAUpdater.cpp

Phi->unorderedDeleteIncomingIf([&](MemoryAccess *MA, BasicBlock *B) {
       return false;
     });
     Phi->addIncoming(NewPhi, New);
     if (onlySingleValue(NewPhi))
       removeMemoryAccess(NewPhi);
   }
 }
 
-void MemorySSAUpdater::removeMemoryAccess(MemoryAccess *MA) {
+void MemorySSAUpdater::removeMemoryAccess(MemoryAccess *MA, bool OptimizePhis) {
   assert(!MSSA->isLiveOnEntryDef(MA) &&
          "Trying to remove the live on entry def");
   // We can only delete phi nodes if they have no uses, or we can replace all
   // uses with a single definition.
   MemoryAccess *NewDefTarget = nullptr;
   if (MemoryPhi *MP = dyn_cast<MemoryPhi>(MA)) {
     // Note that it is sufficient to know that all edges of the phi node have
     // the same argument.  If they do, by the definition of dominance frontiers
     // (which we used to place this phi), that argument must dominate this phi,
     // and thus, must dominate the phi's uses, and so we will not hit the assert
     // below.
     NewDefTarget = onlySingleValue(MP);
     assert((NewDefTarget || MP->use_empty()) &&
            "We can't delete this memory phi");
   } else {
     NewDefTarget = cast<MemoryUseOrDef>(MA)->getDefiningAccess();
   }
 
+  SmallSetVector<MemoryPhi *, 4> PhisToCheck;
+
   // Re-point the uses at our defining access
   if (!isa<MemoryUse>(MA) && !MA->use_empty()) {
     // Reset optimized on users of this store, and reset the uses.
     // A few notes:
     // 1. This is a slightly modified version of RAUW to avoid walking the
     // uses twice here.
     // 2. If we wanted to be complete, we would have to reset the optimized
     // flags on users of phi nodes if doing the below makes a phi node have all
     // the same arguments. Instead, we prefer users to removeMemoryAccess those
     // phi nodes, because doing it here would be N^3.
     if (MA->hasValueHandle())
       ValueHandleBase::ValueIsRAUWd(MA, NewDefTarget);
     // Note: We assume MemorySSA is not used in metadata since it's not really
     // part of the IR.
 
     while (!MA->use_empty()) {
       Use &U = *MA->use_begin();
       if (auto *MUD = dyn_cast<MemoryUseOrDef>(U.getUser()))
         MUD->resetOptimized();
+      if (OptimizePhis)
+        if (MemoryPhi *MP = dyn_cast<MemoryPhi>(U.getUser()))
+          PhisToCheck.insert(MP);
       U.set(NewDefTarget);
     }
   }
 
   // The call below to erase will destroy MA, so we can't change the order we
   // are doing things here
   MSSA->removeFromLookups(MA);
   MSSA->removeFromLists(MA);
+
+  // Optionally optimize Phi uses. This will recursively remove trivial phis.
+  if (PhisToCheck.size()) {
+    SmallVector<TrackingVH<MemoryPhi>, 16> PhisToOptimize;
+    for (MemoryPhi *MP : PhisToCheck)
+      if (onlySingleValue(MP))
+        PhisToOptimize.push_back(MP);
+    PhisToCheck.clear();
+
+    unsigned PhisSize = PhisToOptimize.size();
+    while (PhisSize-- > 0)
+      if (MemoryPhi *MP = PhisToOptimize.pop_back_val()) {
+        auto OperRange = MP->operands();
+        tryRemoveTrivialPhi(MP, OperRange);
+      }
+  }
 }
 
 void MemorySSAUpdater::removeBlocks(
     const SmallPtrSetImpl<BasicBlock *> &DeadBlocks) {
   // First delete all uses of BB in MemoryPhis.
   for (BasicBlock *BB : DeadBlocks) {
     Instruction *TI = BB->getTerminator();
     assert(TI && "Basic block expected to have a terminator instruction");