Fix MSSA update in MergeBlockIntoPredecessor

[apilipenko]

Summary

There are a couple of issues around MergeBlockIntoPredecessor and MSSA update.

Both issues show up when loop-simplify-cfg is used after simple-loop-unswitch.

1. moveAllAfterMergeBlocks(BasicBlock *From, BasicBlock *To, Instruction *Start) expects to be called after all instructions from From were moved to To. Specifically, it assumes that Start belongs to To basic block and iterates over make_range(Start->getIterator(), To->end()) range (see moveAllAccesses). MergeBlockIntoPredecessor currently calls moveAllAfterMergeBlocks before instructions are moved into To and passes Start which belongs to From.

The fix is to call moveAllAfterMergeBlocks after the instructions are moved from From to To.

2. After simple-loop-unswitch some blocks might end up with trivial phis with one incoming value. When such a block is merged with its predecessor (by loop-simplify-cfg) the phi node is left intact after the block is removed. This results in a dangling reference to a phi without a block. unswitch_simplify_mssa_update.ll test case fails with VerifyMemorySSA && EXPENSIVE_CHECKS on this assertion in verifyPrevDefInPhis: assert(LastAcc == IncAcc && "Incorrect incoming access into phi.");

I fixed this by calling tryRemoveTrivialPhi for From block in moveAllAfterMergeBlocks. I'm not certain that this is the right approach. May be we shouldn't have trivial phis to begin with? Please advise.

[asbirlea]

I'm seeing the attached test pass at ToT now.
I may have resolved it in rL373380, which was failing to remove trivial Phis when inserting a backedge block.
Please let me know if you are still seeing issues!

lib/Analysis/MemorySSAUpdater.cpp

for (BasicBlock *Succ : successors(To))
     if (MemoryPhi *MPhi = MSSA->getMemoryAccess(Succ))
       MPhi->setIncomingBlock(MPhi->getBasicBlockIndex(From), To);
 }
 
 void MemorySSAUpdater::moveAllAfterMergeBlocks(BasicBlock *From, BasicBlock *To,
                                                Instruction *Start) {
   assert(From->getSinglePredecessor() == To &&
          "From block is expected to have a single predecessor (To).");
+  if (MemoryPhi *MP = MSSA->getMemoryAccess(From)) {
+    // `From` might have a trivial PHI with a single input. Optimize it
+    // now before `From` is removed.
+    assert(MP.getNumIncomingValues() == 1 && "From can only have one pred");
+    auto *ReplacedWith = tryRemoveTrivialPhi(MP);
+    assert(ReplacedWith != MP &&
+           "If there is a Phi it must be trivially removable");
+  }
   moveAllAccesses(From, To, Start);
   for (BasicBlock *Succ : successors(From))
     if (MemoryPhi *MPhi = MSSA->getMemoryAccess(Succ))
       MPhi->setIncomingBlock(MPhi->getBasicBlockIndex(From), To);
 }
 
 /// If all arguments of a MemoryPHI are defined by the same incoming
 /// argument, return that argument.

lib/Transforms/Utils/BasicBlockUtils.cpp

for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I)
       // This successor of BB may already have PredBB as a predecessor.
       if (llvm::find(successors(PredBB), *I) == succ_end(PredBB))
         Updates.push_back({DominatorTree::Insert, PredBB, *I});
     for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I)
       Updates.push_back({DominatorTree::Delete, BB, *I});
     Updates.push_back({DominatorTree::Delete, PredBB, BB});
   }
 
-  if (MSSAU)
-    MSSAU->moveAllAfterMergeBlocks(BB, PredBB, &*(BB->begin()));
-
-  // Delete the unconditional branch from the predecessor...
-  PredBB->getInstList().pop_back();
+  // We need to update MemorySSA if MSSAU is provided. MemorySSAUpdater expects
+  // a specifc state when notified: `PredBB` still branches to `BB`, but all
+  // instructions were moved from `BB` to `PredBB` and `BB` is now an empty
+  // block. Splice the instructions, but keep the terminators intact.
+  Instruction *PredBBTerm = PredBB->getTerminator();
+  Instruction *BBTerm = BB->getTerminator();
+  Instruction *FirstBBInst = &*(BB->begin());
+  PredBB->getInstList().splice(PredBBTerm->getIterator(), BB->getInstList(),
+                               BB->begin(), BBTerm->getIterator());
 
-  // Make all PHI nodes that referred to BB now refer to Pred as their
-  // source...
-  BB->replaceAllUsesWith(PredBB);
+  if (MSSAU)
+    MSSAU->moveAllAfterMergeBlocks(BB, PredBB, FirstBBInst);
 
-  // Move all definitions in the successor to the predecessor...
-  PredBB->getInstList().splice(PredBB->end(), BB->getInstList());
+  // Finish updating the CFG after MSSAU is notified.
+  BB->replaceSuccessorsPhiUsesWith(PredBB);
+  BBTerm->removeFromParent();
+  BBTerm->insertBefore(PredBBTerm);
+  PredBBTerm->eraseFromParent();
   new UnreachableInst(BB->getContext(), BB);
 
   // Eliminate duplicate dbg.values describing the entry PHI node post-splice.
   for (auto Incoming : IncomingValues) {
     if (isa<Instruction>(*Incoming)) {
       SmallVector<DbgValueInst *, 2> DbgValues;
       SmallDenseSet<std::pair<DILocalVariable *, DIExpression *>, 2>
           DbgValueSet;

test/Transforms/LoopSimplifyCFG/unswitch_simplify_mssa_update.ll

+; RUN: opt -passes='loop-mssa(unswitch,simplify-cfg),verify<memoryssa>' -S < %s | FileCheck %s
+
+declare void @a()
+declare void @b()
+declare void @x()
+
+; CHECK: test_unswitch_simplify
+define void @test_unswitch_simplify(i1* %ptr, i1 %cond, i1 %f, i32 %n) {
+entry:
+  br i1 %f, label %loop_begin, label %loop_exit
+
+loop_begin:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop_latch]
+  call void @x()
+  br i1 %cond, label %loop_a, label %loop_b
+
+loop_a:
+  call void @a()
+  br label %loop_latch
+
+loop_b:
+  call void @b()
+  br label %loop_latch
+
+loop_latch:
+  %iv.next = add i32 %iv, 1
+  %v = icmp eq i32 %iv.next, %n
+  br i1 %v, label %loop_begin, label %loop_exit
+
+loop_exit:
+  ret void
+}
 No newline at end of file