[DSE] Support traversing MemoryPhis.

[fhahn]

Summary

For MemoryPhis, we have to avoid that the MemoryPhi may be executed
before before the access we are currently looking at.

To do this we do a post-order numbering of the basic blocks in the
function and bail out once we reach a MemoryPhi with a larger (or equal)
post-order block number than the current MemoryAccess.
This changes the order in which we visit stores for elimination.

This patch also adds support for exploring multiple paths. We keep a worklist (ToCheck) of memory accesses that might be eliminated by our starting MemoryDef or MemoryPhis for further exploration. For MemoryPhis, we add the incoming values to the worklist, for MemoryDefs we add the defining access.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[fhahn]

Rebase

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[fhahn]

Needs updating to work bottom-up.

[fhahn]

Updated to work with bottom up approach, iteratively.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[fhahn]

Rebased after updates to D72700.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[fhahn]

Rebased after latest changes to earlier patch.

[fhahn]

Ping. Rebased after D72700 landed.

[fhahn]

ping

[asbirlea]

One general comment: this seems to do more than just lifting the MemoryPhi restriction. It's adding some cases that were were not added to the initial version, and that's great, but it would be good to add this in the patch description at least.
(e.g. continue upward search in DSE.cpp:1871, adding removal of partially overlapping stores)

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

1620–1621

cast and remove null check below?

1671–1672

Can you clarify this?
KillingDef is a MemoryDef, UseAccess is a MemoryPhi inside this if. How can they be equal?

[fhahn]

In D72148#1885668, @asbirlea wrote:

j One general comment: this seems to do more than just lifting the MemoryPhi restriction. It's adding some cases that were were not added to the initial version, and that's great, but it would be good to add this in the patch description at least.

(e.g. continue upward search in DSE.cpp:1871, adding removal of partially overlapping stores)

It looks like the overlapping store handling slipped in at some point while rebasing the series. I've removed the changes and will update the separate patch. Sorry about that.

I've extended the description to include the change to how candidates for elimination are explored now. It's quite tied to handling MemoryPhis I think, but I can split it off, if you prefer.

[fhahn]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

1671–1672

I think this was a left-over before switching to a SetVector worklist. With the SetVector, we will never visit the node twice. I've added an assertion to be sure to check for cycles.

[asbirlea]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

1671–1672

Sure, but this equality cannot happen, right? They're different types, and one if not a subtype of the other. One is a MemoryPhi and one a MemoryDef (KillingDef, the store that started the search), they cannot be equal even if you didn't check for cycles.

Sounds like you're thinking of not seeing DomAccess again as the worklist is a set, but you may see it AFAICT. If at :1617 DomAccess uses itself, it will be added to the worklist. I'm inclined to believe the change you meant is:

if (DomAccess == UseAccess)
  continue;

But just as well not having anything is fine. The PushMemUses will be a noop, as all uses have already been added.

1690–1694

Should this be Current == UseAccess ?

[fhahn]

Remove unnecessary assertion, add comment to check.

[fhahn]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

1671–1672

Sure, but this equality cannot happen, right?

Ah yes! I thought it was needed for a test case I added, but after the refactoring it is not necessary any longer. I've removed it

1690–1694

I've added a comment why this check is there. We have to make sure the KillingDef does not read the memory location. The easiest way is to just consider it as regular use, but skip the overwrite check. It probably makes more sense to just check for self reads before getDomMemoryDef. I've added a TODO.

On a related note, I've put up D75025, which should help with avoiding a number of unnecessary duplicated checks for the same killing def.

[fhahn]

Rebase after small fix in b8d638d337e76a632d07d61f4cef59e243b961a8

[fhahn]

ping

[fhahn]

ping.

Add a missing isa<MemoryPhi> check after getting walker result + test case.

[fhahn]

Ping.

@asbirlea I think all your comments should be addressed in the latest version.

Also rebased and fixed debug printing to account for MemoryPhis.

[fhahn]

ping

[asbirlea]

Thank you!

[ebrevnov]

@fhahn is this patch supposed to support the following case?

define i8* @foo(i1 %cond, i8* %p1, i8* %p2) {
entry:

%0 = icmp ne i1 %cond, 0
br i1 %0, label %dead_store, label %merge

dead_store:

store i8 1, i8* %p1
br label %merge

merge:

%p3 = phi i8* [ %p1, %dead_store ], [ %p2, %entry ]
store i8 2, i8* %p3
ret i8* %p3

}

If not, is there any plans or prototype which can handle it? Thanks.

[fhahn]

In D72148#2239379, @ebrevnov wrote:

@fhahn is this patch supposed to support the following case?

define i8* @foo(i1 %cond, i8* %p1, i8* %p2) {
entry:

%0 = icmp ne i1 %cond, 0
br i1 %0, label %dead_store, label %merge

dead_store:

store i8 1, i8* %p1
br label %merge

merge:

%p3 = phi i8* [ %p1, %dead_store ], [ %p2, %entry ]
store i8 2, i8* %p3
ret i8* %p3

}

