Update MemorySSA in SimpleLoopUnswitch.

[asbirlea]

Summary

Teach SimpleLoopUnswitch to preserve MemorySSA.
Enable tests for correctness, dependency disabled by default.

[asbirlea]

Cleanup and clang-format.

[asbirlea]

Rebase after changes in SimpleLoopUnswitch and MemorySSA API changes.

[asbirlea]

Rebase and cleanup. Ready to review.

[asbirlea]

Gentle ping.

[asbirlea]

Update MemorySSA when unswitching guard intrinsics.
And ping.

[chandlerc]

LGTM

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

1266–1271

It would be really nice to avoid building the set here. Can the API be changed to support a list of unique blocks? We already have uniqued them.

(Happy for this to be a follow-up of course.)

[asbirlea]

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

1266–1271

The implementation of removeBlocks checks existence of a basic block in this list(.count(BB)), and I think that's best done with a set. So not sure if we can avoid building the set anyway.
I can push the set creation in an API (i.e. a wrapper taking a list and building the set).

[asbirlea]

Thank you for the review!

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/MemorySSA.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/Utils/Local.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"

 ///
 /// Requires that the loop exit and unswitched basic block are the same, and
 /// that the exiting block was a unique predecessor of that block. Rewrites the
 /// PHI nodes in that block such that what were LCSSA PHI nodes become trivial
 /// PHI nodes from the old preheader that now contains the unswitched
 /// terminator.
 static void rewritePHINodesForUnswitchedExitBlock(BasicBlock &UnswitchedBB,
                                                   BasicBlock &OldExitingBB,
-                                                  BasicBlock &OldPH) {
+                                                  BasicBlock &OldPH,
+                                                  MemorySSA *MSSA) {
   for (PHINode &PN : UnswitchedBB.phis()) {
     // When the loop exit is directly unswitched we just need to update the
     // incoming basic block. We loop to handle weird cases with repeated
     // incoming blocks, but expect to typically only have one operand here.
     for (auto i : seq<int>(0, PN.getNumOperands())) {
       assert(PN.getIncomingBlock(i) == &OldExitingBB &&
              "Found incoming block different from unique predecessor!");
       PN.setIncomingBlock(i, &OldPH);
     }
   }
+
+  if (MSSA)
+    if (MemoryPhi *MPhi = MSSA->getMemoryAccess(&UnswitchedBB)) {
+      // Only updates incoming block. Incoming value will be replaced in all
+      // blocks no longer dominated by OldExitingBB, after DT is updated.
+      MPhi->setIncomingBlock(MPhi->getBasicBlockIndex(&UnswitchedBB), &OldPH);
+    }
 }
 
 /// Rewrite the PHI nodes in the loop exit basic block and the split off
 /// unswitched block.
 ///
 /// Because the exit block remains an exit from the loop, this rewrites the
 /// LCSSA PHI nodes in it to remove the unswitched edge and introduces PHI
 /// nodes into the unswitched basic block to select between the value in the
 /// old preheader and the loop exit.
