[DSE] Basic global DSE, backed by MemorySSA

[bryant]

new dse doesn't depend on domtree.

[bryant]

Reposting this with llvm-commits attached.

include/llvm/InitializePasses.h

 void initializeConstantMergeLegacyPassPass(PassRegistry &);
 void initializeConstantPropagationPass(PassRegistry&);
 void initializeCorrelatedValuePropagationPass(PassRegistry&);
 void initializeCostModelAnalysisPass(PassRegistry&);
 void initializeCrossDSOCFIPass(PassRegistry&);
 void initializeDAEPass(PassRegistry&);
 void initializeDAHPass(PassRegistry&);
 void initializeDCELegacyPassPass(PassRegistry&);
-void initializeDSELegacyPassPass(PassRegistry&);
+void initializeDSELegacyMemDepPass(PassRegistry&);
+void initializeDSELegacyMSSAPass(PassRegistry&);
 void initializeDataFlowSanitizerPass(PassRegistry&);
 void initializeDeadInstEliminationPass(PassRegistry&);
 void initializeDeadMachineInstructionElimPass(PassRegistry&);
 void initializeDelinearizationPass(PassRegistry &);
 void initializeDemandedBitsWrapperPassPass(PassRegistry&);
 void initializeDependenceAnalysisPass(PassRegistry&);
 void initializeDependenceAnalysisWrapperPassPass(PassRegistry&);
 void initializeDetectDeadLanesPass(PassRegistry&);

include/llvm/Transforms/Scalar/DeadStoreElimination.h

 
 #include "llvm/IR/Function.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// This class implements a trivial dead store elimination. We consider
 /// only the redundant stores that are local to a single Basic Block.
-class DSEPass : public PassInfoMixin<DSEPass> {
+template <bool UseMSSA> class DSEPass : public PassInfoMixin<DSEPass<UseMSSA>> {
 public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
 };
 }
 
 #endif // LLVM_TRANSFORMS_SCALAR_DSE_H

lib/Passes/PassBuilder.cpp

#endif
   // Run instcombine after redundancy and dead bit elimination to exploit
   // opportunities opened up by them.
   FPM.addPass(InstCombinePass());
 
   // Re-consider control flow based optimizations after redundancy elimination,
   // redo DCE, etc.
   FPM.addPass(JumpThreadingPass());
   FPM.addPass(CorrelatedValuePropagationPass());
-  FPM.addPass(DSEPass());
+  FPM.addPass(DSEPass<false>());
+  FPM.addPass(DSEPass<true>());
   // FIXME: Enable this when the loop pass manager can support enforcing loop
   // simplified and LCSSA form as well as updating the loop nest after
   // transformations and we finsih porting the loop passes.
 #if 0
   FPM.addPass(createFunctionToLoopPassAdaptor(LICMPass()));
 #endif
 
   // Finally, do an expensive DCE pass to catch all the dead code exposed by

lib/Passes/PassRegistry.def

 FUNCTION_PASS("adce", ADCEPass())
 FUNCTION_PASS("add-discriminators", AddDiscriminatorsPass())
 FUNCTION_PASS("alignment-from-assumptions", AlignmentFromAssumptionsPass())
 FUNCTION_PASS("bdce", BDCEPass())
 FUNCTION_PASS("break-crit-edges", BreakCriticalEdgesPass())
 FUNCTION_PASS("consthoist", ConstantHoistingPass())
 FUNCTION_PASS("correlated-propagation", CorrelatedValuePropagationPass())
 FUNCTION_PASS("dce", DCEPass())
