From fc5c89900f2a4b50e0f3a88ef7c89115d93684f4 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Wed, 27 Nov 2024 11:47:22 +0100
Subject: [PATCH] [SimpleLoopUnswitch] Fix LCSSA phi node invalidation

Fixes https://github.com/llvm/llvm-project/issues/117537.
---
 .../Transforms/Scalar/SimpleLoopUnswitch.cpp  |  5 +-
 .../Transforms/SimpleLoopUnswitch/pr117537.ll | 92 +++++++++++++++++++
 2 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/SimpleLoopUnswitch/pr117537.ll

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index aa3cbc5e4bddc2..d8ef450eeb9a15 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1248,8 +1248,9 @@ static BasicBlock *buildClonedLoopBlocks(
       assert(VMap.lookup(&I) == &ClonedI && "Mismatch in the value map!");
 
       // Forget SCEVs based on exit phis in case SCEV looked through the phi.
-      if (SE && isa<PHINode>(I))
-        SE->forgetValue(&I);
+      if (SE)
+        if (auto *PN = dyn_cast<PHINode>(&I))
+          SE->forgetLcssaPhiWithNewPredecessor(&L, PN);
 
       BasicBlock::iterator InsertPt = MergeBB->getFirstInsertionPt();
 
diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/pr117537.ll b/llvm/test/Transforms/SimpleLoopUnswitch/pr117537.ll
new file mode 100644
index 00000000000000..fd61cfab164d3b
--- /dev/null
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/pr117537.ll
@@ -0,0 +1,92 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes='print<scalar-evolution>,simple-loop-unswitch<nontrivial>,print<scalar-evolution>' -verify-scev < %s 2>/dev/null | FileCheck %s
+
+; Make sure we don't assert due to insufficient SCEV invalidation.
+
+define void @test(ptr %p) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CHECK:%.*]] = icmp eq ptr [[P]], null
+; CHECK-NEXT:    br i1 [[CHECK]], label %[[ENTRY_SPLIT_US:.*]], label %[[ENTRY_SPLIT:.*]]
+; CHECK:       [[ENTRY_SPLIT_US]]:
+; CHECK-NEXT:    br label %[[BB0_US:.*]]
+; CHECK:       [[BB0_US]]:
+; CHECK-NEXT:    br label %[[LOOP0_US:.*]]
+; CHECK:       [[LOOP0_US]]:
+; CHECK-NEXT:    [[V_US:%.*]] = load atomic i32, ptr [[P]] unordered, align 8
+; CHECK-NEXT:    [[ADD_US:%.*]] = add i32 [[V_US]], 3
+; CHECK-NEXT:    br i1 true, label %[[PREHEADER_SPLIT_US:.*]], label %[[BB0_US]]
+; CHECK:       [[PREHEADER_SPLIT_US]]:
+; CHECK-NEXT:    [[ADD_LCSSA_US:%.*]] = phi i32 [ [[ADD_US]], %[[LOOP0_US]] ]
+; CHECK-NEXT:    br label %[[PREHEADER:.*]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    br label %[[BB0:.*]]
+; CHECK:       [[BB0]]:
+; CHECK-NEXT:    br label %[[LATCH:.*]]
+; CHECK:       [[LATCH]]:
+; CHECK-NEXT:    br i1 false, label %[[EXIT0:.*]], label %[[LOOP0:.*]]
+; CHECK:       [[EXIT0]]:
+; CHECK-NEXT:    ret void
+; CHECK:       [[LOOP0]]:
+; CHECK-NEXT:    [[V:%.*]] = load atomic i32, ptr [[P]] unordered, align 8
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[V]], 3
+; CHECK-NEXT:    br i1 true, label %[[PREHEADER_SPLIT:.*]], label %[[BB0]]
+; CHECK:       [[PREHEADER_SPLIT]]:
+; CHECK-NEXT:    [[ADD_LCSSA:%.*]] = phi i32 [ [[ADD]], %[[LOOP0]] ]
+; CHECK-NEXT:    br label %[[PREHEADER]]
+; CHECK:       [[PREHEADER]]:
+; CHECK-NEXT:    [[DOTUS_PHI:%.*]] = phi i32 [ [[ADD_LCSSA]], %[[PREHEADER_SPLIT]] ], [ [[ADD_LCSSA_US]], %[[PREHEADER_SPLIT_US]] ]
+; CHECK-NEXT:    br label %[[LOOP1:.*]]
+; CHECK:       [[LOOP1]]:
+; CHECK-NEXT:    [[IV1:%.*]] = phi i32 [ [[DOTUS_PHI]], %[[PREHEADER]] ], [ [[IV1_NEXT:%.*]], %[[BACKEDGE:.*]] ]
+; CHECK-NEXT:    [[IV1_NEXT]] = add i32 [[IV1]], -33
+; CHECK-NEXT:    br label %[[LOOP2:.*]]
+; CHECK:       [[BACKEDGE]]:
+; CHECK-NEXT:    br i1 true, label %[[EXIT1:.*]], label %[[LOOP1]]
+; CHECK:       [[LOOP2]]:
+; CHECK-NEXT:    [[IV0:%.*]] = phi i32 [ [[IV1]], %[[LOOP1]] ], [ [[IV0_NEXT:%.*]], %[[LOOP2]] ]
+; CHECK-NEXT:    [[IV0_NEXT]] = add nsw i32 [[IV0]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[IV0_NEXT]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label %[[BACKEDGE]], label %[[LOOP2]]
+; CHECK:       [[EXIT1]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %check = icmp eq ptr %p, null
+  br label %bb0
+
+bb0:                                              ; preds = %loop0, %entry
+  br i1 %check, label %loop0, label %latch
+
+latch:                                            ; preds = %bb0
+  br i1 %check, label %exit0, label %loop0
+
+exit0:                                            ; preds = %latch
+  ret void
+
+loop0:                                            ; preds = %latch, %bb0
+  %v = load atomic i32, ptr %p unordered, align 8
+  %add = add i32 %v, 3
+  br i1 true, label %preheader, label %bb0
+
+preheader:                                        ; preds = %loop0
+  br label %loop1
+
+loop1:                                            ; preds = %backedge, %preheader
+  %iv1 = phi i32 [ %add, %preheader ], [ %iv1.next, %backedge ]
+  %iv1.next = add i32 %iv1, -33
+  br label %loop2
+
+backedge:                                         ; preds = %loop2
+  br i1 true, label %exit1, label %loop1
+
+loop2:                                            ; preds = %loop2, %loop1
+  %iv0 = phi i32 [ %iv1, %loop1 ], [ %iv0.next, %loop2 ]
+  %iv0.next = add nsw i32 %iv0, 1
+  %cmp = icmp sgt i32 %iv0.next, 0
+  br i1 %cmp, label %backedge, label %loop2
+
+exit1:                                            ; preds = %backedge
+  ret void
+}