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

1095–1100

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

1095–1100

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/ADT/Twine.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CodeMetrics.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/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"
 #include "llvm/IR/Instructions.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) {
+                                                      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);
+  }
 }
 
 /// 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");
 
   Value *LoopCond = BI.getCondition();
 
   // Need a trivial loop condition to unswitch.
   if (!L.isLoopInvariant(LoopCond))
     return false;
static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
   }
   auto *ContinueBB = BI.getSuccessor(1 - LoopExitSuccIdx);
   auto *ParentBB = BI.getParent();
   if (!areLoopExitPHIsLoopInvariant(L, *ParentBB, *LoopExitBB))
     return false;
 
   LLVM_DEBUG(dbgs() << "    unswitching trivial branch when: " << CondVal
                     << " == " << LoopCond << "\n");
+  if (MSSA) {
+    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 (BasicBlock *PredBB = LoopExitBB->getUniquePredecessor()) {
     (void)PredBB;
     assert(PredBB == 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);
   }
 
   // Now splice the branch to gate reaching the new preheader and re-point its
   // successors.
   OldPH->getInstList().splice(std::prev(OldPH->end()),
                               BI.getParent()->getInstList(), BI);
   OldPH->getTerminator()->eraseFromParent();
   BI.setSuccessor(LoopExitSuccIdx, UnswitchedBB);
   BI.setSuccessor(1 - LoopExitSuccIdx, NewPH);
 
   // Create a new unconditional branch that will continue the loop as a new
   // terminator.
   BranchInst::Create(ContinueBB, ParentBB);
 
   // Rewrite the relevant PHI nodes.
   if (UnswitchedBB == LoopExitBB)
-    rewritePHINodesForUnswitchedExitBlock(*UnswitchedBB, *ParentBB, *OldPH);
+    rewritePHINodesForUnswitchedExitBlock(*UnswitchedBB, *ParentBB, *OldPH, MSSA);
   else
     rewritePHINodesForExitAndUnswitchedBlocks(*LoopExitBB, *UnswitchedBB,
-                                              *ParentBB, *OldPH);
+                                              *ParentBB, *OldPH, MSSA,
+                                              MSSAUpdater);
 
   // Now we need to update the dominator tree.
   DT.applyUpdates(
       {{DT.Delete, ParentBB, UnswitchedBB}, {DT.Insert, OldPH, UnswitchedBB}});
 
   // Since this is an i1 condition we can also trivially replace uses of it
   // within the loop with a constant.
   replaceLoopUsesWithConstant(L, *LoopCond, *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->updatePhisAfterUnswitch(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) {
+    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);
+          SplitBlock(DefaultExitBB, &DefaultExitBB->front(), &DT, &LI, MSSA,
+                     MSSAUpdater);
       rewritePHINodesForExitAndUnswitchedBlocks(*DefaultExitBB, *SplitBB,
-                                                *ParentBB, *OldPH);
+                                                *ParentBB, *OldPH, MSSA,
+                                                MSSAUpdater);
       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);
+                                                *ParentBB, *OldPH, MSSA,
+                                                MSSAUpdater);
     }
     // 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->updatePhisAfterUnswitch(ParentBB, NewPH, &DT);
