Make MemorySSA::dominates/locallydominates constant time

[• dberlin]

Summary

Make MemorySSA::dominates/locallydominates constant time

[• dberlin]

• Update comment

[george.burgess.iv]

LGTM with a few comments; thanks for the patch!

include/llvm/Transforms/Utils/MemorySSA.h

635

Given that we only ever store true for the value, can this be a DenseSet instead?

lib/Transforms/Utils/MemorySSA.cpp

1205

Do we need to invalidate numbering here, as well? (Alternatively, we could start all counts from 2, and always give a newly-created phi a value of 1, but that could also be done later as a part of the "number with a stride of N" optimization)

1582

Nit: Please remove braces

[• dberlin]

include/llvm/Transforms/Utils/MemorySSA.h

635

Replaced with SmallPtrSet

lib/Transforms/Utils/MemorySSA.cpp

1205

Fixed.
We can actually just answer PHI queries without ever numbering, them ever.
There is only ever one phi. It always must be first :)

so after the equality/etc checks in dominates, you could do

if (isa<MemoryPhi>(Dominator))
  return true;
if (isa<MemoryPhi>(Dominatee))
  return false;

But we can do this in a followup if it matters.

include/llvm/Transforms/Utils/MemorySSA.h

private:
   MemoryUseOrDef *createDefinedAccess(Instruction *, MemoryAccess *);
   MemoryAccess *findDominatingDef(BasicBlock *, enum InsertionPlace);
   void removeFromLookups(MemoryAccess *);
 
   MemoryAccess *renameBlock(BasicBlock *, MemoryAccess *);
   void renamePass(DomTreeNode *, MemoryAccess *IncomingVal,
                   SmallPtrSet<BasicBlock *, 16> &Visited);
   AccessList *getOrCreateAccessList(const BasicBlock *);
+  void renumberBlock(const BasicBlock *) const;
+
   AliasAnalysis *AA;
   DominatorTree *DT;
   Function &F;
 
   // Memory SSA mappings
   DenseMap<const Value *, MemoryAccess *> ValueToMemoryAccess;
   AccessMap PerBlockAccesses;
   std::unique_ptr<MemoryAccess> LiveOnEntryDef;
 
+  // Domination mappings
+  // Note that the numbering is local to a block, even though the map is
+  // global.
+  mutable DenseMap<const BasicBlock *, bool> BlockNumberingValid;
+  mutable DenseMap<const MemoryAccess *, unsigned long> BlockNumbering;
+
   // Memory SSA building info
   std::unique_ptr<CachingWalker> Walker;
   unsigned NextID;
 };
 
 // This pass does eager building and then printing of MemorySSA. It is used by
 // the tests to be able to build, dump, and verify Memory SSA.
 class MemorySSAPrinterLegacyPass : public FunctionPass {

lib/Transforms/Utils/MemorySSA.cpp

class ClobberWalker {
   bool shouldIgnoreCache() const {
     // UseCache will only be false when we're debugging, or when expensive
     // checks are enabled. In either case, we don't care deeply about speed.
     return LLVM_UNLIKELY(!UseCache);
   }
 
   void addCacheEntry(const MemoryAccess *What, MemoryAccess *To,
                      const MemoryLocation &Loc) const {
-// EXPENSIVE_CHECKS because most of these queries are redundant, and if What
-// and To are in the same BB, that gives us n^2 behavior.
+// EXPENSIVE_CHECKS because most of these queries are redundant.
 #ifdef EXPENSIVE_CHECKS
     assert(MSSA.dominates(To, What));
 #endif
     if (shouldIgnoreCache())
       return;
     WC.insert(What, To, Loc, Query->IsCall);
   }
 
OptznResult tryOptimizePhi(MemoryPhi *Phi, MemoryAccess *Start,
     SmallVector<TerminatedPath, 4> TerminatedPaths;
 
     addSearches(Phi, PausedSearches, 0);
 
     // Moves the TerminatedPath with the "most dominated" Clobber to the end of
     // Paths.
     auto MoveDominatedPathToEnd = [&](SmallVectorImpl<TerminatedPath> &Paths) {
       assert(!Paths.empty() && "Need a path to move");
-      // FIXME: This is technically n^2 (n = distance(DefPath.First,
-      // DefPath.Last)) because of local dominance checks.
       auto Dom = Paths.begin();
       for (auto I = std::next(Dom), E = Paths.end(); I != E; ++I)
         if (!MSSA.dominates(I->Clobber, Dom->Clobber))
           Dom = I;
       auto Last = Paths.end() - 1;
       if (Last != Dom)
         std::iter_swap(Last, Dom);
     };

 MemorySSA::CachingWalker *MemorySSA::getWalkerImpl() {
   if (Walker)
     return Walker.get();
 
   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);
   MemoryPhi *Phi = new MemoryPhi(BB->getContext(), BB, NextID++);
   ValueToMemoryAccess.insert(std::make_pair(BB, Phi));
   // Phi's always are placed at the front of the block.
   Accesses->push_front(Phi);
   return Phi;
 }
if (Point == Beginning) {
     auto AI = std::find_if(
         Accesses->begin(), Accesses->end(),
         [](const MemoryAccess &MA) { return !isa<MemoryPhi>(MA); });
 
     Accesses->insert(AI, NewAccess);
   } else {
     Accesses->push_back(NewAccess);
   }
-
+  BlockNumberingValid.erase(BB);
   return NewAccess;
 }
 MemoryAccess *MemorySSA::createMemoryAccessBefore(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->insert(AccessList::iterator(InsertPt), NewAccess);
+  BlockNumberingValid.erase(InsertPt->getBlock());
   return NewAccess;
 }
 
 MemoryAccess *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);