-static void rewritePHINodesForExitAndUnswitchedBlocks(BasicBlock &ExitBB,
-                                                      BasicBlock &UnswitchedBB,
-                                                      BasicBlock &OldExitingBB,
-                                                      BasicBlock &OldPH,
-                                                      bool FullUnswitch) {
+static void rewritePHINodesForExitAndUnswitchedBlocks(
+    BasicBlock &ExitBB, BasicBlock &UnswitchedBB, BasicBlock &OldExitingBB,
+    BasicBlock &OldPH, MemorySSA *MSSA, MemorySSAUpdater *MSSAUpdater,
   assert(&ExitBB != &UnswitchedBB &&
          "Must have different loop exit and unswitched blocks!");
   Instruction *InsertPt = &*UnswitchedBB.begin();
   for (PHINode &PN : ExitBB.phis()) {
     auto *NewPN = PHINode::Create(PN.getType(), /*NumReservedValues*/ 2,
                                   PN.getName() + ".split", InsertPt);
 
     // Walk backwards over the old PHI node's inputs to minimize the cost of
for (int i = PN.getNumIncomingValues() - 1; i >= 0; --i) {
       NewPN->addIncoming(Incoming, &OldPH);
     }
 
     // Now replace the old PHI with the new one and wire the old one in as an
     // input to the new one.
     PN.replaceAllUsesWith(NewPN);
     NewPN->addIncoming(&PN, &ExitBB);
   }
+
+  if (MSSA) {
+    // Only updates incoming block. Incoming value will be replaced in all
+    // blocks no longer dominated by OldExitingBB, after DT is updated.
+    MSSAUpdater->wireOldPhiIntoNewPhiAfterUnswitch(
+        &ExitBB, &UnswitchedBB, &OldExitingBB, &OldPH, FullUnswitch);
+  }
 }
 
 /// Unswitch a trivial branch if the condition is loop invariant.
 ///
 /// This routine should only be called when loop code leading to the branch has
 /// been validated as trivial (no side effects). This routine checks if the
 /// condition is invariant and one of the successors is a loop exit. This
 /// allows us to unswitch without duplicating the loop, making it trivial.
 ///
 /// If this routine fails to unswitch the branch it returns false.
 ///
 /// If the branch can be unswitched, this routine splits the preheader and
 /// hoists the branch above that split. Preserves loop simplified form
 /// (splitting the exit block as necessary). It simplifies the branch within
 /// the loop to an unconditional branch but doesn't remove it entirely. Further
 /// cleanup can be done with some simplify-cfg like pass.
 static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
-                                  LoopInfo &LI) {
+                                  LoopInfo &LI, MemorySSA *MSSA,
+                                  MemorySSAUpdater *MSSAUpdater) {
   assert(BI.isConditional() && "Can only unswitch a conditional branch!");
   LLVM_DEBUG(dbgs() << "  Trying to unswitch branch: " << BI << "\n");
 
   // The loop invariant values that we want to unswitch.
   TinyPtrVector<Value *> Invariants;
 
   // When true, we're fully unswitching the branch rather than just unswitching
   // some input conditions to the branch.
LLVM_DEBUG({
     for (Value *Invariant : Invariants) {
       dbgs() << "      " << *Invariant << " == true";
       if (Invariant != Invariants.back())
         dbgs() << " ||";
       dbgs() << "\n";
     }
   });
 
+  if (MSSA) {
+    /*
+    dbgs()<<"Start unswitchTrivialBranch...";
+    if (FullUnswitch)
+      dbgs()<<"FullUnswitch\n";
+    else
+      dbgs()<<"PartialUnswitch\n";
+    MSSA->dump();
+    */
+    MSSA->verifyMemorySSA();
+  }
+
   // Split the preheader, so that we know that there is a safe place to insert
   // the conditional branch. We will change the preheader to have a conditional
   // branch on LoopCond.
   BasicBlock *OldPH = L.getLoopPreheader();
-  BasicBlock *NewPH = SplitEdge(OldPH, L.getHeader(), &DT, &LI);
+  BasicBlock *NewPH =
+      SplitEdge(OldPH, L.getHeader(), &DT, &LI, MSSA, MSSAUpdater);
 
   // Now that we have a place to insert the conditional branch, create a place
   // to branch to: this is the exit block out of the loop that we are
   // unswitching. We need to split this if there are other loop predecessors.
   // Because the loop is in simplified form, *any* other predecessor is enough.
   BasicBlock *UnswitchedBB;
   if (FullUnswitch && LoopExitBB->getUniquePredecessor()) {
     assert(LoopExitBB->getUniquePredecessor() == BI.getParent() &&
            "A branch's parent isn't a predecessor!");
     UnswitchedBB = LoopExitBB;
   } else {
-    UnswitchedBB = SplitBlock(LoopExitBB, &LoopExitBB->front(), &DT, &LI);
+    UnswitchedBB = SplitBlock(LoopExitBB, &LoopExitBB->front(), &DT, &LI, MSSA,
+                              MSSAUpdater);
   }
 
   // Actually move the invariant uses into the unswitched position. If possible,
   // we do this by moving the instructions, but when doing partial unswitching
   // we do it by building a new merge of the values in the unswitched position.
   OldPH->getTerminator()->eraseFromParent();
   if (FullUnswitch) {
     // If fully unswitching, we can use the existing branch instruction.
else
                  Instruction::And &&
              "Must have an `and` of `i1`s for the condition!");
     buildPartialUnswitchConditionalBranch(*OldPH, Invariants, ExitDirection,
                                           *UnswitchedBB, *NewPH);
   }
 
   // Rewrite the relevant PHI nodes.
   if (UnswitchedBB == LoopExitBB)
-    rewritePHINodesForUnswitchedExitBlock(*UnswitchedBB, *ParentBB, *OldPH);
+    rewritePHINodesForUnswitchedExitBlock(*UnswitchedBB, *ParentBB, *OldPH,
+                                          MSSA);
   else
     rewritePHINodesForExitAndUnswitchedBlocks(*LoopExitBB, *UnswitchedBB,
-                                              *ParentBB, *OldPH, FullUnswitch);
+                                              *ParentBB, *OldPH, MSSA,
+                                              MSSAUpdater, FullUnswitch);
 
   // Now we need to update the dominator tree.
   DT.insertEdge(OldPH, UnswitchedBB);
   if (FullUnswitch)
     DT.deleteEdge(ParentBB, UnswitchedBB);
 
   // The constant we can replace all of our invariants with inside the loop
   // body. If any of the invariants have a value other than this the loop won't
   // be entered.
   ConstantInt *Replacement = ExitDirection
                                  ? ConstantInt::getFalse(BI.getContext())
                                  : ConstantInt::getTrue(BI.getContext());
 
   // Since this is an i1 condition we can also trivially replace uses of it
   // within the loop with a constant.
   for (Value *Invariant : Invariants)
     replaceLoopInvariantUses(L, Invariant, *Replacement);
 
+  // After the dominator tree was updated, update values in phis no longer
+  // dominated by ParentBB, to use value incoming from NewPh (incoming BB was
+  // already set to OldPh for those value in the rewritePhi calls above).
+  if (MSSA) {
+    MSSAUpdater->updateHeaderPhisUsesAfterLoopUnswitch(ParentBB, NewPH, &DT);
+    MSSA->verifyMemorySSA();
+  }
+
   ++NumTrivial;
   ++NumBranches;
   return true;
 }
 
 /// Unswitch a trivial switch if the condition is loop invariant.
 ///
 /// This routine should only be called when loop code leading to the switch has

 /// the loop by removing now-dead cases. If the default case is one of those
 /// unswitched, it replaces its destination with a new basic block containing
 /// only unreachable. Such basic blocks, while technically loop exits, are not
 /// considered for unswitching so this is a stable transform and the same
 /// switch will not be revisited. If after unswitching there is only a single
 /// in-loop successor, the switch is further simplified to an unconditional
 /// branch. Still more cleanup can be done with some simplify-cfg like pass.
 static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