If not, is there any plans or prototype which can handle it? Thanks.

Currently we do not phi-translate the location of the 'killing def' and hence do not eliminate the store in this case. I haven't thought too much about it so far, but I think doing the phi translation of the 'killing loc' should not be too hard. It might be as simple as adjusting the ToCheck set to contain a pair (Access to check, Phi-translated DefLoc) and do phi translation when we queue the incoming defs for MemoryPhis (https://github.com/llvm/llvm-project/blob/master/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp#L2284)

I won't have time to look into this in the near future unfortunately, as I am mainly focused on getting things in good shape to enable DSE & MemorySSA by default. It might be good to file a bug report so this opportunity does not drop off our radar. If anyone is interested in giving implementing it a try, even better :)

[ebrevnov]

I won't have time to look into this in the near future unfortunately, as I am mainly focused on getting things in good shape to enable DSE & MemorySSA by default. It might be good to file a bug report so this opportunity does not drop off our radar. If anyone is interested in giving implementing it a try, even better :)

Florian, thank you for quick feedback. I prototyped phi translation approach and it does work on the posted test case. Unfortunately, it fails on a bit more complicated case which I will take a look tomorrow. If I have something working reasonably well in the end I will share it.

[fhahn]

In D72148#2242301, @ebrevnov wrote:

I won't have time to look into this in the near future unfortunately, as I am mainly focused on getting things in good shape to enable DSE & MemorySSA by default. It might be good to file a bug report so this opportunity does not drop off our radar. If anyone is interested in giving implementing it a try, even better :)

Florian, thank you for quick feedback. I prototyped phi translation approach and it does work on the posted test case. Unfortunately, it fails on a bit more complicated case which I will take a look tomorrow. If I have something working reasonably well in the end I will share it.

That's great! There are some changes to how we find MemoryDefs that may be killed by a killing def, rGe717fdb0f155 and D86487. This might actually make things a bit easier, as we now fully control the traversal of the defining accesses.

[fhahn]

In D72148#2242386, @fhahn wrote:

In D72148#2242301, @ebrevnov wrote:

I won't have time to look into this in the near future unfortunately, as I am mainly focused on getting things in good shape to enable DSE & MemorySSA by default. It might be good to file a bug report so this opportunity does not drop off our radar. If anyone is interested in giving implementing it a try, even better :)

Florian, thank you for quick feedback. I prototyped phi translation approach and it does work on the posted test case. Unfortunately, it fails on a bit more complicated case which I will take a look tomorrow. If I have something working reasonably well in the end I will share it.

That's great! There are some changes to how we find MemoryDefs that may be killed by a killing def, rGe717fdb0f155 and D86487. This might actually make things a bit easier, as we now fully control the traversal of the defining accesses.

Hi @ebrevnov,

I just wanted to briefly check in to see if you made any progress on this or if you think it is worth adding a bug report to keep track of the improvement?

Cheers,
Florian

[ebrevnov]

Hi @ebrevnov,

I just wanted to briefly check in to see if you made any progress on this or if you think it is worth adding a bug report to keep track of the improvement?

Hi Florian,

I have a prototype which is able to handle my case. I can upload it to the phabricator for the reference. It's based on about month old sources because if I rebase to newer things stop working. Another missing piece is caching results of phi translation. In the current version phi translation is done from scratch for each new memory access. It shouldn't be hard to restructure the code to reuse phi translated address from "previous" memory access.

Thanks
Evgeniy

[fhahn]

In D72148#2286883, @ebrevnov wrote:

Hi @ebrevnov,

I just wanted to briefly check in to see if you made any progress on this or if you think it is worth adding a bug report to keep track of the improvement?

Hi Florian,

I have a prototype which is able to handle my case. I can upload it to the phabricator for the reference. It's based on about month old sources because if I rebase to newer things stop working. Another missing piece is caching results of phi translation. In the current version phi translation is done from scratch for each new memory access. It shouldn't be hard to restructure the code to reuse phi translated address from "previous" memory access.

Yeah there have been some substantial changes to the way we find the candidates to remove, but I think the same idea should still apply. It would be ideal to update & rebase the patch, but if you won't have time to do so it would be great if you could share the older version regardless.

[ebrevnov]

Yeah there have been some substantial changes to the way we find the candidates to remove, but I think the same idea should still apply. It would be ideal to update & rebase the patch, but if you won't have time to do so it would be great if you could share the older version regardless.

Rebased and uploaded here https://reviews.llvm.org/D90095

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

 // traversal.  Doing so would be pretty trivial.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/DeadStoreElimination.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MemoryDependenceAnalysis.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