+  BlockNumberingValid.erase(InsertPt->getBlock());
   return NewAccess;
 }
 
 /// \brief Helper function to create new memory accesses
 MemoryUseOrDef *MemorySSA::createNewAccess(Instruction *I) {
   // The assume intrinsic has a control dependency which we model by claiming
   // that it writes arbitrarily. Ignore that fake memory dependency here.
   // FIXME: Replace this special casing with a more accurate modelling of

 
 /// \brief Properly remove \p MA from all of MemorySSA's lookup tables.
 ///
 /// Because of the way the intrusive list and use lists work, it is important to
 /// do removal in the right order.
 void MemorySSA::removeFromLookups(MemoryAccess *MA) {
   assert(MA->use_empty() &&
          "Trying to remove memory access that still has uses");
+  BlockNumbering.erase(MA);
   if (MemoryUseOrDef *MUD = dyn_cast<MemoryUseOrDef>(MA))
     MUD->setDefiningAccess(nullptr);
   // Invalidate our walker's cache if necessary
   if (!isa<MemoryUse>(MA))
     Walker->invalidateInfo(MA);
   // The call below to erase will destroy MA, so we can't change the order we
   // are doing things here
   Value *MemoryInst;

 MemoryAccess *MemorySSA::getMemoryAccess(const Value *I) const {
   return ValueToMemoryAccess.lookup(I);
 }
 
 MemoryPhi *MemorySSA::getMemoryAccess(const BasicBlock *BB) const {
   return cast_or_null<MemoryPhi>(getMemoryAccess((const Value *)BB));
 }
 
+/// Perform a local numbering on blocks so that instruction ordering can be
+/// determined in constant time.
+/// TODO: We currently just number in order.  If we numbered by N, we could
+/// allow N-1 sequences of insertBefore/insertAfter without needing to update
+/// the numbering.
+void MemorySSA::renumberBlock(const BasicBlock *B) const {
+  // The pre-increment ensures the numbers really start at 1.
+  unsigned long CurrentNumber = 0;
+  const AccessList *AL = getBlockAccesses(B);
+  assert(AL != nullptr && "Asking to renumber an empty block");
+  for (const auto &I : *AL) {
+    BlockNumbering[&I] = ++CurrentNumber;
+  }
+  BlockNumberingValid[B] = true;
+}
+
 /// \brief Determine, for two memory accesses in the same block,
 /// whether \p Dominator dominates \p Dominatee.
 /// \returns True if \p Dominator dominates \p Dominatee.
 bool MemorySSA::locallyDominates(const MemoryAccess *Dominator,
                                  const MemoryAccess *Dominatee) const {
-  assert((Dominator->getBlock() == Dominatee->getBlock()) &&
-         "Asking for local domination when accesses are in different blocks!");
 
+  const BasicBlock *DominatorBlock = Dominator->getBlock();
+  const BasicBlock *DominateeBlock = Dominatee->getBlock();
+
+  assert((DominatorBlock == DominateeBlock) &&
+         "Asking for local domination when accesses are in different blocks!");
   // A node dominates itself.
   if (Dominatee == Dominator)
     return true;
 
   // When Dominatee is defined on function entry, it is not dominated by another
   // memory access.
   if (isLiveOnEntryDef(Dominatee))
     return false;
 
   // When Dominator is defined on function entry, it dominates the other memory
   // access.
   if (isLiveOnEntryDef(Dominator))
     return true;
 
-  // Get the access list for the block
-  const AccessList *AccessList = getBlockAccesses(Dominator->getBlock());
-  AccessList::const_reverse_iterator It(Dominator->getIterator());
-
-  // If we hit the beginning of the access list before we hit dominatee, we must
-  // dominate it
-  return std::none_of(It, AccessList->rend(),
-                      [&](const MemoryAccess &MA) { return &MA == Dominatee; });
+  if (!BlockNumberingValid.lookup(DominatorBlock))
+    renumberBlock(DominatorBlock);
+
+  unsigned long DominatorNum = BlockNumbering.lookup(Dominator);
+  // All numbers start with 1
+  assert(DominatorNum != 0 && "Block was not numbered properly");
+  unsigned long DominateeNum = BlockNumbering.lookup(Dominatee);
+  assert(DominateeNum != 0 && "Block was not numbered properly");
+  return DominatorNum < DominateeNum;
 }
 
 bool MemorySSA::dominates(const MemoryAccess *Dominator,
                           const MemoryAccess *Dominatee) const {
   if (Dominator == Dominatee)
     return true;
 
   if (isLiveOnEntryDef(Dominatee))

 void MemorySSAWrapperPass::print(raw_ostream &OS, const Module *M) const {
   MSSA->print(OS);
 }
 
 MemorySSAWalker::MemorySSAWalker(MemorySSA *M) : MSSA(M) {}
 
 MemorySSA::CachingWalker::CachingWalker(MemorySSA *M, AliasAnalysis *A,
                                         DominatorTree *D)
-    : MemorySSAWalker(M), Walker(*M, *A, *D, Cache),
-      AutoResetWalker(true) {}
+    : MemorySSAWalker(M), Walker(*M, *A, *D, Cache), AutoResetWalker(true) {}
 
 MemorySSA::CachingWalker::~CachingWalker() {}
 
 void MemorySSA::CachingWalker::invalidateInfo(MemoryAccess *MA) {
   // TODO: We can do much better cache invalidation with differently stored
   // caches.  For now, for MemoryUses, we simply remove them
   // from the cache, and kill the entire call/non-call cache for everything
   // else.  The problem is for phis or defs, currently we'd need to follow use