-                                  LoopInfo &LI) {
+                                  LoopInfo &LI, MemorySSA *MSSA,
+                                  MemorySSAUpdater *MSSAUpdater) {
   LLVM_DEBUG(dbgs() << "  Trying to unswitch switch: " << SI << "\n");
   Value *LoopCond = SI.getCondition();
 
   // If this isn't switching on an invariant condition, we can't unswitch it.
   if (!L.isLoopInvariant(LoopCond))
     return false;
 
   auto *ParentBB = SI.getParent();
if (!L.contains(SI.getDefaultDest()) &&
       areLoopExitPHIsLoopInvariant(L, *ParentBB, *SI.getDefaultDest()) &&
       !isa<UnreachableInst>(SI.getDefaultDest()->getTerminator()))
     DefaultExitBB = SI.getDefaultDest();
   else if (ExitCaseIndices.empty())
     return false;
 
   LLVM_DEBUG(dbgs() << "    unswitching trivial cases...\n");
 
+  if (MSSA) {
+    dbgs() << "Start unswitchTrivialSwitch\n";
+    MSSA->verifyMemorySSA();
+  }
+
   SmallVector<std::pair<ConstantInt *, BasicBlock *>, 4> ExitCases;
   ExitCases.reserve(ExitCaseIndices.size());
   // We walk the case indices backwards so that we remove the last case first
   // and don't disrupt the earlier indices.
   for (unsigned Index : reverse(ExitCaseIndices)) {
     auto CaseI = SI.case_begin() + Index;
     // Save the value of this case.
     ExitCases.push_back({CaseI->getCaseValue(), CaseI->getCaseSuccessor()});
if (SI.getNumCases() == 0)
       CommonSuccBB = SI.getDefaultDest();
     else if (SI.getDefaultDest() != CommonSuccBB)
       CommonSuccBB = nullptr;
   }
 
   // Split the preheader, so that we know that there is a safe place to insert
   // the switch.
   BasicBlock *OldPH = L.getLoopPreheader();
-  BasicBlock *NewPH = SplitEdge(OldPH, L.getHeader(), &DT, &LI);
+  BasicBlock *NewPH =
+      SplitEdge(OldPH, L.getHeader(), &DT, &LI, MSSA, MSSAUpdater);
   OldPH->getTerminator()->eraseFromParent();
 
   // Now add the unswitched switch.
   auto *NewSI = SwitchInst::Create(LoopCond, NewPH, ExitCases.size(), OldPH);
 
   // Rewrite the IR for the unswitched basic blocks. This requires two steps.
   // First, we split any exit blocks with remaining in-loop predecessors. Then
   // we update the PHIs in one of two ways depending on if there was a split.
   // We walk in reverse so that we split in the same order as the cases
   // appeared. This is purely for convenience of reading the resulting IR, but
   // it doesn't cost anything really.
   SmallPtrSet<BasicBlock *, 2> UnswitchedExitBBs;
   SmallDenseMap<BasicBlock *, BasicBlock *, 2> SplitExitBBMap;
   // Handle the default exit if necessary.
   // FIXME: It'd be great if we could merge this with the loop below but LLVM's
   // ranges aren't quite powerful enough yet.
   if (DefaultExitBB) {
     if (pred_empty(DefaultExitBB)) {
       UnswitchedExitBBs.insert(DefaultExitBB);
-      rewritePHINodesForUnswitchedExitBlock(*DefaultExitBB, *ParentBB, *OldPH);
+      rewritePHINodesForUnswitchedExitBlock(*DefaultExitBB, *ParentBB, *OldPH,
+                                            MSSA);
     } else {
-      auto *SplitBB =
-          SplitBlock(DefaultExitBB, &DefaultExitBB->front(), &DT, &LI);
+      auto *SplitBB = SplitBlock(DefaultExitBB, &DefaultExitBB->front(), &DT,
+                                 &LI, MSSA, MSSAUpdater);
       rewritePHINodesForExitAndUnswitchedBlocks(
-          *DefaultExitBB, *SplitBB, *ParentBB, *OldPH, /*FullUnswitch*/ true);
+          *DefaultExitBB, *SplitBB, *ParentBB, *OldPH, MSSA, MSSAUpdater,
+          /*FullUnswitch*/ true);
       DefaultExitBB = SplitExitBBMap[DefaultExitBB] = SplitBB;
     }
   }
   // Note that we must use a reference in the for loop so that we update the
   // container.
   for (auto &CasePair : reverse(ExitCases)) {
     // Grab a reference to the exit block in the pair so that we can update it.
     BasicBlock *ExitBB = CasePair.second;
 
     // If this case is the last edge into the exit block, we can simply reuse it
     // as it will no longer be a loop exit. No mapping necessary.
     if (pred_empty(ExitBB)) {
       // Only rewrite once.
       if (UnswitchedExitBBs.insert(ExitBB).second)
-        rewritePHINodesForUnswitchedExitBlock(*ExitBB, *ParentBB, *OldPH);
+        rewritePHINodesForUnswitchedExitBlock(*ExitBB, *ParentBB, *OldPH, MSSA);
       continue;
     }
 
     // Otherwise we need to split the exit block so that we retain an exit
     // block from the loop and a target for the unswitched condition.
     BasicBlock *&SplitExitBB = SplitExitBBMap[ExitBB];
     if (!SplitExitBB) {
       // If this is the first time we see this, do the split and remember it.
-      SplitExitBB = SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI);
+      SplitExitBB =
+          SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI, MSSA, MSSAUpdater);
       rewritePHINodesForExitAndUnswitchedBlocks(
-          *ExitBB, *SplitExitBB, *ParentBB, *OldPH, /*FullUnswitch*/ true);
+          *ExitBB, *SplitExitBB, *ParentBB, *OldPH, MSSA, MSSAUpdater,
+          /*FullUnswitch*/ true);
     }
     // Update the case pair to point to the split block.
     CasePair.second = SplitExitBB;
   }
 
   // Now add the unswitched cases. We do this in reverse order as we built them
   // in reverse order.
   for (auto CasePair : reverse(ExitCases)) {
static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
   }
   for (auto SplitUnswitchedPair : SplitExitBBMap) {
     auto *UnswitchedBB = SplitUnswitchedPair.second;
     DTUpdates.push_back({DT.Delete, ParentBB, UnswitchedBB});
     DTUpdates.push_back({DT.Insert, OldPH, UnswitchedBB});
   }
   DT.applyUpdates(DTUpdates);
 
