Update MemorySSA in LoopUnswitch.

[asbirlea]

Summary

Update MemorySSA in old LoopUnswitch pass.
Actual dependency and update is disabled by default.

[asbirlea]

Rebase on dependent patch.

[asbirlea]

Update to use new API in dependent patch.

[asbirlea]

Rebase and update after API change.

[asbirlea]

Rebase. Ready to review.

[chandlerc]

LGTM

lib/Transforms/Scalar/LoopUnswitch.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/Analysis/DivergenceAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/MemorySSA.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/Utils/Local.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"

 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
 #include <cassert>
 #include <map>
 #include <set>
class LoopUnswitch : public LoopPass {
 
     LUAnalysisCache BranchesInfo;
 
     bool OptimizeForSize;
     bool redoLoop = false;
 
     Loop *currentLoop = nullptr;
     DominatorTree *DT = nullptr;
+    MemorySSA *MSSA = nullptr;
+    std::unique_ptr<MemorySSAUpdater> MSSAUpdater;
     BasicBlock *loopHeader = nullptr;
     BasicBlock *loopPreheader = nullptr;
 
     bool SanitizeMemory;
     LoopSafetyInfo SafetyInfo;
 
     // LoopBlocks contains all of the basic blocks of the loop, including the
     // preheader of the loop, the body of the loop, and the exit blocks of the
public:
     bool isUnreachableDueToPreviousUnswitching(BasicBlock *);
 
