MemorySSA: Link all defs together into an intrusive defslist, to make updater easier

[• dberlin]

Summary

This is the first in a series of patches to add a simple, generalized updater to MemorySSA.

For MemorySSA, every def is may-def, instead of the normal must-def.
(the best way to think of memoryssa is "everything is really one variable, with different versions of that variable at different points in the program).
This means when updating, we end up having to do a bunch of work to touch defs below and above us.

In order to support this quickly, i have ilist'd all the defs for each block. ilist supports tags, so this is quite easy. the only slightly messy part is that you can't have two iplists for the same type that differ only whether they have the ownership part enabled or not, because the traits are for the value type.

The verifiers have been updated to test that the def order is correct.

[• dberlin]

• Update with more changes from my branch. SpliceAbove is completely broken right now for defs spliced into blocks only with uses, as it does not update downards. I've disabled the test because i'll delete the API as part of the updater

[• dberlin]

An example where spliceabove is broken:

MemoryUse(0)
use(a)
2= MemoryDef(1)
store b
3 = MemoryDef(2)
store a

Splice store a above use a produces:

3 = MemoryDef(0)
store a
MemoryUse(3)
use(a)
2= MemoryDef(1)
store b

This would work in SSA, because you would just have two variables live at the same time.
However, it does not work for MemorySSA, because only one variable may be live at a time, and therefore, you need to update downwards as well.

The only cases where the effects are truly local is when you splice into the middle of a block that has defs above or below you.

[george.burgess.iv]

Thanks for the patch! The overall approach here seems solid to me.

include/llvm/Transforms/Utils/MemorySSA.h

114

I dunno how cluttered llvm:: is with random pass/analysis-specific things, but it may be nice to scope these somehow (e.g. in a MSSAHelpers struct or namespace or ...)

573

"since iplist's own if you don't change the traits" sounds strange to me

689

Less newlines, please

700

Can this just be a map of const BasicBlock * to unique_ptr<struct { AccessList; DefsList; }> (excuse the not-actual-c++)? Looks like we look both up together a lot, and their lifetimes are sorta-related.

(If it ends up being a big refactor, I'm fine if we do this as a follow-up)

lib/Transforms/Utils/MemorySSA.cpp

1523

Nit: if we're not going to "cache" our old value, I think this can can be sunk to the if.

1533

InsertIntoDef = true

1582

It looks like we have this pattern:

void foo(MemoryAccess *MA) {
  getOrCreateAccessList(MA->getBlock())->add_thing(MA);
  if (!isa<MemoryUse>(MA)) {
    getOrCreateDefsList(MA->getBlock())->add_thing(MA);
  }
}
// Where add_thing is one of {insert, push_back, push_front}.

a lot. Can we maybe factor that out into a insertMemoryAccessBefore/insertMemoryAccessAfter(AccessList::iterator, MemoryAccess *) thing that handles the DefList bits for us?

1613

Looks like we put Phis in this list, too. Is this supposed to push Defs in front of phis?

1614

Please delete this newline (and below)

1642

Is Defs->end() correct here? What if I insert a store before the MemoryUse in

define void @foo() {
  %1 = alloca i8
  ; MemoryUse(liveOnEntry)
  %2 = load i8, i8* %1
  ; 1 = MemoryDef(liveOnEntry)
  store i8 0, i8* %1
}

?

1660

Please remove the extra parens

1661

Same "is Defs->end() correct?" comment

1927

Please remove the extra parens

2384

Is this from clang-format?

unittests/Transforms/Utils/MemorySSA.cpp

488

I'm assuming this is temporary

[• dberlin]

include/llvm/Transforms/Utils/MemorySSA.h

114

Fixed :)

573

I tried to rewrite this into something more understandable.

700

So, i'm not strongly opposed, but note:

We currently guarantee two invariants:

1. If PerBlockAccesses has your block, then the access list for the block is not empty (and vice versa. If it does not have the block, there are no accesses)
2. If PerBlockDefs has your block, then the defs list for the block is not empty.

If we put them both together, then obviously we can still maintain invariant 1, but we will break invariant 2.

Breaking that invariant means things will now have to check that the defs list is not null everywhere.

Personally, i don't think it's worth it, but happy to hear another opinion :)
(and obviously, happy to document these invariants)

lib/Transforms/Utils/MemorySSA.cpp

1523