-#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
struct DSEState {
   SmallVector<MemoryDef *, 64> MemDefs;
   // Any that should be skipped as they are already deleted
   SmallPtrSet<MemoryAccess *, 4> SkipStores;
   // Keep track of all of the objects that are invisible to the caller until the
   // function returns.
   SmallPtrSet<const Value *, 16> InvisibleToCaller;
   // Keep track of blocks with throwing instructions not modeled in MemorySSA.
   SmallPtrSet<BasicBlock *, 16> ThrowingBlocks;
+  // Post-order numbers for each basic block. Used to figure out if memory
+  // accesses are executed before another access.
+  DenseMap<BasicBlock *, unsigned> PostOrderNumbers;
 
   /// Keep track of instructions (partly) overlapping with killing MemoryDefs per
   /// basic block.
   DenseMap<BasicBlock *, InstOverlapIntervalsTy> IOLs;
 
   DSEState(Function &F, AliasAnalysis &AA, MemorySSA &MSSA, DominatorTree &DT,
            PostDominatorTree &PDT, const TargetLibraryInfo &TLI)
       : F(F), AA(AA), MSSA(MSSA), DT(DT), PDT(PDT), TLI(TLI) {}
 
   static DSEState get(Function &F, AliasAnalysis &AA, MemorySSA &MSSA,
                       DominatorTree &DT, PostDominatorTree &PDT,
                       const TargetLibraryInfo &TLI) {
     DSEState State(F, AA, MSSA, DT, PDT, TLI);
     // Collect blocks with throwing instructions not modeled in MemorySSA and
     // alloc-like objects.
-    for (Instruction &I : instructions(F)) {
+    unsigned PO = 0;
+    for (BasicBlock *BB : post_order(&F)) {
+      State.PostOrderNumbers[BB] = PO++;
+      for (Instruction &I : *BB) {
         if (I.mayThrow() && !MSSA.getMemoryAccess(&I))
           State.ThrowingBlocks.insert(I.getParent());
 
         auto *MD = dyn_cast_or_null<MemoryDef>(MSSA.getMemoryAccess(&I));
         if (MD && State.MemDefs.size() < MemorySSADefsPerBlockLimit &&
             hasAnalyzableMemoryWrite(&I, TLI) && isRemovable(&I))
           State.MemDefs.push_back(MD);
 
         // Track alloca and alloca-like objects. Here we care about objects not
         // visible to the caller during function execution. Alloca objects are
-      // invalid in the caller, for alloca-like objects we ensure that they are
-      // not captured throughout the function.
+        // invalid in the caller, for alloca-like objects we ensure that they
+        // are not captured throughout the function.
         if (isa<AllocaInst>(&I) ||
             (isAllocLikeFn(&I, &TLI) && !PointerMayBeCaptured(&I, false, true)))
           State.InvisibleToCaller.insert(&I);
       }
+    }
+
     // Treat byval or inalloca arguments the same as Allocas, stores to them are
     // dead at the end of the function.
     for (Argument &AI : F.args())
       if (AI.hasByValOrInAllocaAttr())
         State.InvisibleToCaller.insert(&AI);
     return State;
   }
 
bool isReadClobber(MemoryLocation DefLoc, Instruction *UseInst) const {
     if (isRefSet(MR))
       MR = AA.callCapturesBefore(UseInst, DefLoc, &DT);
     return isRefSet(MR);
   }
 
   // Find a MemoryDef writing to \p DefLoc and dominating \p Current, with no
   // read access in between or return None otherwise. The returned value may not
   // (completely) overwrite \p DefLoc. Currently we bail out when we encounter
-  // any of the following
-  //  * An aliasing MemoryUse (read).
-  //  * A MemoryPHI.
+  // an aliasing MemoryUse (read).
   Optional<MemoryAccess *> getDomMemoryDef(MemoryDef *KillingDef,
                                            MemoryAccess *Current,
                                            MemoryLocation DefLoc,
                                            bool DefVisibleToCaller,
                                            int &ScanLimit) const {
-    MemoryDef *DomDef;
-    MemoryAccess *StartDef = Current;
+    MemoryAccess *DomAccess;
     bool StepAgain;
     LLVM_DEBUG(dbgs() << "  trying to get dominating access for " << *Current
                       << "\n");
     // Find the next clobbering Mod access for DefLoc, starting at Current.
     do {
       StepAgain = false;
       // Reached TOP.
       if (MSSA.isLiveOnEntryDef(Current))
         return None;
 
-      MemoryUseOrDef *CurrentUD = dyn_cast<MemoryUseOrDef>(Current);
-      if (!CurrentUD)
-        return None;
-
+      if (isa<MemoryPhi>(Current)) {
+        DomAccess = Current;
+        break;
+      }
+      MemoryUseOrDef *CurrentUD = cast<MemoryUseOrDef>(Current);
       // Look for access that clobber DefLoc.
-      MemoryAccess *DomAccess =
-          MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(
-              CurrentUD->getDefiningAccess(), DefLoc);
-      DomDef = dyn_cast<MemoryDef>(DomAccess);
-      if (!DomDef || MSSA.isLiveOnEntryDef(DomDef))
+      DomAccess = MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(CurrentUD,
+                                                                      DefLoc);
+      if (MSSA.isLiveOnEntryDef(DomAccess))
         return None;
 
+      if (isa<MemoryPhi>(DomAccess))
+        break;
+
       // Check if we can skip DomDef for DSE. We also require the KillingDef
       // execute whenever DomDef executes and use post-dominance to ensure that.
-      if (canSkipDef(DomDef, DefVisibleToCaller) ||
+
+      MemoryDef *DomDef = dyn_cast<MemoryDef>(DomAccess);
+      if ((DomDef && canSkipDef(DomDef, DefVisibleToCaller)) ||
           !PDT.dominates(KillingDef->getBlock(), DomDef->getBlock())) {
         StepAgain = true;
-        Current = DomDef;
+        Current = DomDef->getDefiningAccess();
       }
 
     } while (StepAgain);
 
-    LLVM_DEBUG(dbgs() << "  Checking for reads of " << *DomDef << " ("
-                      << *DomDef->getMemoryInst() << ")\n");
+    LLVM_DEBUG({
+      dbgs() << "  Checking for reads of " << *DomAccess;
+      if (isa<MemoryDef>(DomAccess))
+        dbgs() << " (" << *cast<MemoryDef>(DomAccess)->getMemoryInst() << ")\n";
+    });
 
     SmallSetVector<MemoryAccess *, 32> WorkList;
     auto PushMemUses = [&WorkList](MemoryAccess *Acc) {
       for (Use &U : Acc->uses())
         WorkList.insert(cast<MemoryAccess>(U.getUser()));
     };
-    PushMemUses(DomDef);
+    PushMemUses(DomAccess);
 
     // Check if DomDef may be read.
     for (unsigned I = 0; I < WorkList.size(); I++) {
       MemoryAccess *UseAccess = WorkList[I];
 
       LLVM_DEBUG(dbgs() << "   Checking use " << *UseAccess);
       if (--ScanLimit == 0) {
         LLVM_DEBUG(dbgs() << "  ...  hit scan limit\n");
         return None;
       }
 
-      // Bail out on MemoryPhis for now.
       if (isa<MemoryPhi>(UseAccess)) {
-        LLVM_DEBUG(dbgs() << "  ...  hit MemoryPhi\n");
-        return None;
+        PushMemUses(UseAccess);
+        continue;
       }
 
       Instruction *UseInst = cast<MemoryUseOrDef>(UseAccess)->getMemoryInst();
       LLVM_DEBUG(dbgs() << " (" << *UseInst << ")\n");
 
       if (isNoopIntrinsic(cast<MemoryUseOrDef>(UseAccess))) {
         PushMemUses(UseAccess);
         continue;
       }
 
       // Uses which may read the original MemoryDef mean we cannot eliminate the
       // original MD. Stop walk.
       if (isReadClobber(DefLoc, UseInst)) {
         LLVM_DEBUG(dbgs() << "  ... found read clobber\n");
         return None;
       }
 
-      if (StartDef == UseAccess)
+      // For the KillingDef we only have to check if it reads the memory
+      // location.
+      // TODO: It would probably be better to check for self-reads before
+      // calling the function.
+      if (KillingDef == UseAccess)
         continue;
 
       // Check all uses for MemoryDefs, except for defs completely overwriting
       // the original location. Otherwise we have to check uses of *all*
       // MemoryDefs we discover, including non-aliasing ones. Otherwise we might
       // miss cases like the following
       //   1 = Def(LoE) ; <----- DomDef stores [0,1]
       //   2 = Def(1)   ; (2, 1) = NoAlias,   stores [2,3]
       //   Use(2)       ; MayAlias 2 *and* 1, loads [0, 3].
       //                  (The Use points to the *first* Def it may alias)
       //   3 = Def(1)   ; <---- Current  (3, 2) = NoAlias, (3,1) = MayAlias,
       //                  stores [0,1]
       if (MemoryDef *UseDef = dyn_cast<MemoryDef>(UseAccess)) {
         if (!isCompleteOverwrite(DefLoc, UseInst))
           PushMemUses(UseDef);
       }
     }
 
-    // No aliasing MemoryUses of DomDef found, DomDef is potentially dead.
-    return {DomDef};
+    // No aliasing MemoryUses of DomAccess found, DomAccess is potentially dead.
+    return {DomAccess};
   }
 
   // Delete dead memory defs
   void deleteDeadInstruction(Instruction *SI) {
     MemorySSAUpdater Updater(&MSSA);
     SmallVector<Instruction *, 32> NowDeadInsts;
     NowDeadInsts.push_back(SI);
     --NumFastOther;
bool eliminateDeadStoresMemorySSA(Function &F, AliasAnalysis &AA,
                                   PostDominatorTree &PDT,
                                   const TargetLibraryInfo &TLI) {
   const DataLayout &DL = F.getParent()->getDataLayout();
   bool MadeChange = false;
 
   DSEState State = DSEState::get(F, AA, MSSA, DT, PDT, TLI);
   // For each store:
   for (unsigned I = 0; I < State.MemDefs.size(); I++) {
-    MemoryDef *Current = State.MemDefs[I];
-    if (State.SkipStores.count(Current))
+    MemoryDef *KillingDef = State.MemDefs[I];
+    if (State.SkipStores.count(KillingDef))
       continue;
-    Instruction *SI = cast<MemoryDef>(Current)->getMemoryInst();
+    Instruction *SI = KillingDef->getMemoryInst();
     auto MaybeSILoc = State.getLocForWriteEx(SI);
     if (!MaybeSILoc) {
       LLVM_DEBUG(dbgs() << "Failed to find analyzable write location for "
                         << *SI << "\n");
       continue;
     }
     MemoryLocation SILoc = *MaybeSILoc;
     assert(SILoc.Ptr && "SILoc should not be null");
     const Value *SILocUnd = GetUnderlyingObject(SILoc.Ptr, DL);
     Instruction *DefObj =
         const_cast<Instruction *>(dyn_cast<Instruction>(SILocUnd));
     bool DefVisibleToCaller = !State.InvisibleToCaller.count(SILocUnd);
     if (DefObj && ((isAllocLikeFn(DefObj, &TLI) &&
                     !PointerMayBeCapturedBefore(DefObj, false, true, SI, &DT))))
       DefVisibleToCaller = false;
 
-    LLVM_DEBUG(dbgs() << "Trying to eliminate MemoryDefs killed by " << *SI
-                      << "\n");
+    MemoryAccess *Current = KillingDef;
+    LLVM_DEBUG(dbgs() << "Trying to eliminate MemoryDefs killed by "
+                      << *KillingDef << " (" << *SI << ")\n");
 
     int ScanLimit = MemorySSAScanLimit;
-    MemoryDef *StartDef = Current;
-    // Walk MemorySSA upward to find MemoryDefs that might be killed by SI.
-    while (Optional<MemoryAccess *> Next = State.getDomMemoryDef(
-               StartDef, Current, SILoc, DefVisibleToCaller, ScanLimit)) {
+    // Worklist of MemoryAccesses that may be killed by KillingDef.
+    SetVector<MemoryAccess *> ToCheck;
+    ToCheck.insert(KillingDef->getDefiningAccess());
+
+    // Check if MemoryAccesses in the worklist are killed by KillingDef.
+    for (unsigned I = 0; I < ToCheck.size(); I++) {
+      Current = ToCheck[I];
+      if (State.SkipStores.count(Current))
+        continue;
+
+      Optional<MemoryAccess *> Next = State.getDomMemoryDef(
+          KillingDef, Current, SILoc, DefVisibleToCaller, ScanLimit);
+
+      if (!Next) {
+        LLVM_DEBUG(dbgs() << "  finished walk\n");
+        continue;
+      }
+
       MemoryAccess *DomAccess = *Next;
       LLVM_DEBUG(dbgs() << " Checking if we can kill " << *DomAccess << "\n");
+      if (isa<MemoryPhi>(DomAccess)) {
+        for (Value *V : cast<MemoryPhi>(DomAccess)->incoming_values()) {
+          MemoryAccess *IncomingAccess = cast<MemoryAccess>(V);
+          BasicBlock *IncomingBlock = IncomingAccess->getBlock();
+          BasicBlock *PhiBlock = DomAccess->getBlock();
+
+          // We only consider incoming MemoryAccesses that come before the
+          // MemoryPhi. Otherwise we could discover candidates that do not
+          // strictly dominate our starting def.
+          if (State.PostOrderNumbers[IncomingBlock] >
+              State.PostOrderNumbers[PhiBlock])
+            ToCheck.insert(IncomingAccess);
+        }
+        continue;
+      }
       MemoryDef *NextDef = dyn_cast<MemoryDef>(DomAccess);
       Instruction *NI = NextDef->getMemoryInst();
       LLVM_DEBUG(dbgs() << "  def " << *NI << "\n");
 
-      if (!hasAnalyzableMemoryWrite(NI, TLI))
-        break;
+      if (!hasAnalyzableMemoryWrite(NI, TLI)) {
+        LLVM_DEBUG(dbgs() << " skip, cannot analyze def\n");
+        continue;
+      }
 
       if (!isRemovable(NI)) {
         LLVM_DEBUG(dbgs() << " skip, cannot remove def\n");
         continue;
       }
 
       MemoryLocation NILoc = *State.getLocForWriteEx(NI);
       // Check for anything that looks like it will be a barrier to further
       // removal
       if (State.isDSEBarrier(SI, SILoc, SILocUnd, NI, NILoc)) {
-        LLVM_DEBUG(dbgs() << "  stop, barrier\n");
-        break;
+        LLVM_DEBUG(dbgs() << "  skip, barrier\n");
+        continue;
       }
 
       // Before we try to remove anything, check for any extra throwing
       // instructions that block us from DSEing
       if (State.mayThrowBetween(SI, NI, SILocUnd)) {
-        LLVM_DEBUG(dbgs() << " stop, may throw!\n");
+        LLVM_DEBUG(dbgs() << " skip, may throw!\n");
         break;
       }
 
       if (!DebugCounter::shouldExecute(MemorySSACounter))
         break;
 
       // Check if NI overwrites SI.
       int64_t InstWriteOffset, DepWriteOffset;
       auto Iter = State.IOLs.insert(
           std::make_pair<BasicBlock *, InstOverlapIntervalsTy>(
               NI->getParent(), InstOverlapIntervalsTy()));
       auto &IOL = Iter.first->second;
       OverwriteResult OR = isOverwrite(SILoc, NILoc, DL, TLI, DepWriteOffset,
                                        InstWriteOffset, NI, IOL, AA, &F);
 
+      ToCheck.insert(NextDef->getDefiningAccess());
       if (OR == OW_Complete) {
         LLVM_DEBUG(dbgs() << "DSE: Remove Dead Store:\n  DEAD: " << *NI
                           << "\n  KILLER: " << *SI << '\n');
         State.deleteDeadInstruction(NI);
         ++NumFastStores;
         MadeChange = true;
-      } else
-        Current = NextDef;
+      }
     }
   }
 
   if (EnablePartialOverwriteTracking)
     for (auto &KV : State.IOLs)
       MadeChange |= removePartiallyOverlappedStores(&AA, DL, KV.second);
 
   return MadeChange;

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-exceptions.ll

 ; CHECK-NEXT:  block1:
 ; CHECK-NEXT:    [[SV:%.*]] = alloca i32
 ; CHECK-NEXT:    br label [[BLOCK2:%.*]]
 ; CHECK:       block2:
 ; CHECK-NEXT:    store i32 20, i32* [[SV]]
 ; CHECK-NEXT:    invoke void @f()
 ; CHECK-NEXT:    to label [[BLOCK3:%.*]] unwind label [[CATCH_DISPATCH:%.*]]
 ; CHECK:       block3:
-; CHECK-NEXT:    store i32 30, i32* [[SV]]
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       catch.dispatch:
 ; CHECK-NEXT:    [[CS1:%.*]] = catchswitch within none [label %catch] unwind label [[CLEANUP:%.*]]
 ; CHECK:       catch:
 ; CHECK-NEXT:    [[C:%.*]] = catchpad within [[CS1]] []
 ; CHECK-NEXT:    [[LV:%.*]] = load i32, i32* [[SV]]
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       cleanup:

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-loops.ll

 end:
   ret void
 }
 
 
 define void @test14(i32* noalias %P) {
 ; CHECK-LABEL: @test14(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    store i32 1, i32* [[P:%.*]]
 ; CHECK-NEXT:    br label [[FOR:%.*]]
 ; CHECK:       for:
-; CHECK-NEXT:    store i32 0, i32* [[P]]
+; CHECK-NEXT:    store i32 0, i32* [[P:%.*]]
 ; CHECK-NEXT:    br i1 false, label [[FOR]], label [[END:%.*]]
 ; CHECK:       end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   store i32 1, i32* %P
   br label %for
 for:
end:
   ret void
 }
 
 define void @test21(i32* noalias %P) {
 ; CHECK-LABEL: @test21(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ARRAYIDX0:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 1
 ; CHECK-NEXT:    [[P3:%.*]] = bitcast i32* [[ARRAYIDX0]] to i8*
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 4 [[P3]], i8 0, i64 28, i1 false)
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i8, i8* [[P3]], i64 4
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 4 [[TMP0]], i8 0, i64 24, i1 false)
 ; CHECK-NEXT:    br label [[FOR:%.*]]
 ; CHECK:       for:
 ; CHECK-NEXT:    [[ARRAYIDX1:%.*]] = getelementptr inbounds i32, i32* [[P]], i64 1
 ; CHECK-NEXT:    store i32 1, i32* [[ARRAYIDX1]], align 4
 ; CHECK-NEXT:    br i1 false, label [[FOR]], label [[END:%.*]]
 ; CHECK:       end:
 ; CHECK-NEXT:    ret void
 ;
for.body:
   %lv = load i32, i32* %p
   br label %for.header
 
 end:
   store i32 3, i32* %p, align 4
   ret void
 }
 
+%struct.hoge = type { i32, i32 }
+
+@global = external local_unnamed_addr global %struct.hoge*, align 8
+
+define void @widget(i8* %tmp) {
+; CHECK-LABEL: @widget(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 1 [[TMP:%.*]], i8* nonnull align 16 undef, i64 64, i1 false)
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[TMP2:%.*]] = load %struct.hoge*, %struct.hoge** @global, align 8
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds [[STRUCT_HOGE:%.*]], %struct.hoge* [[TMP2]], i64 undef, i32 1
+; CHECK-NEXT:    store i32 0, i32* [[TMP3]], align 4
+; CHECK-NEXT:    [[TMP4:%.*]] = load %struct.hoge*, %struct.hoge** @global, align 8
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds [[STRUCT_HOGE]], %struct.hoge* [[TMP4]], i64 undef, i32 1
+; CHECK-NEXT:    store i32 10, i32* [[TMP5]], align 4
+; CHECK-NEXT:    br label [[BB1]]
+;
+bb:
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 1 %tmp, i8* nonnull align 16 undef, i64 64, i1 false)
+  br label %bb1
+
+bb1:                                              ; preds = %bb1, %bb
+  %tmp2 = load %struct.hoge*, %struct.hoge** @global, align 8
+  %tmp3 = getelementptr inbounds %struct.hoge, %struct.hoge* %tmp2, i64 undef, i32 1
+  store i32 0, i32* %tmp3, align 4
+  %tmp4 = load %struct.hoge*, %struct.hoge** @global, align 8
+  %tmp5 = getelementptr inbounds %struct.hoge, %struct.hoge* %tmp4, i64 undef, i32 1
+  store i32 10, i32* %tmp5, align 4
+  br label %bb1
+}
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg)

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-memoryphis.ll

bb3:
   store i32 0, i32* %P
   ret void
 }
 
 define void @test5(i32* noalias %P) {
 ; CHECK-LABEL: @test5(
 ; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb1:
-; CHECK-NEXT:    store i32 1, i32* [[P:%.*]]
 ; CHECK-NEXT:    br label [[BB3:%.*]]
 ; CHECK:       bb2:
-; CHECK-NEXT:    store i32 1, i32* [[P]]
 ; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
-; CHECK-NEXT:    store i32 0, i32* [[P]]
+; CHECK-NEXT:    store i32 0, i32* [[P:%.*]]
 ; CHECK-NEXT:    ret void
 ;
   br i1 true, label %bb1, label %bb2
 bb1:
   store i32 1, i32* %P
   br label %bb3
 bb2:
   store i32 1, i32* %P
   br label %bb3
 bb3:
   store i32 0, i32* %P
   ret void
 }
 
 define void @test8(i32* %P, i32* %Q) {
 ; CHECK-LABEL: @test8(
 ; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb1:
-; CHECK-NEXT:    store i32 1, i32* [[P:%.*]]
 ; CHECK-NEXT:    br label [[BB3:%.*]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    store i32 1, i32* [[Q:%.*]]
 ; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
-; CHECK-NEXT:    store i32 0, i32* [[P]]
+; CHECK-NEXT:    store i32 0, i32* [[P:%.*]]
 ; CHECK-NEXT:    ret void
 ;
   br i1 true, label %bb1, label %bb2
 bb1:
   store i32 1, i32* %P
   br label %bb3
 bb2:
   store i32 1, i32* %Q

 define void @widget(i32* %Ptr, i1 %c1, i1 %c2, i32 %v1, i32 %v2, i32 %v3) {
 ; CHECK-LABEL: @widget(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    tail call void @hoge()
 ; CHECK-NEXT:    br i1 [[C1:%.*]], label [[BB3:%.*]], label [[BB1:%.*]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br i1 [[C2:%.*]], label [[BB2:%.*]], label [[BB3]]
 ; CHECK:       bb2:
-; CHECK-NEXT:    store i32 -1, i32* [[PTR:%.*]], align 4
 ; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB4:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    switch i32 [[V1:%.*]], label [[BB8:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[BB5:%.*]]
 ; CHECK-NEXT:    i32 1, label [[BB6:%.*]]
 ; CHECK-NEXT:    i32 2, label [[BB7:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       bb5:
-; CHECK-NEXT:    store i32 0, i32* [[PTR]], align 4
+; CHECK-NEXT:    store i32 0, i32* [[PTR:%.*]], align 4
 ; CHECK-NEXT:    br label [[BB8]]
 ; CHECK:       bb6:
 ; CHECK-NEXT:    store i32 1, i32* [[PTR]], align 4
 ; CHECK-NEXT:    br label [[BB8]]
 ; CHECK:       bb7:
 ; CHECK-NEXT:    store i32 2, i32* [[PTR]], align 4
 ; CHECK-NEXT:    br label [[BB8]]
 ; CHECK:       bb8:

 
 bb7:                                              ; preds = %bb4
   store i32 2, i32* %Ptr, align 4
   br label %bb8
 
 bb8:                                              ; preds = %bb7, %bb6, %bb5, %bb4
   br label %bb4
 }
+
+
+declare void @fn1_test11()
+declare void @fn2_test11()
+
+define void @test11(i1 %c, i8** %ptr.1) {
+; CHECK-LABEL: @test11(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF_THEN:%.*]], label [[EXIT:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    tail call void @fn2_test11() #0
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    store i8* null, i8** [[PTR_1:%.*]], align 8
+; CHECK-NEXT:    tail call void @fn2_test11() #0
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %c, label %if.then, label %exit
+
+if.then:                                      ; preds = %entry
+  tail call void @fn2_test11() #1
+  br label %exit
+
+exit:
+  store i8* null, i8** %ptr.1, align 8
+  tail call void @fn2_test11() #1
+  ret void
+}
+
+attributes #1 = { nounwind }

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-simple.ll

 ; CHECK-NEXT:    [[P_1:%.*]] = getelementptr i32, i32* [[P]], i32 1
 ; CHECK-NEXT:    store i32 1, i32* [[P_1]]
 ; CHECK-NEXT:    [[P_64:%.*]] = bitcast i32* [[P]] to i64*
 ; CHECK-NEXT:    [[LV:%.*]] = load i64, i64* [[P_64]]
 ; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB3:%.*]]
 ; CHECK:       bb2:
-; CHECK-NEXT:    ret void
+; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
 ; CHECK-NEXT:    store i32 2, i32* [[P]]
 ; CHECK-NEXT:    ret void
 ;
   store i32 0, i32* %P
   %P.1 = getelementptr i32, i32* %P, i32 1
   store i32 1, i32* %P.1
 
   %P.64 = bitcast i32* %P to i64*
   %lv = load i64, i64* %P.64
   br i1 true, label %bb1, label %bb2
 bb1:
   br label %bb3
 bb2:
-  ret void
+  br label %bb3
 bb3:
   store i32 2, i32* %P
   ret void
 }
+
+define void @test10(i32* %P) {
+; CHECK-LABEL: @test10(
+; CHECK-NEXT:    store i32 0, i32* [[P:%.*]]
+; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    store i32 0, i32* [[P]]
+; CHECK-NEXT:    br label [[BB3:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    ret void
+; CHECK:       bb3:
+; CHECK-NEXT:    ret void
+;
+  store i32 0, i32* %P
+  br i1 true, label %bb1, label %bb2
+bb1:
+  store i32 0, i32* %P
+  br label %bb3
+bb2:
+  ret void
+bb3:
+  ret void
+}
+
+
+define void @test11() {
+; CHECK-LABEL: @test11(
+; CHECK-NEXT:    [[P:%.*]] = alloca i32
+; CHECK-NEXT:    store i32 0, i32* [[P]]
+; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    store i32 0, i32* [[P]]
+; CHECK-NEXT:    br label [[BB3:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    ret void
+; CHECK:       bb3:
+; CHECK-NEXT:    ret void
+;
+  %P = alloca i32
+  store i32 0, i32* %P
+  br i1 true, label %bb1, label %bb2
+bb1:
+  store i32 0, i32* %P
+  br label %bb3
+bb2:
+  ret void
+bb3:
+  ret void
+}
+
+
+define void @test12(i32* %P) {
+; CHECK-LABEL: @test12(
+; CHECK-NEXT:    store i32 0, i32* [[P:%.*]]
+; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    store i32 1, i32* [[P]]
+; CHECK-NEXT:    br label [[BB3:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    store i32 1, i32* [[P]]
+; CHECK-NEXT:    ret void
+; CHECK:       bb3:
+; CHECK-NEXT:    ret void
+;
+  store i32 0, i32* %P
+  br i1 true, label %bb1, label %bb2
+bb1:
+  store i32 1, i32* %P
+  br label %bb3
+bb2:
+  store i32 1, i32* %P
+  ret void
+bb3:
+  ret void
+}
+
+
+define void @test13(i32* %P) {
+; CHECK-LABEL: @test13(
+; CHECK-NEXT:    store i32 0, i32* [[P:%.*]]
+; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    store i32 1, i32* [[P]]
+; CHECK-NEXT:    br label [[BB3:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    store i32 1, i32* [[P]]
+; CHECK-NEXT:    br label [[BB3]]
+; CHECK:       bb3:
+; CHECK-NEXT:    ret void
+;
+  store i32 0, i32* %P
+  br i1 true, label %bb1, label %bb2
+bb1:
+  store i32 1, i32* %P
+  br label %bb3
+bb2:
+  store i32 1, i32* %P
+  br label %bb3
+bb3:
+  ret void
+}