[llvm] [SimpleLoopUnswitch] Preserve one PHI when removing a predecessor of a BB (PR #116813)

[via llvm-commits]

https://github.com/DianQK created https://github.com/llvm/llvm-project/pull/116813

Fixes #116809.

>From b2eb1f4d3f536fe82812943ded40f8e94cb0be40 Mon Sep 17 00:00:00 2001
From: DianQK <dianqk at dianqk.net>
Date: Tue, 19 Nov 2024 22:01:12 +0800
Subject: [PATCH 1/3] Pre-commit test cases

---
 .../nontrivial-unswitch-pred-removed.ll       | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)
 create mode 100644 llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll

diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll
new file mode 100644
index 00000000000000..a5fa6c8f2acdf2
--- /dev/null
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll
@@ -0,0 +1,20 @@
+; RUN: opt -passes='loop-mssa(simple-loop-unswitch<nontrivial>),print<memoryssa>' -verify-memoryssa -disable-output -S < %s 2>&1 | FileCheck %s
+
+; CHECK: preds = %bb2{{$}}
+; CHECK-NEXT: MemoryDef
+; CHECK-NEXT: call i32 @bar()
+
+define i32 @foo(i1 %arg, ptr %arg1) {
+bb:
+  br label %bb2
+
+bb2:                                              ; preds = %bb2, %bb
+  %i = select i1 %arg, ptr %arg1, ptr @bar
+  %i3 = call i32 %i()
+  br i1 %arg, label %bb2, label %bb4
+
+bb4:                                              ; preds = %bb2
+  ret i32 %i3
+}
+
+declare i32 @bar() nounwind willreturn memory(none)

>From b164143bd0cf537cd17958628cda31acdeb86a8e Mon Sep 17 00:00:00 2001
From: DianQK <dianqk at dianqk.net>
Date: Tue, 19 Nov 2024 22:11:46 +0800
Subject: [PATCH 2/3] [SimpleLoopUnswitch] Preserve one PHI when removing a
 predecessor of a BB

---
 llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp         | 2 +-
 .../nontrivial-unswitch-pred-removed.ll                   | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index aa3cbc5e4bddc2..d6fcf57f2b0647 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1716,7 +1716,7 @@ static void deleteDeadBlocksFromLoop(Loop &L,
     auto *BB = DeathCandidates.pop_back_val();
     if (!DeadBlockSet.count(BB) && !DT.isReachableFromEntry(BB)) {
       for (BasicBlock *SuccBB : successors(BB)) {
-        SuccBB->removePredecessor(BB);
+        SuccBB->removePredecessor(BB, /*KeepOneInputPHIs*/ true);
         DeathCandidates.push_back(SuccBB);
       }
       DeadBlockSet.insert(BB);
diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll
index a5fa6c8f2acdf2..3545492f1e5e23 100644
--- a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-pred-removed.ll
@@ -1,8 +1,12 @@
 ; RUN: opt -passes='loop-mssa(simple-loop-unswitch<nontrivial>),print<memoryssa>' -verify-memoryssa -disable-output -S < %s 2>&1 | FileCheck %s
+; RUN: opt -passes='loop-mssa(simple-loop-unswitch<nontrivial>,licm)' -verify-memoryssa -disable-output -S 
 
-; CHECK: preds = %bb2{{$}}
+; Check that SimpleLoopUnswitch preserves the MemoryDef of `call i32 @bar()` by preserving the PHI node.
+; Also, check that executing LICM after SimpleLoopUnswitch does not result in a crash.
+
+; CHECK: %unswitched.select = phi ptr [ @bar, %bb2 ]
 ; CHECK-NEXT: MemoryDef
-; CHECK-NEXT: call i32 @bar()
+; CHECK-NEXT: call i32 %unswitched.select()
 
 define i32 @foo(i1 %arg, ptr %arg1) {
 bb:

>From 6b80bdb95627e4047c98e28991c503e3e90fddfa Mon Sep 17 00:00:00 2001
From: DianQK <dianqk at dianqk.net>
Date: Tue, 19 Nov 2024 22:21:25 +0800
Subject: [PATCH 3/3] [SimpleLoopUnswitch] Preserve one PHI when removing a
 predecessor of a cloned BB

---
 llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index d6fcf57f2b0647..d807c66bd96195 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1678,7 +1678,7 @@ deleteDeadClonedBlocks(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
       if (BasicBlock *ClonedBB = cast_or_null<BasicBlock>(VMap->lookup(BB)))
         if (!DT.isReachableFromEntry(ClonedBB)) {
           for (BasicBlock *SuccBB : successors(ClonedBB))
-            SuccBB->removePredecessor(ClonedBB);
+            SuccBB->removePredecessor(ClonedBB, /*KeepOneInputPHIs*/ true);
           DeadBlocks.push_back(ClonedBB);
         }