[LoopDeletion] Teach LoopDeletion to preserve MemorySSA if available.

If MemorySSA analysis is analysis, LoopDeletion now preserves it.

llvm/include/llvm/Transforms/Utils/LoopUtils.h

@@ -41,6 +41,7 @@ class DataLayout;
 class Loop;
 class LoopInfo;
 class MemoryAccess;
+class MemorySSA;
 class MemorySSAUpdater;
 class OptimizationRemarkEmitter;
 class PredicatedScalarEvolution;
@@ -144,12 +145,12 @@ bool hoistRegion(DomTreeNode *, AliasAnalysis *, LoopInfo *, DominatorTree *,
 ///   - The loop needs to have a Preheader
 ///   - A unique dedicated exit block must exist
 ///
-/// This also updates the relevant analysis information in \p DT, \p SE, and \p
-/// LI if pointers to those are provided.
+/// This also updates the relevant analysis information in \p DT, \p SE, \p LI
+/// and \p MSSA if pointers to those are provided.
 /// It also updates the loop PM if an updater struct is provided.
 
 void deleteDeadLoop(Loop *L, DominatorTree *DT, ScalarEvolution *SE,
-                    LoopInfo *LI);
+                    LoopInfo *LI, MemorySSA *MSSA = nullptr);
 
 /// Try to promote memory values to scalars by sinking stores out of
 /// the loop and moving loads to before the loop.  We do this by looping over

llvm/lib/Transforms/Scalar/LoopDeletion.cpp

@@ -18,6 +18,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/MemorySSA.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/InitializePasses.h"
@@ -134,7 +135,8 @@ static bool isLoopNeverExecuted(Loop *L) {
 /// is unable to delete it due to hoisting trivially loop invariant
 /// instructions out of the loop.
 static LoopDeletionResult deleteLoopIfDead(Loop *L, DominatorTree &DT,
-                                           ScalarEvolution &SE, LoopInfo &LI) {
+                                           ScalarEvolution &SE, LoopInfo &LI,
+                                           MemorySSA *MSSA) {
   assert(L->isLCSSAForm(DT) && "Expected LCSSA!");
 
   // We can only remove the loop if there is a preheader that we can branch from
@@ -164,7 +166,7 @@ static LoopDeletionResult deleteLoopIfDead(Loop *L, DominatorTree &DT,
       std::fill(P.incoming_values().begin(), P.incoming_values().end(),
                 UndefValue::get(P.getType()));
     }
-    deleteDeadLoop(L, &DT, &SE, &LI);
+    deleteDeadLoop(L, &DT, &SE, &LI, MSSA);
     ++NumDeleted;
     return LoopDeletionResult::Deleted;
   }
@@ -200,7 +202,7 @@ static LoopDeletionResult deleteLoopIfDead(Loop *L, DominatorTree &DT,
   }
 
   LLVM_DEBUG(dbgs() << "Loop is invariant, delete it!");
-  deleteDeadLoop(L, &DT, &SE, &LI);
+  deleteDeadLoop(L, &DT, &SE, &LI, MSSA);
   ++NumDeleted;
 
   return LoopDeletionResult::Deleted;
@@ -213,14 +215,17 @@ PreservedAnalyses LoopDeletionPass::run(Loop &L, LoopAnalysisManager &AM,
   LLVM_DEBUG(dbgs() << "Analyzing Loop for deletion: ");
   LLVM_DEBUG(L.dump());
   std::string LoopName = L.getName();
-  auto Result = deleteLoopIfDead(&L, AR.DT, AR.SE, AR.LI);
+  auto Result = deleteLoopIfDead(&L, AR.DT, AR.SE, AR.LI, AR.MSSA);
   if (Result == LoopDeletionResult::Unmodified)
     return PreservedAnalyses::all();
 
   if (Result == LoopDeletionResult::Deleted)
     Updater.markLoopAsDeleted(L, LoopName);
 
-  return getLoopPassPreservedAnalyses();
+  auto PA = getLoopPassPreservedAnalyses();
+  if (AR.MSSA)
+    PA.preserve<MemorySSAAnalysis>();
+  return PA;
 }
 
 namespace {
@@ -235,6 +240,7 @@ public:
   bool runOnLoop(Loop *L, LPPassManager &) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addPreserved<MemorySSAWrapperPass>();
     getLoopAnalysisUsage(AU);
   }
 };
