[MemorySSAUpdater] Fix iterator invalidation bug in applyInsertUpdates

[dtcxzyw]

This patch defers resetting optimized accesses until all uses are replaced, to avoid invalidating the iterator.

Closes #139103.
Closes #139289.
Closes #139308.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

Changes

This patch defers resetting optimized accesses until all uses are replaced, to avoid invalidating the iterator.

Closes #139103.
Closes #139289.
Closes #139308.

---

Full diff: https://github.com/llvm/llvm-project/pull/139370.diff

2 Files Affected:

• (modified) llvm/lib/Analysis/MemorySSAUpdater.cpp (+7-1)
• (added) llvm/test/Analysis/MemorySSA/pr139103.ll (+47)

diff --git a/llvm/lib/Analysis/MemorySSAUpdater.cpp b/llvm/lib/Analysis/MemorySSAUpdater.cpp
index aa9f0b6e100c4..ecfecb03c3757 100644
--- a/llvm/lib/Analysis/MemorySSAUpdater.cpp
+++ b/llvm/lib/Analysis/MemorySSAUpdater.cpp
@@ -1119,6 +1119,9 @@ void MemorySSAUpdater::applyInsertUpdates(ArrayRef<CFGUpdate> Updates,
     if (auto DefsList = MSSA->getWritableBlockDefs(BlockWithDefsToReplace)) {
       for (auto &DefToReplaceUses : *DefsList) {
         BasicBlock *DominatingBlock = DefToReplaceUses.getBlock();
+        // We defer resetting optimized accesses until all uses are replaced, to
+        // avoid invalidating the iterator.
+        SmallVector<MemoryUseOrDef *, 4> ResetOptimized;
         for (Use &U : llvm::make_early_inc_range(DefToReplaceUses.uses())) {
           MemoryAccess *Usr = cast<MemoryAccess>(U.getUser());
           if (MemoryPhi *UsrPhi = dyn_cast<MemoryPhi>(Usr)) {
@@ -1135,10 +1138,13 @@ void MemorySSAUpdater::applyInsertUpdates(ArrayRef<CFGUpdate> Updates,
                 assert(IDom && "Block must have a valid IDom.");
                 U.set(GetLastDef(IDom->getBlock()));
               }
-              cast<MemoryUseOrDef>(Usr)->resetOptimized();
+              ResetOptimized.push_back(cast<MemoryUseOrDef>(Usr));
             }
           }
         }
+
+        for (auto *Usr : ResetOptimized)
+          Usr->resetOptimized();
       }
     }
   }
diff --git a/llvm/test/Analysis/MemorySSA/pr139103.ll b/llvm/test/Analysis/MemorySSA/pr139103.ll
new file mode 100644
index 0000000000000..60a8c8bea09c7
--- /dev/null
+++ b/llvm/test/Analysis/MemorySSA/pr139103.ll
@@ -0,0 +1,47 @@
+; RUN: opt -disable-output -passes="loop-mssa(licm,loop-rotate,licm,simple-loop-unswitch<nontrivial>),print<memoryssa>" -verify-memoryssa < %s 2>&1 | FileCheck %s
+
+; Make sure that we update MSSA correctly in this case.
+
+; CHECK-LABEL: MemorySSA for function: test
+; CHECK: for.header2.preheader:
+; CHECK-NEXT: 11 = MemoryPhi({entry.split,liveOnEntry},{for.header,9})
+; CHECK: for.body.us:
+; CHECK-NEXT: 7 = MemoryPhi({for.header2.preheader.split.us,11},{for.header2.us,9})
+; CHECK-NEXT: 8 = MemoryDef(7)->7
+; CHECK-NEXT: store i32 0, ptr %p, align 4
+; CHECK-NEXT: 9 = MemoryDef(8)->8
+; CHECK-NEXT: store i8 0, ptr %p, align 1
+
+define void @test(ptr %p, i1 %cond) {
+entry:
+  br label %for.header
+
+for.header:
+  br i1 false, label %exit.loopexit1, label %for.header2.preheader
+
+for.header2.preheader:
+  br label %for.body
+
+for.header2:
+  br i1 false, label %for.latch, label %for.body
+
+for.body:
+  store i32 0, ptr %p, align 4
+  store i8 0, ptr %p, align 1
+  br i1 %cond, label %for.header2, label %exit.loopexit
+
+for.latch:
+  br i1 false, label %for.inc, label %exit.loopexit1
+
+for.inc:
+  br label %for.header
+
+exit.loopexit:
+  br label %exit
+
+exit.loopexit1:
+  br label %exit
+
+exit:
+  ret void
+}

[dtcxzyw]

Note: I found this bug by commenting out Next = nullptr;. We may need some assertion check to avoid this kind of bugs.

[nikic]

LGTM

[nikic]

I'd be curious though on how #138962 exposed this issue. Given that MSSA doesn't involve ConstantData, I wouldn't have expected it to change anything.