MemorySSA: Add support for caching clobbering access in stores

[• dberlin]

Summary

This enables us to cache the clobbering access for stores, despite the
fact that we can't rewrite the use-def chains themselves.

Early testing shows that, after this change, for larger testcases, it will be a significant net positive (memory and time) to remove the walker caching.

[davide]

LG

lib/Transforms/Utils/MemorySSA.cpp

2202 ↗

(On Diff #93743)

auto. I know it was already there, but while you're at it :)

2220 ↗

(On Diff #93743)

auto

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

public:
 
   /// \brief Get the access that produces the memory state used by this Use.
   MemoryAccess *getDefiningAccess() const { return getOperand(0); }
 
   static inline bool classof(const Value *MA) {
     return MA->getValueID() == MemoryUseVal || MA->getValueID() == MemoryDefVal;
   }
 
+  // Sadly, these have to be public because they are needed in some of the iterators.
+  virtual bool isOptimized() const = 0;
+  virtual MemoryAccess *getOptimized() const = 0;
+  virtual void setOptimized(MemoryAccess *) = 0;
+
+  /// \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.
+  virtual void resetOptimized() = 0;
+
 protected:
   friend class MemorySSA;
   friend class MemorySSAUpdater;
   MemoryUseOrDef(LLVMContext &C, MemoryAccess *DMA, unsigned Vty,
                  Instruction *MI, BasicBlock *BB)
       : MemoryAccess(C, Vty, BB, 1), MemoryInst(MI) {
     setDefiningAccess(DMA);
   }
-
-  void setDefiningAccess(MemoryAccess *DMA) { setOperand(0, DMA); }
+  void setDefiningAccess(MemoryAccess *DMA, bool Optimized = false) {
+    setOperand(0, DMA);
+    if (!Optimized)
+      return;
+    setOptimized(DMA);
+  }
 
 private:
   Instruction *MemoryInst;
 };
 
 template <>
 struct OperandTraits<MemoryUseOrDef>
     : public FixedNumOperandTraits<MemoryUseOrDef, 1> {};
public:
   void *operator new(size_t, unsigned) = delete;
 
   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);
+  virtual void setOptimized(MemoryAccess *DMA) override {
   }
 
-  bool isOptimized() const {
+  virtual bool isOptimized() const override {
     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; }
+  virtual MemoryAccess *getOptimized() const override {
+    return getDefiningAccess();
+  }
+  virtual void resetOptimized() override { OptimizedID = INVALID_MEMORYACCESS_ID; }
 
 protected:
   friend class MemorySSA;
 
   unsigned getID() const override {
     llvm_unreachable("MemoryUses do not have IDs");
   }
 

 /// associated with them. This use points to the nearest reaching
 /// MemoryDef/MemoryPhi.
 class MemoryDef final : public MemoryUseOrDef {
 public:
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(MemoryAccess);
 
   MemoryDef(LLVMContext &C, MemoryAccess *DMA, Instruction *MI, BasicBlock *BB,
             unsigned Ver)
-      : MemoryUseOrDef(C, DMA, MemoryDefVal, MI, BB), ID(Ver) {}
+      : MemoryUseOrDef(C, DMA, MemoryDefVal, MI, BB), ID(Ver),
+        Optimized(nullptr), OptimizedID(INVALID_MEMORYACCESS_ID) {}
 
   // allocate space for exactly one operand
   void *operator new(size_t s) { return User::operator new(s, 1); }
   void *operator new(size_t, unsigned) = delete;
 
   static inline bool classof(const Value *MA) {
     return MA->getValueID() == MemoryDefVal;
   }
 
+  virtual void setOptimized(MemoryAccess *MA) override {
+    Optimized = MA;
+    OptimizedID = getDefiningAccess()->getID();
+  }
+  virtual MemoryAccess *getOptimized() const override { return Optimized; }
+  virtual bool isOptimized() const override {
+    return getOptimized() && OptimizedID == getDefiningAccess()->getID();
+  }
+  virtual void resetOptimized() override { OptimizedID = INVALID_MEMORYACCESS_ID; }
+
+
   void print(raw_ostream &OS) const override;
 
 protected:
   friend class MemorySSA;
 
   unsigned getID() const override { return ID; }
 
 private:
   const unsigned ID;
+  MemoryAccess *Optimized;
+  unsigned int OptimizedID;
 };
 
 template <>
 struct OperandTraits<MemoryDef> : public FixedNumOperandTraits<MemoryDef, 1> {};
 DEFINE_TRANSPARENT_OPERAND_ACCESSORS(MemoryDef, MemoryAccess)
 
 /// \brief Represents phi nodes for memory accesses.
 ///
: public iterator_facade_base<def_chain_iterator<T>,
                                   std::forward_iterator_tag, MemoryAccess *> {
   def_chain_iterator() : MA(nullptr) {}
   def_chain_iterator(T MA) : MA(MA) {}
 
   T operator*() const { return MA; }
 
   def_chain_iterator &operator++() {
     // N.B. liveOnEntry has a null defining access.
-    if (auto *MUD = dyn_cast<MemoryUseOrDef>(MA))
-      MA = MUD->getDefiningAccess();
+    if (auto *MUD = dyn_cast<MemoryUseOrDef>(MA)) {
+      if (MUD->isOptimized())
+        MA = MUD->getOptimized();
       else
+        MA = MUD->getDefiningAccess();
+    } else {
       MA = nullptr;
+    }
+
     return *this;
   }
 
   bool operator==(const def_chain_iterator &O) const { return MA == O.MA; }
 
 private:
   T MA;
 };

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

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();
+  if (auto *MUD = dyn_cast<MemoryUseOrDef>(StartingAccess))
+    if (MUD->isOptimized())
+      return MUD->getOptimized();
 
   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);
+    if (auto *MUD = dyn_cast<MemoryUseOrDef>(StartingAccess))
+      MUD->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);
+  if (auto *MUD = dyn_cast<MemoryUseOrDef>(StartingAccess))
+    MUD->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/lib/Transforms/Utils/MemorySSAUpdater.cpp

if (!isa<MemoryUse>(MA) && !MA->use_empty()) {
     // 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();
+      if (auto *MUD = dyn_cast<MemoryUseOrDef>(U.getUser()))
+        MUD->resetOptimized();
       U.set(NewDefTarget);
     }
   }
 
   // The call below to erase will destroy MA, so we can't change the order we
   // are doing things here
   MSSA->removeFromLookups(MA);
   MSSA->removeFromLists(MA);