From f661ce209f190c993bc8635adba5476edbe90a4a Mon Sep 17 00:00:00 2001
From: Anna Thomas <anna@azul.com>
Date: Thu, 2 Sep 2021 17:26:45 -0400
Subject: [PATCH] [LoopPredication] Fix MemorySSA crash in predicateLoopExits

The attached testcase crashes without the patch (Not the same accesses
in the same order).

When we move instructions before another instruction, we also need to
update the memory accesses corresponding to it.

Reviewed-By: asbirlea
Differential Revision: https://reviews.llvm.org/D109197
---
 llvm/lib/Transforms/Scalar/LoopPredication.cpp     |  4 ++++
 .../Transforms/LoopPredication/predicate-exits.ll  | 28 ++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/LoopPredication.cpp b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
index 49c5b96..0d5d7e7 100644
--- a/llvm/lib/Transforms/Scalar/LoopPredication.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
@@ -1127,6 +1127,10 @@ bool LoopPredication::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
   // modifying the operand is legal.
   auto *IP = cast<Instruction>(WidenableBR->getCondition());
   IP->moveBefore(WidenableBR);
+  if (MSSAU)
+    if (auto *MUD = MSSAU->getMemorySSA()->getMemoryAccess(IP))
+       MSSAU->moveToPlace(MUD, WidenableBR->getParent(),
+                          MemorySSA::BeforeTerminator);
   Rewriter.setInsertPoint(IP);
   IRBuilder<> B(IP);
 
diff --git a/llvm/test/Transforms/LoopPredication/predicate-exits.ll b/llvm/test/Transforms/LoopPredication/predicate-exits.ll
index 326c967..74f32ae 100644
--- a/llvm/test/Transforms/LoopPredication/predicate-exits.ll
+++ b/llvm/test/Transforms/LoopPredication/predicate-exits.ll
@@ -1082,6 +1082,34 @@ exit:
   ret i32 %result
 }
 
+define void @test_memssa() {
+bb:
+  %tmp = call i1 @llvm.experimental.widenable.condition()
+  %tmp1 = call i1 @llvm.experimental.widenable.condition()
+  br i1 %tmp, label %bb3, label %bb2
+
+bb2:                                              ; preds = %bb
+  unreachable
+
+bb3:                                              ; preds = %bb
+  br label %bb4
+
+bb4:                                              ; preds = %bb6, %bb3
+  %tmp5 = phi i32 [ %tmp7, %bb6 ], [ 0, %bb3 ]
+  br i1 undef, label %bb10, label %bb6
+
+bb6:                                              ; preds = %bb4
+  %tmp7 = add nuw nsw i32 %tmp5, 1
+  %tmp8 = icmp ult i32 %tmp7, undef
+  br i1 %tmp8, label %bb4, label %bb9
+
+bb9:                                              ; preds = %bb6
+  ret void
+
+bb10:                                             ; preds = %bb4
+  ret void
+}
+
 
 
 declare void @unknown()
-- 
2.7.4