MemorySSA: Allow movement to arbitrary places

[• dberlin]

Summary

Extend the MemorySSAUpdater API to allow movement to arbitrary places

[• dberlin]

Sorry, i split two commits up wrong, and it was easier to roll forward than backwards.
Happy to address any comments folks have

include/llvm/Transforms/Utils/MemorySSA.h

DefsList *getWritableBlockDefs(const BasicBlock *BB) const {
     auto It = PerBlockDefs.find(BB);
     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);
 
 private:
   class CachingWalker;
   class OptimizeUses;
 
   CachingWalker *getWalkerImpl();
   void buildMemorySSA();
   void optimizeUses();

include/llvm/Transforms/Utils/MemorySSAUpdater.h

 // Create the memory access you want for the instruction (this is mainly so
 // we know where it is, without having to duplicate the entire set of create
 // functions MemorySSA supports).
 // Call insertDef or insertUse depending on whether it's a MemoryUse or a
 // MemoryDef.
 // That's it.
 //
 // For moving, first, move the instruction itself using the normal SSA
-// instruction moving API, then just call moveBefore or moveAfter with the right
-// arguments.
+// instruction moving API, then just call moveBefore, moveAfter,or moveTo with
+// the right arguments.
 //
-// walk memory instructions using a use/def graph.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_MEMORYSSAUPDATER_H
 #define LLVM_TRANSFORMS_UTILS_MEMORYSSAUPDATER_H
 
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/BasicBlock.h"
private:
   SmallPtrSet<BasicBlock *, 8> VisitedBlocks;
 
 public:
   MemorySSAUpdater(MemorySSA *MSSA) : MSSA(MSSA) {}
   void insertDef(MemoryDef *Def);
   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,
-              MemorySSA::AccessList::iterator Where);
+              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/MemorySSAUpdater.cpp

while (!Worklist.empty()) {
           Worklist.push_back(S);
         }
       }
     }
   }
 }
 
 // Move What before Where in the MemorySSA IR.
+template <class WhereType>
 void MemorySSAUpdater::moveTo(MemoryUseOrDef *What, BasicBlock *BB,
-                              MemorySSA::AccessList::iterator Where) {
+                              WhereType Where) {
   // Replace all our users with our defining access.
   What->replaceAllUsesWith(What->getDefiningAccess());
 
   // Let MemorySSA take care of moving it around in the lists.
   MSSA->moveTo(What, BB, Where);
 
   // Now reinsert it into the IR and do whatever fixups needed.
   if (auto *MD = dyn_cast<MemoryDef>(What))
     insertDef(MD);
   else
     insertUse(cast<MemoryUse>(What));
 }
+
 // Move What before Where in the MemorySSA IR.
 void MemorySSAUpdater::moveBefore(MemoryUseOrDef *What, MemoryUseOrDef *Where) {
   moveTo(What, Where->getBlock(), Where->getIterator());
 }
 
 // Move What after Where in the MemorySSA IR.
 void MemorySSAUpdater::moveAfter(MemoryUseOrDef *What, MemoryUseOrDef *Where) {
   moveTo(What, Where->getBlock(), ++Where->getIterator());
 }
 
+void MemorySSAUpdater::moveToPlace(MemoryUseOrDef *What, BasicBlock *BB,
+                                   MemorySSA::InsertionPlace Where) {
+  return moveTo(What, BB, Where);
+}
 } // namespace llvm

unittests/Transforms/Utils/MemorySSA.cpp

TEST_F(MemorySSATest, MoveAStoreUpdaterMove) {
   Updater.moveAfter(SideStoreAccess, EntryStoreAccess);
   // After, it's a phi of the side store.
   EXPECT_EQ(MergePhi->getIncomingValue(0), SideStoreAccess);
   EXPECT_EQ(MergePhi->getIncomingValue(1), SideStoreAccess);
 
   MSSA.verifyMemorySSA();
 }
 
+TEST_F(MemorySSATest, MoveAStoreAllAround) {
+  // We create a diamond where there is a in the entry, a store on one side, and
+  // a load at the end.  After building MemorySSA, we test updating by moving
+  // the store from the side block to the entry block, then to the other side
+  // block, then to before the load.  This does not destroy the old access.
+  F = Function::Create(
+      FunctionType::get(B.getVoidTy(), {B.getInt8PtrTy()}, false),
+      GlobalValue::ExternalLinkage, "F", &M);
+  BasicBlock *Entry(BasicBlock::Create(C, "", F));
+  BasicBlock *Left(BasicBlock::Create(C, "", F));
+  BasicBlock *Right(BasicBlock::Create(C, "", F));
+  BasicBlock *Merge(BasicBlock::Create(C, "", F));
+  B.SetInsertPoint(Entry);
+  Argument *PointerArg = &*F->arg_begin();
+  StoreInst *EntryStore = B.CreateStore(B.getInt8(16), PointerArg);
+  B.CreateCondBr(B.getTrue(), Left, Right);
+  B.SetInsertPoint(Left);
+  auto *SideStore = B.CreateStore(B.getInt8(16), PointerArg);
+  BranchInst::Create(Merge, Left);
+  BranchInst::Create(Merge, Right);
+  B.SetInsertPoint(Merge);
+  auto *MergeLoad = B.CreateLoad(PointerArg);
+  setupAnalyses();
+  MemorySSA &MSSA = *Analyses->MSSA;
+  MemorySSAUpdater Updater(&MSSA);
+
+  // Move the store
+  auto *EntryStoreAccess = MSSA.getMemoryAccess(EntryStore);
+  auto *SideStoreAccess = MSSA.getMemoryAccess(SideStore);
+  // Before, the load will point to a phi of the EntryStore and SideStore.
+  auto *LoadAccess = cast<MemoryUse>(MSSA.getMemoryAccess(MergeLoad));
+  EXPECT_TRUE(isa<MemoryPhi>(LoadAccess->getDefiningAccess()));
+  MemoryPhi *MergePhi = cast<MemoryPhi>(LoadAccess->getDefiningAccess());
+  EXPECT_EQ(MergePhi->getIncomingValue(1), EntryStoreAccess);
+  EXPECT_EQ(MergePhi->getIncomingValue(0), SideStoreAccess);
+  // Move the store before the entry store
+  SideStore->moveBefore(*EntryStore->getParent(), EntryStore->getIterator());
+  Updater.moveBefore(SideStoreAccess, EntryStoreAccess);
+  // After, it's a phi of the entry store.
+  EXPECT_EQ(MergePhi->getIncomingValue(0), EntryStoreAccess);
+  EXPECT_EQ(MergePhi->getIncomingValue(1), EntryStoreAccess);
+  MSSA.verifyMemorySSA();
+  // Now move the store to the right branch
+  SideStore->moveBefore(*Right, Right->begin());
+  Updater.moveToPlace(SideStoreAccess, Right, MemorySSA::Beginning);
+  MSSA.verifyMemorySSA();
+  EXPECT_EQ(MergePhi->getIncomingValue(0), EntryStoreAccess);
+  EXPECT_EQ(MergePhi->getIncomingValue(1), SideStoreAccess);
+  // Now move it before the load
+  SideStore->moveBefore(MergeLoad);
+  Updater.moveBefore(SideStoreAccess, LoadAccess);
+  EXPECT_EQ(MergePhi->getIncomingValue(0), EntryStoreAccess);
+  EXPECT_EQ(MergePhi->getIncomingValue(1), EntryStoreAccess);
+  MSSA.verifyMemorySSA();
+}
+
 TEST_F(MemorySSATest, RemoveAPhi) {
   // We create a diamond where there is a store on one side, and then a load
   // after the merge point.  This enables us to test a bunch of different
   // removal cases.
   F = Function::Create(
       FunctionType::get(B.getVoidTy(), {B.getInt8PtrTy()}, false),
       GlobalValue::ExternalLinkage, "F", &M);
   BasicBlock *Entry(BasicBlock::Create(C, "", F));