     /// This transformation requires natural loop information & requires that
     /// loop preheaders be inserted into the CFG.
     ///
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.addRequired<AssumptionCacheTracker>();
       AU.addRequired<TargetTransformInfoWrapperPass>();
+      if (EnableMSSALoopDependency)
+        AU.addRequired<MemorySSAWrapperPass>();
       if (hasBranchDivergence)
         AU.addRequired<DivergenceAnalysis>();
       getLoopAnalysisUsage(AU);
     }
 
   private:
     void releaseMemory() override {
       BranchesInfo.forgetLoop(currentLoop);

 
 char LoopUnswitch::ID = 0;
 
 INITIALIZE_PASS_BEGIN(LoopUnswitch, "loop-unswitch", "Unswitch loops",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DivergenceAnalysis)
 INITIALIZE_PASS_END(LoopUnswitch, "loop-unswitch", "Unswitch loops",
                       false, false)
 
 Pass *llvm::createLoopUnswitchPass(bool Os, bool hasBranchDivergence) {
   return new LoopUnswitch(Os, hasBranchDivergence);
 }
 
bool LoopUnswitch::runOnLoop(Loop *L, LPPassManager &LPM_Ref) {
   if (skipLoop(L))
     return false;
 
   AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(
       *L->getHeader()->getParent());
   LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   LPM = &LPM_Ref;
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  if (EnableMSSALoopDependency) {
+    MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
+    MSSAUpdater = make_unique<MemorySSAUpdater>(MSSA);
+  }
   currentLoop = L;
   Function *F = currentLoop->getHeader()->getParent();
 
   SanitizeMemory = F->hasFnAttribute(Attribute::SanitizeMemory);
   if (SanitizeMemory)
     computeLoopSafetyInfo(&SafetyInfo, L);
 
+  if (MSSA)
+    MSSA->verifyMemorySSA();
+
   bool Changed = false;
   do {
     assert(currentLoop->isLCSSAForm(*DT));
     redoLoop = false;
     Changed |= processCurrentLoop();
   } while(redoLoop);
 
+  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();
+  }
+
   return Changed;
 }
 
 // Return true if the BasicBlock BB is unreachable from the loop header.
 // Return false, otherwise.
 bool LoopUnswitch::isUnreachableDueToPreviousUnswitching(BasicBlock *BB) {
   auto *Node = DT->getNode(BB)->getIDom();
   BasicBlock *DomBB = Node->getBlock();
LLVM_DEBUG(dbgs() << "loop-unswitch: Trivial-Unswitch loop %"
                     << loopHeader->getName() << " [" << L->getBlocks().size()
                     << " blocks] in Function "
                     << L->getHeader()->getParent()->getName()
                     << " on cond: " << *Val << " == " << *Cond << "\n");
 
   // First step, split the preheader, so that we know that there is a safe place
   // to insert the conditional branch.  We will change loopPreheader to have a
   // conditional branch on Cond.
-  BasicBlock *NewPH = SplitEdge(loopPreheader, loopHeader, DT, LI);
+  BasicBlock *NewPH =
+      SplitEdge(loopPreheader, loopHeader, DT, LI, MSSA, MSSAUpdater.get());
 
   // 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 should
   // short-circuit to.
 
   // Split this block now, so that the loop maintains its exit block, and so
   // that the jump from the preheader can execute the contents of the exit block
   // without actually branching to it (the exit block should be dominated by the
   // loop header, not the preheader).
   assert(!L->contains(ExitBlock) && "Exit block is in the loop?");
-  BasicBlock *NewExit = SplitBlock(ExitBlock, &ExitBlock->front(), DT, LI);
+  BasicBlock *NewExit = SplitBlock(ExitBlock, &ExitBlock->front(), DT, LI, MSSA,
+                                   MSSAUpdater.get());
 
   // Okay, now we have a position to branch from and a position to branch to,
   // insert the new conditional branch.
   auto *OldBranch = dyn_cast<BranchInst>(loopPreheader->getTerminator());
   assert(OldBranch && "Failed to split the preheader");
   EmitPreheaderBranchOnCondition(Cond, Val, NewExit, NewPH, OldBranch, TI);
   LPM->deleteSimpleAnalysisValue(OldBranch, L);
 
   // EmitPreheaderBranchOnCondition removed the OldBranch from the function.
   // Delete it, as it is no longer needed.
   delete OldBranch;
 
+  if (MSSA)
+    MSSAUpdater->updatePhisWhenAddingBBPredecessors(NewExit, ExitBlock, DT);
+
   // We need to reprocess this loop, it could be unswitched again.
   redoLoop = true;
 
   // Now that we know that the loop is never entered when this condition is a
   // particular value, rewrite the loop with this info.  We know that this will
   // at least eliminate the old branch.
   RewriteLoopBodyWithConditionConstant(L, Cond, Val, false);
   ++NumTrivial;
void LoopUnswitch::SplitExitEdges(Loop *L,
 
   for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i) {
     BasicBlock *ExitBlock = ExitBlocks[i];
     SmallVector<BasicBlock *, 4> Preds(pred_begin(ExitBlock),
                                        pred_end(ExitBlock));
 
     // Although SplitBlockPredecessors doesn't preserve loop-simplify in
     // general, if we call it on all predecessors of all exits then it does.
-    SplitBlockPredecessors(ExitBlock, Preds, ".us-lcssa", DT, LI, nullptr,
-                           nullptr, /*PreserveLCSSA*/ true);
+    SplitBlockPredecessors(ExitBlock, Preds, ".us-lcssa", DT, LI, MSSA,
+                           MSSAUpdater.get(), /*PreserveLCSSA*/ true);
   }
 }
 
 /// We determined that the loop is profitable to unswitch when LIC equal Val.
 /// Split it into loop versions and test the condition outside of either loop.
 /// Return the loops created as Out1/Out2.
 void LoopUnswitch::UnswitchNontrivialCondition(Value *LIC, Constant *Val,
                                                Loop *L, TerminatorInst *TI) {
   Function *F = loopHeader->getParent();
   LLVM_DEBUG(dbgs() << "loop-unswitch: Unswitching loop %"
                     << loopHeader->getName() << " [" << L->getBlocks().size()
                     << " blocks] in Function " << F->getName() << " when '"
                     << *Val << "' == " << *LIC << "\n");
 
   if (auto *SEWP = getAnalysisIfAvailable<ScalarEvolutionWrapperPass>())
     SEWP->getSE().forgetLoop(L);
 
   LoopBlocks.clear();
   NewBlocks.clear();
 
   // First step, split the preheader and exit blocks, and add these blocks to
   // the LoopBlocks list.
-  BasicBlock *NewPreheader = SplitEdge(loopPreheader, loopHeader, DT, LI);
+  BasicBlock *NewPreheader =
+      SplitEdge(loopPreheader, loopHeader, DT, LI, MSSA, MSSAUpdater.get());
   LoopBlocks.push_back(NewPreheader);
 
   // We want the loop to come after the preheader, but before the exit blocks.
   LoopBlocks.insert(LoopBlocks.end(), L->block_begin(), L->block_end());
 
   SmallVector<BasicBlock*, 8> ExitBlocks;
   L->getUniqueExitBlocks(ExitBlocks);
 
void LoopUnswitch::UnswitchNontrivialCondition(Value *LIC, Constant *Val,
 
   // Keep a WeakTrackingVH holding onto LIC.  If the first call to
   // RewriteLoopBody
   // deletes the instruction (for example by simplifying a PHI that feeds into
   // the condition that we're unswitching on), we don't rewrite the second
   // iteration.
   WeakTrackingVH LICHandle(LIC);
 
+  if (MSSA) {
+    // Update MemorySSA after cloning, and before splitting to unreachables,
+    // since that invalidates the 1:1 mapping of clones in VMap.
+    SmallPtrSet<BasicBlock *, 4> subLoopHeaders;
+    subLoopHeaders.insert(loopHeader);
+    for (Loop *InnerL : L->getSubLoops())
+      subLoopHeaders.insert(InnerL->getHeader());
+    MSSAUpdater->updateForClonedLoop(loopHeader, ExitBlocks, VMap,
+                                     subLoopHeaders, DT);
+    MSSA->verifyMemorySSA();
+  }
+
   // Now we rewrite the original code to know that the condition is true and the
   // new code to know that the condition is false.
   RewriteLoopBodyWithConditionConstant(L, LIC, Val, false);
 
   // It's possible that simplifying one loop could cause the other to be
   // changed to another value or a constant.  If its a constant, don't simplify
   // it.
   if (!LoopProcessWorklist.empty() && LoopProcessWorklist.back() == NewLoop &&
       LICHandle && !isa<Constant>(LICHandle))
     RewriteLoopBodyWithConditionConstant(NewLoop, LICHandle, Val, true);
+
+  if (MSSA) {
+    MSSA->verifyMemorySSA();
+  }
 }
 
 /// Remove all instances of I from the worklist vector specified.
 static void RemoveFromWorklist(Instruction *I,
                                std::vector<Instruction*> &Worklist) {
 
   Worklist.erase(std::remove(Worklist.begin(), Worklist.end(), I),
                  Worklist.end());
for (User *U : LIC->users()) {
     // to the latch.
     if (Latch && DT->dominates(SISucc, Latch))
       continue;
 
     // FIXME: This is a hack.  We need to keep the successor around
     // and hooked up so as to preserve the loop structure, because
     // trying to update it is complicated.  So instead we preserve the
     // loop structure and put the block on a dead code path.
-    SplitEdge(Switch, SISucc, DT, LI);
+    SplitEdge(Switch, SISucc, DT, LI, MSSA, MSSAUpdater.get());
     // Compute the successors instead of relying on the return value
     // of SplitEdge, since it may have split the switch successor
     // after PHI nodes.
     BasicBlock *NewSISucc = DeadCase.getCaseSuccessor();
     BasicBlock *OldSISucc = *succ_begin(NewSISucc);
     // Create an "unreachable" destination.
     BasicBlock *Abort = BasicBlock::Create(Context, "us-unreachable",
                                            Switch->getParent(),
if (isInstructionTriviallyDead(I)) {
       LLVM_DEBUG(dbgs() << "Remove dead instruction '" << *I << "\n");
 
       // Add uses to the worklist, which may be dead now.
       for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i)
         if (Instruction *Use = dyn_cast<Instruction>(I->getOperand(i)))
           Worklist.push_back(Use);
       LPM->deleteSimpleAnalysisValue(I, L);
       RemoveFromWorklist(I, Worklist);
+      if (MSSA)
+        if (MemoryAccess *MA = MSSA->getMemoryAccess(I))
+          MSSAUpdater->removeMemoryAccess(MA);
       I->eraseFromParent();
       ++NumSimplify;
       continue;
     }
 
     // See if instruction simplification can hack this up.  This is common for
     // things like "select false, X, Y" after unswitching made the condition be
     // 'false'.  TODO: update the domtree properly so we can pass it here.
if (BranchInst *BI = dyn_cast<BranchInst>(I)) {
 
         // If Succ has any successors with PHI nodes, update them to have
         // entries coming from Pred instead of Succ.
         Succ->replaceAllUsesWith(Pred);
 
         // Move all of the successor contents from Succ to Pred.
         Pred->getInstList().splice(BI->getIterator(), Succ->getInstList(),
                                    Succ->begin(), Succ->end());
+        if (MSSA)
+          MSSAUpdater->moveAllAfterMergeBlocks(Succ, Pred, BI);
         LPM->deleteSimpleAnalysisValue(BI, L);
         RemoveFromWorklist(BI, Worklist);
         BI->eraseFromParent();
 
         // Remove Succ from the loop tree.
         LI->removeBlock(Succ);
         LPM->deleteSimpleAnalysisValue(Succ, L);
         Succ->eraseFromParent();