Update MemorySSA in LoopInstSimplify.

[asbirlea]

Summary

Add MemorySSA as a depency to LoopInstInstSimplify and preserve it.
Disabled by default until all passes preserve MemorySSA.

[chandlerc]

lib/Transforms/Scalar/LoopInstSimplify.cpp

80–82

Can be written simpler as:

MemorySSA *MSSA = MSSAU ? MSSAU->getMemorySSA() : nullptr;

139

IR seems weird as a variable name here?

Maybe SimplifiedI or SimpleI? Also, does it make sense to test this before getMemoryAccess?

198

We we really want a unique_ptr instead of an Optional?

[asbirlea]

Address comments.

[asbirlea]

My bad, previous update won't work, updater needs to be allocated. Please comment on the use of make_unique (= operator in Optional does an std::move).

[chandlerc]

In D50906#1206772, @asbirlea wrote:

My bad, previous update won't work, updater needs to be allocated. Please comment on the use of make_unique (= operator in Optional does an std::move).

Can we make MemorySSAUpdater movable? This seems likely to be a pretty common pattern...

[asbirlea]

Update to normal constructor.

[chandlerc]

LGTM

[asbirlea]

Thank you for the review!

lib/Transforms/Scalar/LoopInstSimplify.cpp

 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InstructionSimplify.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/TargetLibraryInfo.h"
