From 57ae9bb9323548b2ad4ba8274c3910bf9c764983 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Sun, 20 Sep 2020 12:40:35 +0100
Subject: [PATCH] [LSR] Preserve MSSA when using SplitCriticalEdge.

LSR claims to MemorySSA, but we also have to make sure it is preserved
when splitting critical edges. This can be done by passing MSSAU to
SplitCriticalEdge.

Fixes PR47557.
---
 .../Transforms/Scalar/LoopStrengthReduce.cpp  |  9 ++--
 .../MemorySSA/update-remove-dead-blocks.ll    | 46 +++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Analysis/MemorySSA/update-remove-dead-blocks.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 47329fa1f043..936ab8e823df 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -5382,10 +5382,11 @@ void LSRInstance::RewriteForPHI(
           // Split the critical edge.
           BasicBlock *NewBB = nullptr;
           if (!Parent->isLandingPad()) {
-            NewBB = SplitCriticalEdge(BB, Parent,
-                                      CriticalEdgeSplittingOptions(&DT, &LI)
-                                          .setMergeIdenticalEdges()
-                                          .setKeepOneInputPHIs());
+            NewBB =
+                SplitCriticalEdge(BB, Parent,
+                                  CriticalEdgeSplittingOptions(&DT, &LI, MSSAU)
+                                      .setMergeIdenticalEdges()
+                                      .setKeepOneInputPHIs());
           } else {
             SmallVector<BasicBlock*, 2> NewBBs;
             SplitLandingPadPredecessors(Parent, BB, "", "", NewBBs, &DT, &LI);
diff --git a/llvm/test/Analysis/MemorySSA/update-remove-dead-blocks.ll b/llvm/test/Analysis/MemorySSA/update-remove-dead-blocks.ll
new file mode 100644
index 000000000000..fab278784583
--- /dev/null
+++ b/llvm/test/Analysis/MemorySSA/update-remove-dead-blocks.ll
@@ -0,0 +1,46 @@
+; 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
+}