[DSE] Enable MSSA DSE to optimize across PHIs

[ebrevnov]

Summary

Currently DSE is unable to optimize stores if there is a phi in between. This patch adds basic support by "phi translating" killing location through phis.

[ebrevnov]

Rename DefLoc to ResDefLoc to fix name conflict.

[fhahn]

Great, thanks for sharing!

This currently seems to lead to a bit less stores removed overall on MultiSource/SPEC2000/SPEC2006 with -O3 -flto, but that might be related to some of the limits needing adjustments.

As a first step, it would probably be good to add a set of tests for the new functionality.

[ebrevnov]

Thanks for giving it a try!

Did you use latest update? If not would be nice to remeasure.
Unfortunately, I don't have access to mentioned Spec and not aware of Multisource benchmarks to look on my side. Is it feasible for you to share IR's for several cases so I could take a look? Would be nice to have them in regression test suite anyway.
On the other side I can try to run similar experiment on our internal tests.

[ebrevnov]

This currently seems to lead to a bit less stores removed overall on MultiSource/SPEC2000/SPEC2006 with -O3 -flto, but that might be related to some of the limits needing adjustments.

I found PartialLimit preventing DSE to optimize one of my simple cases. Here is the fix https://reviews.llvm.org/D90371 which I believe is beneficial regardless.

I also noted in case of earlier access is overwritten by the later one from the beginning and the end we DSE only beginning. Did you meet such behavior? Any plans fixing that?

Thanks
Evgeniy

[fhahn]

Hi! I'm not sure what the current status of this patch is, but I put up an initial patch that just does PHI translation when stepping across memory phis: D98288

Do you think it would be possible to port/add the logic to translate phis for any access?

[ebrevnov]

Extracted small NFC out and rebase

[vdsered]

@ebrevnov
Hi, any progress on the patch? I'd like to work on it if you agreed to that.

[ebrevnov]

Rebase

[ebrevnov]

Extracted more code into separate patches

[ebrevnov]

In D90095#2843496, @vdsered wrote:

@ebrevnov
Hi, any progress on the patch? I'd like to work on it if you agreed to that.

Hi Daniil. The help is more than welcome. I've extracted couple of patches which I hope to be able to land as ground work (see dependencies). Next step is to extract minimal change to introduce basic phi translation support. The challenge here is to take best of D98288 and this one as they are doing essentially the same things.

[ebrevnov]

Extracted more stuff to a separate patches.

[ebrevnov]

This patch is now close to what I wanted it to be (though I still need to work on comments) and you are welcome to take a look.

[ebrevnov]

Rebase + Formatting

[vdsered]

@ebrevnov I see that in memoryIsNotModifiedBetween function and there is some sort of a local phi-translation (I mean it is used only inside this function) as I understand. Do we need to update this function somehow?

[ebrevnov]

Rebase

[ebrevnov]

In D90095#2992286, @vdsered wrote:

@ebrevnov I see that in memoryIsNotModifiedBetween function and there is some sort of a local phi-translation (I mean it is used only inside this function) as I understand. Do we need to update this function somehow?

Looks like this function does phi translation along the way as well. I hope there is an easy way to merge them. Let me try...

[fhahn]

IIUC the dependencies for the patch landed by now. Please let me know when this is ready for review. It would also be good to add a brief explanation of the approach both the the description and in code comments.

Also, I think this might need a bit more additional testing by building larger programs. For example, the following example crashes at the moment:

@global = external global i64, align 8

define i32 @widget(i1 %arg, i8* %arg1, i8* %arg2) {
bb:
  store i64 777, i64* @global, align 8
  br i1 %arg, label %bb3, label %bb4

bb3:                                              ; preds = %bb
  %tmp = getelementptr i8, i8* %arg1, i64 1
  store i8 11, i8* %tmp, align 8
  br label %bb4

bb4:                                              ; preds = %bb3, %bb
  %tmp5 = phi i8* [ %arg1, %bb3 ], [ %arg2, %bb ]
  %tmp6 = getelementptr i8, i8* %tmp5, i64 1
  store i8 99, i8* %tmp6, align 8
  ret i32 10
}

[xbolva00]

Status of this patch?

[ebrevnov]

The ball is on my side but I can't return to it due to other higher priority tasks. Please let me know if there is a demand from your side so I could re-prioritize this one. Also, as mentioned earlier, the help is more than welcome :-)

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

 //===- DeadStoreElimination.cpp - MemorySSA Backed Dead Store Elimination -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // The code below implements dead store elimination using MemorySSA. It uses

 #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/AssumptionCache.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/MustExecute.h"
bool canSkipDef(MemoryDef *D, bool DefVisibleToCaller,
 
   // Skip intrinsics that do not really read or modify memory.
   if (isNoopIntrinsic(DI))
     return true;
 
   return false;
 }
 
