[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

2193

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.

include/llvm/Transforms/Utils/MemorySSA.h

 
 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();
+  }
 
 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.

lib/Transforms/Utils/MemorySSA.cpp

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() {

 
 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.");
 }

unittests/Transforms/Utils/MemorySSA.cpp

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);
 }