1 files changed, 49 insertions, 12 deletions

diff --git a/llvm/lib/Transforms/Scalar/GVNHoist.cpp b/llvm/lib/Transforms/Scalar/GVNHoist.cpp
index 19b8dd76942..8dd9eea7a9b 100644
--- a/llvm/lib/Transforms/Scalar/GVNHoist.cpp
+++ b/llvm/lib/Transforms/Scalar/GVNHoist.cpp
@@ -193,14 +193,11 @@ public:
     VN.setAliasAnalysis(AA);
     VN.setMemDep(MD);
     bool Res = false;
+    MemorySSA M(F, AA, DT);
+    MSSA = &M;
 
     // FIXME: use lazy evaluation of VN to avoid the fix-point computation.
     while (1) {
-      // FIXME: only compute MemorySSA once. We need to update the analysis in
-      // the same time as transforming the code.
-      MemorySSA M(F, AA, DT);
-      MSSA = &M;
-
       // Perform DFS Numbering of instructions.
       unsigned I = 0;
       for (const BasicBlock *BB : depth_first(&F.getEntryBlock()))
@@ -208,15 +205,15 @@ public:
           DFSNumber.insert({&Inst, ++I});
 
       auto HoistStat = hoistExpressions(F);
-      if (HoistStat.first + HoistStat.second == 0) {
+      if (HoistStat.first + HoistStat.second == 0)
         return Res;
-      }
-      if (HoistStat.second > 0) {
+
+      if (HoistStat.second > 0)
         // To address a limitation of the current GVN, we need to rerun the
-        // hoisting after we hoisted loads in order to be able to hoist all
-        // scalars dependent on the hoisted loads. Same for stores.
+        // hoisting after we hoisted loads or stores in order to be able to
+        // hoist all scalars dependent on the hoisted ld/st.
         VN.clear();
-      }
+
       Res = true;
 
       // DFS numbers change when instructions are hoisted: clear and recompute.
@@ -310,7 +307,8 @@ private:
 
     for (User *U : Def->users())
       if (auto *MU = dyn_cast<MemoryUse>(U)) {
-        BasicBlock *UBB = MU->getBlock();
+        // FIXME: MU->getBlock() does not get updated when we move the instruction.
+        BasicBlock *UBB = MU->getMemoryInst()->getParent();
         // Only analyze uses in BB.
         if (BB != UBB)
           continue;
@@ -742,9 +740,23 @@ private:
             !makeGepOperandsAvailable(Repl, HoistPt, InstructionsToHoist))
           continue;
 
+        // Move the instruction at the end of HoistPt.
         Repl->moveBefore(HoistPt->getTerminator());
       }
 
+      MemoryAccess *NewMemAcc = nullptr;
+      if (MemoryAccess *MA = MSSA->getMemoryAccess(Repl)) {
+        if (MemoryUseOrDef *OldMemAcc = dyn_cast<MemoryUseOrDef>(MA)) {
+          // The definition of this ld/st will not change: ld/st hoisting is
+          // legal when the ld/st is not moved past its current definition.
+          MemoryAccess *Def = OldMemAcc->getDefiningAccess();
+          NewMemAcc = MSSA->createMemoryAccessInBB(Repl, Def, HoistPt,
+                                                   MemorySSA::End);
+          OldMemAcc->replaceAllUsesWith(NewMemAcc);
+          MSSA->removeMemoryAccess(OldMemAcc);
+        }
+      }
+
       if (isa<LoadInst>(Repl))
         ++NL;
       else if (isa<StoreInst>(Repl))
@@ -775,11 +787,36 @@ private:
           } else if (isa<CallInst>(Repl)) {
             ++NumCallsRemoved;
           }
+
+          if (NewMemAcc) {
+            // Update the uses of the old MSSA access with NewMemAcc.
+            MemoryAccess *OldMA = MSSA->getMemoryAccess(I);
+            OldMA->replaceAllUsesWith(NewMemAcc);
+            MSSA->removeMemoryAccess(OldMA);
+          }
+
           Repl->intersectOptionalDataWith(I);
           combineKnownMetadata(Repl, I);
           I->replaceAllUsesWith(Repl);
           I->eraseFromParent();
         }
+
+      // Remove MemorySSA phi nodes with the same arguments.
+      if (NewMemAcc) {
+        SmallPtrSet<MemoryPhi *, 4> UsePhis;
+        for (User *U : NewMemAcc->users())
+          if (MemoryPhi *Phi = dyn_cast<MemoryPhi>(U))
+            UsePhis.insert(Phi);
+
+        for (auto *Phi : UsePhis) {
+          auto In = Phi->incoming_values();
+          if (std::all_of(In.begin(), In.end(),
+                          [&](Use &U){return U == NewMemAcc;})) {
+            Phi->replaceAllUsesWith(NewMemAcc);
+            MSSA->removeMemoryAccess(Phi);
+          }
+        }
+      }
     }
 
     NumHoisted += NL + NS + NC + NI;