+  // Update MemorySSA only after the dominator tree was updated.
+  if (MSSA) {
+    MSSAUpdater->updateHeaderPhisUsesAfterLoopUnswitch(ParentBB, NewPH, &DT);
+    dbgs() << "End unswitchTrivialSwitch\n";
+    MSSA->verifyMemorySSA();
+  }
+
   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
   ++NumTrivial;
   ++NumSwitches;
   return true;
 }
 
 /// This routine scans the loop to find a branch or switch which occurs before
 /// any side effects occur. These can potentially be unswitched without
 /// duplicating the loop. If a branch or switch is successfully unswitched the
 /// scanning continues to see if subsequent branches or switches have become
 /// trivial. Once all trivial candidates have been unswitched, this routine
 /// returns.
 ///
 /// The return value indicates whether anything was unswitched (and therefore
 /// changed).
 static bool unswitchAllTrivialConditions(Loop &L, DominatorTree &DT,
-                                         LoopInfo &LI) {
+                                         LoopInfo &LI, MemorySSA *MSSA,
+                                         MemorySSAUpdater *MSSAUpdater) {
   bool Changed = false;
 
+  // dbgs() << "Start unswitchAllTrivialConditions\n";
   // If loop header has only one reachable successor we should keep looking for
   // trivial condition candidates in the successor as well. An alternative is
   // to constant fold conditions and merge successors into loop header (then we
   // only need to check header's terminator). The reason for not doing this in
   // LoopUnswitch pass is that it could potentially break LoopPassManager's
   // invariants. Folding dead branches could either eliminate the current loop
   // or make other loops unreachable. LCSSA form might also not be preserved
   // after deleting branches. The following code keeps traversing loop header's
do {
 
     if (auto *SI = dyn_cast<SwitchInst>(CurrentTerm)) {
       // Don't bother trying to unswitch past a switch with a constant
       // condition. This should be removed prior to running this pass by
       // simplify-cfg.
       if (isa<Constant>(SI->getCondition()))
         return Changed;
 
-      if (!unswitchTrivialSwitch(L, *SI, DT, LI))
+      if (!unswitchTrivialSwitch(L, *SI, DT, LI, MSSA, MSSAUpdater))
         // Couldn't unswitch this one so we're done.
         return Changed;
 
       // Mark that we managed to unswitch something.
       Changed = true;
 
       // If unswitching turned the terminator into an unconditional branch then
       // we can continue. The unswitching logic specifically works to fold any
do {
     // Don't bother trying to unswitch past an unconditional branch or a branch
     // with a constant value. These should be removed by simplify-cfg prior to
     // running this pass.
     if (!BI->isConditional() || isa<Constant>(BI->getCondition()))
       return Changed;
 
     // Found a trivial condition candidate: non-foldable conditional branch. If
     // we fail to unswitch this, we can't do anything else that is trivial.
-    if (!unswitchTrivialBranch(L, *BI, DT, LI))
+    if (!unswitchTrivialBranch(L, *BI, DT, LI, MSSA, MSSAUpdater))
       return Changed;
 
     // Mark that we managed to unswitch something.
     Changed = true;
 
     // If we only unswitched some of the conditions feeding the branch, we won't
     // have collapsed it to a single successor.
     BI = cast<BranchInst>(CurrentBB->getTerminator());

 /// update the `AssumptionCache` provided in `AC`.
 static BasicBlock *buildClonedLoopBlocks(
     Loop &L, BasicBlock *LoopPH, BasicBlock *SplitBB,
     ArrayRef<BasicBlock *> ExitBlocks, BasicBlock *ParentBB,
     BasicBlock *UnswitchedSuccBB, BasicBlock *ContinueSuccBB,
     const SmallDenseMap<BasicBlock *, BasicBlock *, 16> &DominatingSucc,
     ValueToValueMapTy &VMap,
     SmallVectorImpl<DominatorTree::UpdateType> &DTUpdates, AssumptionCache &AC,
-    DominatorTree &DT, LoopInfo &LI) {
+    DominatorTree &DT, LoopInfo &LI, MemorySSA *MSSA,
+    MemorySSAUpdater *MSSAUpdater) {
   SmallVector<BasicBlock *, 4> NewBlocks;
   NewBlocks.reserve(L.getNumBlocks() + ExitBlocks.size());
 
   // We will need to clone a bunch of blocks, wrap up the clone operation in
   // a helper.
   auto CloneBlock = [&](BasicBlock *OldBB) {
     // Clone the basic block and insert it before the new preheader.
     BasicBlock *NewBB = CloneBasicBlock(OldBB, VMap, ".us", OldBB->getParent());
for (auto *ExitBB : ExitBlocks) {
     if (SkipBlock(ExitBB))
       continue;
 
     // When we are going to clone an exit, we don't need to clone all the
     // instructions in the exit block and we want to ensure we have an easy
     // place to merge the CFG, so split the exit first. This is always safe to
     // do because there cannot be any non-loop predecessors of a loop exit in
     // loop simplified form.
-    auto *MergeBB = SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI);
+    auto *MergeBB =
+        SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI, MSSA, MSSAUpdater);
 
     // Rearrange the names to make it easier to write test cases by having the
     // exit block carry the suffix rather than the merge block carrying the
     // suffix.
     MergeBB->takeName(ExitBB);
     ExitBB->setName(Twine(MergeBB->getName()) + ".split");
 
     // Now clone the original exit block.
#endif
     NonChildClonedLoops.push_back(cloneLoopNest(
         *ChildL, ExitLoopMap.lookup(ClonedChildHeader), VMap, LI));
   }
 }
 
 static void
 deleteDeadClonedBlocks(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
                        ArrayRef<std::unique_ptr<ValueToValueMapTy>> VMaps,
-                       DominatorTree &DT) {
+                       DominatorTree &DT, MemorySSAUpdater *MSSAUpdater) {
   // Find all the dead clones, and remove them from their successors.
   SmallVector<BasicBlock *, 16> DeadBlocks;
   for (BasicBlock *BB : llvm::concat<BasicBlock *const>(L.blocks(), ExitBlocks))
     for (auto &VMap : VMaps)
       if (BasicBlock *ClonedBB = cast_or_null<BasicBlock>(VMap->lookup(BB)))
         if (!DT.isReachableFromEntry(ClonedBB)) {
           for (BasicBlock *SuccBB : successors(ClonedBB))
             SuccBB->removePredecessor(ClonedBB);
           DeadBlocks.push_back(ClonedBB);
         }
 
+  // Remove all MemorySSA in the dead blocks
+  if (MSSAUpdater) {
+    SmallPtrSet<BasicBlock *, 16> DeadBlockSet(DeadBlocks.begin(),
+                                               DeadBlocks.end());
+    MSSAUpdater->removeBlocks(DeadBlockSet);
+  }
+
   // Drop any remaining references to break cycles.
   for (BasicBlock *BB : DeadBlocks)
     BB->dropAllReferences();
   // Erase them from the IR.
   for (BasicBlock *BB : DeadBlocks)
     BB->eraseFromParent();
 }
 
-static void
-deleteDeadBlocksFromLoop(Loop &L,
+static void deleteDeadBlocksFromLoop(Loop &L,
                                      SmallVectorImpl<BasicBlock *> &ExitBlocks,
-                         DominatorTree &DT, LoopInfo &LI) {
+                                     DominatorTree &DT, LoopInfo &LI,
+                                     MemorySSAUpdater *MSSAUpdater) {
   // Find all the dead blocks, and remove them from their successors.
   SmallVector<BasicBlock *, 16> DeadBlocks;
   for (BasicBlock *BB : llvm::concat<BasicBlock *const>(L.blocks(), ExitBlocks))
     if (!DT.isReachableFromEntry(BB)) {
       for (BasicBlock *SuccBB : successors(BB))
         SuccBB->removePredecessor(BB);
       DeadBlocks.push_back(BB);
     }
 
   SmallPtrSet<BasicBlock *, 16> DeadBlockSet(DeadBlocks.begin(),
                                              DeadBlocks.end());
 
+  // Remove all MemorySSA in the dead blocks
+  if (MSSAUpdater)
+    MSSAUpdater->removeBlocks(DeadBlockSet);
+
   // Filter out the dead blocks from the exit blocks list so that it can be
   // used in the caller.
   llvm::erase_if(ExitBlocks,
                  [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
 
   // Walk from this loop up through its parents removing all of the dead blocks.
   for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop()) {
     for (auto *BB : DeadBlocks)
for (DomTreeNode *ChildN : *N) {
              "Cannot visit a node twice when walking a tree!");
       DomWorklist.push_back(ChildN);
     }
   } while (!DomWorklist.empty());
 }
 
 static bool unswitchNontrivialInvariants(
     Loop &L, TerminatorInst &TI, ArrayRef<Value *> Invariants,
-    DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
+    DominatorTree &DT, LoopInfo &LI, MemorySSA *MSSA,
+    MemorySSAUpdater *MSSAUpdater, AssumptionCache &AC,
     function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB) {
   auto *ParentBB = TI.getParent();
   BranchInst *BI = dyn_cast<BranchInst>(&TI);
   SwitchInst *SI = BI ? nullptr : cast<SwitchInst>(&TI);
 
   // We can only unswitch switches, conditional branches with an invariant
   // condition, or combining invariant conditions with an instruction.
   assert((SI || BI->isConditional()) &&
          "Can only unswitch switches and conditional branch!");
   bool FullUnswitch = SI || BI->getCondition() == Invariants[0];
   if (FullUnswitch)
     assert(Invariants.size() == 1 &&
            "Cannot have other invariants with full unswitching!");
   else
     assert(isa<Instruction>(BI->getCondition()) &&
            "Partial unswitching requires an instruction as the condition!");
 
+  if (MSSA) {
+    // dbgs() << "Start unswitchInvariantBranch, is full? "<<FullUnswitch<<"\n";
+    MSSA->verifyMemorySSA();
+  }
+
   // Constant and BBs tracking the cloned and continuing successor. When we are
   // unswitching the entire condition, this can just be trivially chosen to
   // unswitch towards `true`. However, when we are unswitching a set of
   // invariants combined with `and` or `or`, the combining operation determines
   // the best direction to unswitch: we want to unswitch the direction that will
   // collapse the branch.
   bool Direction = true;
   int ClonedSucc = 0;
static bool unswitchNontrivialInvariants(
   // FIXME: We should teach SplitBlock to handle this and remove this
   // restriction.
   for (auto *ExitBB : ExitBlocks)
     if (isa<CleanupPadInst>(ExitBB->getFirstNonPHI()))
       return false;
 
   // Compute the parent loop now before we start hacking on things.
   Loop *ParentL = L.getParentLoop();
+  // Get blocks in RPO order for MSSA update, also before changing CFG.
+  LoopBlocksRPO LBRPO(&L);
+  if (MSSA)
+    LBRPO.perform(&LI);
 
   // Compute the outer-most loop containing one of our exit blocks. This is the
   // furthest up our loopnest which can be mutated, which we will use below to
   // update things.
   Loop *OuterExitL = &L;
   for (auto *ExitBB : ExitBlocks) {
     Loop *NewOuterExitL = LI.getLoopFor(ExitBB);
     if (!NewOuterExitL) {
if (SuccBB->getUniquePredecessor() ||
       });
 
   // Split the preheader, so that we know that there is a safe place to insert
   // the conditional branch. We will change the preheader to have a conditional
   // branch on LoopCond. The original preheader will become the split point
   // between the unswitched versions, and we will have a new preheader for the
   // original loop.
   BasicBlock *SplitBB = L.getLoopPreheader();
-  BasicBlock *LoopPH = SplitEdge(SplitBB, L.getHeader(), &DT, &LI);
+  BasicBlock *LoopPH =
+      SplitEdge(SplitBB, L.getHeader(), &DT, &LI, MSSA, MSSAUpdater);
 
   // Keep track of the dominator tree updates needed.
   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
 
   // Clone the loop for each unswitched successor.
   SmallVector<std::unique_ptr<ValueToValueMapTy>, 4> VMaps;
   VMaps.reserve(UnswitchedSuccBBs.size());
   SmallDenseMap<BasicBlock *, BasicBlock *, 4> ClonedPHs;
   for (auto *SuccBB : UnswitchedSuccBBs) {
     VMaps.emplace_back(new ValueToValueMapTy());
-    ClonedPHs[SuccBB] = buildClonedLoopBlocks(
-        L, LoopPH, SplitBB, ExitBlocks, ParentBB, SuccBB, RetainedSuccBB,
-        DominatingSucc, *VMaps.back(), DTUpdates, AC, DT, LI);
+    ClonedPHs[SuccBB] =
+        buildClonedLoopBlocks(L, LoopPH, SplitBB, ExitBlocks, ParentBB, SuccBB,
+                              RetainedSuccBB, DominatingSucc, *VMaps.back(),
+                              DTUpdates, AC, DT, LI, MSSA, MSSAUpdater);
   }
 
   // The stitching of the branched code back together depends on whether we're
   // doing full unswitching or not with the exception that we always want to
   // nuke the initial terminator placed in the split block.
   SplitBB->getTerminator()->eraseFromParent();
   if (FullUnswitch) {
-    for (BasicBlock *SuccBB : UnswitchedSuccBBs) {
-      // Remove the parent as a predecessor of the unswitched successor.
-      SuccBB->removePredecessor(ParentBB,
-                                /*DontDeleteUselessPHIs*/ true);
-      DTUpdates.push_back({DominatorTree::Delete, ParentBB, SuccBB});
-    }
-
-    // Now splice the terminator from the original loop and rewrite its
+    // Splice the terminator from the original loop and rewrite its
     // successors.
     SplitBB->getInstList().splice(SplitBB->end(), ParentBB->getInstList(), TI);
+
+    // Keep a clone of the terminator for now.
+    Instruction *NewTI = TI.clone();
+    if (TI.hasName())
+      NewTI->setName(TI.getName() + ".us");
+    ParentBB->getInstList().push_back(NewTI);
+
     if (BI) {
       assert(UnswitchedSuccBBs.size() == 1 &&
              "Only one possible unswitched block for a branch!");
       BasicBlock *ClonedPH = ClonedPHs.begin()->second;
       BI->setSuccessor(ClonedSucc, ClonedPH);
       BI->setSuccessor(1 - ClonedSucc, LoopPH);
       DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
     } else {
       assert(SI && "Must either be a branch or switch!");
 
       // Walk the cases and directly update their successors.
       for (auto &Case : SI->cases())
         Case.setSuccessor(ClonedPHs.find(Case.getCaseSuccessor())->second);
       // We need to use the set to populate domtree updates as even when there
       // are multiple cases pointing at the same successor we only want to
       // insert one edge in the domtree.
       for (BasicBlock *SuccBB : UnswitchedSuccBBs)
         DTUpdates.push_back(
             {DominatorTree::Insert, SplitBB, ClonedPHs.find(SuccBB)->second});
 
       SI->setDefaultDest(LoopPH);
     }
 
+    if (MSSA) {
+      DT.applyUpdates(DTUpdates);
+      DTUpdates.clear();
+      for (auto &VMap : VMaps)
+        MSSAUpdater->updateForClonedLoop(LBRPO, ExitBlocks, *VMap,
+                                         /*IgnoreIncomingWithNoClones=*/true);
+      MSSAUpdater->updateExitBlocksForClonedLoop(ExitBlocks, VMaps, &DT);
+      for (BasicBlock *SuccBB : UnswitchedSuccBBs)
+        MSSAUpdater->removeEdge(ParentBB, SuccBB);
+    }
+
+    for (BasicBlock *SuccBB : UnswitchedSuccBBs) {
+      // Remove the parent as a predecessor of the unswitched successor.
+      SuccBB->removePredecessor(ParentBB,
+                                /*DontDeleteUselessPHIs*/ true);
+      DTUpdates.push_back({DominatorTree::Delete, ParentBB, SuccBB});
+    }
+    ParentBB->getTerminator()->eraseFromParent();
     // Create a new unconditional branch to the continuing block (as opposed to
     // the one cloned).
     BranchInst::Create(RetainedSuccBB, ParentBB);
   } else {
     assert(BI && "Only branches have partial unswitching.");
     assert(UnswitchedSuccBBs.size() == 1 &&
            "Only one possible unswitched block for a branch!");
     BasicBlock *ClonedPH = ClonedPHs.begin()->second;
     // When doing a partial unswitch, we have to do a bit more work to build up
     // the branch in the split block.
     buildPartialUnswitchConditionalBranch(*SplitBB, Invariants, Direction,
                                           *ClonedPH, *LoopPH);
     DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
   }
 
   // Apply the updates accumulated above to get an up-to-date dominator tree.
   DT.applyUpdates(DTUpdates);
 
   // Now that we have an accurate dominator tree, first delete the dead cloned
   // blocks so that we can accurately build any cloned loops. It is important to
   // not delete the blocks from the original loop yet because we still want to
   // reference the original loop to understand the cloned loop's structure.
-  deleteDeadClonedBlocks(L, ExitBlocks, VMaps, DT);
+  deleteDeadClonedBlocks(L, ExitBlocks, VMaps, DT, MSSAUpdater);
 
   // Build the cloned loop structure itself. This may be substantially
   // different from the original structure due to the simplified CFG. This also
   // handles inserting all the cloned blocks into the correct loops.
   SmallVector<Loop *, 4> NonChildClonedLoops;
   for (std::unique_ptr<ValueToValueMapTy> &VMap : VMaps)
     buildClonedLoops(L, ExitBlocks, *VMap, LI, NonChildClonedLoops);
 
   // Now that our cloned loops have been built, we can update the original loop.
   // First we delete the dead blocks from it and then we rebuild the loop
   // structure taking these deletions into account.
-  deleteDeadBlocksFromLoop(L, ExitBlocks, DT, LI);
+  deleteDeadBlocksFromLoop(L, ExitBlocks, DT, LI, MSSAUpdater);
+
+  if (MSSA) {
+    // dbgs()<<"After deleting dead blocks\n";
+    // MSSA->dump();
+    MSSA->verifyMemorySSA();
+  }
+
   SmallVector<Loop *, 4> HoistedLoops;
   bool IsStillLoop = rebuildLoopAfterUnswitch(L, ExitBlocks, LI, HoistedLoops);
 
+  if (MSSA) {
+    // dbgs()<<"After rebuilding loop\n";
+    // MSSA->dump();
+    MSSA->verifyMemorySSA();
+  }
+
   // This transformation has a high risk of corrupting the dominator tree, and
   // the below steps to rebuild loop structures will result in hard to debug
   // errors in that case so verify that the dominator tree is sane first.
   // FIXME: Remove this when the bugs stop showing up and rely on existing
   // verification steps.
   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 
   if (BI) {
#endif
   // For that we need to merge together the updated loops and the cloned loops
   // and check whether the original loop survived.
   SmallVector<Loop *, 4> SibLoops;
   for (Loop *UpdatedL : llvm::concat<Loop *>(NonChildClonedLoops, HoistedLoops))
     if (UpdatedL->getParentLoop() == ParentL)
       SibLoops.push_back(UpdatedL);
   UnswitchCB(IsStillLoop, SibLoops);
 
+  if (MSSA) {
+    MSSA->verifyMemorySSA();
+  }
+
   ++NumBranches;
   return true;
 }
 
 /// Recursively compute the cost of a dominator subtree based on the per-block
 /// cost map provided.
 ///
 /// The recursive computation is memozied into the provided DT-indexed cost map
int Cost = std::accumulate(
         return Sum + computeDomSubtreeCost(*ChildN, BBCostMap, DTCostMap);
       });
   bool Inserted = DTCostMap.insert({&N, Cost}).second;
   (void)Inserted;
   assert(Inserted && "Should not insert a node while visiting children!");
   return Cost;
 }
 
-static bool unswitchBestCondition(
-    Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
+static bool
+unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI, MemorySSA *MSSA,
+                      MemorySSAUpdater *MSSAUpdater, AssumptionCache &AC,
                       TargetTransformInfo &TTI,
                       function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB) {
   // Collect all invariant conditions within this loop (as opposed to an inner
   // loop which would be handled when visiting that inner loop).
   SmallVector<std::pair<TerminatorInst *, TinyPtrVector<Value *>>, 4>
       UnswitchCandidates;
   for (auto *BB : L.blocks()) {
     if (LI.getLoopFor(BB) != &L)
       continue;
 
if (BestUnswitchCost >= UnswitchThreshold) {
     LLVM_DEBUG(dbgs() << "Cannot unswitch, lowest cost found: "
                       << BestUnswitchCost << "\n");
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "  Trying to unswitch non-trivial (cost = "
                     << BestUnswitchCost << ") terminator: " << *BestUnswitchTI
                     << "\n");
-  return unswitchNontrivialInvariants(
-      L, *BestUnswitchTI, BestUnswitchInvariants, DT, LI, AC, UnswitchCB);
+  return unswitchNontrivialInvariants(L, *BestUnswitchTI,
+                                      BestUnswitchInvariants, DT, LI, MSSA,
+                                      MSSAUpdater, AC, UnswitchCB);
 }
 
 /// Unswitch control flow predicated on loop invariant conditions.
 ///
 /// This first hoists all branches or switches which are trivial (IE, do not
 /// require duplicating any part of the loop) out of the loop body. It then
 /// looks at other loop invariant control flows and tries to unswitch those as
 /// well by cloning the loop if the result is small enough.
 static bool
 unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
-             TargetTransformInfo &TTI, bool NonTrivial,
+             TargetTransformInfo &TTI, MemorySSA *MSSA,
+             MemorySSAUpdater *MSSAUpdater, bool NonTrivial,
              function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB) {
   assert(L.isRecursivelyLCSSAForm(DT, LI) &&
          "Loops must be in LCSSA form before unswitching.");
   bool Changed = false;
 
   // Must be in loop simplified form: we need a preheader and dedicated exits.
   if (!L.isLoopSimplifyForm())
     return false;
 
   // Try trivial unswitch first before loop over other basic blocks in the loop.
-  if (unswitchAllTrivialConditions(L, DT, LI)) {
+  if (unswitchAllTrivialConditions(L, DT, LI, MSSA, MSSAUpdater)) {
     // If we unswitched successfully we will want to clean up the loop before
     // processing it further so just mark it as unswitched and return.
     UnswitchCB(/*CurrentLoopValid*/ true, {});
     return true;
   }
 
   // If we're not doing non-trivial unswitching, we're done. We both accept
   // a parameter but also check a local flag that can be used for testing
   // a debugging.
   if (!NonTrivial && !EnableNonTrivialUnswitch)
     return false;
 
   // For non-trivial unswitching, because it often creates new loops, we rely on
   // the pass manager to iterate on the loops rather than trying to immediately
   // reach a fixed point. There is no substantial advantage to iterating
   // internally, and if any of the new loops are simplified enough to contain
   // trivial unswitching we want to prefer those.
 
   // Try to unswitch the best invariant condition. We prefer this full unswitch to
   // a partial unswitch when possible below the threshold.
-  if (unswitchBestCondition(L, DT, LI, AC, TTI, UnswitchCB))
+  if (unswitchBestCondition(L, DT, LI, MSSA, MSSAUpdater, AC, TTI, UnswitchCB))
     return true;
 
   // No other opportunities to unswitch.
   return Changed;
 }
 
 PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
                                               LoopStandardAnalysisResults &AR,
auto UnswitchCB = [&L, &U, &LoopName](bool CurrentLoopValid,
     // If the current loop remains valid, we should revisit it to catch any
     // other unswitch opportunities. Otherwise, we need to mark it as deleted.
     if (CurrentLoopValid)
       U.revisitCurrentLoop();
     else
       U.markLoopAsDeleted(L, LoopName);
   };
 
-  if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.TTI, NonTrivial,
-                    UnswitchCB))
+  std::unique_ptr<MemorySSAUpdater> MSSAUpdater;
+  if (AR.MSSA) {
+    MSSAUpdater = make_unique<MemorySSAUpdater>(AR.MSSA);
+    // dbgs()<<"[SimpleLoopUnswitch]=====Before unswitchLoop call\n";
+    // AR.MSSA->dump();
+    AR.MSSA->verifyMemorySSA();
+  }
+  if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.TTI, AR.MSSA, MSSAUpdater.get(),
+                    NonTrivial, UnswitchCB))
     return PreservedAnalyses::all();
 
+  if (AR.MSSA) {
+    // dbgs()<<"[SimpleLoopUnswitch]=====After unswitchLoop call\n";
+    // AR.MSSA->dump();
+    AR.MSSA->verifyMemorySSA();
+  }
+
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
   return getLoopPassPreservedAnalyses();
 }
 
 namespace {
 
initializeSimpleLoopUnswitchLegacyPassPass(
         *PassRegistry::getPassRegistry());
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
+    if (EnableMSSALoopDependency)
+      AU.addRequired<MemorySSAWrapperPass>();
     getLoopAnalysisUsage(AU);
   }
 };
 
 } // end anonymous namespace
 
 bool SimpleLoopUnswitchLegacyPass::runOnLoop(Loop *L, LPPassManager &LPM) {
   if (skipLoop(L))
     return false;
 
   Function &F = *L->getHeader()->getParent();
 
   LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << *L
                     << "\n");
 
   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  MemorySSA *MSSA = nullptr;
