[MemorySSA] Fix removing dead accesses from MemoryPhis.

[fhahn]

Summary

The current logic to remove all uses of MemoryAccess defined in dead
blocks from MemoryPhis outside the set of dead blocks. MemoryPhis using
a definition from a dead block do not necessarily have to be in a direct
successor of a dead block, for example if the successor does not contain
any memory accesses.

To catch those additional MemoryPhis that need updating we need to check
all uses of all memory defs in the dead blocks and remove them from any
MemoryPhi outside the set of dead blocks.

Note that we sill need the existing code where we check the successors
for memory phis to catch the case where a MemoryPhi has an incoming
value from a dead block but is not defined in any dead block.

This can happen during unswitching, where a region between the memory
def and a memory phi becomes unreachable.

I *think* there's still a case we might miss in the above case: when the
memory phi is not a direct successor of a dead block. Not sure if this
can happen in practice.

Fixes PR47557.

Note that I did not manage to reduce the test cases so far that it only
requires a single pass to crash.

[fhahn]

Turns out that things actually go slightly wrong earlier during LSR. LSR did not pass MSSU to SplitCriticalEdge, which caused the MemoryPhi to have slightly wrong incoming blocks. I pushed a fix for that, as it is straight-forward to fix (57ae9bb9323548b2ad4ba8274c3910bf9c764983). With that, this patch is no longer needed for the test case.

llvm/lib/Analysis/MemorySSAUpdater.cpp

for (BasicBlock *BB : DeadBlocks) {
     Instruction *TI = BB->getTerminator();
     assert(TI && "Basic block expected to have a terminator instruction");
     for (BasicBlock *Succ : successors(TI))
       if (!DeadBlocks.count(Succ))
         if (MemoryPhi *MP = MSSA->getMemoryAccess(Succ)) {
           MP->unorderedDeleteIncomingBlock(BB);
           tryRemoveTrivialPhi(MP);
         }
-    // Drop all references of all accesses in BB
+
+    // Drop all references of all accesses in BB and remove those accesses from
+    // MemoryPhis outside of the dead blocks.
     if (MemorySSA::AccessList *Acc = MSSA->getWritableBlockAccesses(BB))
-      for (MemoryAccess &MA : *Acc)
+      for (MemoryAccess &MA : *Acc) {
         MA.dropAllReferences();
+
+        for (auto UI = MA.use_begin(), UE = MA.use_end(); UI != UE;) {
+          auto &U = *UI++;
+          auto *MemPhi = dyn_cast_or_null<MemoryPhi>(U.getUser());
+          if (!MemPhi || DeadBlocks.contains(MemPhi->getBlock()))
+            continue;
+          MemPhi->unorderedDeleteIncomingIf(
+              [&MA](MemoryAccess *Acc, BasicBlock *BB) { return &MA == Acc; });
+          tryRemoveTrivialPhi(MemPhi);
+        }
+      }
   }
 
   // Next, delete all memory accesses in each block
   for (BasicBlock *BB : DeadBlocks) {
     MemorySSA::AccessList *Acc = MSSA->getWritableBlockAccesses(BB);
     if (!Acc)
       continue;
     for (auto AB = Acc->begin(), AE = Acc->end(); AB != AE;) {

llvm/test/Analysis/MemorySSA/update-remove-dead-blocks.ll

+; RUN: opt -loop-unswitch -loop-reduce -loop-simplifycfg -verify-memoryssa -S %s | FileCheck %s
+
+; TODO: also run with NPM, but currently LSR does not preserve LCSSA, causing a verification failure on the test.
+;   opt -passes='loop-mssa(unswitch<nontrivial>,loop-reduce,simplify-cfg)' -verify-memoryssa -S %s | FileCheck %s
+
+; Test case for PR47557.
+
+; REQUIRES: x86-registered-target
+
+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"
+
+@a = external global i32, align 4
+@c = external global [1 x i32], align 4
+
+define i32* @test() {
+; CHECK-LABEL: @test
+;
+entry:                                                ; preds = %entry
+  br label %for.cond
+
+for.cond:                                         ; preds = %cleanup, %entry
+  %storemerge = phi i64 [ 0, %entry ], [ %inc7, %cleanup ]
+  br label %for.cond2.1
+
+for.body3:                                        ; preds = %for.cond2.2, %for.cond2.1
+  %arrayidx = getelementptr inbounds [1 x i32], [1 x i32]* @c, i64 0, i64 %storemerge
+  ret i32* %arrayidx
+
+cleanup:                                          ; preds = %for.end5, %if.then
+  %inc7 = add nsw i64 %storemerge, 1
+  br label %for.cond
+
+for.cond2.1:                                      ; preds = %for.cond
+  br i1 true, label %for.inc.1, label %for.body3
+
+for.inc.1:                                        ; preds = %for.end.1
+  br i1 false, label %for.body.2, label %cleanup
+
+for.body.2:                                       ; preds = %for.inc.1
+  store i32 0, i32* @a, align 4
+  br label %for.cond2.2
+
+for.cond2.2:                                      ; preds = %for.body.2
+  br i1 true, label %cleanup, label %for.body3
+}