@@ -255,11 +261,15 @@ bool LoopDeletionLegacyPass::runOnLoop(Loop *L, LPPassManager &LPM) {
   DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   ScalarEvolution &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
   LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+  auto *MSSAAnalysis = getAnalysisIfAvailable<MemorySSAWrapperPass>();
+  MemorySSA *MSSA = nullptr;
+  if (MSSAAnalysis)
+    MSSA = &MSSAAnalysis->getMSSA();
 
   LLVM_DEBUG(dbgs() << "Analyzing Loop for deletion: ");
   LLVM_DEBUG(L->dump());
 
-  LoopDeletionResult Result = deleteLoopIfDead(L, DT, SE, LI);
+  LoopDeletionResult Result = deleteLoopIfDead(L, DT, SE, LI, MSSA);
 
   if (Result == LoopDeletionResult::Deleted)
     LPM.markLoopAsDeleted(*L);

llvm/lib/Transforms/Utils/LoopUtils.cpp

@@ -506,13 +506,16 @@ llvm::collectChildrenInLoop(DomTreeNode *N, const Loop *CurLoop) {
   return Worklist;
 }
 
-void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT = nullptr,
-                          ScalarEvolution *SE = nullptr,
-                          LoopInfo *LI = nullptr) {
+void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT, ScalarEvolution *SE,
+                          LoopInfo *LI, MemorySSA *MSSA) {
   assert((!DT || L->isLCSSAForm(*DT)) && "Expected LCSSA!");
   auto *Preheader = L->getLoopPreheader();
   assert(Preheader && "Preheader should exist!");
 
+  std::unique_ptr<MemorySSAUpdater> MSSAU;
+  if (MSSA)
+    MSSAU = std::make_unique<MemorySSAUpdater>(MSSA);
+
   // Now that we know the removal is safe, remove the loop by changing the
   // branch from the preheader to go to the single exit block.
   //
@@ -585,18 +588,33 @@ void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT = nullptr,
            "Should have exactly one value and that's from the preheader!");
   }
 
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+  if (DT) {
+    DTU.applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock}});
+    if (MSSA) {
+      MSSAU->applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock}}, *DT);
+      if (VerifyMemorySSA)
+        MSSA->verifyMemorySSA();
+    }
+  }
+
   // Disconnect the loop body by branching directly to its exit.
   Builder.SetInsertPoint(Preheader->getTerminator());
   Builder.CreateBr(ExitBlock);
   // Remove the old branch.
   Preheader->getTerminator()->eraseFromParent();
 
-  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
   if (DT) {
-    // Update the dominator tree by informing it about the new edge from the
-    // preheader to the exit and the removed edge.
-    DTU.applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock},
-                      {DominatorTree::Delete, Preheader, L->getHeader()}});
+    DTU.applyUpdates({{DominatorTree::Delete, Preheader, L->getHeader()}});
+    if (MSSA) {
+      MSSAU->applyUpdates({{DominatorTree::Delete, Preheader, L->getHeader()}},
+                          *DT);
+      if (VerifyMemorySSA)
+        MSSA->verifyMemorySSA();
+      SmallSetVector<BasicBlock *, 8> DeadBlockSet(L->block_begin(),
+                                                   L->block_end());
+      MSSAU->removeBlocks(DeadBlockSet);
+    }
   }
 
   // Use a map to unique and a vector to guarantee deterministic ordering.
@@ -657,6 +675,9 @@ void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT = nullptr,
   for (auto *Block : L->blocks())
     Block->dropAllReferences();
 
+  if (MSSA && VerifyMemorySSA)
+    MSSA->verifyMemorySSA();
+
   if (LI) {
     // Erase the instructions and the blocks without having to worry
     // about ordering because we already dropped the references.