+  std::unique_ptr<MemorySSAUpdater> MSSAUpdater;
+  if (EnableMSSALoopDependency) {
+    MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
+    MSSAUpdater = make_unique<MemorySSAUpdater>(MSSA);
+  }
 
   auto UnswitchCB = [&L, &LPM](bool CurrentLoopValid,
                                ArrayRef<Loop *> NewLoops) {
     // If we did a non-trivial unswitch, we have added new (cloned) loops.
     for (auto *NewL : NewLoops)
       LPM.addLoop(*NewL);
 
     // If the current loop remains valid, re-add it to the queue. This is
     // a little wasteful as we'll finish processing the current loop as well,
     // but it is the best we can do in the old PM.
     if (CurrentLoopValid)
       LPM.addLoop(*L);
     else
       LPM.markLoopAsDeleted(*L);
   };
 
-  bool Changed =
-      unswitchLoop(*L, DT, LI, AC, TTI, NonTrivial, UnswitchCB);
+  if (MSSA) {
+    // dbgs()<<"Before unswitchLoop call\n";
+    MSSA->verifyMemorySSA();
+  }
+
+  bool Changed = unswitchLoop(*L, DT, LI, AC, TTI, MSSA, MSSAUpdater.get(),
+                              NonTrivial, UnswitchCB);
+
+  if (MSSA) {
+    // TODO: Compare performance of update vs rebuild MSSA. Hack:
+    // auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+    // MSSA->~MemorySSA();
+    // new (MSSA) MemorySSA(*F, &AA, DT);
+
+    // dbgs()<<"After unswitchLoop call\n";
+    MSSA->verifyMemorySSA();
+  }
 
   // If anything was unswitched, also clear any cached information about this
   // loop.
   LPM.deleteSimpleAnalysisLoop(L);
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 
   return Changed;
 }
 
 char SimpleLoopUnswitchLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(SimpleLoopUnswitchLegacyPass, "simple-loop-unswitch",
                       "Simple unswitch loops", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
+INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(SimpleLoopUnswitchLegacyPass, "simple-loop-unswitch",
                     "Simple unswitch loops", false, false)
 
 Pass *llvm::createSimpleLoopUnswitchLegacyPass(bool NonTrivial) {
   return new SimpleLoopUnswitchLegacyPass(NonTrivial);
 }