[llvm] [MemorySSAUpdater] Fix iterator invalidation bug in `applyInsertUpdates` (PR #139370)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

<details>
<summary>Changes</summary>

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

Closes https://github.com/llvm/llvm-project/issues/139103.
Closes https://github.com/llvm/llvm-project/issues/139289.
Closes https://github.com/llvm/llvm-project/issues/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
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
+}

``````````

</details>


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