[MemorySSA] Update MemoryPhi wiring for block splitting to consider if identical edges were merged.

[asbirlea]

Summary

Block splitting is done with either identical edges being merged, or not.
Only critical edges can be split without merging identical edges based on an option.
Teach the memoryssa updater to take this into account: for the same edge between two blocks only move one entry from the Phi in Old to the new Phi in New.

[george.burgess.iv]

Thanks!

lib/Analysis/MemorySSAUpdater.cpp

501–502

(I see that practically, the answer is "no". Just curious how strongly we're tailoring this function for the user)

Is it theoretically reasonable for Preds to contain duplicates of the same BasicBlock if !IdenticalEdgesWereMerged? The concern being that if I, as the user, call wireOldPredecessorsToNewImmediatePredecessor(Old, New, {BB, BB}, false);, we'll only end up moving one instance of BB.

If this concern is founded, please add an assert that PredsSet.size() == Preds.size() if !IdenticalEdgesWereMerged, so it's obvious that said case is broken if someone tries to depend on it in the future (which I'm totally cool with, given that, like said, this can't happen at the moment).

[asbirlea]

Address comment. Thank you for the review!

[asbirlea]

lib/Analysis/MemorySSAUpdater.cpp

501–502

Agreed, added the assert!

[asbirlea]

Nit: remove braces.

include/llvm/Analysis/MemorySSAUpdater.h

public:
   /// |------|        |------|
   void moveAllAfterMergeBlocks(BasicBlock *From, BasicBlock *To,
                                Instruction *Start);
   /// A new empty BasicBlock (New) now branches directly to Old. Some of
   /// Old's predecessors (Preds) are now branching to New instead of Old.
   /// If New is the only predecessor, move Old's Phi, if present, to New.
   /// Otherwise, add a new Phi in New with appropriate incoming values, and
   /// update the incoming values in Old's Phi node too, if present.
-  void
-  wireOldPredecessorsToNewImmediatePredecessor(BasicBlock *Old, BasicBlock *New,
-                                               ArrayRef<BasicBlock *> Preds);
-
+  void wireOldPredecessorsToNewImmediatePredecessor(
+      BasicBlock *Old, BasicBlock *New, ArrayRef<BasicBlock *> Preds,
+      bool IdenticalEdgesWereMerged = true);
   // The below are utility functions. Other than creation of accesses to pass
   // to insertDef, and removeAccess to remove accesses, you should generally
   // not attempt to update memoryssa yourself. It is very non-trivial to get
   // the edge cases right, and the above calls already operate in near-optimal
   // time bounds.
 
   /// Create a MemoryAccess in MemorySSA at a specified point in a block,
   /// with a specified clobbering definition.

lib/Analysis/MemorySSAUpdater.cpp

if (!MA)
       MA = cast<MemoryAccess>(Arg);
     else if (MA != Arg)
       return nullptr;
   }
   return MA;
 }
 
 void MemorySSAUpdater::wireOldPredecessorsToNewImmediatePredecessor(
-    BasicBlock *Old, BasicBlock *New, ArrayRef<BasicBlock *> Preds) {
+    BasicBlock *Old, BasicBlock *New, ArrayRef<BasicBlock *> Preds,
+    bool IdenticalEdgesWereMerged) {
   assert(!MSSA->getWritableBlockAccesses(New) &&
          "Access list should be null for a new block.");
   MemoryPhi *Phi = MSSA->getMemoryAccess(Old);
   if (!Phi)
     return;
   if (pred_size(Old) == 1) {
     assert(pred_size(New) == Preds.size() &&
            "Should have moved all predecessors.");
     MSSA->moveTo(Phi, New, MemorySSA::Beginning);
   } else {
     assert(!Preds.empty() && "Must be moving at least one predecessor to the "
                              "new immediate predecessor.");
     MemoryPhi *NewPhi = MSSA->createMemoryPhi(New);
     SmallPtrSet<BasicBlock *, 16> PredsSet(Preds.begin(), Preds.end());
+    // Currently only support the case of removing a single incoming edge when
+    // identical edges were not merged.
+    if (!IdenticalEdgesWereMerged)
+      assert(PredsSet.size() == Preds.size() &&
+             "If identical edges were not merged, we cannot have duplicate "
+             "blocks in the predecessors");
     Phi->unorderedDeleteIncomingIf([&](MemoryAccess *MA, BasicBlock *B) {
       if (PredsSet.count(B)) {
         NewPhi->addIncoming(MA, B);
+        if (!IdenticalEdgesWereMerged)
+          PredsSet.erase(B);
         return true;
       }
       return false;
     });
     Phi->addIncoming(NewPhi, New);
     if (onlySingleValue(NewPhi))
       removeMemoryAccess(NewPhi);
   }

lib/Transforms/Utils/BreakCriticalEdges.cpp

if (Options.MergeIdenticalEdges) {
     }
   }
 
   // If we have nothing to update, just return.
   auto *DT = Options.DT;
   auto *LI = Options.LI;
   auto *MSSAU = Options.MSSAU;
   if (MSSAU)
-    MSSAU->wireOldPredecessorsToNewImmediatePredecessor(DestBB, NewBB, {TIBB});
+    MSSAU->wireOldPredecessorsToNewImmediatePredecessor(
+        DestBB, NewBB, {TIBB}, Options.MergeIdenticalEdges);
 
   if (!DT && !LI)
     return NewBB;
 
   if (DT) {
     // Update the DominatorTree.
     //       ---> NewBB -----\
     //      /                 V