+    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;
 
   // 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
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))
         // Coludn'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;
 
     // We unswitched the branch. This should always leave us with an
     // unconditional branch that we can follow now.
     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 SmallPtrSetImpl<BasicBlock *> &SkippedLoopAndExitBlocks,
     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 (SkippedLoopAndExitBlocks.count(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.
for (auto *ChildLoopBB : ChildL->blocks())
       assert(VMap.count(ChildLoopBB) &&
              "Cloned a child loop header but not all of that loops blocks!");
 #endif
 
     NonChildClonedLoops.push_back(cloneLoopNest(
         *ChildL, ExitLoopMap.lookup(ClonedChildHeader), VMap, LI));
   }
 
   // Return the main cloned loop if any.
   return ClonedL;
 }
 
 static void
 deleteDeadBlocksFromLoop(Loop &L,
                          const SmallVectorImpl<BasicBlock *> &DeadBlocks,
                          SmallVectorImpl<BasicBlock *> &ExitBlocks,
-                         DominatorTree &DT, LoopInfo &LI) {
+                         DominatorTree &DT, LoopInfo &LI,
+                         MemorySSAUpdater *MSSAUpdater) {
   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); });
 
   // Remove these blocks from their successors.
   for (auto *BB : DeadBlocks)
     for (BasicBlock *SuccBB : successors(BB))

 /// clone the necessary parts of the loop to maintain behavior.
 ///
 /// It also updates both dominator tree and loopinfo based on the unswitching.
 ///
 /// Once unswitching has been performed it runs the provided callback to report
 /// the new loops and no-longer valid loops to the caller.
 static bool unswitchInvariantBranch(
     Loop &L, BranchInst &BI, DominatorTree &DT, LoopInfo &LI,
-    AssumptionCache &AC,
+    MemorySSA *MSSA, MemorySSAUpdater *MSSAUpdater, AssumptionCache &AC,
     function_ref<void(bool, ArrayRef<Loop *>)> NonTrivialUnswitchCB) {
   assert(BI.isConditional() && "Can only unswitch a conditional branch!");
   assert(L.isLoopInvariant(BI.getCondition()) &&
          "Can only unswitch an invariant branch condition!");
 
+  if (MSSA) {
+    MSSA->verifyMemorySSA();
+  }
+
   // Constant and BBs tracking the cloned and continuing successor.
   const int ClonedSucc = 0;
   auto *ParentBB = BI.getParent();
   auto *UnswitchedSuccBB = BI.getSuccessor(ClonedSucc);
   auto *ContinueSuccBB = BI.getSuccessor(1 - ClonedSucc);
 
   assert(UnswitchedSuccBB != ContinueSuccBB &&
          "Should not unswitch a branch that always goes to the same place!");
bool DeleteUnswitchedSucc =
       });
 
   // 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 a mapping for the cloned values.
   ValueToValueMapTy VMap;
 
   // Keep track of the dominator tree updates needed.
   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
 
   // Build the cloned blocks from the loop.
   auto *ClonedPH = buildClonedLoopBlocks(
       L, LoopPH, SplitBB, ExitBlocks, ParentBB, UnswitchedSuccBB,
-      ContinueSuccBB, SkippedLoopAndExitBlocks, VMap, DTUpdates, AC, DT, LI);
+      ContinueSuccBB, SkippedLoopAndExitBlocks, VMap, DTUpdates, AC, DT, LI,
+      MSSA, MSSAUpdater);
 
   // Remove the parent as a predecessor of the unswitched successor.
   UnswitchedSuccBB->removePredecessor(ParentBB, /*DontDeleteUselessPHIs*/ true);
 
   // Now splice the branch from the original loop and use it to select between
   // the two loops.
   SplitBB->getTerminator()->eraseFromParent();
   SplitBB->getInstList().splice(SplitBB->end(), ParentBB->getInstList(), BI);
   BI.setSuccessor(ClonedSucc, ClonedPH);
   BI.setSuccessor(1 - ClonedSucc, LoopPH);
 
   // Create a new unconditional branch to the continuing block (as opposed to
   // the one cloned).
-  BranchInst::Create(ContinueSuccBB, ParentBB);
+  BranchInst *GoodBranch = BranchInst::Create(ContinueSuccBB, ParentBB);
 
   // Before we update the dominator tree, collect the dead blocks if we're going
   // to end up deleting the unswitched successor.
   SmallVector<BasicBlock *, 16> DeadBlocks;
   if (DeleteUnswitchedSucc) {
     DeadBlocks.push_back(UnswitchedSuccBB);
     for (int i = 0; i < (int)DeadBlocks.size(); ++i) {
       // If we reach an exit block, stop recursing as the unswitched loop will
       // end up reaching the merge block which we make the successor of the
       // exit.
       if (ExitBlockSet.count(DeadBlocks[i]))
         continue;
 
       // Insert the children that are within the loop or exit block set. Other
       // children may reach out of the loop. While we don't expect these to be
       // dead (as the unswitched clone should reach them) we don't try to prove
       // that here.
       for (DomTreeNode *ChildN : *DT[DeadBlocks[i]])
         if (L.contains(ChildN->getBlock()) ||
             ExitBlockSet.count(ChildN->getBlock()))
           DeadBlocks.push_back(ChildN->getBlock());
     }
   }
 
