MemorySSA: Add support for renaming uses in the updater.

[• dberlin]

Summary

This lets one add aliasing stores to the updater.
(i'm next going to move the creation/etc functions to the updater)

[• dberlin]

• Remove errant MSSA.dump added to unit test

[davide]

The code logic looks fine. Some minor comments inline.

include/llvm/Transforms/Utils/MemorySSAUpdater.h

89–93

Unrelated formatting, I guess.

lib/Transforms/Utils/MemorySSA.cpp

1187–1198

Full stop at the ned of the comment (here and everywhere else)

unittests/Transforms/Utils/MemorySSA.cpp

197

IIRC dump() is supposed to be used only in a debugger and print() should be used instead, but I might be wrong.

[davide]

Oh, just realized I commented while you're updating. Oh well.

[davide]

lib/Transforms/Utils/MemorySSA.cpp

1914–1915

unrelated formatting I think (also I think the #if is redundant if you mark dump() as LLVM_DUMP_METHOD?)

[davide]

LGTM

[Prazek]

include/llvm/Transforms/Utils/MemorySSA.h

692

why 2 last params are not used?

unittests/Transforms/Utils/MemorySSA.cpp

186–190

Shouldn't it have some EXPECTs checking that renaming happened?

[george.burgess.iv]

handful o'nits to supplement piotr's comments, then this LGTM.

thanks!

include/llvm/Transforms/Utils/MemorySSAUpdater.h

66–83

should this be ///?

lib/Transforms/Utils/MemorySSA.cpp

1188

missing space between "multiple blocks"

1191

please sink this into the if condition

1195

please remove redundant parens

lib/Transforms/Utils/MemorySSAUpdater.cpp

299

please do if (auto ... = dyn_cast<MemoryDef>(FirstDef)) instead

unittests/Transforms/Utils/MemorySSA.cpp

197

please remove

[• dberlin]

include/llvm/Transforms/Utils/MemorySSA.h

692

There is only one caller that should be using it.
I removed the params here.

include/llvm/Transforms/Utils/MemorySSAUpdater.h

66–83

My understanding is that this isn't required anymore, but maybe i'm wrong?

[• dberlin]

• Update for review comments

[george.burgess.iv]

LGTM. Thanks again!

include/llvm/Transforms/Utils/MemorySSAUpdater.h

66–83

no clue. I'm happy to assume you're right; we can fix it later if doxygen doesn't like it. ¯\_(ツ)_/¯

[Prazek]

include/llvm/Transforms/Utils/MemorySSAUpdater.h

66–83

I am not sure if doxygen will work, but I am used to doxygen comments for definitions because my ide color them differently :p

include/llvm/Transforms/Utils/MemorySSA.h

DefsList *getWritableBlockDefs(const BasicBlock *BB) const {
     return It == PerBlockDefs.end() ? nullptr : It->second.get();
   }
 
   // This is used by the updater to perform the internal memoryssa machinations
   // for moves.  It does not always leave the IR in a correct state, and relies
   // on the updater to fixup what it breaks, so it is not public.
   void moveTo(MemoryUseOrDef *What, BasicBlock *BB, AccessList::iterator Where);
   void moveTo(MemoryUseOrDef *What, BasicBlock *BB, InsertionPlace Point);
+  // Rename the dominator tree branch rooted at BB.
+  void renamePass(BasicBlock *BB, MemoryAccess *IncomingVal,
+                  SmallPtrSet<BasicBlock *, 16> &Visited) {
+    renamePass(DT->getNode(BB), IncomingVal, Visited, true, true);
+  }
 
 private:
   class CachingWalker;
   class OptimizeUses;
 
   CachingWalker *getWalkerImpl();
   void buildMemorySSA();
   void optimizeUses();
 
   void verifyUseInDefs(MemoryAccess *, MemoryAccess *) const;
   using AccessMap = DenseMap<const BasicBlock *, std::unique_ptr<AccessList>>;
   using DefsMap = DenseMap<const BasicBlock *, std::unique_ptr<DefsList>>;
 
   void
   determineInsertionPoint(const SmallPtrSetImpl<BasicBlock *> &DefiningBlocks);
   void computeDomLevels(DenseMap<DomTreeNode *, unsigned> &DomLevels);
   void markUnreachableAsLiveOnEntry(BasicBlock *BB);
   bool dominatesUse(const MemoryAccess *, const MemoryAccess *) const;
   MemoryUseOrDef *createNewAccess(Instruction *);
   MemoryUseOrDef *createDefinedAccess(Instruction *, MemoryAccess *);
   MemoryAccess *findDominatingDef(BasicBlock *, enum InsertionPlace);
   void removeFromLookups(MemoryAccess *);
   void removeFromLists(MemoryAccess *, bool ShouldDelete = true);
-
   void placePHINodes(const SmallPtrSetImpl<BasicBlock *> &,
                      const DenseMap<const BasicBlock *, unsigned int> &);
-  MemoryAccess *renameBlock(BasicBlock *, MemoryAccess *);
+  MemoryAccess *renameBlock(BasicBlock *, MemoryAccess *, bool);
+  void renameSuccessorPhis(BasicBlock *, MemoryAccess *, bool);
   void renamePass(DomTreeNode *, MemoryAccess *IncomingVal,
-                  SmallPtrSet<BasicBlock *, 16> &Visited);
+                  SmallPtrSet<BasicBlock *, 16> &Visited,
+                  bool SkipVisited = false, bool RenameAllUses = false);
   AccessList *getOrCreateAccessList(const BasicBlock *);
   DefsList *getOrCreateDefsList(const BasicBlock *);
   void renumberBlock(const BasicBlock *) const;
   void insertIntoListsForBlock(MemoryAccess *, const BasicBlock *,
                                InsertionPlace);
   void insertIntoListsBefore(MemoryAccess *, const BasicBlock *,
                              AccessList::iterator);
   AliasAnalysis *AA;

include/llvm/Transforms/Utils/MemorySSAUpdater.h

 class MemorySSAUpdater {
 private:
   MemorySSA *MSSA;
   SmallVector<MemoryPhi *, 8> InsertedPHIs;
   SmallPtrSet<BasicBlock *, 8> VisitedBlocks;
 
 public:
   MemorySSAUpdater(MemorySSA *MSSA) : MSSA(MSSA) {}
-  void insertDef(MemoryDef *Def);
+  // Insert a definition into the MemorySSA IR.  RenameUses will rename any use
+  // below the new def block (and any inserted phis).  RenameUses should be set
+  // to true if the definition may cause new aliases for loads below it.  This
+  // is not the case for hoisting or sinking or other forms of code *movement*.
+  // It *is* the case for straight code insertion.
+  // For example:
+  // store a
+  // if (foo) { }
+  // load a
+  //
+  // Moving the store into the if block, and calling insertDef, does not
+  // require RenameUses.
+  // However, changing it to:
+  // store a
+  // if (foo) { store b }
+  // load a
+  // Where a mayalias b, *does* require RenameUses be set to true.
+  void insertDef(MemoryDef *Def, bool RenameUses = false);
   void insertUse(MemoryUse *Use);
   void moveBefore(MemoryUseOrDef *What, MemoryUseOrDef *Where);
   void moveAfter(MemoryUseOrDef *What, MemoryUseOrDef *Where);
   void moveToPlace(MemoryUseOrDef *What, BasicBlock *BB,
                    MemorySSA::InsertionPlace Where);
+
 private:
   // Move What before Where in the MemorySSA IR.
   template <class WhereType>
   void moveTo(MemoryUseOrDef *What, BasicBlock *BB,
               WhereType Where);
   MemoryAccess *getPreviousDef(MemoryAccess *);
   MemoryAccess *getPreviousDefInBlock(MemoryAccess *);
   MemoryAccess *getPreviousDefFromEnd(BasicBlock *);
   MemoryAccess *getPreviousDefRecursive(BasicBlock *);
   MemoryAccess *recursePhi(MemoryAccess *Phi);
   template <class RangeType>
   MemoryAccess *tryRemoveTrivialPhi(MemoryPhi *Phi, RangeType &Operands);
   void fixupDefs(const SmallVectorImpl<MemoryAccess *> &);
 };
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_MEMORYSSAUPDATER_H

lib/Transforms/Utils/MemorySSA.cpp

public:
   void resetClobberWalker() { Walker.reset(); }
 
   void verify(const MemorySSA *MSSA) override {
     MemorySSAWalker::verify(MSSA);
     Walker.verify(MSSA);
   }
 };
 
