[MemorySSA] Fix bug in CachingMemorySSAWalker::invalidateInfo

[gberry]

Summary

CachingMemorySSAWalker::invalidateInfo was using IsCall to determine
which cache map needed to be cleared of entries referring to the invalidated
MemoryAccess, but there could also be entries referring to it in the
other cache map (value entries, not key entries). This change just
clears both tables to be conservatively correct.

Also add a verifyRemoved() function, called when expensive
checks (i.e. XDEBUG) are enabled to verify that the invalidated
MemoryAccess object is not referenced in any of the caches.

[george.burgess.iv]

LGTM with one nit. Thank you!

include/llvm/Transforms/Utils/MemorySSA.h

759 ↗

(On Diff #54570)

It might be useful to be able to call this e.g. directly from GDB, so can we relax this to if NDEBUG isn't defined? (To be clear, I agree that the automatic check on removal should happen iff XDEBUG is given)

[• dberlin]

Thank you for fixing this!

[• dberlin]

(It also occurs to me we can probably just use WeakVH since these are all
Value's now, and then check for null)

llvm/trunk/include/llvm/Transforms/Utils/MemorySSA.h

void doCacheRemove(const MemoryAccess *, const UpwardsMemoryQuery &,
                      const MemoryLocation &);
 
 private:
   MemoryAccessPair UpwardsDFSWalk(MemoryAccess *, const MemoryLocation &,
                                   UpwardsMemoryQuery &, bool);
   MemoryAccess *getClobberingMemoryAccess(MemoryAccess *, UpwardsMemoryQuery &);
   bool instructionClobbersQuery(const MemoryDef *, UpwardsMemoryQuery &,
                                 const MemoryLocation &Loc) const;
+  void verifyRemoved(MemoryAccess *);
   SmallDenseMap<ConstMemoryAccessPair, MemoryAccess *>
       CachedUpwardsClobberingAccess;
   DenseMap<const MemoryAccess *, MemoryAccess *> CachedUpwardsClobberingCall;
   AliasAnalysis *AA;
   DominatorTree *DT;
 };
 
 /// \brief Iterator base class used to implement const and non-const iterators

llvm/trunk/lib/Transforms/Utils/MemorySSA.cpp

void CachingMemorySSAWalker::invalidateInfo(MemoryAccess *MA) {
   if (MemoryUse *MU = dyn_cast<MemoryUse>(MA)) {
     UpwardsMemoryQuery Q;
     Instruction *I = MU->getMemoryInst();
     Q.IsCall = bool(ImmutableCallSite(I));
     Q.Inst = I;
     if (!Q.IsCall)
       Q.StartingLoc = MemoryLocation::get(I);
     doCacheRemove(MA, Q, Q.StartingLoc);
-    return;
-  }
+  } else {
     // If it is not a use, the best we can do right now is destroy the cache.
-  bool IsCall = false;
-
-  if (auto *MUD = dyn_cast<MemoryUseOrDef>(MA)) {
-    Instruction *I = MUD->getMemoryInst();
-    IsCall = bool(ImmutableCallSite(I));
-  }
-  if (IsCall)
     CachedUpwardsClobberingCall.clear();
-  else
     CachedUpwardsClobberingAccess.clear();
   }
 
+#ifdef XDEBUG
+  // Run this only when expensive checks are enabled.
+  verifyRemoved(MA);
+#endif
+}
+
 void CachingMemorySSAWalker::doCacheRemove(const MemoryAccess *M,
                                            const UpwardsMemoryQuery &Q,
                                            const MemoryLocation &Loc) {
   if (Q.IsCall)
     CachedUpwardsClobberingCall.erase(M);
   else
     CachedUpwardsClobberingAccess.erase({M, Loc});
 }
CachingMemorySSAWalker::getClobberingMemoryAccess(const Instruction *I) {
   DEBUG(dbgs() << "Starting Memory SSA clobber for " << *I << " is ");
   DEBUG(dbgs() << *DefiningAccess << "\n");
   DEBUG(dbgs() << "Final Memory SSA clobber for " << *I << " is ");
   DEBUG(dbgs() << *Result << "\n");
 
   return Result;
 }
 
+// Verify that MA doesn't exist in any of the caches.
+void CachingMemorySSAWalker::verifyRemoved(MemoryAccess *MA) {
+#ifndef NDEBUG
+  for (auto &P : CachedUpwardsClobberingAccess)
+    assert(P.first.first != MA && P.second != MA &&
+           "Found removed MemoryAccess in cache.");
+  for (auto &P : CachedUpwardsClobberingCall)
+    assert(P.first != MA && P.second != MA &&
+           "Found removed MemoryAccess in cache.");
+#endif // !NDEBUG
+}
+
 MemoryAccess *
 DoNothingMemorySSAWalker::getClobberingMemoryAccess(const Instruction *I) {
   MemoryAccess *MA = MSSA->getMemoryAccess(I);
   if (auto *Use = dyn_cast<MemoryUseOrDef>(MA))
     return Use->getDefiningAccess();
   return MA;
 }
 
 MemoryAccess *DoNothingMemorySSAWalker::getClobberingMemoryAccess(
     MemoryAccess *StartingAccess, MemoryLocation &) {
   if (auto *Use = dyn_cast<MemoryUseOrDef>(StartingAccess))
     return Use->getDefiningAccess();
   return StartingAccess;
 }
 }