-FUNCTION_PASS("dse", DSEPass())
+FUNCTION_PASS("dse", DSEPass<false>())
+FUNCTION_PASS("dsem", DSEPass<true>())
 FUNCTION_PASS("dot-cfg", CFGPrinterPass())
 FUNCTION_PASS("dot-cfg-only", CFGOnlyPrinterPass())
 FUNCTION_PASS("early-cse", EarlyCSEPass(/*UseMemorySSA=*/false))
 FUNCTION_PASS("early-cse-memssa", EarlyCSEPass(/*UseMemorySSA=*/true))
 FUNCTION_PASS("gvn-hoist", GVNHoistPass())
 FUNCTION_PASS("instcombine", InstCombinePass())
 FUNCTION_PASS("instsimplify", InstSimplifierPass())
 FUNCTION_PASS("invalidate<all>", InvalidateAllAnalysesPass())

lib/Transforms/Scalar/DeadStoreElimination.cpp

 // basic-block local redundant stores.
 //
 // FIXME: This should eventually be extended to be a post-dominator tree
 // traversal.  Doing so would be pretty trivial.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.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/PostDominators.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Pass.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/DeadStoreElimination.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/MemorySSA.h"
 #include <map>
 using namespace llvm;
 
 #define DEBUG_TYPE "dse"
 
 STATISTIC(NumRedundantStores, "Number of redundant stores deleted");
 STATISTIC(NumFastStores, "Number of stores deleted");
 STATISTIC(NumFastOther, "Number of other instrs removed");
 STATISTIC(NumCompletePartials, "Number of stores dead by later partials");
 
 static cl::opt<bool> EnablePartialOverwriteTracking(
     "enable-dse-partial-overwrite-tracking", cl::init(true), cl::Hidden,
     cl::desc("Enable partial-overwrite tracking in DSE"));
 
+static cl::opt<bool> EnableMSSA("use-dse-mssa", cl::init(false), cl::Hidden,
+                                cl::desc("Use the new MemorySSA-backed DSE."));
 
 //===----------------------------------------------------------------------===//
 // Helper functions
 //===----------------------------------------------------------------------===//
 typedef std::map<int64_t, int64_t> OverlapIntervalsTy;
 typedef DenseMap<Instruction *, OverlapIntervalsTy> InstOverlapIntervalsTy;
 
 /// Delete this instruction.  Before we do, go through and zero out all the
for (BasicBlock &BB : F)
     // Only check non-dead blocks.  Dead blocks may have strange pointer
     // cycles that will confuse alias analysis.
     if (DT->isReachableFromEntry(&BB))
       MadeChange |= eliminateDeadStores(BB, AA, MD, DT, TLI);
 
   return MadeChange;
 }
 