+using PhiTransPtr = std::pair<WeakTrackingVH, int64_t>;
+using ExMemoryAccess = std::pair<MemoryAccess *, PhiTransPtr>;
+
+static PhiTransPtr phiTranslatePtr(const PhiTransPtr &Ptr, BasicBlock *FromBB,
+                                   BasicBlock *ToBB, const DataLayout &DL,
+                                   DominatorTree &DT, AssumptionCache *AC) {
+  PhiTransPtr ResPtr;
+  DenseMap<BasicBlock *, PhiTransPtr> Visited;
+  SmallVector<std::pair<BasicBlock *, const PhiTransPtr>, 16> WorkList;
+
+  if (FromBB == ToBB) {
+    return Ptr;
+  }
+
+  WorkList.push_back({FromBB, Ptr});
+
+  while (!WorkList.empty()) {
+    auto CurrNode = WorkList.pop_back_val();
+    BasicBlock *CurrBB = CurrNode.first;
+    const PhiTransPtr &CurrPtr = CurrNode.second;
+
+    for (pred_iterator PI = pred_begin(CurrBB), E = pred_end(CurrBB); PI != E;
+         ++PI) {
+      BasicBlock *PredBB = *PI;
+      PHITransAddr TransAddr{CurrPtr.first, DL, AC};
+
+      // TODO:
+      if (!DT.dominates(ToBB, PredBB))
+        continue;
+
+      if (TransAddr.NeedsPHITranslationFromBlock(CurrBB) &&
+          (!TransAddr.IsPotentiallyPHITranslatable() ||
+           TransAddr.PHITranslateValue(CurrBB, PredBB, &DT, false)))
+        return PhiTransPtr{};
+
+      auto Inserted =
+          Visited.try_emplace(PredBB, TransAddr.getAddr(), CurrPtr.second);
+      auto &TransPtr = Inserted.first->second;
+      if (!Inserted.second) {
+        if (TransAddr.getAddr() != TransPtr.first ||
+            CurrPtr.second != TransPtr.second)
+          // We already visited this block before. If it was with a different
+          // address - bail out!
+          return PhiTransPtr{};
+        continue;
+      }
+
+      if (PredBB == ToBB) {
+        ResPtr = TransPtr;
+        continue;
+      }
+
+      WorkList.push_back({PredBB, TransPtr});
+    }
+  }
+
+  assert(ResPtr.first.pointsToAliveValue() &&
+         "PHI translation is expected to complete successfully");
+  assert(ResPtr.second == 0 &&
+         "Unexpected non-zero offset during phi translation");
+  return ResPtr;
+}
+
+static ExMemoryAccess
+phiTransFromMemoryAccessTo(const ExMemoryAccess &FromAccess,
+                           MemoryAccess *ToAccess, const DataLayout &DL,
+                           DominatorTree &DT, AssumptionCache *AC) {
+  PhiTransPtr ResAddr;
+
+  if (FromAccess.second.first.pointsToAliveValue()) {
+
+    ResAddr = phiTranslatePtr(FromAccess.second, FromAccess.first->getBlock(),
+                              ToAccess->getBlock(), DL, DT, AC);
+  }
+
+  return std::make_pair(ToAccess, ResAddr);
+}
+
+static ExMemoryAccess getExDefiningAccess(const ExMemoryAccess &CurrAccess,
+                                          const DataLayout &DL,
+                                          DominatorTree &DT,
+                                          AssumptionCache *AC) {
+  assert(isa<MemoryUseOrDef>(CurrAccess.first) && "TODO");
+  MemoryAccess *Next =
+      cast<MemoryUseOrDef>(CurrAccess.first)->getDefiningAccess();
+  return phiTransFromMemoryAccessTo(CurrAccess, Next, DL, DT, AC);
+}
+
 struct DSEState {
   Function &F;
   AliasAnalysis &AA;
 
   /// The single BatchAA instance that is used to cache AA queries. It will
   /// not be invalidated over the whole run. This is safe, because:
   /// 1. Only memory writes are removed, so the alias cache for memory
   ///    locations remains valid.
   /// 2. No new instructions are added (only instructions removed), so cached
   ///    information for a deleted value cannot be accessed by a re-used new
   ///    value pointer.
   BatchAAResults BatchAA;
 
   MemorySSA &MSSA;
   DominatorTree &DT;
   PostDominatorTree &PDT;
   const TargetLibraryInfo &TLI;
   const DataLayout &DL;
   const LoopInfo &LI;
+  AssumptionCache *AC;
 
   // Whether the function contains any irreducible control flow, useful for
   // being accurately able to detect loops.
   bool ContainsIrreducibleLoops;
 
   // All MemoryDefs that potentially could kill other MemDefs.
   SmallVector<MemoryDef *, 64> MemDefs;
   // Any that should be skipped as they are already deleted
struct DSEState {
   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,
-           const LoopInfo &LI)
+           const LoopInfo &LI, AssumptionCache *AC)
       : F(F), AA(AA), BatchAA(AA), MSSA(MSSA), DT(DT), PDT(PDT), TLI(TLI),