+  if (MSSA) {
+    GoodBranch->eraseFromParent();
+    // Create a new branch to facilitate MemorySSA update.
+    BranchInst *RevBranch =
+        BranchInst::Create(UnswitchedSuccBB, ContinueSuccBB,
+                           UndefValue::get(
+                               llvm::Type::getInt1Ty(ParentBB->getContext())),
+                           ParentBB);
+
+    //DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
+    //DT.applyUpdates(DTUpdates);
+
+    SmallPtrSet<BasicBlock *, 4> subLoopHeaders;
+    subLoopHeaders.insert(L.getHeader());
+    for (Loop *InnerL : L.getSubLoops())
+      subLoopHeaders.insert(InnerL->getHeader());
+
+    // Note: this will work with a partially invalid DT, but needs the reverse
+    // branch set for correct updates.
+    MSSAUpdater->updateForClonedLoop(L.getHeader(), ExitBlocks, VMap,
+                                     subLoopHeaders, &DT,
+                                     /*IgnoreIncomingWithNoClones=*/true);
+
+    // Now remove the branch and create the right one below.
+    RevBranch->eraseFromParent();
+    BranchInst::Create(ContinueSuccBB, ParentBB);
+    //DTUpdates.clear();
+    //DTUpdates.push_back({DominatorTree::Delete, ParentBB, UnswitchedSuccBB});
+    //DT.applyUpdates(DTUpdates);
+  }
+
   // Add the remaining edges to our updates and apply them to get an up-to-date
   // dominator tree. Note that this will cause the dead blocks above to be
   // unreachable and no longer in the dominator tree.
   DTUpdates.push_back({DominatorTree::Delete, ParentBB, UnswitchedSuccBB});
   DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
   DT.applyUpdates(DTUpdates);
 
   // 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;
   Loop *ClonedL =
       buildClonedLoops(L, ExitBlocks, VMap, LI, NonChildClonedLoops);
 
   // Delete anything that was made dead in the original loop due to
   // unswitching.
   if (!DeadBlocks.empty())
-    deleteDeadBlocksFromLoop(L, DeadBlocks, ExitBlocks, DT, LI);
+    deleteDeadBlocksFromLoop(L, DeadBlocks, ExitBlocks, DT, LI, MSSAUpdater);
+
+  if (MSSA) {
+    MSSA->verifyMemorySSA();
+  }
 
   SmallVector<Loop *, 4> HoistedLoops;
   bool IsStillLoop = rebuildLoopAfterUnswitch(L, ExitBlocks, LI, HoistedLoops);
 
   // 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
#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);
   NonTrivialUnswitchCB(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

 /// 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 *>)> NonTrivialUnswitchCB) {
   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.
-  Changed |= unswitchAllTrivialConditions(L, DT, LI);
+  Changed |= unswitchAllTrivialConditions(L, DT, LI, MSSA, MSSAUpdater);
 
   // 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 Changed;
 
   // Collect all remaining invariant branch conditions within this loop (as
for (TerminatorInst *CandidateTI : UnswitchCandidates) {
     }
   }
 
   if (BestUnswitchCost < UnswitchThreshold) {
     LLVM_DEBUG(dbgs() << "  Trying to unswitch non-trivial (cost = "
                       << BestUnswitchCost << ") branch: " << *BestUnswitchTI
                       << "\n");
     Changed |= unswitchInvariantBranch(L, cast<BranchInst>(*BestUnswitchTI), DT,
-                                       LI, AC, NonTrivialUnswitchCB);
+                                       LI, MSSA, MSSAUpdater, AC,
+                                       NonTrivialUnswitchCB);
   } else {
     LLVM_DEBUG(dbgs() << "Cannot unswitch, lowest cost found: "
                       << BestUnswitchCost << "\n");
   }
 
   return Changed;
 }
 
auto NonTrivialUnswitchCB = [&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,
-                    NonTrivialUnswitchCB))
+  std::unique_ptr<MemorySSAUpdater> MSSAUpdater;
+  if (AR.MSSA)
+    MSSAUpdater = make_unique<MemorySSAUpdater>(AR.MSSA);
+  if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.TTI, AR.MSSA, MSSAUpdater.get(),
+                    NonTrivial, NonTrivialUnswitchCB))
     return PreservedAnalyses::all();
 
+  if (AR.MSSA) {
+    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 NonTrivialUnswitchCB = [&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);
   };
 
+  if (MSSA) {
+    MSSA->verifyMemorySSA();
+  }
+
   bool Changed =
-      unswitchLoop(*L, DT, LI, AC, TTI, NonTrivial, NonTrivialUnswitchCB);
+      unswitchLoop(*L, DT, LI, AC, TTI, MSSA, MSSAUpdater.get(), NonTrivial,
+                   NonTrivialUnswitchCB);
+
+  if (MSSA) {
+    // TODO: Compare performance of update vs rebuild MSSA. Hack:
+    // auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+    // MSSA->~MemorySSA();
+    // new (MSSA) MemorySSA(*F, &AA, DT);
+
+    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);
 }