+struct WalkResult {
+  enum {
+    NextDef,
+    NextPhi,
+    KilledByUse,
+    SplitPoint,
+    ReachedEnd,
+  } State;
+  MemoryAccess *MA;
+};
+
+class DSEWalker {
+  Function *F;
+
+  unsigned InstNum;
+  DenseMap<const Value *, unsigned> InstNums;
+  // ^ Post-order numbering of MemoryPhis and instructions, a) that throw, or b)
+  // are DSE candidates. Used to detect intervening MayThrows or loop latches
+  SmallVector<unsigned, 32> MayThrows;
+  SmallPtrSet<const Value *, 16> Returns;
+  // ^ Values returned by F.
+  DenseSet<const Value *> NonEscapes;
+  // ^ Args and insts that don't escape on unwind of F.
+  InstOverlapIntervalsTy IOL;
+  // ^ isOverwrite needs this.
+public:
+  SmallVector<MemoryDef *, 32> Stores;
+  // ^ Stores to visit, post-ordered
+
+private:
+  AliasAnalysis *AA;
+  MemorySSA *MSSA;
+  const PostDominatorTree *PDT;
+  const TargetLibraryInfo *TLI;
+
+  unsigned addPO(Instruction &I) {
+    InstNums[&I] = InstNum++;
+    DEBUG(dbgs() << "PO: " << I << ", " << InstNums[&I] << "\n");
+    return InstNums[&I];
+  }
+
+  unsigned addPO(MemoryPhi &P) {
+    InstNums[&P] = InstNum++;
+    DEBUG(dbgs() << "PO: " << P << ", " << InstNums[&P] << "\n");
+    return InstNums[&P];
+  }
+
+  bool nonEscapingOnUnwind(Instruction &I) {
+    return isa<AllocaInst>(&I) ||
+           (isAllocLikeFn(&I, TLI) && !PointerMayBeCaptured(&I, false, true));
+  }
+
+  MemoryLocation getWriteLoc(Instruction *I) const {
+    MemoryLocation Loc = getLocForWrite(I, *AA);
+    if (!Loc.Ptr) {
+      if (auto CS = CallSite(I)) {
+        if (Function *F = CS.getCalledFunction()) {
+          StringRef FnName = F->getName();
+          if ((TLI->has(LibFunc::strcpy) &&
+               FnName == TLI->getName(LibFunc::strcpy)) ||
+              (TLI->has(LibFunc::strncpy) &&
+               FnName == TLI->getName(LibFunc::strncpy)) ||
+              (TLI->has(LibFunc::strcat) &&
+               FnName == TLI->getName(LibFunc::strcat)) ||
+              (TLI->has(LibFunc::strncat) &&
+               FnName == TLI->getName(LibFunc::strncat)))
+            return MemoryLocation::getForArgument(ImmutableCallSite(I), 0,
+                                                  *TLI);
+        }
+      }
+    }
+    return Loc;
+  }
+
+  // For free and lifetime_end, the size of MemLoc doesn't matter.
+  Value *getLifetimeEndish(Instruction &I) {
+    const DataLayout &DL = F->getParent()->getDataLayout();
+    if (isFreeCall(&I, TLI))
+      return GetUnderlyingObject(cast<CallInst>(&I)->getArgOperand(0), DL);
+    if (auto *II = dyn_cast<IntrinsicInst>(&I))
+      if (II->getIntrinsicID() == Intrinsic::lifetime_end)
+        // TODO: this isn't entirely  correct, since lifetime_end contains a
+        // size param.
+        return GetUnderlyingObject(cast<CallInst>(&I)->getArgOperand(1), DL);
+
+    return nullptr;
+  }
+
+public:
+  // Represents a DSE candidate and its cached escape info.
+  struct Candidate {
+    MemoryDef *D;
+    MemoryLocation Loc;
+    const Value *Und;
+    bool Returned;
+    // ^ Is this candidate's memory location returned by the function?
+    bool Escapes;
+    // ^ If execution in the function unwinds after this candidate writes, will
+    // its memory location escapes?
+  };
+
+  void deleteDead(MemoryDef &D) {
+    Instruction &I = *D.getMemoryInst();
+    DEBUG(dbgs() << "DSE-ing:\n\t" << D << "\n\t" << I << "\n");
+
+    SmallVector<Value *, 32> DeadPool(I.value_op_begin(), I.value_op_end());
+    DenseSet<Instruction *> Deleted; // Prevents double deletion.
+    MSSA->removeMemoryAccess(&D);
+    I.eraseFromParent();
+    IOL.erase(&I);
+
+    while (!DeadPool.empty()) {
+      auto *Cand = dyn_cast<Instruction>(DeadPool.pop_back_val());
+      if (Cand && !Deleted.count(Cand) &&
+          // Check if safe to delete, accounting for possible volatile or
+          // atomic.
+          isInstructionTriviallyDead(Cand, TLI)) {
+        DeadPool.insert(DeadPool.end(), Cand->value_op_begin(),
+                        Cand->value_op_end());
+        if (MemoryAccess *MA = MSSA->getMemoryAccess(Cand))
+          MSSA->removeMemoryAccess(MA);
+        Cand->eraseFromParent();
+        IOL.erase(Cand);
+        Deleted.insert(Cand);
+        ++NumFastOther;
+      }
+    }
+  }
+
+  // TODO: this fails to catch the following case:
+  //
+  // p = load x;
+  // while (...) { store p, x; store _, x; ... }
+  //
+  // memoryIsNotModifiedBetween is able to detect this.
+  bool isNoop(const MemoryDef &D) {
+    if (auto *SI = dyn_cast<StoreInst>(D.getMemoryInst())) {
+      MemoryAccess *Clob = MSSA->getWalker()->getClobberingMemoryAccess(SI);
+      if (auto *LI = dyn_cast<LoadInst>(SI->getValueOperand())) {
+        if (auto *U = dyn_cast<MemoryUse>(MSSA->getMemoryAccess(LI)))
+          // If LI is a MemoryUse, then MSSA guarantees that its defining access
+          // is accurate.
+          return U->getDefiningAccess() == Clob &&
+                 AA->isMustAlias(MemoryLocation::get(LI),
+                                 MemoryLocation::get(SI));
+        // Otherwise, LI is an abnormal load, i.e. atomic and/or volatile.
+        return Clob == MSSA->getMemoryAccess(LI) &&
+               AA->isMustAlias(MemoryLocation::get(LI),
+                               MemoryLocation::get(SI));
+      } else if (auto *MUD = dyn_cast<MemoryUseOrDef>(Clob)) {
+        // Storing zero to calloc-ed memory counts as no-op.
+        if (!MSSA->isLiveOnEntryDef(MUD)) {
+          Constant *C;
+          return (C = dyn_cast<Constant>(SI->getValueOperand())) &&
+                 C->isNullValue() && isCallocLikeFn(MUD->getMemoryInst(), TLI);
+        }
+      }
+    }
+    // TODO: Can also handle memmove(a <- a) and memcpy(a <- a), though the
+    // latter is technically UB.
+    return false;
+  }
+
+  unsigned po(Instruction &I) {
+    assert(InstNums.find(&I) != InstNums.end() && "Unindexed instruction.");
+    return InstNums.find(&I)->second;
+  }
+
+  unsigned po(MemoryPhi &P) {
+    assert(InstNums.find(&P) != InstNums.end() && "Unindexed MemoryPhi.");
+    return InstNums.find(&P)->second;
+  }
+
+  bool throwInRange(unsigned Earlier, unsigned Later) const {
+    DEBUG(dbgs() << "checking for throws between " << Earlier << " and "
+                 << Later << "\n");
+    assert(Earlier >= Later &&
+           "Larger number == later in post-order == earlier in rpo.");
+    return std::upper_bound(MayThrows.begin(), MayThrows.end(), Earlier) !=
+           std::lower_bound(MayThrows.begin(), MayThrows.end(), Later);
+  }
+
+  // Could Def prevent the DSE of a candidate store to Loc (which necessarily
+  // isn't volatile or atomic)? Example reasons for a yes:
+  // - Def's atomicity is stronger than monotonic
+  // - Def itself uses Loc, e.g., memcpy(... <- Loc)
+  bool isDSEBarrier(MemoryDef &D, const Candidate &Cand) {
+    Instruction *I = D.getMemoryInst();
+
+    if (getLifetimeEndish(*I))
+      return false;
+
+    if (I->isAtomic()) {
+      auto F = [](AtomicOrdering A) {
+        return A == AtomicOrdering::Monotonic || A == AtomicOrdering::Unordered;
+      };
+      if (auto *LI = dyn_cast<LoadInst>(I))
+        // Compensate for AA's skittishness around atomics.
+        return !(F(LI->getOrdering()) &&
+                 AA->isNoAlias(MemoryLocation::get(LI), Cand.Loc));
+      else if (auto *SI = dyn_cast<StoreInst>(I))
+        return !F(SI->getOrdering());
+      else if (isa<FenceInst>(I) && !Cand.Escapes && !Cand.Returned)
+        // From the old DSE: "Fences only constraints the ordering of already
+        // visible stores, it does not make a store visible to other threads.
+        // So, skipping over a fence does not change a store from being dead."
+        return false;
+    }
+
+    return AA->getModRefInfo(I, Cand.Loc) & MRI_Ref;
+  }
+
+  // Given the current walk location DefOrPhi, attempt to move downwards to the
+  // next MemoryDef that could DSE Cand.
+  WalkResult walkNext(MemoryAccess *DefOrPhi, const Candidate &Cand) {
+    DEBUG(dbgs() << "descending past " << *DefOrPhi << "\n");
+    WalkResult Res = {WalkResult::ReachedEnd, nullptr};
+
+    // Ensure that the uselist 1) doesn't MRI_Ref Cand.Loc, and 2) contains at
+    // most one MemoryDef or MemoryPhi.
+    for (Use &U : DefOrPhi->uses()) {
+      if (auto *Phi = dyn_cast<MemoryPhi>(U.getUser())) {
+        unsigned EarlierNum;
+        if (auto *D = dyn_cast<MemoryDef>(DefOrPhi))
+          EarlierNum = InstNums.find(D->getMemoryInst())->second;
+        else
+          EarlierNum = InstNums.find(DefOrPhi)->second;
+
+        DEBUG(dbgs() << "found phi: " << *Phi << ", "
+                     << InstNums.find(Phi)->second << ", " << EarlierNum
+                     << "\n");
+
+        if (Res.MA || InstNums.find(Phi)->second > EarlierNum)
+          // More than one MemoryDef or phi in the uselist implies a split point
+          // in the MSSA graph. A Phi with a lower (higher) RPO (PO) number
+          // means that we've encountered a loop latch.
+          return {WalkResult::SplitPoint, nullptr};
+        Res = {WalkResult::NextPhi, Phi};
+      } else if (auto *Load = dyn_cast<MemoryUse>(U.getUser())) {
+        if (AA->getModRefInfo(Load->getMemoryInst(), Cand.Loc) & MRI_Ref) {
+          // For a pair of stores to DSE, there can't be any intervening uses of
+          // the stored-to memory.
+          DEBUG(dbgs() << "used by " << *Load << "\n");
+          return {WalkResult::KilledByUse, Load};
+        }
+      } else if (auto *D = dyn_cast<MemoryDef>(U.getUser())) {
+        if (isDSEBarrier(*D, Cand)) {
+          DEBUG(dbgs() << "used by " << *D << "\n");
+          return {WalkResult::KilledByUse, D};
+        } else if (Res.MA)
+          return {WalkResult::SplitPoint, nullptr};
+        Res = {WalkResult::NextDef, D};
+      } else
+        llvm_unreachable("Unexpected MemorySSA node type.");
+    }
+
+    if (Res.State == WalkResult::ReachedEnd)
+      Res.MA = DefOrPhi;
+    return Res;
+  }
+
+  Candidate makeCand(MemoryDef &D) const {
+    assert(hasMemoryWrite(D.getMemoryInst(), *TLI) &&
+           "DSE candidates must write to an analyzable memory location.");
+    const DataLayout &DL = F->getParent()->getDataLayout();
+    MemoryLocation Loc = getWriteLoc(D.getMemoryInst());
+    assert(Loc.Ptr && "Expected a Loc!");
+    const Value *Und = GetUnderlyingObject(Loc.Ptr, DL);
+    bool Returned = Returns.count(Und);
+    DEBUG(dbgs() << "is " << D << " returned? " << Returned << "\n");
+    bool Escapes =
+        !(NonEscapes.count(Und) || ([&]() {
+            SmallVector<Value *, 4> Unds;
+            GetUnderlyingObjects(const_cast<Value *>(Und), Unds, DL);
+            return all_of(Unds, [&](Value *V) { return NonEscapes.count(V); });
+          })());
+    DEBUG(dbgs() << "does " << D << " escape? " << Escapes << "\n");
+    return {&D, Loc, Und, Returned, Escapes};
+  }
+
+  bool canDSE(const Candidate &Earlier, const MemoryDef &Later,
+              bool NonLocal = false) {
+    DEBUG(dbgs() << "can dse " << Later << "?\n");
+    Instruction &EarlierI = *Earlier.D->getMemoryInst();
+    Instruction &LaterI = *Later.getMemoryInst();
+
+    const DataLayout &DL = F->getParent()->getDataLayout();
+    if (Value *Ptr = getLifetimeEndish(LaterI)) {
+      // TODO: Possibly avoid repeated calls to GetUnderlyingObject through
+      // getLifetimeEndish.
+      DEBUG(dbgs() << "checking alias with free:\n\t" << *Earlier.Und << "\n\t"
+                   << *Ptr << "\n");
+      return AA->isMustAlias(Earlier.Und, Ptr);
+    }
+
+    MemoryLocation LLoc = getWriteLoc(&LaterI);
+    if (!LLoc.Ptr ||
+        (NonLocal &&
+         !PDT->dominates(Later.getBlock(), Earlier.D->getBlock())) ||
+        (Earlier.Escapes && throwInRange(po(EarlierI), po(LaterI))))
+      return false;
+
+    int64_t EOff, LOff;
+    OverwriteResult O =
+        isOverwrite(LLoc, Earlier.Loc, DL, *TLI, EOff, LOff, &EarlierI, IOL);
+    DEBUG(dbgs() << "Overwrite: " << O << "\n");
+    return O == OverwriteComplete;
+  }
+
+  DSEWalker(Function &F_, AliasAnalysis &AA_, MemorySSA &MSSA_,
+            const PostDominatorTree &PDT_, const TargetLibraryInfo &TLI_)
+      : F(&F_), InstNum(0), AA(&AA_), MSSA(&MSSA_), PDT(&PDT_), TLI(&TLI_) {
+    // Record non-escaping args.
+    for (Argument &Arg : F->args())
+      if (Arg.hasByValOrInAllocaAttr())
+        NonEscapes.insert(&Arg);
+
+    // Number instructions of interest by post-order
+    for (BasicBlock *BB : post_order(F)) {
+      for (Instruction &I : reverse(*BB)) {
+        if (I.mayThrow()) {
+          MayThrows.push_back(addPO(I));
+        } else if (auto *Def =
+                       dyn_cast_or_null<MemoryDef>(MSSA->getMemoryAccess(&I))) {
+          addPO(I);
+          // Only consider true DSE candidates. Volatile/atomic stores don't
+          // qualify, for instance.
+          if (hasMemoryWrite(&I, *TLI) && isRemovable(&I)) {
+            DEBUG(dbgs() << "candidate: " << I << "\n");
+            Stores.push_back(Def);
+          }
+        }
+
+        if (nonEscapingOnUnwind(I)) {
+          DEBUG(dbgs() << "found non-escaping mem: " << I << "\n");
+          NonEscapes.insert(&I);
+        }
+      }
+
+      if (MemoryPhi *Phi = MSSA->getMemoryAccess(BB))
+        // Phis numbers are used to recognize loop latches.
+        addPO(*Phi);
+
+      if (auto *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {
+        const DataLayout &DL = F->getParent()->getDataLayout();
+        // TODO: Possibly ignore unreachable blocks, which would require DT.
+        if (Value *RetVal = RI->getReturnValue())
+          Returns.insert(GetUnderlyingObject(RetVal, DL));
+      }
+    }
+  }
+};
+
+static bool eliminateDeadStoresMSSA(Function &F, AliasAnalysis &AA,
+                                    MemorySSA &MSSA,
+                                    const PostDominatorTree &PDT,
+                                    const TargetLibraryInfo &TLI) {
+  DEBUG(MSSA.print(dbgs()));
+  DSEWalker Walker(F, AA, MSSA, PDT, TLI);
+
+  using Candidate = DSEWalker::Candidate;
+
+  bool Changed = false;
+  for (MemoryDef *D : Walker.Stores) {
+    DEBUG(dbgs() << "inspecting " << *D->getMemoryInst() << "\n");
+
+    if (Walker.isNoop(*D)) {
+      DEBUG(dbgs() << "deleting no-op store " << *D << "\n");
+      Walker.deleteDead(*D);
+      Changed = true;
+      continue;
+    }
+
+    Candidate Cand = Walker.makeCand(*D);
+    WalkResult Res;
+    // Attempt to DSE within the same block because post-dom checks are limited
+    // to BasicBlock granularity.
+    bool next = false;
+    for (Res = Walker.walkNext(D, Cand); Res.State == WalkResult::NextDef &&
+                                         Res.MA->getBlock() == D->getBlock();
+         Res = Walker.walkNext(Res.MA, Cand)) {
+      if (Walker.canDSE(Cand, *cast<MemoryDef>(Res.MA), /* NonLocal */ false)) {
+        Walker.deleteDead(*D);
+        Changed = true;
+        next = true;
+        break;
+      }
+    }
+
+    if (next)
+      continue;
+
+    // Non-local search.
+    for (; Res.State <= WalkResult::NextPhi;
+         Res = Walker.walkNext(Res.MA, Cand)) {
+      if (auto *Def = dyn_cast<MemoryDef>(Res.MA)) {
+        if (Walker.canDSE(Cand, *Def, /* NonLocal */ true)) {
+          Walker.deleteDead(*D);
+          Changed = true;
+          break;
+        }
+      }
+    }
+
+    DEBUG(dbgs() << "Finished walking. " << Res.State << "\n");
+    if (Res.State == WalkResult::ReachedEnd && !Cand.Escapes &&
+        !Cand.Returned) {
+      DEBUG(dbgs() << "Caught unused write to non-escaping memory.\n");
+      Walker.deleteDead(*D);
+      Changed = true;
+    }
+  }
+  return Changed;
+}
+
 //===----------------------------------------------------------------------===//
 // DSE Pass
 //===----------------------------------------------------------------------===//
-PreservedAnalyses DSEPass::run(Function &F, FunctionAnalysisManager &AM) {
+template <>
+PreservedAnalyses DSEPass<false>::run(Function &F,
+                                      FunctionAnalysisManager &AM) {
   AliasAnalysis *AA = &AM.getResult<AAManager>(F);
   DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);
   MemoryDependenceResults *MD = &AM.getResult<MemoryDependenceAnalysis>(F);
   const TargetLibraryInfo *TLI = &AM.getResult<TargetLibraryAnalysis>(F);
 
   if (!eliminateDeadStores(F, AA, MD, DT, TLI))
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserve<DominatorTreeAnalysis>();
   PA.preserve<GlobalsAA>();
   PA.preserve<MemoryDependenceAnalysis>();
   return PA;
 }
 
+template <>
+PreservedAnalyses DSEPass<true>::run(Function &F, FunctionAnalysisManager &AM) {
+  if (!eliminateDeadStoresMSSA(F, AM.getResult<AAManager>(F),
+                               AM.getResult<MemorySSAAnalysis>(F).getMSSA(),
+                               AM.getResult<PostDominatorTreeAnalysis>(F),
+                               AM.getResult<TargetLibraryAnalysis>(F)))
+    return PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  PA.preserve<PostDominatorTreeAnalysis>();
+  PA.preserve<MemorySSAAnalysis>();
+  PA.preserve<GlobalsAA>();
+  return PA;
+}
+
 namespace {
 /// A legacy pass for the legacy pass manager that wraps \c DSEPass.
-class DSELegacyPass : public FunctionPass {
+template <bool UseMSSA> class DSELegacyPass : public FunctionPass {
 public:
   DSELegacyPass() : FunctionPass(ID) {
-    initializeDSELegacyPassPass(*PassRegistry::getPassRegistry());
+    if (UseMSSA)
+      initializeDSELegacyMSSAPass(*PassRegistry::getPassRegistry());
+    else
+      initializeDSELegacyMemDepPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
 
-    DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     AliasAnalysis *AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
-    MemoryDependenceResults *MD =
-        &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
     const TargetLibraryInfo *TLI =
         &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
-
+    if (UseMSSA) {
+      PostDominatorTree &PDT =
+          getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
+      MemorySSA &MSSA = getAnalysis<MemorySSAWrapperPass>().getMSSA();
+      return eliminateDeadStoresMSSA(F, *AA, MSSA, PDT, *TLI);
+    } else {
+      DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+      MemoryDependenceResults *MD =
+          &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
       return eliminateDeadStores(F, AA, MD, DT, TLI);
     }
+  }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
-    AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<AAResultsWrapperPass>();
-    AU.addRequired<MemoryDependenceWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
-    AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addPreserved<GlobalsAAWrapperPass>();
+
+    if (UseMSSA) {
+      AU.addRequired<PostDominatorTreeWrapperPass>();
+      AU.addRequired<MemorySSAWrapperPass>();
+      AU.addPreserved<PostDominatorTreeWrapperPass>();
+      AU.addPreserved<MemorySSAWrapperPass>();
+    } else {
+      AU.addRequired<DominatorTreeWrapperPass>();
+      AU.addRequired<MemoryDependenceWrapperPass>();
+      AU.addPreserved<DominatorTreeWrapperPass>();
       AU.addPreserved<MemoryDependenceWrapperPass>();
     }
+  }
 
   static char ID; // Pass identification, replacement for typeid
 };
 } // end anonymous namespace
 
-char DSELegacyPass::ID = 0;
-INITIALIZE_PASS_BEGIN(DSELegacyPass, "dse", "Dead Store Elimination", false,
+using DSELegacyMemDep = DSELegacyPass<false>;
+template <> char DSELegacyPass<false>::ID = 0;
+INITIALIZE_PASS_BEGIN(DSELegacyMemDep, "dse", "Dead Store Elimination", false,
                       false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(GlobalsAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
-INITIALIZE_PASS_END(DSELegacyPass, "dse", "Dead Store Elimination", false,
+INITIALIZE_PASS_END(DSELegacyMemDep, "dse", "Dead Store Elimination", false,
                     false)
 
+using DSELegacyMSSA = DSELegacyPass<true>;
+template <> char DSELegacyPass<true>::ID = 0;
+INITIALIZE_PASS_BEGIN(DSELegacyMSSA, "dsem",
+                      "Dead Store Elimination (Memory SSA)", false, false)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(GlobalsAAWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_END(DSELegacyMSSA, "dsem",
+                    "Dead Store Elimination (Memory SSA)", false, false)
+
 FunctionPass *llvm::createDeadStoreEliminationPass() {
-  return new DSELegacyPass();
+  if (EnableMSSA)
+    return new DSELegacyPass<true>();
+  return new DSELegacyPass<false>();
 }

lib/Transforms/Scalar/Scalar.cpp

void llvm::initializeScalarOpts(PassRegistry &Registry) {
   initializeBDCELegacyPassPass(Registry);
   initializeAlignmentFromAssumptionsPass(Registry);
   initializeConstantHoistingLegacyPassPass(Registry);
   initializeConstantPropagationPass(Registry);
   initializeCorrelatedValuePropagationPass(Registry);
   initializeDCELegacyPassPass(Registry);
   initializeDeadInstEliminationPass(Registry);
   initializeScalarizerPass(Registry);
-  initializeDSELegacyPassPass(Registry);
+  initializeDSELegacyMemDepPass(Registry);
+  initializeDSELegacyMSSAPass(Registry);
   initializeGuardWideningLegacyPassPass(Registry);
   initializeGVNLegacyPassPass(Registry);
   initializeNewGVNPass(Registry);
   initializeEarlyCSELegacyPassPass(Registry);
   initializeEarlyCSEMemSSALegacyPassPass(Registry);
   initializeGVNHoistLegacyPassPass(Registry);
   initializeFlattenCFGPassPass(Registry);
   initializeInductiveRangeCheckEliminationPass(Registry);