Fixed. We now cache the old value, so we don't do a lookup for every def.

1582

You can't use the AccessList iterator, but we can at least create two helpers

1613

No, fixed.

1642

No, it is not correct.
We have to find the nearest def, above or below, and do before/after.
Writing the helper functions made me realize this as well.

We'll do it the slow way first, then we can make it faster if necessary (we have the block numbering, we could make it note which are defs, and then find the nearest def to the access passed here).
Since this is an update API, i doubt it'll be that slow in practice.

2384

Not sure how it got there

unittests/Transforms/Utils/MemorySSA.cpp

488

Yes, it's going away completely as part of the updater update.

[• dberlin]

• Update for review comments

[• dberlin]

Note that insertIntoListsBefore requires the block because the iterator may point at the end.

include/llvm/Transforms/Utils/MemorySSA.h

85

I have no idea what happened here, i'll fix (but clang format sorted it into this order)

[• dberlin]

• Remove double include of iterator_range.h
• Update formatting

[bryant]

In D29046#653835, @dberlin wrote:

An example where spliceabove is broken:

MemoryUse(0)
use(a)
2= MemoryDef(1)
store b
3 = MemoryDef(2)
store a

Splice store a above use a produces:

3 = MemoryDef(0)
store a
MemoryUse(3)
use(a)
2= MemoryDef(1)
store b

This would work in SSA, because you would just have two variables live at the same time.
However, it does not work for MemorySSA, because only one variable may be live at a time, and therefore, you need to update downwards as well.

The only cases where the effects are truly local is when you splice into the middle of a block that has defs above or below you.

This isn't of much consequence, but I wanted to point out that your example already violates the pre-conditions to spliceMemoryAccessAbove:

// \brief Splice \p What to just before \p Where.
//
// In order to be efficient, the following conditions must be met:
//   - \p Where  dominates \p What,
//   - All memory accesses in [\p Where, \p What) are no-alias with \p What.
//
// TODO: relax the MemoryDef requirement on Where.

Still, a generalized updater is way more useful than a tightly-scoped splicing API.

[george.burgess.iv]

LGTM after 3 more small nits. Thanks again!

include/llvm/Transforms/Utils/MemorySSA.h

573

Much better, thanks! (Looks like clang-format broke the comment, though. :P)

700

Then I'm happy with our current approach. Documenting the invariants would be nice, too. :)

lib/Transforms/Utils/MemorySSA.cpp

1490

insertIntoListsForBlock?

[• dberlin]

lib/Transforms/Utils/MemorySSA.cpp

1490

I just used createMemoryPhi instead

[george.burgess.iv]

lib/Transforms/Utils/MemorySSA.cpp

1490

WFM