-        DL(F.getParent()->getDataLayout()), LI(LI) {}
+        DL(F.getParent()->getDataLayout()), LI(LI), AC(AC),
+        ContainsIrreducibleLoops(false) {}
 
   static DSEState get(Function &F, AliasAnalysis &AA, MemorySSA &MSSA,
                       DominatorTree &DT, PostDominatorTree &PDT,
-                      const TargetLibraryInfo &TLI, const LoopInfo &LI) {
-    DSEState State(F, AA, MSSA, DT, PDT, TLI, LI);
+                      const TargetLibraryInfo &TLI, const LoopInfo &LI,
+                      AssumptionCache *AC) {
+    DSEState State(F, AA, MSSA, DT, PDT, TLI, LI, AC);
     // Collect blocks with throwing instructions not modeled in MemorySSA and
     // alloc-like objects.
     unsigned PO = 0;
     for (BasicBlock *BB : post_order(&F)) {
       State.PostOrderNumbers[BB] = PO++;
       for (Instruction &I : *BB) {
         MemoryAccess *MA = MSSA.getMemoryAccess(&I);
         if (I.mayThrow() && !MA)
bool isGuaranteedLoopInvariant(const Value *Ptr) {
     }
     if (auto *GEP = dyn_cast<GEPOperator>(Ptr)) {
       return IsGuaranteedLoopInvariantBase(GEP->getPointerOperand()) &&
              GEP->hasAllConstantIndices();
     }
     return IsGuaranteedLoopInvariantBase(Ptr);
   }
 
-  // Find a MemoryDef writing to \p KillingLoc and dominating \p StartAccess,
+  // Find a MemoryDef writing to \p KillingLoc and dominating \p ExStartAccess,
   // with no read access between them or on any other path to a function exit
   // block if \p KillingLoc is not accessible after the function returns. If
   // there is no such MemoryDef, return None. The returned value may not
   // (completely) overwrite \p KillingLoc. Currently we bail out when we
   // encounter an aliasing MemoryUse (read).
-  Optional<MemoryAccess *>
-  getDomMemoryDef(MemoryDef *KillingDef, MemoryAccess *StartAccess,
+  Optional<ExMemoryAccess>
+  getDomMemoryDef(MemoryDef *KillingDef, const ExMemoryAccess &ExStartAccess,
                   const MemoryLocation &KillingLoc, const Value *KillingUndObj,
                   unsigned &ScanLimit, unsigned &WalkerStepLimit,
                   bool IsMemTerm, unsigned &PartialLimit) {
     if (ScanLimit == 0 || WalkerStepLimit == 0) {
       LLVM_DEBUG(dbgs() << "\n    ...  hit scan limit\n");
       return None;
     }
 
-    MemoryAccess *Current = StartAccess;
+    MemoryAccess *const StartAccess = ExStartAccess.first;
     Instruction *KillingI = KillingDef->getMemoryInst();
     LLVM_DEBUG(dbgs() << "  trying to get dominating access\n");
 
-    // Find the next clobbering Mod access for DefLoc, starting at StartAccess.
+    ExMemoryAccess ExCurrentAccess = ExStartAccess;
     Optional<MemoryLocation> CurrentLoc;
-    for (;; Current = cast<MemoryDef>(Current)->getDefiningAccess()) {
+    Optional<std::pair<MemoryLocation, uint64_t>> PhiTransKillingLocAndOffset;
+
+    // Find the next clobbering Mod access for KillingLoc, starting at
+    // StartAccess.
+    for (;;
+         ExCurrentAccess = getExDefiningAccess(ExCurrentAccess, DL, DT, AC)) {
+      MemoryAccess *CurrentAccess = ExCurrentAccess.first;
+
+      PhiTransKillingLocAndOffset.reset();
+      if (ExCurrentAccess.second.first.pointsToAliveValue()) {
+        PhiTransKillingLocAndOffset = {
+            KillingLoc.getWithNewPtr(ExCurrentAccess.second.first)
+                .getWithoutAATags(),
+            ExCurrentAccess.second.second};
+      }
+
       LLVM_DEBUG({
-        dbgs() << "   visiting " << *Current;
-        if (!MSSA.isLiveOnEntryDef(Current) && isa<MemoryUseOrDef>(Current))
-          dbgs() << " (" << *cast<MemoryUseOrDef>(Current)->getMemoryInst()
+        dbgs() << "   visiting " << *CurrentAccess;
+        if (!MSSA.isLiveOnEntryDef(CurrentAccess) &&
+            isa<MemoryUseOrDef>(CurrentAccess))
+          dbgs() << " ("
+                 << *cast<MemoryUseOrDef>(CurrentAccess)->getMemoryInst()
                  << ")";
         dbgs() << "\n";
       });
 
       // Reached TOP.
-      if (MSSA.isLiveOnEntryDef(Current)) {
+      if (MSSA.isLiveOnEntryDef(CurrentAccess)) {
         LLVM_DEBUG(dbgs() << "   ...  found LiveOnEntryDef\n");
         return None;
       }
 
       // Cost of a step. Accesses in the same block are more likely to be valid
       // candidates for elimination, hence consider them cheaper.
-      unsigned StepCost = KillingDef->getBlock() == Current->getBlock()
+      unsigned StepCost = KillingDef->getBlock() == CurrentAccess->getBlock()
                               ? MemorySSASameBBStepCost
                               : MemorySSAOtherBBStepCost;
       if (WalkerStepLimit <= StepCost) {
         LLVM_DEBUG(dbgs() << "   ...  hit walker step limit\n");
         return None;
       }
       WalkerStepLimit -= StepCost;
 
       // Return for MemoryPhis. They cannot be eliminated directly and the
       // caller is responsible for traversing them.
-      if (isa<MemoryPhi>(Current)) {
+      if (isa<MemoryPhi>(CurrentAccess)) {
         LLVM_DEBUG(dbgs() << "   ...  found MemoryPhi\n");
-        return Current;
+        return ExCurrentAccess;
       }
 
       // Below, check if CurrentDef is a valid candidate to be eliminated by
       // KillingDef. If it is not, check the next candidate.
-      MemoryDef *CurrentDef = cast<MemoryDef>(Current);
+      MemoryDef *CurrentDef = cast<MemoryDef>(CurrentAccess);
       Instruction *CurrentI = CurrentDef->getMemoryInst();
 
       if (canSkipDef(CurrentDef, !isInvisibleToCallerBeforeRet(KillingUndObj),
                      TLI))
         continue;
 
       // Before we try to remove anything, check for any extra throwing
       // instructions that block us from DSEing
for (;;
       // If Current is known to be on path that reads DefLoc or is a read
       // clobber, bail out, as the path is not profitable. We skip this check
       // for intrinsic calls, because the code knows how to handle memcpy
       // intrinsics.
       if (!isa<IntrinsicInst>(CurrentI) && isReadClobber(KillingLoc, CurrentI))
         return None;
 
       // Quick check if there are direct uses that are read-clobbers.
-      if (any_of(Current->uses(), [this, &KillingLoc, StartAccess](Use &U) {
+      if (any_of(
+              CurrentAccess->uses(), [this, &KillingLoc, StartAccess](Use &U) {
                 if (auto *UseOrDef = dyn_cast<MemoryUseOrDef>(U.getUser()))
                   return !MSSA.dominates(StartAccess, UseOrDef) &&
                          isReadClobber(KillingLoc, UseOrDef->getMemoryInst());
                 return false;
               })) {
         LLVM_DEBUG(dbgs() << "   ...  found a read clobber\n");
         return None;
       }
 
       // If Current cannot be analyzed or is not removable, check the next
       // candidate.
       if (!hasAnalyzableMemoryWrite(CurrentI, TLI) || !isRemovable(CurrentI))
         continue;
 
       // If Current does not have an analyzable write location, skip it
+      int64_t CurrOffset = 0;
       CurrentLoc = getLocForWriteEx(CurrentI);
       if (!CurrentLoc)
         continue;
 
       // AliasAnalysis does not account for loops. Limit elimination to
       // candidates for which we can guarantee they always store to the same
       // memory location and not located in different loops.
       if (!isGuaranteedLoopIndependent(CurrentI, KillingI, *CurrentLoc)) {
         LLVM_DEBUG(dbgs() << "  ... not guaranteed loop independent\n");
         WalkerStepLimit -= 1;
         continue;
       }
 
       if (IsMemTerm) {
         // If the killing def is a memory terminator (e.g. lifetime.end), check
         // the next candidate if the current Current does not write the same
         // underlying object as the terminator.
         if (!isMemTerminator(*CurrentLoc, CurrentI, KillingI))
           continue;
       } else {
-        int64_t KillingOffset = 0;
-        int64_t DeadOffset = 0;
-        auto OR = isOverwrite(KillingI, CurrentI, KillingLoc, *CurrentLoc,
-                              KillingOffset, DeadOffset);
+        auto &KillingMemLoc = PhiTransKillingLocAndOffset
+                                  ? (*PhiTransKillingLocAndOffset).first
+                                  : KillingLoc;
+        int64_t KillingMemOffset = PhiTransKillingLocAndOffset
+                                       ? (*PhiTransKillingLocAndOffset).second
+                                       : 0;
+        assert(KillingMemOffset == 0 && "Unexpected non-zero offset");
+        auto OR = isOverwrite(KillingI, CurrentI, KillingMemLoc, *CurrentLoc,
+                              KillingMemOffset, CurrOffset);
         // If Current does not write to the same object as KillingDef, check
         // the next candidate.
         if (OR == OW_Unknown)
           continue;
         else if (OR == OW_MaybePartial) {
           // If KillingDef only partially overwrites Current, check the next
           // candidate if the partial step limit is exceeded. This aggressively
           // limits the number of candidates for partial store elimination,
getDomMemoryDef(MemoryDef *KillingDef, const ExMemoryAccess &ExStartAccess,
     };
 
     // Accesses to objects accessible after the function returns can only be
     // eliminated if the access is dead along all paths to the exit. Collect
     // the blocks with killing (=completely overwriting MemoryDefs) and check if
     // they cover all paths from MaybeDeadAccess to any function exit.
     SmallPtrSet<Instruction *, 16> KillingDefs;
     KillingDefs.insert(KillingDef->getMemoryInst());
-    MemoryAccess *MaybeDeadAccess = Current;
+    MemoryAccess *MaybeDeadAccess = ExCurrentAccess.first;
     MemoryLocation MaybeDeadLoc = *CurrentLoc;
     Instruction *MaybeDeadI = cast<MemoryDef>(MaybeDeadAccess)->getMemoryInst();
     LLVM_DEBUG(dbgs() << "  Checking for reads of " << *MaybeDeadAccess << " ("
                       << *MaybeDeadI << ")\n");
 
     SmallSetVector<MemoryAccess *, 32> WorkList;
     auto PushMemUses = [&WorkList](MemoryAccess *Acc) {
       for (Use &U : Acc->uses())
if (!isInvisibleToCallerAfterRet(KillingUndObj)) {
           break;
         CommonPred = PDT.findNearestCommonDominator(CommonPred, BB);
       }
 
       // If CommonPred is in the set of killing blocks, just check if it
       // post-dominates MaybeDeadAccess.
       if (KillingBlocks.count(CommonPred)) {
         if (PDT.dominates(CommonPred, MaybeDeadAccess->getBlock()))
-          return {MaybeDeadAccess};
+          return {ExCurrentAccess};
         return None;
       }
 
       // If the common post-dominator does not post-dominate MaybeDeadAccess,
       // there is a path from MaybeDeadAccess to an exit not going through a
       // killing block.
       if (PDT.dominates(CommonPred, MaybeDeadAccess->getBlock())) {
         SetVector<BasicBlock *> WorkList;
 
         // If CommonPred is null, there are multiple exits from the function.
         // They all have to be added to the worklist.
         if (CommonPred)
           WorkList.insert(CommonPred);
         else
           for (BasicBlock *R : PDT.roots())
             WorkList.insert(R);
 
         NumCFGTries++;
         // Check if all paths starting from an exit node go through one of the
         // killing blocks before reaching MaybeDeadAccess.
         for (unsigned I = 0; I < WorkList.size(); I++) {
           NumCFGChecks++;
-          BasicBlock *Current = WorkList[I];
-          if (KillingBlocks.count(Current))
+          BasicBlock *CurrentBB = WorkList[I];
+          if (KillingBlocks.count(CurrentBB))
             continue;
-          if (Current == MaybeDeadAccess->getBlock())
+          if (CurrentBB == MaybeDeadAccess->getBlock())
             return None;
 
           // MaybeDeadAccess is reachable from the entry, so we don't have to
           // explore unreachable blocks further.
-          if (!DT.isReachableFromEntry(Current))
+          if (!DT.isReachableFromEntry(CurrentBB))
             continue;
 
-          for (BasicBlock *Pred : predecessors(Current))
+          for (BasicBlock *Pred : predecessors(CurrentBB))
             WorkList.insert(Pred);
 
           if (WorkList.size() >= MemorySSAPathCheckLimit)
             return None;
         }
         NumCFGSuccess++;
-        return {MaybeDeadAccess};
+        return {ExCurrentAccess};
       }
       return None;
     }
 
     // No aliasing MemoryUses of MaybeDeadAccess found, MaybeDeadAccess is
     // potentially dead.
-    return {MaybeDeadAccess};
+    return {ExCurrentAccess};
   }
 
   // Delete dead memory defs
   void deleteDeadInstruction(Instruction *SI) {
     MemorySSAUpdater Updater(&MSSA);
     SmallVector<Instruction *, 32> NowDeadInsts;
     NowDeadInsts.push_back(SI);
     --NumFastOther;
bool storeIsNoop(MemoryDef *Def, const Value *DefUO) {
 
     return false;
   }
 };
 
 static bool eliminateDeadStores(Function &F, AliasAnalysis &AA, MemorySSA &MSSA,
                                 DominatorTree &DT, PostDominatorTree &PDT,
                                 const TargetLibraryInfo &TLI,
-                                const LoopInfo &LI) {
+                                const LoopInfo &LI, AssumptionCache *AC) {
   bool MadeChange = false;
 
-  DSEState State = DSEState::get(F, AA, MSSA, DT, PDT, TLI, LI);
+  const DataLayout &DL = F.getParent()->getDataLayout();
+  DSEState State = DSEState::get(F, AA, MSSA, DT, PDT, TLI, LI, AC);
   // For each store:
   for (unsigned I = 0; I < State.MemDefs.size(); I++) {
     MemoryDef *KillingDef = State.MemDefs[I];
     if (State.SkipStores.count(KillingDef))
       continue;
     Instruction *KillingI = KillingDef->getMemoryInst();
 
     Optional<MemoryLocation> MaybeKillingLoc;
     if (State.isMemTerminatorInst(KillingI))
       MaybeKillingLoc = State.getLocForTerminator(KillingI).map(
           [](const std::pair<MemoryLocation, bool> &P) { return P.first; });
     else
       MaybeKillingLoc = State.getLocForWriteEx(KillingI);
 
     if (!MaybeKillingLoc) {
       LLVM_DEBUG(dbgs() << "Failed to find analyzable write location for "
                         << *KillingI << "\n");
       continue;
     }
     MemoryLocation KillingLoc = *MaybeKillingLoc;
     assert(KillingLoc.Ptr && "KillingLoc should not be null");
     const Value *KillingUndObj = getUnderlyingObject(KillingLoc.Ptr);
 
-    MemoryAccess *Current = KillingDef;
     LLVM_DEBUG(dbgs() << "Trying to eliminate MemoryDefs killed by "
                       << *KillingDef << " (" << *KillingI << ")\n");
 
     unsigned ScanLimit = MemorySSAScanLimit;
     unsigned WalkerStepLimit = MemorySSAUpwardsStepLimit;
     unsigned PartialLimit = MemorySSAPartialStoreLimit;
     // Worklist of MemoryAccesses that may be killed by KillingDef.
-    SetVector<MemoryAccess *> ToCheck;
+    SmallVector<ExMemoryAccess> ToCheck;
 
-    if (KillingUndObj)
-      ToCheck.insert(KillingDef->getDefiningAccess());
+    if (KillingUndObj) {
+      ExMemoryAccess ExKillingDef{KillingDef,
+                                  {const_cast<Value *>(KillingLoc.Ptr), 0}};
+      ToCheck.push_back(getExDefiningAccess(ExKillingDef, DL, DT, AC));
+    }
 
     bool Shortend = false;
     bool IsMemTerm = State.isMemTerminatorInst(KillingI);
     // Check if MemoryAccesses in the worklist are killed by KillingDef.
-    for (unsigned I = 0; I < ToCheck.size(); I++) {
-      Current = ToCheck[I];
+    while (!ToCheck.empty()) {
+      const ExMemoryAccess CurrentAccess = ToCheck.pop_back_val();
+      MemoryAccess *Current = CurrentAccess.first;
       if (State.SkipStores.count(Current))
         continue;
 
-      Optional<MemoryAccess *> MaybeDeadAccess = State.getDomMemoryDef(
-          KillingDef, Current, KillingLoc, KillingUndObj, ScanLimit,
+      Optional<ExMemoryAccess> MaybeDeadAccess = State.getDomMemoryDef(
+          KillingDef, CurrentAccess, KillingLoc, KillingUndObj, ScanLimit,
           WalkerStepLimit, IsMemTerm, PartialLimit);
 
       if (!MaybeDeadAccess) {
         LLVM_DEBUG(dbgs() << "  finished walk\n");
         continue;
       }
 
-      MemoryAccess *DeadAccess = *MaybeDeadAccess;
+      ExMemoryAccess &ExDeadAccess = *MaybeDeadAccess;
+      MemoryAccess *DeadAccess = ExDeadAccess.first;
       LLVM_DEBUG(dbgs() << " Checking if we can kill " << *DeadAccess);
       if (isa<MemoryPhi>(DeadAccess)) {
         LLVM_DEBUG(dbgs() << "\n  ... adding incoming values to worklist\n");
         for (Value *V : cast<MemoryPhi>(DeadAccess)->incoming_values()) {
           MemoryAccess *IncomingAccess = cast<MemoryAccess>(V);
           BasicBlock *IncomingBlock = IncomingAccess->getBlock();
           BasicBlock *PhiBlock = DeadAccess->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);
+            ToCheck.push_back(phiTransFromMemoryAccessTo(
+                ExDeadAccess, IncomingAccess, DL, DT, AC));
         }
         continue;
       }
       auto *DeadDefAccess = cast<MemoryDef>(DeadAccess);
       Instruction *DeadI = DeadDefAccess->getMemoryInst();
       LLVM_DEBUG(dbgs() << " (" << *DeadI << ")\n");
-      ToCheck.insert(DeadDefAccess->getDefiningAccess());
+      ToCheck.push_back(getExDefiningAccess(ExDeadAccess, DL, DT, AC));
       NumGetDomMemoryDefPassed++;
 
       if (!DebugCounter::shouldExecute(MemorySSACounter))
         continue;
 
       MemoryLocation DeadLoc = *State.getLocForWriteEx(DeadI);
 
       if (IsMemTerm) {
         const Value *DeadUndObj = getUnderlyingObject(DeadLoc.Ptr);
         if (KillingUndObj != DeadUndObj)
           continue;
         LLVM_DEBUG(dbgs() << "DSE: Remove Dead Store:\n  DEAD: " << *DeadI
                           << "\n  KILLER: " << *KillingI << '\n');
         State.deleteDeadInstruction(DeadI);
         ++NumFastStores;
         MadeChange = true;
       } else {
         // Check if DeadI overwrites KillingI.
         int64_t KillingOffset = 0;
         int64_t DeadOffset = 0;
+
+        if (ExDeadAccess.second.first.pointsToAliveValue()) {
+          KillingLoc = KillingLoc.getWithNewPtr(ExDeadAccess.second.first)
+                      .getWithoutAATags();
+          KillingOffset = ExDeadAccess.second.second;
+        }
+        assert(KillingOffset == 0 && "Unexpected non-zero offset");
+
         OverwriteResult OR = State.isOverwrite(
             KillingI, DeadI, KillingLoc, DeadLoc, KillingOffset, DeadOffset);
         if (OR == OW_MaybePartial) {
           auto Iter = State.IOLs.insert(
               std::make_pair<BasicBlock *, InstOverlapIntervalsTy>(
                   DeadI->getParent(), InstOverlapIntervalsTy()));
           auto &IOL = Iter.first->second;
           OR = isPartialOverwrite(KillingLoc, DeadLoc, KillingOffset,

 //===----------------------------------------------------------------------===//
 PreservedAnalyses DSEPass::run(Function &F, FunctionAnalysisManager &AM) {
   AliasAnalysis &AA = AM.getResult<AAManager>(F);
   const TargetLibraryInfo &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
   MemorySSA &MSSA = AM.getResult<MemorySSAAnalysis>(F).getMSSA();
   PostDominatorTree &PDT = AM.getResult<PostDominatorTreeAnalysis>(F);
   LoopInfo &LI = AM.getResult<LoopAnalysis>(F);
+  AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
 
-  bool Changed = eliminateDeadStores(F, AA, MSSA, DT, PDT, TLI, LI);
+  bool Changed = eliminateDeadStores(F, AA, MSSA, DT, PDT, TLI, LI, &AC);
 
 #ifdef LLVM_ENABLE_STATS
   if (AreStatisticsEnabled())
     for (auto &I : instructions(F))
       NumRemainingStores += isa<StoreInst>(&I);
 #endif
 
   if (!Changed)
bool runOnFunction(Function &F) override {
     AliasAnalysis &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
     DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     const TargetLibraryInfo &TLI =
         getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
     MemorySSA &MSSA = getAnalysis<MemorySSAWrapperPass>().getMSSA();
     PostDominatorTree &PDT =
         getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
     LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+    AssumptionCache &AC =
+        getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
 
-    bool Changed = eliminateDeadStores(F, AA, MSSA, DT, PDT, TLI, LI);
+    bool Changed = eliminateDeadStores(F, AA, MSSA, DT, PDT, TLI, LI, &AC);
 
 #ifdef LLVM_ENABLE_STATS
     if (AreStatisticsEnabled())
       for (auto &I : instructions(F))
         NumRemainingStores += isa<StoreInst>(&I);
 #endif
 
     return Changed;
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     AU.addRequired<AAResultsWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addPreserved<GlobalsAAWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addRequired<PostDominatorTreeWrapperPass>();
     AU.addRequired<MemorySSAWrapperPass>();
     AU.addPreserved<PostDominatorTreeWrapperPass>();
     AU.addPreserved<MemorySSAWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addPreserved<LoopInfoWrapperPass>();
+    AU.addRequired<AssumptionCacheTracker>();
+    AU.addPreserved<AssumptionCacheTracker>();
   }
 };
 
 } // end anonymous namespace
 
 char DSELegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(DSELegacyPass, "dse", "Dead Store Elimination", false,

llvm/test/Transforms/DeadStoreElimination/phi-translation.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -dse -S %s | FileCheck %s
 
-; TODO: Both the stores in %then and %else can be eliminated by translating %p
+; Both the stores in %then and %else can be eliminated by translating %p
 ; through the phi.
 define void @memoryphi_translate_1(i1 %c) {
 ; CHECK-LABEL: @memoryphi_translate_1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A_1:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    [[A_2:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[THEN:%.*]], label [[ELSE:%.*]]
 ; CHECK:       then:
-; CHECK-NEXT:    store i8 0, i8* [[A_1]], align 1
 ; CHECK-NEXT:    br label [[END:%.*]]
 ; CHECK:       else:
-; CHECK-NEXT:    store i8 9, i8* [[A_2]], align 1
 ; CHECK-NEXT:    br label [[END]]
 ; CHECK:       end:
 ; CHECK-NEXT:    [[P:%.*]] = phi i8* [ [[A_1]], [[THEN]] ], [ [[A_2]], [[ELSE]] ]
 ; CHECK-NEXT:    store i8 10, i8* [[P]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %a.1 = alloca i8
else:
   br label %end
 
 end:
   %p = phi i8* [ %a.1, %then ], [ %a.2, %else ]
   store i8 10, i8* %p
   ret void
 }
 
-; TODO: The store in %else can be eliminated by translating %p through the phi.
+; The store in %else can be eliminated by translating %p through the phi.
 ; The store in %then cannot be eliminated, because %a.1 is read before the final
 ; store.
 define i8 @memoryphi_translate_2(i1 %c) {
 ; CHECK-LABEL: @memoryphi_translate_2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A_1:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    [[A_2:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[THEN:%.*]], label [[ELSE:%.*]]
 ; CHECK:       then:
 ; CHECK-NEXT:    store i8 0, i8* [[A_1]], align 1
 ; CHECK-NEXT:    br label [[END:%.*]]
 ; CHECK:       else:
-; CHECK-NEXT:    store i8 9, i8* [[A_2]], align 1
 ; CHECK-NEXT:    br label [[END]]
 ; CHECK:       end:
 ; CHECK-NEXT:    [[P:%.*]] = phi i8* [ [[A_1]], [[THEN]] ], [ [[A_2]], [[ELSE]] ]
 ; CHECK-NEXT:    [[L:%.*]] = load i8, i8* [[A_1]], align 1
 ; CHECK-NEXT:    store i8 10, i8* [[P]], align 1
 ; CHECK-NEXT:    ret i8 [[L]]
 ;
 entry:

 
 end:
   %p = phi i8* [ %a.1, %then ], [ %a.2, %else ]
   %l = load i8, i8* %a.1
   store i8 10, i8* %p
   ret i8 %l
 }
 
-; TODO: The store in %then can be eliminated by translating %p through the phi.
+; The store in %then can be eliminated by translating %p through the phi.
 ; The store in %else cannot be eliminated, because %a.2 is read before the final
 ; store.
 define i8 @memoryphi_translate_3(i1 %c) {
 ; CHECK-LABEL: @memoryphi_translate_3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A_1:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    [[A_2:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[THEN:%.*]], label [[ELSE:%.*]]
 ; CHECK:       then:
-; CHECK-NEXT:    store i8 0, i8* [[A_1]], align 1
 ; CHECK-NEXT:    br label [[END:%.*]]
 ; CHECK:       else:
 ; CHECK-NEXT:    store i8 9, i8* [[A_2]], align 1
 ; CHECK-NEXT:    br label [[END]]
 ; CHECK:       end:
 ; CHECK-NEXT:    [[P:%.*]] = phi i8* [ [[A_1]], [[THEN]] ], [ [[A_2]], [[ELSE]] ]
 ; CHECK-NEXT:    [[L:%.*]] = load i8, i8* [[A_2]], align 1
 ; CHECK-NEXT:    store i8 10, i8* [[P]], align 1

 ; CHECK-LABEL: @test_trans_null(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[THEN:%.*]], label [[ELSE:%.*]]
 ; CHECK:       then:
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       else:
 ; CHECK-NEXT:    call void @fn()
 ; CHECK-NEXT:    [[BC:%.*]] = bitcast i8* undef to i16*
-; CHECK-NEXT:    [[GEP_1:%.*]] = getelementptr inbounds i16, i16* [[BC]], i64 2
-; CHECK-NEXT:    store i16 8, i16* [[GEP_1]], align 2
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[P:%.*]] = phi i16* [ [[PTR:%.*]], [[THEN]] ], [ [[BC]], [[ELSE]] ]
 ; CHECK-NEXT:    [[GEP_2:%.*]] = getelementptr inbounds i16, i16* [[P]], i64 2
 ; CHECK-NEXT:    store i16 8, i16* [[GEP_2]], align 2
 ; CHECK-NEXT:    ret void
 ;
 entry: