[MSSA] Initial version of avoiding unnecessary use walks when calling getClobberingMemoryAccess

[• dberlin]

Summary

This allows us to mark when uses have been optimized.
This lets us avoid rewalking (IE when people call getClobberingAccess on everything), and also
enables us to later relax the requirement of use optimization during updates with less cost.

[• dberlin]

(Unittest that resets the defining access, and calls getClobberingAccess, coming)

[george.burgess.iv]

lib/Transforms/Utils/MemorySSA.cpp

2185 ↗

(On Diff #73232)

RE: Note: ...: can you bring me through why having this on a Def is would be useful? Is the cache just invalidated too aggressively?

[• dberlin]

• Move optimized flag to uses
• Update to work in all update cases

[george.burgess.iv]

LGTM, thanks!

[• dberlin]

I know you LGTM'd this, but i updated it t#o show the final state. The
only case we don't handle automatically right now is if removing
memory accesses causes single valued phis.

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

 namespace llvm {
 
 class DominatorTree;
 class Function;
 class Instruction;
 class MemoryAccess;
 class LLVMContext;
 class raw_ostream;
+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.

 
 protected:
   friend class MemorySSA;
   friend class MemoryUseOrDef;
   friend class MemoryUse;
   friend class MemoryDef;
   friend class MemoryPhi;
 
-  /// \brief Used internally to give IDs to MemoryAccesses for printing
+  /// \brief Used for debugging and tracking things about MemoryAccesses.
+  /// Guaranteed unique among MemoryAccesses, no guarantees otherwise.
   virtual unsigned getID() const = 0;
 
   MemoryAccess(LLVMContext &C, unsigned Vty, BasicBlock *BB,
                unsigned NumOperands)
       : User(Type::getVoidTy(C), Vty, nullptr, NumOperands), Block(BB) {}
 
 private:
   MemoryAccess(const MemoryAccess &);

 
 public:
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(MemoryAccess);
 
   // allocate space for exactly one operand
   void *operator new(size_t s) { return User::operator new(s, 1); }
 
   MemoryUse(LLVMContext &C, MemoryAccess *DMA, Instruction *MI, BasicBlock *BB)
-      : MemoryUseOrDef(C, DMA, MemoryUseVal, MI, BB) {}
+      : MemoryUseOrDef(C, DMA, MemoryUseVal, MI, BB), OptimizedID(0) {}
 
   static inline bool classof(const MemoryUse *) { return true; }
   static inline bool classof(const Value *MA) {
     return MA->getValueID() == MemoryUseVal;
   }
 
   void print(raw_ostream &OS) const override;
+  void setDefiningAccess(MemoryAccess *DMA, bool Optimized = false) {
+    if (Optimized)
+      OptimizedID = DMA->getID();
+    MemoryUseOrDef::setDefiningAccess(DMA);
+  }
+  bool isOptimized() const {
+    return getDefiningAccess() && OptimizedID == getDefiningAccess()->getID();
+  }
+  /// \brief Reset the ID of what this MemoryUse was optimized to, causing it to
+  /// be rewalked by the walker if necessary.
+  /// This really should only be called by tests.
+  void resetOptimized() { OptimizedID = INVALID_MEMORYACCESS_ID; }
 
 protected:
   friend class MemorySSA;
 
   unsigned getID() const override {
     llvm_unreachable("MemoryUses do not have IDs");
   }
+
+private:
+  unsigned int OptimizedID;
 };
 
 template <>
 struct OperandTraits<MemoryUse> : public FixedNumOperandTraits<MemoryUse, 1> {};
 DEFINE_TRANSPARENT_OPERAND_ACCESSORS(MemoryUse, MemoryAccess)
 
 /// \brief Represents a read-write access to memory, whether it is a must-alias,
 /// or a may-alias.
static inline bool classof(const Value *MA) {
     return MA->getValueID() == MemoryDefVal;
   }
 
   void print(raw_ostream &OS) const override;
 
 protected:
   friend class MemorySSA;
 
-  // For debugging only. This gets used to give memory accesses pretty numbers
-  // when printing them out
   unsigned getID() const override { return ID; }
 
 private:
   const unsigned ID;
 };
 
 template <>
 struct OperandTraits<MemoryDef> : public FixedNumOperandTraits<MemoryDef, 1> {};
protected:
   friend class MemorySSA;
   /// \brief this is more complicated than the generic
   /// User::allocHungoffUses, because we have to allocate Uses for the incoming
   /// values and pointers to the incoming blocks, all in one allocation.
   void allocHungoffUses(unsigned N) {
     User::allocHungoffUses(N, /* IsPhi */ true);
   }
 
-  /// For debugging only. This gets used to give memory accesses pretty numbers
-  /// when printing them out
   unsigned getID() const final { return ID; }
 
 private:
   // For debugging only
   const unsigned ID;
   unsigned ReservedSpace;
 
   /// \brief This grows the operand list in response to a push_back style of

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

for (auto AI = Accesses->begin(), AE = Accesses->end(); AI != AE;) {
     else
       Accesses->erase(AI);
     AI = Next;
   }
 }
 
 MemorySSA::MemorySSA(Function &Func, AliasAnalysis *AA, DominatorTree *DT)
     : AA(AA), DT(DT), F(Func), LiveOnEntryDef(nullptr), Walker(nullptr),
-      NextID(0) {
+      NextID(INVALID_MEMORYACCESS_ID) {
   buildMemorySSA();
 }
 
 MemorySSA::~MemorySSA() {
   // Drop all our references
   for (const auto &Pair : PerBlockAccesses)
     for (MemoryAccess &MA : *Pair.second)
       MA.dropAllReferences();
for (MemoryAccess &MA : *Accesses) {
     auto *MU = dyn_cast<MemoryUse>(&MA);
     if (!MU) {
       VersionStack.push_back(&MA);
       ++StackEpoch;
       continue;
     }
 
     if (isUseTriviallyOptimizableToLiveOnEntry(*AA, MU->getMemoryInst())) {
-      MU->setDefiningAccess(MSSA->getLiveOnEntryDef());
+      MU->setDefiningAccess(MSSA->getLiveOnEntryDef(), true);
       continue;
     }
 
     MemoryLocOrCall UseMLOC(MU);
     auto &LocInfo = LocStackInfo[UseMLOC];
     // If the pop epoch changed, it means we've removed stuff from top of
     // stack due to changing blocks. We may have to reset the lower bound or
     // last kill info.
while (UpperBound > LocInfo.LowerBound) {
         FoundClobberResult = true;
         break;
       }
       --UpperBound;
     }
     // At the end of this loop, UpperBound is either a clobber, or lower bound
     // PHI walking may cause it to be < LowerBound, and in fact, < LastKill.
     if (FoundClobberResult || UpperBound < LocInfo.LastKill) {
-      MU->setDefiningAccess(VersionStack[UpperBound]);
+      MU->setDefiningAccess(VersionStack[UpperBound], true);
       // We were last killed now by where we got to
       LocInfo.LastKill = UpperBound;
     } else {
       // Otherwise, we checked all the new ones, and now we know we can get to
       // LastKill.
-      MU->setDefiningAccess(VersionStack[LocInfo.LastKill]);
+      MU->setDefiningAccess(VersionStack[LocInfo.LastKill], true);
     }
     LocInfo.LowerBound = VersionStack.size() - 1;
     LocInfo.LowerBoundBlock = BB;
   }
 }
 
 /// Optimize uses to point to their actual clobbering definitions.
 void MemorySSA::OptimizeUses::optimizeUses() {
if (MemoryPhi *MP = dyn_cast<MemoryPhi>(MA)) {
     NewDefTarget = onlySingleValue(MP);
     assert((NewDefTarget || MP->use_empty()) &&
            "We can't delete this memory phi");
   } else {
     NewDefTarget = cast<MemoryUseOrDef>(MA)->getDefiningAccess();
   }
 
   // Re-point the uses at our defining access
-  if (!MA->use_empty())
-    MA->replaceAllUsesWith(NewDefTarget);
+  if (!MA->use_empty()) {
+    // Reset optimized on users of this store, and reset the uses.
+    // A few notes:
+    // 1. This is a slightly modified version of RAUW to avoid walking the
+    // uses twice here.
+    // 2. If we wanted to be complete, we would have to reset the optimized
+    // flags on users of phi nodes if doing the below makes a phi node have all
+    // the same arguments. Instead, we prefer users to removeMemoryAccess those
+    // phi nodes, because doing it here would be N^3.
+    if (MA->hasValueHandle())
+      ValueHandleBase::ValueIsRAUWd(MA, NewDefTarget);
+    // Note: We assume MemorySSA is not used in metadata since it's not really
+    // part of the IR.
+
+    while (!MA->use_empty()) {
+      Use &U = *MA->use_begin();
+      if (MemoryUse *MU = dyn_cast<MemoryUse>(U.getUser()))
+        MU->resetOptimized();
+      U.set(NewDefTarget);
+    }
+  }
 
   // The call below to erase will destroy MA, so we can't change the order we
   // are doing things here
   removeFromLookups(MA);
 }
 
 void MemorySSA::print(raw_ostream &OS) const {
   MemorySSAAnnotatedWriter Writer(this);
void MemorySSA::CachingWalker::invalidateInfo(MemoryAccess *MA) {
   // See if this is a MemoryUse, if so, just remove the cached info. MemoryUse
   // is by definition never a barrier, so nothing in the cache could point to
   // this use. In that case, we only need invalidate the info for the use
   // itself.
 
   if (MemoryUse *MU = dyn_cast<MemoryUse>(MA)) {
     UpwardsMemoryQuery Q(MU->getMemoryInst(), MU);
     Cache.remove(MU, Q.StartingLoc, Q.IsCall);
+    MU->resetOptimized();
   } else {
     // If it is not a use, the best we can do right now is destroy the cache.
     Cache.clear();
   }
 
 #ifdef EXPENSIVE_CHECKS
   verifyRemoved(MA);
 #endif

 
 MemoryAccess *
 MemorySSA::CachingWalker::getClobberingMemoryAccess(MemoryAccess *MA) {
   auto *StartingAccess = dyn_cast<MemoryUseOrDef>(MA);
   // If this is a MemoryPhi, we can't do anything.
   if (!StartingAccess)
     return MA;
 
+  // If this is an already optimized use or def, return the optimized result.
+  // Note: Currently, we do not store the optimized def result because we'd need
+  // a separate field, since we can't use it as the defining access.
+  if (MemoryUse *MU = dyn_cast<MemoryUse>(StartingAccess))
+    if (MU->isOptimized())
+      return MU->getDefiningAccess();
+
   const Instruction *I = StartingAccess->getMemoryInst();
   UpwardsMemoryQuery Q(I, StartingAccess);
   // We can't sanely do anything with a fences, they conservatively
   // clobber all memory, and have no locations to get pointers from to
   // try to disambiguate.
   if (!Q.IsCall && I->isFenceLike())
     return StartingAccess;
 
   if (auto *CacheResult = Cache.lookup(StartingAccess, Q.StartingLoc, Q.IsCall))
     return CacheResult;
 
   if (isUseTriviallyOptimizableToLiveOnEntry(*MSSA->AA, I)) {
     MemoryAccess *LiveOnEntry = MSSA->getLiveOnEntryDef();
     Cache.insert(StartingAccess, LiveOnEntry, Q.StartingLoc, Q.IsCall);
+    if (MemoryUse *MU = dyn_cast<MemoryUse>(StartingAccess))
+      MU->setDefiningAccess(LiveOnEntry, true);
     return LiveOnEntry;
   }
 
   // Start with the thing we already think clobbers this location
   MemoryAccess *DefiningAccess = StartingAccess->getDefiningAccess();
 
   // At this point, DefiningAccess may be the live on entry def.
   // If it is, we will not get a better result.
   if (MSSA->isLiveOnEntryDef(DefiningAccess))
     return DefiningAccess;
 
   MemoryAccess *Result = getClobberingMemoryAccess(DefiningAccess, Q);
   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");
+  if (MemoryUse *MU = dyn_cast<MemoryUse>(StartingAccess))
+    MU->setDefiningAccess(Result, true);
 
   return Result;
 }
 
 // Verify that MA doesn't exist in any of the caches.
 void MemorySSA::CachingWalker::verifyRemoved(MemoryAccess *MA) {
   assert(!Cache.contains(MA) && "Found removed MemoryAccess in cache.");
 }

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

TEST_F(MemorySSATest, RemoveMemoryAccess) {
   MSSA.removeMemoryAccess(StoreAccess);
   MSSA.verifyMemorySSA();
   // After the removeaccess, let's see if we got the right accesses
   // The load should still point to the phi ...
   EXPECT_EQ(DefiningAccess, LoadAccess->getDefiningAccess());
   // but we should now get live on entry for the clobbering definition of the
   // load, since it will walk past the phi node since every argument is the
   // same.
+  // XXX: This currently requires either removing the phi or resetting optimized
+  // on the load
+
+  EXPECT_FALSE(
+      MSSA.isLiveOnEntryDef(Walker->getClobberingMemoryAccess(LoadInst)));
+  // If we reset optimized, we get live on entry.
+  LoadAccess->resetOptimized();
   EXPECT_TRUE(
       MSSA.isLiveOnEntryDef(Walker->getClobberingMemoryAccess(LoadInst)));
-
   // The phi should now be a two entry phi with two live on entry defs.
   for (const auto &Op : DefiningAccess->operands()) {
     MemoryAccess *Operand = cast<MemoryAccess>(&*Op);
     EXPECT_TRUE(MSSA.isLiveOnEntryDef(Operand));
   }
 
   // Now we try to remove the single valued phi
   MSSA.removeMemoryAccess(DefiningAccess);
TEST_F(MemorySSATest, WalkerInvariantLoadOpt) {
   // flexible to future changes.
   Walker->invalidateInfo(LoadMA);
   Load->setMetadata(LLVMContext::MD_invariant_load, MDNode::get(C, {}));
 
   MemoryAccess *LoadClobber = Walker->getClobberingMemoryAccess(LoadMA);
   EXPECT_EQ(LoadClobber, MSSA.getLiveOnEntryDef());
 }
 
-// At one point, we were building MSSA with 0 AA passes. This ensures that we
-// actually use BasicAA.
-TEST_F(MemorySSATest, AAIsPresentAtBuildTime) {
+// Test loads get reoptimized properly by the walker.
+TEST_F(MemorySSATest, WalkerReopt) {
   F = Function::Create(FunctionType::get(B.getVoidTy(), {}, false),
                        GlobalValue::ExternalLinkage, "F", &M);
   B.SetInsertPoint(BasicBlock::Create(C, "", F));
-
   Type *Int8 = Type::getInt8Ty(C);
-  Constant *One = ConstantInt::get(Int8, 1);
-  Value *AllocaA = B.CreateAlloca(Int8);
-  Instruction *StoreA = B.CreateStore(One, AllocaA);
-
-  Value *AllocaB = B.CreateAlloca(Int8);
-  B.CreateStore(One, AllocaB);
-  Instruction *LoadA = B.CreateLoad(AllocaA);
+  Value *AllocaA = B.CreateAlloca(Int8, ConstantInt::get(Int8, 1), "A");
+  Instruction *SIA = B.CreateStore(ConstantInt::get(Int8, 0), AllocaA);
+  Value *AllocaB = B.CreateAlloca(Int8, ConstantInt::get(Int8, 1), "B");
+  Instruction *SIB = B.CreateStore(ConstantInt::get(Int8, 0), AllocaB);
+  Instruction *LIA = B.CreateLoad(AllocaA);
 
   setupAnalyses();
   MemorySSA &MSSA = *Analyses->MSSA;
-  auto *MU = cast<MemoryUse>(MSSA.getMemoryAccess(LoadA));
-  EXPECT_EQ(MU->getDefiningAccess(), MSSA.getMemoryAccess(StoreA));
+  MemorySSAWalker *Walker = Analyses->Walker;
+
+  MemoryAccess *LoadClobber = Walker->getClobberingMemoryAccess(LIA);
+  MemoryUse *LoadAccess = cast<MemoryUse>(MSSA.getMemoryAccess(LIA));
+  EXPECT_EQ(LoadClobber, MSSA.getMemoryAccess(SIA));
+  EXPECT_TRUE(MSSA.isLiveOnEntryDef(Walker->getClobberingMemoryAccess(SIA)));
+  MSSA.removeMemoryAccess(LoadAccess);
+
+  // 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);
 }