include/llvm/Transforms/Utils/MemorySSA.h

 // definitions, which would require multiple phis, and multiple memoryaccesses
 // per instruction.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_MEMORYSSA_H
 #define LLVM_TRANSFORMS_UTILS_MEMORYSSA_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/ilist.h"
 #include "llvm/ADT/ilist_node.h"
 #include "llvm/ADT/iterator.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/PHITransAddr.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/OperandTraits.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"

 namespace llvm {
 
 class Function;
 class Instruction;
 class MemoryAccess;
 class LLVMContext;
 class raw_ostream;
 
+struct AllAccessTag {};
+struct DefsOnlyTag {};
+
 enum {
   // Used to signify what the default invalid ID is for MemoryAccess's
   // getID()
   INVALID_MEMORYACCESS_ID = 0
 };
 
 template <class T> class memoryaccess_def_iterator_base;
 using memoryaccess_def_iterator = memoryaccess_def_iterator_base<MemoryAccess>;
 using const_memoryaccess_def_iterator =
     memoryaccess_def_iterator_base<const MemoryAccess>;
 
 // \brief The base for all memory accesses. All memory accesses in a block are
 // linked together using an intrusive list.
-class MemoryAccess : public User, public ilist_node<MemoryAccess> {
+class MemoryAccess : public User,
+                     public ilist_node<MemoryAccess, ilist_tag<AllAccessTag>>,
+                     public ilist_node<MemoryAccess, ilist_tag<DefsOnlyTag>> {
 public:
+  using AllAccessType = ilist_node<MemoryAccess, ilist_tag<AllAccessTag>>;
+  using DefsOnlyType = ilist_node<MemoryAccess, ilist_tag<DefsOnlyTag>>;
+
   // Methods for support type inquiry through isa, cast, and
   // dyn_cast
   static inline bool classof(const MemoryAccess *) { return true; }
   static inline bool classof(const Value *V) {
     unsigned ID = V->getValueID();
     return ID == MemoryUseVal || ID == MemoryPhiVal || ID == MemoryDefVal;
   }
 
public:
   /// \brief This iterator walks over all of the defs in a given
   /// MemoryAccess. For MemoryPhi nodes, this walks arguments. For
   /// MemoryUse/MemoryDef, this walks the defining access.
   memoryaccess_def_iterator defs_begin();
   const_memoryaccess_def_iterator defs_begin() const;
   memoryaccess_def_iterator defs_end();
   const_memoryaccess_def_iterator defs_end() const;
 
+  /// \brief Get the iterators for the all access list and the defs only list
+  /// We default to the all access list.
+  AllAccessType::self_iterator getIterator() {
+    return this->AllAccessType::getIterator();
+  }
+  AllAccessType::const_self_iterator getIterator() const {
+    return this->AllAccessType::getIterator();
+  }
+  AllAccessType::reverse_self_iterator getReverseIterator() {
+    return this->AllAccessType::getReverseIterator();
+  }
+  AllAccessType::const_reverse_self_iterator getReverseIterator() const {
+    return this->AllAccessType::getReverseIterator();
+  }
+  DefsOnlyType::self_iterator getDefsIterator() {
+    return this->DefsOnlyType::getIterator();
+  }
+  DefsOnlyType::const_self_iterator getDefsIterator() const {
+    return this->DefsOnlyType::getIterator();
+  }
+  DefsOnlyType::reverse_self_iterator getReverseDefsIterator() {
+    return this->DefsOnlyType::getReverseIterator();
+  }
+  DefsOnlyType::const_reverse_self_iterator getReverseDefsIterator() const {
+    return this->DefsOnlyType::getReverseIterator();
+  }
+
 protected:
   friend class MemorySSA;
   friend class MemoryUseOrDef;
   friend class MemoryUse;
   friend class MemoryDef;
   friend class MemoryPhi;
 
   /// \brief Used for debugging and tracking things about MemoryAccesses.
public:
   inline bool isLiveOnEntryDef(const MemoryAccess *MA) const {
     return MA == LiveOnEntryDef.get();
   }
 
   inline MemoryAccess *getLiveOnEntryDef() const {
     return LiveOnEntryDef.get();
   }
 
-  using AccessList = iplist<MemoryAccess>;
+  // Sadly, ilist_traits controls the ownership properties of the list, and we
+  // can't differentiate the ilist traits based on tags, only on the value
+  // type.  So we can't use iplist for both since iplist's own if you don't
+  // change the traits.  Thus, we use simple_ilist for one, and iplist for the
+  // other.  This means the deflist usage is a bit uglier, but c'est la vie.
+  using AccessList = iplist<MemoryAccess, ilist_tag<AllAccessTag>>;
+  using DefsList = simple_ilist<MemoryAccess, ilist_tag<DefsOnlyTag>>;
 
   /// \brief Return the list of MemoryAccess's for a given basic block.
   ///
   /// This list is not modifiable by the user.
   const AccessList *getBlockAccesses(const BasicBlock *BB) const {
     return getWritableBlockAccesses(BB);
   }
 
+  /// \brief Return the list of MemoryDef's and MemoryPhi's for a given basic
+  /// block.
+  ///
+  /// This list is not modifiable by the user.
+  const DefsList *getBlockDefs(const BasicBlock *BB) const {
+    return getWritableBlockDefs(BB);
+  }
+
   /// \brief Create an empty MemoryPhi in MemorySSA for a given basic block.
   /// Only one MemoryPhi for a block exists at a time, so this function will
   /// assert if you try to create one where it already exists.
   MemoryPhi *createMemoryPhi(BasicBlock *BB);
 
   enum InsertionPlace { Beginning, End };
 
   /// \brief Create a MemoryAccess in MemorySSA at a specified point in a block,
protected:
   // Used by Memory SSA annotater, dumpers, and wrapper pass
   friend class MemorySSAAnnotatedWriter;
   friend class MemorySSAPrinterLegacyPass;
 
   void verifyDefUses(Function &F) const;
   void verifyDomination(Function &F) const;
   void verifyOrdering(Function &F) const;
 
-  // This is used by the use optimizer class
+  // This is used by the use optimizer and updater.
   AccessList *getWritableBlockAccesses(const BasicBlock *BB) const {
     auto It = PerBlockAccesses.find(BB);
     return It == PerBlockAccesses.end() ? nullptr : It->second.get();
-  }
 
+  // This is used by the use optimizer and updater.
+  DefsList *getWritableBlockDefs(const BasicBlock *BB) const {
+    auto It = PerBlockDefs.find(BB);
+    return It == PerBlockDefs.end() ? nullptr : It->second.get();
+  }
+
+  
+
 private:
   class CachingWalker;
   class OptimizeUses;
 
   CachingWalker *getWalkerImpl();
   void buildMemorySSA();
   void optimizeUses();
 
   void verifyUseInDefs(MemoryAccess *, MemoryAccess *) const;
   using AccessMap = DenseMap<const BasicBlock *, std::unique_ptr<AccessList>>;
+  using DefsMap = DenseMap<const BasicBlock *, std::unique_ptr<DefsList>>;
 
   void
   determineInsertionPoint(const SmallPtrSetImpl<BasicBlock *> &DefiningBlocks);
   void computeDomLevels(DenseMap<DomTreeNode *, unsigned> &DomLevels);
   void markUnreachableAsLiveOnEntry(BasicBlock *BB);
   bool dominatesUse(const MemoryAccess *, const MemoryAccess *) const;
   MemoryUseOrDef *createNewAccess(Instruction *);
   MemoryUseOrDef *createDefinedAccess(Instruction *, MemoryAccess *);
   MemoryAccess *findDominatingDef(BasicBlock *, enum InsertionPlace);
   void removeFromLookups(MemoryAccess *);
 
   void placePHINodes(const SmallPtrSetImpl<BasicBlock *> &,
                      const DenseMap<const BasicBlock *, unsigned int> &);
   MemoryAccess *renameBlock(BasicBlock *, MemoryAccess *);
   void renamePass(DomTreeNode *, MemoryAccess *IncomingVal,
                   SmallPtrSet<BasicBlock *, 16> &Visited);
   AccessList *getOrCreateAccessList(const BasicBlock *);
+  DefsList *getOrCreateDefsList(const BasicBlock *);
   void renumberBlock(const BasicBlock *) const;
 
   AliasAnalysis *AA;
   DominatorTree *DT;
   Function &F;
 
   // Memory SSA mappings
   DenseMap<const Value *, MemoryAccess *> ValueToMemoryAccess;
   AccessMap PerBlockAccesses;
+  DefsMap PerBlockDefs;
   std::unique_ptr<MemoryAccess> LiveOnEntryDef;
 
   // Domination mappings
   // Note that the numbering is local to a block, even though the map is
   // global.
   mutable SmallPtrSet<const BasicBlock *, 16> BlockNumberingValid;
   mutable DenseMap<const MemoryAccess *, unsigned long> BlockNumbering;
 

lib/Transforms/Utils/MemorySSA.cpp

 
 MemorySSA::AccessList *MemorySSA::getOrCreateAccessList(const BasicBlock *BB) {
   auto Res = PerBlockAccesses.insert(std::make_pair(BB, nullptr));
 
   if (Res.second)
     Res.first->second = make_unique<AccessList>();
   return Res.first->second.get();
 }
+MemorySSA::DefsList *MemorySSA::getOrCreateDefsList(const BasicBlock *BB) {
+  auto Res = PerBlockDefs.insert(std::make_pair(BB, nullptr));
+
+  if (Res.second)
+    Res.first->second = make_unique<DefsList>();
+  return Res.first->second.get();
+}
 
 /// This class is a batch walker of all MemoryUse's in the program, and points
 /// their defining access at the thing that actually clobbers them.  Because it
 /// is a batch walker that touches everything, it does not operate like the
 /// other walkers.  This walker is basically performing a top-down SSA renaming
 /// pass, where the version stack is used as the cache.  This enables it to be
 /// significantly more time and memory efficient than using the regular walker,
 /// which is walking bottom-up.
std::sort(IDFBlocks.begin(), IDFBlocks.end(),
             [&BBNumbers](const BasicBlock *A, const BasicBlock *B) {
               return BBNumbers.lookup(A) < BBNumbers.lookup(B);
             });
 
   // Now place MemoryPhi nodes.
   for (auto &BB : IDFBlocks) {
     // Insert phi node
     AccessList *Accesses = getOrCreateAccessList(BB);
+    DefsList *Defs = getOrCreateDefsList(BB);
     MemoryPhi *Phi = new MemoryPhi(BB->getContext(), BB, NextID++);
     ValueToMemoryAccess[BB] = Phi;
     // Phi's always are placed at the front of the block.
     Accesses->push_front(Phi);
+    Defs->push_front(*Phi);
   }
 }
 
 void MemorySSA::buildMemorySSA() {
   // We create an access to represent "live on entry", for things like
   // arguments or users of globals, where the memory they use is defined before
   // the beginning of the function. We do not actually insert it into the IR.
   // We do not define a live on exit for the immediate uses, and thus our
void MemorySSA::buildMemorySSA() {
   SmallPtrSet<BasicBlock *, 32> DefiningBlocks;
   SmallPtrSet<BasicBlock *, 32> DefUseBlocks;
   // Go through each block, figure out where defs occur, and chain together all
   // the accesses.
   for (BasicBlock &B : F) {
     BBNumbers[&B] = NextBBNum++;
     bool InsertIntoDef = false;
     AccessList *Accesses = nullptr;
+    DefsList *Defs = nullptr;
     for (Instruction &I : B) {
       MemoryUseOrDef *MUD = createNewAccess(&I);
       if (!MUD)
         continue;
-      InsertIntoDef |= isa<MemoryDef>(MUD);
 
       if (!Accesses)
         Accesses = getOrCreateAccessList(&B);
       Accesses->push_back(MUD);
+      if (isa<MemoryDef>(MUD)) {
+        InsertIntoDef |= true;
+        Defs = getOrCreateDefsList(&B);
+        Defs->push_back(*MUD);
+      }
     }
     if (InsertIntoDef)
       DefiningBlocks.insert(&B);
     if (Accesses)
       DefUseBlocks.insert(&B);
   }
   placePHINodes(DefiningBlocks, BBNumbers);
 
MemorySSA::CachingWalker *MemorySSA::getWalkerImpl() {
 
   Walker = make_unique<CachingWalker>(this, AA, DT);
   return Walker.get();
 }
 
 MemoryPhi *MemorySSA::createMemoryPhi(BasicBlock *BB) {
   assert(!getMemoryAccess(BB) && "MemoryPhi already exists for this BB");
   AccessList *Accesses = getOrCreateAccessList(BB);
+  DefsList *Defs = getOrCreateDefsList(BB);
   MemoryPhi *Phi = new MemoryPhi(BB->getContext(), BB, NextID++);
   ValueToMemoryAccess[BB] = Phi;
   // Phi's always are placed at the front of the block.
   Accesses->push_front(Phi);
+  Defs->push_front(*Phi);
   BlockNumberingValid.erase(BB);
   return Phi;
 }
 
 MemoryUseOrDef *MemorySSA::createDefinedAccess(Instruction *I,
                                                MemoryAccess *Definition) {
   assert(!isa<PHINode>(I) && "Cannot create a defined access for a PHI");
   MemoryUseOrDef *NewAccess = createNewAccess(I);
MemoryAccess *MemorySSA::createMemoryAccessInBB(Instruction *I,
   MemoryUseOrDef *NewAccess = createDefinedAccess(I, Definition);
   auto *Accesses = getOrCreateAccessList(BB);
   if (Point == Beginning) {
     // It goes after any phi nodes
     auto AI = find_if(
         *Accesses, [](const MemoryAccess &MA) { return !isa<MemoryPhi>(MA); });
 
     Accesses->insert(AI, NewAccess);
+    if (!isa<MemoryUse>(NewAccess)) {
+      auto *Defs = getOrCreateDefsList(BB);
+      Defs->push_front(*NewAccess);
+   
+    }
   } else {
     Accesses->push_back(NewAccess);
+    if (!isa<MemoryUse>(NewAccess)) {
+      auto *Defs = getOrCreateDefsList(BB);
+      Defs->push_back(*NewAccess);
+    
+    }
   }
   BlockNumberingValid.erase(BB);
   return NewAccess;
 }
 
 MemoryUseOrDef *MemorySSA::createMemoryAccessBefore(Instruction *I,
                                                     MemoryAccess *Definition,
                                                     MemoryUseOrDef *InsertPt) {
   assert(I->getParent() == InsertPt->getBlock() &&
          "New and old access must be in the same block");
   MemoryUseOrDef *NewAccess = createDefinedAccess(I, Definition);
-  auto *Accesses = getOrCreateAccessList(InsertPt->getBlock());
+  // Should already exist, since insertpt is in that block.
+  auto *Accesses = getWritableBlockAccesses(InsertPt->getBlock());
   Accesses->insert(AccessList::iterator(InsertPt), NewAccess);
   BlockNumberingValid.erase(InsertPt->getBlock());
+  if (!isa<MemoryUse>(NewAccess)) {
+    // May not exist, since we are not guaranteed insertion point is a def.
+    auto *Defs = getOrCreateDefsList(InsertPt->getBlock());
+    auto DefsPoint =
+        isa<MemoryDef>(InsertPt) ? InsertPt->getDefsIterator() : Defs->end();
+    Defs->insert(DefsPoint, *NewAccess);
+  }
+
   return NewAccess;
 }
 
 MemoryUseOrDef *MemorySSA::createMemoryAccessAfter(Instruction *I,
                                                    MemoryAccess *Definition,
                                                    MemoryAccess *InsertPt) {
   assert(I->getParent() == InsertPt->getBlock() &&
          "New and old access must be in the same block");
   MemoryUseOrDef *NewAccess = createDefinedAccess(I, Definition);
   auto *Accesses = getOrCreateAccessList(InsertPt->getBlock());
   Accesses->insertAfter(AccessList::iterator(InsertPt), NewAccess);
+  if (!isa<MemoryUse>(NewAccess)) {
+    // May not exist, since we are not guaranteed insertion point is a def.
+    auto *Defs = getOrCreateDefsList(InsertPt->getBlock());
+    auto DefsPoint = isa<MemoryDef>(InsertPt) ? ++(InsertPt->getDefsIterator())
+                                              : Defs->end();
+    Defs->insert(DefsPoint, *NewAccess);
+  }
   BlockNumberingValid.erase(InsertPt->getBlock());
+
   return NewAccess;
 }
 
 void MemorySSA::spliceMemoryAccessAbove(MemoryDef *Where,
                                         MemoryUseOrDef *What) {
   assert(What != getLiveOnEntryDef() &&
          Where != getLiveOnEntryDef() && "Can't splice (above) LOE.");
   assert(dominates(Where, What) && "Only upwards splices are permitted.");
if (MemoryUseOrDef *MUD = dyn_cast<MemoryUseOrDef>(MA)) {
     MemoryInst = MUD->getMemoryInst();
   } else {
     MemoryInst = MA->getBlock();
   }
   auto VMA = ValueToMemoryAccess.find(MemoryInst);
   if (VMA->second == MA)
     ValueToMemoryAccess.erase(VMA);
 
+  // The access list owns the reference, so we erase it from the non-owning list
+  // first.
+  if (!isa<MemoryUse>(MA)) {
+    auto DefsIt = PerBlockDefs.find(MA->getBlock());
+    std::unique_ptr<DefsList> &Defs = DefsIt->second;
+    Defs->remove(*MA);
+    if (Defs->empty())
+      PerBlockDefs.erase(DefsIt);
+  }
+
   auto AccessIt = PerBlockAccesses.find(MA->getBlock());
   std::unique_ptr<AccessList> &Accesses = AccessIt->second;
   Accesses->erase(MA);
   if (Accesses->empty())
     PerBlockAccesses.erase(AccessIt);
 }
 
 void MemorySSA::removeMemoryAccess(MemoryAccess *MA) {

 
 /// \brief Verify that the order and existence of MemoryAccesses matches the
 /// order and existence of memory affecting instructions.
 void MemorySSA::verifyOrdering(Function &F) const {
   // Walk all the blocks, comparing what the lookups think and what the access
   // lists think, as well as the order in the blocks vs the order in the access
   // lists.
   SmallVector<MemoryAccess *, 32> ActualAccesses;
+  SmallVector<MemoryAccess *, 32> ActualDefs;
   for (BasicBlock &B : F) {
     const AccessList *AL = getBlockAccesses(&B);
+    const auto *DL = getBlockDefs(&B);
     MemoryAccess *Phi = getMemoryAccess(&B);
-    if (Phi)
+    if (Phi) {
       ActualAccesses.push_back(Phi);
+      ActualDefs.push_back(Phi);
+    }
+
     for (Instruction &I : B) {
       MemoryAccess *MA = getMemoryAccess(&I);
-      assert((!MA || AL) && "We have memory affecting instructions "
-                            "in this block but they are not in the "
-                            "access list");
-      if (MA)
+      assert((!MA || (AL && (isa<MemoryUse>(MA) || DL))) &&
+             "We have memory affecting instructions "
+             "access list or defs list");
+      if (MA) {
         ActualAccesses.push_back(MA);
+        if (isa<MemoryDef>(MA))
+          ActualDefs.push_back(MA);
+      }
     }
     // Either we hit the assert, really have no accesses, or we have both
-    // accesses and an access list
-    if (!AL)
+    // accesses and an access list.
+    // Same with defs.
+    if (!AL && !DL)
       continue;
     assert(AL->size() == ActualAccesses.size() &&
            "We don't have the same number of accesses in the block as on the "
            "access list");
+    assert(DL ||
+           ActualDefs.size() == 0 &&
+               "Either we should have a defs list, or we should have no defs");
+    assert((!DL || (DL->size() == ActualDefs.size())) &&
+           "We don't have the same number of defs in the block as on the "
+           "def list");
     auto ALI = AL->begin();
     auto AAI = ActualAccesses.begin();
     while (ALI != AL->end() && AAI != ActualAccesses.end()) {
       assert(&*ALI == *AAI && "Not the same accesses in the same order");
       ++ALI;
       ++AAI;
     }
     ActualAccesses.clear();
+    if (DL) {
+      auto DLI = DL->begin();
+      auto ADI = ActualDefs.begin();
+      while (DLI != DL->end() && ADI != ActualDefs.end()) {
+        assert(&*DLI == *ADI && "Not the same defs in the same order");
+        ++DLI;
+        ++ADI;
+      }
+    }
+    ActualDefs.clear();
   }
 }
 
 /// \brief Verify the domination properties of MemorySSA by checking that each
 /// definition dominates all of its uses.
 void MemorySSA::verifyDomination(Function &F) const {
 #ifndef NDEBUG
   for (BasicBlock &B : F) {

 }
 
 MemoryAccess *DoNothingMemorySSAWalker::getClobberingMemoryAccess(
     MemoryAccess *StartingAccess, const MemoryLocation &) {
   if (auto *Use = dyn_cast<MemoryUseOrDef>(StartingAccess))
     return Use->getDefiningAccess();
   return StartingAccess;
 }
+
 } // namespace llvm

unittests/Transforms/Utils/MemorySSA.cpp

TEST_F(MemorySSATest, WalkerReopt) {
   // Create the load memory access pointing to an unoptimized place.
   MemoryUse *NewLoadAccess = cast<MemoryUse>(MSSA.createMemoryAccessInBB(
       LIA, MSSA.getMemoryAccess(SIB), LIA->getParent(), MemorySSA::End));
   // This should it cause it to be optimized
   EXPECT_EQ(Walker->getClobberingMemoryAccess(NewLoadAccess), LoadClobber);
   EXPECT_EQ(NewLoadAccess->getDefiningAccess(), LoadClobber);
 }
 
+#if 0
 // Test out MemorySSA::spliceMemoryAccessAbove.
 TEST_F(MemorySSATest, SpliceAboveMemoryDef) {
   F = Function::Create(FunctionType::get(B.getVoidTy(), {}, false),
                        GlobalValue::ExternalLinkage, "F", &M);
   B.SetInsertPoint(BasicBlock::Create(C, "", F));
 
   Type *Int8 = Type::getInt8Ty(C);
   Value *A = B.CreateAlloca(Int8, ConstantInt::get(Int8, 1), "A");
EXPECT_EQ(Walker.getClobberingMemoryAccess(LoadC),
             MSSA.getMemoryAccess(StoreC));
 
   // exercise block numbering
   EXPECT_TRUE(MSSA.locallyDominates(MSSA.getMemoryAccess(StoreC),
                                     MSSA.getMemoryAccess(StoreB)));
   EXPECT_TRUE(MSSA.locallyDominates(MSSA.getMemoryAccess(StoreA1),
                                     MSSA.getMemoryAccess(StoreA2)));
 }
+#endif