-#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/User.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include <algorithm>
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-instsimplify"
 
 STATISTIC(NumSimplified, "Number of redundant instructions simplified");
 
 static bool simplifyLoopInst(Loop &L, DominatorTree &DT, LoopInfo &LI,
-                             AssumptionCache &AC,
-                             const TargetLibraryInfo &TLI) {
+                             AssumptionCache &AC, const TargetLibraryInfo &TLI,
+                             MemorySSAUpdater *MSSAU) {
   const DataLayout &DL = L.getHeader()->getModule()->getDataLayout();
   SimplifyQuery SQ(DL, &TLI, &DT, &AC);
 
   // On the first pass over the loop body we try to simplify every instruction.
   // On subsequent passes, we can restrict this to only simplifying instructions
   // where the inputs have been updated. We end up needing two sets: one
   // containing the instructions we are simplifying in *this* pass, and one for
   // the instructions we will want to simplify in the *next* pass. We use
static bool simplifyLoopInst(Loop &L, DominatorTree &DT, LoopInfo &LI,
 
   // First we want to create an RPO traversal of the loop body. By processing in
   // RPO we can ensure that definitions are processed prior to uses (for non PHI
   // uses) in all cases. This ensures we maximize the simplifications in each
   // iteration over the loop and minimizes the possible causes for continuing to
   // iterate.
   LoopBlocksRPO RPOT(&L);
   RPOT.perform(&LI);
+  MemorySSA *MSSA = MSSAU ? MSSAU->getMemorySSA() : nullptr;
 
   bool Changed = false;
   for (;;) {
+    if (MSSAU && VerifyMemorySSA)
+      MSSA->verifyMemorySSA();
     for (BasicBlock *BB : RPOT) {
       for (Instruction &I : *BB) {
         if (auto *PI = dyn_cast<PHINode>(&I))
           VisitedPHIs.insert(PI);
 
         if (I.use_empty()) {
           if (isInstructionTriviallyDead(&I, &TLI))
             DeadInsts.push_back(&I);
for (BasicBlock *BB : RPOT) {
           // should always be PHI nodes due to LCSSA form, and we don't want to
           // try to simplify those away.
           assert((L.contains(UserI) || isa<PHINode>(UserI)) &&
                  "Uses outside the loop should be PHI nodes due to LCSSA!");
           if (!IsFirstIteration && L.contains(UserI))
             ToSimplify->insert(UserI);
         }
 
+        if (MSSAU)
+          if (Instruction *SimpleI = dyn_cast_or_null<Instruction>(V))
+            if (MemoryAccess *MA = MSSA->getMemoryAccess(&I))
+              if (MemoryAccess *ReplacementMA = MSSA->getMemoryAccess(SimpleI))
+                MA->replaceAllUsesWith(ReplacementMA);
+
         assert(I.use_empty() && "Should always have replaced all uses!");
         if (isInstructionTriviallyDead(&I, &TLI))
           DeadInsts.push_back(&I);
         ++NumSimplified;
         Changed = true;
       }
     }
 
     // Delete any dead instructions found thus far now that we've finished an
     // iteration over all instructions in all the loop blocks.
     if (!DeadInsts.empty()) {
       Changed = true;
-      RecursivelyDeleteTriviallyDeadInstructions(DeadInsts, &TLI);
+      RecursivelyDeleteTriviallyDeadInstructions(DeadInsts, &TLI, MSSAU);
     }
 
+    if (MSSAU && VerifyMemorySSA)
+      MSSA->verifyMemorySSA();
+
     // If we never found a PHI that needs to be simplified in the next
     // iteration, we're done.
     if (Next->empty())
       break;
 
     // Otherwise, put the next set in place for the next iteration and reset it
     // and the visited PHIs for that iteration.
     std::swap(Next, ToSimplify);
if (skipLoop(L))
       return false;
     DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     AssumptionCache &AC =
         getAnalysis<AssumptionCacheTracker>().getAssumptionCache(
             *L->getHeader()->getParent());
     const TargetLibraryInfo &TLI =
         getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
+    MemorySSA *MSSA = nullptr;
+    Optional<MemorySSAUpdater> MSSAUpdater;
+    if (EnableMSSALoopDependency) {
+      MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
+      MSSAUpdater = MemorySSAUpdater(MSSA);
+    }
 
-    return simplifyLoopInst(*L, DT, LI, AC, TLI);
+    return simplifyLoopInst(*L, DT, LI, AC, TLI,
+                            MSSAUpdater.hasValue() ? MSSAUpdater.getPointer()
+                                                   : nullptr);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.setPreservesCFG();
+    if (EnableMSSALoopDependency) {
+      AU.addRequired<MemorySSAWrapperPass>();
+      AU.addPreserved<MemorySSAWrapperPass>();
+    }
     getLoopAnalysisUsage(AU);
   }
 };
 
 } // end anonymous namespace
 
 PreservedAnalyses LoopInstSimplifyPass::run(Loop &L, LoopAnalysisManager &AM,
                                             LoopStandardAnalysisResults &AR,
                                             LPMUpdater &) {
-  if (!simplifyLoopInst(L, AR.DT, AR.LI, AR.AC, AR.TLI))
+  Optional<MemorySSAUpdater> MSSAUpdater;
+  if (AR.MSSA) {
+    MSSAUpdater = MemorySSAUpdater(AR.MSSA);
+    AR.MSSA->verifyMemorySSA();
+  }
+  if (!simplifyLoopInst(L, AR.DT, AR.LI, AR.AC, AR.TLI,
+                        MSSAUpdater.hasValue() ? MSSAUpdater.getPointer()
+                                               : nullptr))
     return PreservedAnalyses::all();
 
   auto PA = getLoopPassPreservedAnalyses();
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }
 
 char LoopInstSimplifyLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(LoopInstSimplifyLegacyPass, "loop-instsimplify",
                       "Simplify instructions in loops", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
+INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(LoopInstSimplifyLegacyPass, "loop-instsimplify",
                     "Simplify instructions in loops", false, false)
 
 Pass *llvm::createLoopInstSimplifyPass() {
   return new LoopInstSimplifyLegacyPass();
 }

test/Transforms/LoopInstSimplify/basic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S %s -passes=loop-instsimplify | FileCheck %s
+; RUN: opt -S %s -passes=loop-instsimplify -enable-mssa-loop-dependency=true -verify-memoryssa | FileCheck %s
 
 ; Test very basic folding and propagation occurs within a loop body. This should
 ; collapse to the loop iteration structure and the LCSSA PHI node.
 define i32 @test1(i32 %n, i32 %x) {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop: