[MemorySSA] Make the use of moveAllAfterMergeBlocks consistent.

[asbirlea]

Summary

The rule for the moveAllAfterMergeBlocks API si for all instructions
from From to have been moved to To, while keeping the CFG edges (and
block terminators) unchanged.
Update all the callsites for moveAllAfterMergeBlocks to follow this.

Pending follow-up: since the same behavior is needed everytime, merge
all callsites into one. The common denominator may be the call to
MergeBlockIntoPredecessor.

Resolves PR43569.

[george.burgess.iv]

Thanks for this!

LGTM with a few nits

lib/Analysis/MemorySSAUpdater.cpp

1185 ↗

(On Diff #223920)

nit: !Defs->empty()?

(Though I'm surprised this check is required -- shouldn't we be removing entries from the map as they become empty?)

1186 ↗

(On Diff #223920)

If I'm reading the description properly, this function assumes that everything in From is now in To, no?

If so, cast<MemoryPhi>() seems more appropriate here, since the only thing that can remain is a single Phi

test/Analysis/MemorySSA/pr43569.ll

1 ↗

(On Diff #223920)

; REQUIRES: asserts?

[asbirlea]

Address comments.

lib/Analysis/MemorySSAUpdater.cpp

1186 ↗

(On Diff #223920)

This method (moveAllAccesses) is also called from moveAllAfterSpliceBlocks, when From is not necessarily empty. I'll update the comment above.

llvm/lib/Analysis/MemorySSAUpdater.cpp

 // All accesses in To used to be in From. Move to end and update access lists.
 void MemorySSAUpdater::moveAllAccesses(BasicBlock *From, BasicBlock *To,
                                        Instruction *Start) {
 
   MemorySSA::AccessList *Accs = MSSA->getWritableBlockAccesses(From);
   if (!Accs)
     return;
 
+  assert(Start->getParent() == To && "Incorrect Start instruction");
   MemoryAccess *FirstInNew = nullptr;
   for (Instruction &I : make_range(Start->getIterator(), To->end()))
     if ((FirstInNew = MSSA->getMemoryAccess(&I)))
       break;
-  if (!FirstInNew)
-    return;
-
+  if (FirstInNew) {
     auto *MUD = cast<MemoryUseOrDef>(FirstInNew);
     do {
       auto NextIt = ++MUD->getIterator();
       MemoryUseOrDef *NextMUD = (!Accs || NextIt == Accs->end())
                                     ? nullptr
                                     : cast<MemoryUseOrDef>(&*NextIt);
       MSSA->moveTo(MUD, To, MemorySSA::End);
-    // Moving MUD from Accs in the moveTo above, may delete Accs, so we need to
-    // retrieve it again.
+      // Moving MUD from Accs in the moveTo above, may delete Accs, so we need
+      // to retrieve it again.
       Accs = MSSA->getWritableBlockAccesses(From);
       MUD = NextMUD;
     } while (MUD);
   }
 
+  // If all accesses were moved and only a trivial Phi remains, we try to remove
+  // that Phi. This is needed when From is going to be deleted.
+  auto *Defs = MSSA->getWritableBlockDefs(From);
+  if (Defs && !Defs->empty())
+    if (auto *Phi = dyn_cast<MemoryPhi>(&*Defs->begin()))
+      tryRemoveTrivialPhi(Phi);
+}
+
 void MemorySSAUpdater::moveAllAfterSpliceBlocks(BasicBlock *From,
                                                 BasicBlock *To,
                                                 Instruction *Start) {
   assert(MSSA->getBlockAccesses(To) == nullptr &&
          "To block is expected to be free of MemoryAccesses.");
   moveAllAccesses(From, To, Start);
   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 &&
+  assert(From->getUniquePredecessor() == To &&
          "From block is expected to have a single predecessor (To).");
   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

llvm/lib/Transforms/Scalar/LoopUnswitch.cpp

if (BranchInst *BI = dyn_cast<BranchInst>(I)) {
         LLVM_DEBUG(dbgs() << "Merging blocks: " << Pred->getName() << " <- "
                           << Succ->getName() << "\n");
 
         // Resolve any single entry PHI nodes in Succ.
         while (PHINode *PN = dyn_cast<PHINode>(Succ->begin()))
           ReplaceUsesOfWith(PN, PN->getIncomingValue(0), Worklist, L, LPM,
                             MSSAU.get());
 
-        // If Succ has any successors with PHI nodes, update them to have
-        // entries coming from Pred instead of Succ.
-        Succ->replaceAllUsesWith(Pred);
+        Instruction *STI = Succ->getTerminator();
+        Instruction *Start = &*Succ->begin();
+        // If there's nothing to move, mark the starting instruction as the last
+        // instruction in the block.
+        if (Start == STI)
+          Start = BI;
 
         // Move all of the successor contents from Succ to Pred.
         Pred->getInstList().splice(BI->getIterator(), Succ->getInstList(),
-                                   Succ->begin(), Succ->end());
+                                   Succ->begin(), STI->getIterator());
         if (MSSAU)
-          MSSAU->moveAllAfterMergeBlocks(Succ, Pred, BI);
+          MSSAU->moveAllAfterMergeBlocks(Succ, Pred, Start);
+
+        // Move terminator instruction from Succ now, we're deleting BI below.
+        // FIXME: remove BI first might be more intuitive.
+        Pred->getInstList().splice(Pred->end(), Succ->getInstList());
+
+        // If Succ has any successors with PHI nodes, update them to have
+        // entries coming from Pred instead of Succ.
+        Succ->replaceAllUsesWith(Pred);
+
         LPM->deleteSimpleAnalysisValue(BI, L);
         RemoveFromWorklist(BI, Worklist);
         BI->eraseFromParent();
 
         // Remove Succ from the loop tree.
         LI->removeBlock(Succ);
         LPM->deleteSimpleAnalysisValue(Succ, L);
         Succ->eraseFromParent();

llvm/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()));
+  Instruction *PTI = PredBB->getTerminator();
+  Instruction *STI = BB->getTerminator();
+  Instruction *Start = &*BB->begin();
+  // If there's nothing to move, mark the starting instruction as the last
+  // instruction in the block.
+  if (Start == STI)
+    Start = PTI;
 
-  // Delete the unconditional branch from the predecessor...
-  PredBB->getInstList().pop_back();
+  // Move all definitions in the successor to the predecessor...
+  PredBB->getInstList().splice(PTI->getIterator(), BB->getInstList(),
+                               BB->begin(), STI->getIterator());
+
+  if (MSSAU)
+    MSSAU->moveAllAfterMergeBlocks(BB, PredBB, Start);
 
   // Make all PHI nodes that referred to BB now refer to Pred as their
   // source...
   BB->replaceAllUsesWith(PredBB);
 
-  // Move all definitions in the successor to the predecessor...
+  // Delete the unconditional branch from the predecessor...
+  PredBB->getInstList().pop_back();
+
+  // Move terminator instruction and add unreachable to now empty BB.
   PredBB->getInstList().splice(PredBB->end(), BB->getInstList());
   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>
bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
   // Finally, erase the old block and update dominator info.
   if (DTU) {
     assert(BB->getInstList().size() == 1 &&
            isa<UnreachableInst>(BB->getTerminator()) &&
            "The successor list of BB isn't empty before "
            "applying corresponding DTU updates.");
     DTU->applyUpdatesPermissive(Updates);
     DTU->deleteBB(BB);
-  }
-
-  else {
+  } else {
     BB->eraseFromParent(); // Nuke BB if DTU is nullptr.
   }
+
   return true;
 }
 
 void llvm::ReplaceInstWithValue(BasicBlock::InstListType &BIL,
                                 BasicBlock::iterator &BI, Value *V) {
   Instruction &I = *BI;
   // Replaces all of the uses of the instruction with uses of the value
   I.replaceAllUsesWith(V);

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

if (!BI)
     return false;
 
   if (!shouldSpeculateInstrs(Latch->begin(), Jmp->getIterator(), L))
     return false;
 
   LLVM_DEBUG(dbgs() << "Folding loop latch " << Latch->getName() << " into "
                     << LastExit->getName() << "\n");
 
+  Instruction *FirstLatchInst = &*Latch->begin();
+  // If there's nothing to move, mark the starting instruction as the last
+  // instruction in the block.
+  if (FirstLatchInst == Jmp)
+    FirstLatchInst = BI;
+
   // Hoist the instructions from Latch into LastExit.
-  Instruction *FirstLatchInst = &*(Latch->begin());
   LastExit->getInstList().splice(BI->getIterator(), Latch->getInstList(),
                                  Latch->begin(), Jmp->getIterator());
 
   // Update MemorySSA
   if (MSSAU)
     MSSAU->moveAllAfterMergeBlocks(Latch, LastExit, FirstLatchInst);
 
   unsigned FallThruPath = BI->getSuccessor(0) == Latch ? 0 : 1;

llvm/test/Analysis/MemorySSA/pr43569.ll

+; RUN: opt -pgo-kind=pgo-instr-gen-pipeline -aa-pipeline=default -passes="default<O3>" -enable-nontrivial-unswitch -S < %s | FileCheck %s
+; REQUIRES: asserts
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@__profn_c = private constant [1 x i8] c"c"
+@b = common dso_local global i32 0, align 4
+@a = common dso_local global i16 0, align 2
+
+; CHECK-LABEL: @c()
+; Function Attrs: nounwind uwtable
+define dso_local void @c() #0 {
+entry:
+  call void @llvm.instrprof.increment(i8* getelementptr inbounds ([1 x i8], [1 x i8]* @__profn_c, i32 0, i32 0), i64 68269137, i32 3, i32 0)
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.end, %entry
+  call void @llvm.instrprof.increment(i8* getelementptr inbounds ([1 x i8], [1 x i8]* @__profn_c, i32 0, i32 0), i64 68269137, i32 3, i32 1)
+  store i32 0, i32* @b, align 4
+  br label %for.cond1
+
+for.cond1:                                        ; preds = %for.inc, %for.cond
+  %0 = load i32, i32* @b, align 4
+  %1 = load i16, i16* @a, align 2
+  %conv = sext i16 %1 to i32
+  %cmp = icmp slt i32 %0, %conv
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:                                         ; preds = %for.cond1
+  call void @llvm.instrprof.increment(i8* getelementptr inbounds ([1 x i8], [1 x i8]* @__profn_c, i32 0, i32 0), i64 68269137, i32 3, i32 2)
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body
+  %2 = load i32, i32* @b, align 4
+  %inc = add nsw i32 %2, 1
+  store i32 %inc, i32* @b, align 4
+  br label %for.cond1
+
+for.end:                                          ; preds = %for.cond1
+  br label %for.cond
+}
+
+; Function Attrs: nounwind
+declare void @llvm.instrprof.increment(i8*, i64, i32, i32) #1
+
+attributes #0 = { nounwind uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #1 = { nounwind }
+