+void MemorySSA::renameSuccessorPhis(BasicBlock *BB, MemoryAccess *IncomingVal,
+                                    bool RenameAllUses) {
+  // Pass through values to our successors
+  for (const BasicBlock *S : successors(BB)) {
+    auto It = PerBlockAccesses.find(S);
+    // Rename the phi nodes in our successor block
+    if (It == PerBlockAccesses.end() || !isa<MemoryPhi>(It->second->front()))
+      continue;
+    AccessList *Accesses = It->second.get();
+    auto *Phi = cast<MemoryPhi>(&Accesses->front());
+    if (RenameAllUses) {
+      int PhiIndex = Phi->getBasicBlockIndex(BB);
+      assert(PhiIndex != -1 && "Incomplete phi during partial rename");
+      Phi->setIncomingValue(PhiIndex, IncomingVal);
+    } else
+      Phi->addIncoming(IncomingVal, BB);
+  }
+}
+
 /// \brief Rename a single basic block into MemorySSA form.
 /// Uses the standard SSA renaming algorithm.
 /// \returns The new incoming value.
-MemoryAccess *MemorySSA::renameBlock(BasicBlock *BB,
-                                     MemoryAccess *IncomingVal) {
+MemoryAccess *MemorySSA::renameBlock(BasicBlock *BB, MemoryAccess *IncomingVal,
+                                     bool RenameAllUses) {
   auto It = PerBlockAccesses.find(BB);
   // Skip most processing if the list is empty.
   if (It != PerBlockAccesses.end()) {
     AccessList *Accesses = It->second.get();
     for (MemoryAccess &L : *Accesses) {
       if (MemoryUseOrDef *MUD = dyn_cast<MemoryUseOrDef>(&L)) {
-        if (MUD->getDefiningAccess() == nullptr)
+        if (MUD->getDefiningAccess() == nullptr || RenameAllUses)
           MUD->setDefiningAccess(IncomingVal);
         if (isa<MemoryDef>(&L))
           IncomingVal = &L;
       } else {
         IncomingVal = &L;
       }
     }
   }
-
-  // Pass through values to our successors
-  for (const BasicBlock *S : successors(BB)) {
-    auto It = PerBlockAccesses.find(S);
-    // Rename the phi nodes in our successor block
-    if (It == PerBlockAccesses.end() || !isa<MemoryPhi>(It->second->front()))
-      continue;
-    AccessList *Accesses = It->second.get();
-    auto *Phi = cast<MemoryPhi>(&Accesses->front());
-    Phi->addIncoming(IncomingVal, BB);
-  }
-
   return IncomingVal;
 }
 
 /// \brief This is the standard SSA renaming algorithm.
 ///
 /// We walk the dominator tree in preorder, renaming accesses, and then filling
 /// in phi nodes in our successors.
 void MemorySSA::renamePass(DomTreeNode *Root, MemoryAccess *IncomingVal,
-                           SmallPtrSet<BasicBlock *, 16> &Visited) {
+                           SmallPtrSet<BasicBlock *, 16> &Visited,
+                           bool SkipVisited, bool RenameAllUses) {
   SmallVector<RenamePassData, 32> WorkStack;
-  IncomingVal = renameBlock(Root->getBlock(), IncomingVal);
+  // Skip everything if we already renamed this block
+  if (!Visited.insert(Root->getBlock()).second)
+    return;
+
+  IncomingVal = renameBlock(Root->getBlock(), IncomingVal, RenameAllUses);
   WorkStack.push_back({Root, Root->begin(), IncomingVal});
-  Visited.insert(Root->getBlock());
 
   while (!WorkStack.empty()) {
     DomTreeNode *Node = WorkStack.back().DTN;
     DomTreeNode::const_iterator ChildIt = WorkStack.back().ChildIt;
     IncomingVal = WorkStack.back().IncomingVal;
 
     if (ChildIt == Node->end()) {
       WorkStack.pop_back();
     } else {
       DomTreeNode *Child = *ChildIt;
       ++WorkStack.back().ChildIt;
       BasicBlock *BB = Child->getBlock();
-      Visited.insert(BB);
-      IncomingVal = renameBlock(BB, IncomingVal);
+      // We already visited this during our renaming, which can happen when
+      // being asked to rename multiple blocks. Figure out the incoming val,
+      // which is the last def.
+      if (!Visited.insert(BB).second) {
+        // Incoming value can only change if there is a block def, and in that
+        // case, it's the last block def in the list.
+        if (auto *BlockDefs = getWritableBlockDefs(BB))
+          IncomingVal = &*BlockDefs->rbegin();
+      } else
+        IncomingVal = renameBlock(BB, IncomingVal, RenameAllUses);
+      renameSuccessorPhis(BB, IncomingVal, RenameAllUses);
       WorkStack.push_back({Child, Child->begin(), IncomingVal});
     }
   }
 }
 
 /// \brief Compute dominator levels, used by the phi insertion algorithm above.
 void MemorySSA::computeDomLevels(DenseMap<DomTreeNode *, unsigned> &DomLevels) {
   for (auto DFI = df_begin(DT->getRootNode()), DFE = df_end(DT->getRootNode());
        DFI != DFE; ++DFI)

 }
 
 void MemorySSA::print(raw_ostream &OS) const {
   MemorySSAAnnotatedWriter Writer(this);
   F.print(OS, &Writer);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
-LLVM_DUMP_METHOD void MemorySSA::dump() const {
-  print(dbgs());
-}
+LLVM_DUMP_METHOD void MemorySSA::dump() const { print(dbgs()); }
 #endif
 
 void MemorySSA::verifyMemorySSA() const {
   verifyDefUses(F);
   verifyDomination(F);
   verifyOrdering(F);
   Walker->verify(this);
 }
 
void MemoryUse::print(raw_ostream &OS) const {
   if (UO && UO->getID())
     OS << UO->getID();
   else
     OS << LiveOnEntryStr;
   OS << ')';
 }
 
 void MemoryAccess::dump() const {
-  // Cannot completely remove virtual function even in release mode.
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   print(dbgs());
   dbgs() << "\n";
 #endif
 }
 
 char MemorySSAPrinterLegacyPass::ID = 0;
 

lib/Transforms/Utils/MemorySSAUpdater.cpp

 }
 
 // A brief description of the algorithm:
 // First, we compute what should define the new def, using the SSA
 // construction algorithm.
 // Then, we update the defs below us (and any new phi nodes) in the graph to
 // point to the correct new defs, to ensure we only have one variable, and no
 // disconnected stores.
-void MemorySSAUpdater::insertDef(MemoryDef *MD) {
+void MemorySSAUpdater::insertDef(MemoryDef *MD, bool RenameUses) {
   InsertedPHIs.clear();
 
   // See if we had a local def, and if not, go hunting.
   MemoryAccess *DefBefore = getPreviousDefInBlock(MD);
   bool DefBeforeSameBlock = DefBefore != nullptr;
   if (!DefBefore)
     DefBefore = getPreviousDefRecursive(MD->getBlock());
 
void MemorySSAUpdater::insertDef(MemoryDef *MD, bool RenameUses) {
 
   while (!FixupList.empty()) {
     unsigned StartingPHISize = InsertedPHIs.size();
     fixupDefs(FixupList);
     FixupList.clear();
     // Put any new phis on the fixup list, and process them
     FixupList.append(InsertedPHIs.end() - StartingPHISize, InsertedPHIs.end());
   }
+  // Now that all fixups are done, rename all uses if we are asked.
+  if (RenameUses) {
+    SmallPtrSet<BasicBlock *, 16> Visited;
+    BasicBlock *StartBlock = MD->getBlock();
+    // We are guaranteed there is a def in the block, because we just got it
+    // handed to us in this function.
+    MemoryAccess *FirstDef = &*MSSA->getWritableBlockDefs(StartBlock)->begin();
+     // Convert to incoming value if it's a memorydef. A phi *is* already an
+     // incoming value.
+    if (auto *MD = dyn_cast<MemoryDef>(FirstDef))
+      FirstDef = MD->getDefiningAccess();
+    
+    MSSA->renamePass(MD->getBlock(), FirstDef, Visited);
+    // We just inserted a phi into this block, so the incoming value will become
+    // the phi anyway, so it does not matter what we pass.
+    for (auto *MP : InsertedPHIs)
+      MSSA->renamePass(MP->getBlock(), nullptr, Visited);
+  }
 }
 
 void MemorySSAUpdater::fixupDefs(const SmallVectorImpl<MemoryAccess *> &Vars) {
   SmallPtrSet<const BasicBlock *, 8> Seen;
   SmallVector<const BasicBlock *, 16> Worklist;
   for (auto *NewDef : Vars) {
     // First, see if there is a local def after the operand.
     auto *Defs = MSSA->getWritableBlockDefs(NewDef->getBlock());

unittests/Transforms/Utils/MemorySSA.cpp

TEST_F(MemorySSATest, CreateLoadsAndStoreUpdater) {
   EXPECT_EQ(FirstLoadAccess->getDefiningAccess(), EntryStoreAccess);
 
   // Create a store on the left
   // Add the store
   B.SetInsertPoint(Left, Left->begin());
   StoreInst *LeftStore = B.CreateStore(B.getInt8(16), PointerArg);
   MemoryAccess *LeftStoreAccess = MSSA.createMemoryAccessInBB(
       LeftStore, nullptr, Left, MemorySSA::Beginning);
-  Updater.insertDef(cast<MemoryDef>(LeftStoreAccess));
+  Updater.insertDef(cast<MemoryDef>(LeftStoreAccess), false);
   // We don't touch existing loads, so we need to create a new one to get a phi
   // Add the second load
   B.SetInsertPoint(Merge, Merge->begin());
   LoadInst *SecondLoad = B.CreateLoad(PointerArg);
 
   // MemoryPHI should not already exist.
   MP = MSSA.getMemoryAccess(Merge);
   EXPECT_EQ(MP, nullptr);
 
   // Create the load memory access
   MemoryUse *SecondLoadAccess = cast<MemoryUse>(MSSA.createMemoryAccessInBB(
       SecondLoad, nullptr, Merge, MemorySSA::Beginning));
   Updater.insertUse(SecondLoadAccess);
   // Now the load should be a phi of the entry store and the left store
   MemoryPhi *MergePhi =
       dyn_cast<MemoryPhi>(SecondLoadAccess->getDefiningAccess());
   EXPECT_NE(MergePhi, nullptr);
   EXPECT_EQ(MergePhi->getIncomingValue(0), EntryStoreAccess);
   EXPECT_EQ(MergePhi->getIncomingValue(1), LeftStoreAccess);
   // Now create a store below the existing one in the entry
   B.SetInsertPoint(Entry, --Entry->end());
   StoreInst *SecondEntryStore = B.CreateStore(B.getInt8(16), PointerArg);
   MemoryAccess *SecondEntryStoreAccess = MSSA.createMemoryAccessInBB(
       SecondEntryStore, nullptr, Entry, MemorySSA::End);
-  Updater.insertDef(cast<MemoryDef>(SecondEntryStoreAccess));
+  // Insert it twice just to test renaming
+  Updater.insertDef(cast<MemoryDef>(SecondEntryStoreAccess), false);
+  EXPECT_NE(FirstLoadAccess->getDefiningAccess(), MergePhi);
+  Updater.insertDef(cast<MemoryDef>(SecondEntryStoreAccess), true);
+  EXPECT_EQ(FirstLoadAccess->getDefiningAccess(), MergePhi);
   // and make sure the phi below it got updated, despite being blocks away
   MergePhi = dyn_cast<MemoryPhi>(SecondLoadAccess->getDefiningAccess());
   EXPECT_NE(MergePhi, nullptr);
   EXPECT_EQ(MergePhi->getIncomingValue(0), SecondEntryStoreAccess);
   EXPECT_EQ(MergePhi->getIncomingValue(1), LeftStoreAccess);
   MSSA.verifyMemorySSA();
 }
 
 TEST_F(MemorySSATest, CreateALoadUpdater) {
   // We create a diamond, then build memoryssa with no memory accesses, and
   // incrementally update it by inserting a store in one of the branches, and a
   // load in the merge point
   F = Function::Create(
       FunctionType::get(B.getVoidTy(), {B.getInt8PtrTy()}, false),
       GlobalValue::ExternalLinkage, "F", &M);