[llvm] [LoopPeel] LCSSA form is destroyed by LoopPeel, preserve it (PR #78696)

[Vedant Paranjape via llvm-commits]

https://github.com/vedantparanjape-amd updated https://github.com/llvm/llvm-project/pull/78696

>From 31d98719b41a8868865f91e293fe66029b78bf50 Mon Sep 17 00:00:00 2001
From: Vedant Paranjape <vedant.paranjape at amd.com>
Date: Fri, 19 Jan 2024 10:53:36 +0000
Subject: [PATCH] [SimplifyIndVar] LCSSA form is destroyed by simplifyLoopIVs,
 preserve it

In LoopUnroll, peelLoop is called on the loop. After the loop is peeled
it calls simplifyLoopAfterUnroll on the loop. This call to
simplifyLoopAfterUnroll doesn't preserve the LCSSA form of the parent
loop and thus during the next call to peelLoop the LCSSA form is already
broken.

LoopPeel util takes in the PreserveLCSSA argument and it passes
on the same argument to simplifyLoop which checks if the loop is in a
valid LCSSA form, when (PreserveLCSSA = true).

This causes an assert in simplifyLoop when (PreserveLCSSA = true), as
during the last call LCSSA for the loop wasn't preserved, and thus
crashes at the following assert.

assert(L->isRecursivelyLCSSAForm(*DT, *LI) &&
           "Requested to preserve LCSSA, but it's already broken.");

Upon debugging, it is evident that simplifyLoopIVs call inside
simplifyLoopAfterUnroll breaks the LCSSA form. This patch fixes
llvm#77118, it checks if the replacement of IV Users with Loop Invariant
preserves the LCSSA form. If it does not, it simply bails out and
returns false.
---
 llvm/lib/Transforms/Utils/SimplifyIndVar.cpp  |  7 +++
 .../gh-issue77118-broken-lcssa-form.ll        | 58 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 llvm/test/Transforms/LoopUnroll/gh-issue77118-broken-lcssa-form.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
index 5aa6df49e7f676..4f9bbaaa02588a 100644
--- a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
@@ -644,6 +644,13 @@ bool SimplifyIndvar::replaceIVUserWithLoopInvariant(Instruction *I) {
 
   auto *Invariant = Rewriter.expandCodeFor(S, I->getType(), IP);
 
+  if (!LI->replacementPreservesLCSSAForm(I, Invariant)) {
+    LLVM_DEBUG(dbgs() << "INDVARS: Can not replace IV user: " << *I
+                      << " with loop invariant: "  << *S
+                      << " as it breaks LCSSA form " << '\n');
+    return false;
+  }
+
   I->replaceAllUsesWith(Invariant);
   LLVM_DEBUG(dbgs() << "INDVARS: Replace IV user: " << *I
                     << " with loop invariant: " << *S << '\n');
diff --git a/llvm/test/Transforms/LoopUnroll/gh-issue77118-broken-lcssa-form.ll b/llvm/test/Transforms/LoopUnroll/gh-issue77118-broken-lcssa-form.ll
new file mode 100644
index 00000000000000..e530edf68199bb
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/gh-issue77118-broken-lcssa-form.ll
@@ -0,0 +1,58 @@
+; RUN: opt -passes=loop-unroll -unroll-peel-count=2 -S -disable-output -debug-only=loop-unroll < %s 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+define void @test() {
+; CHECK-LABEL: Loop Unroll: F[test] Loop %loop3
+; CHECK-NEXT:    Loop Size = 7
+; CHECK-NEXT:  PEELING loop %loop3 with iteration count 2!
+; CHECK-NEXT:  Loop Unroll: F[test] Loop %loop2
+; CHECK-NEXT:    Loop Size = 28
+; CHECK-NEXT:  PEELING loop %loop2 with iteration count 2!
+; CHECK-NEXT:  Loop Unroll: F[test] Loop %loop4
+; CHECK-NEXT:    Loop Size = 3
+; CHECK-NEXT:  PEELING loop %loop4 with iteration count 2!
+; CHECK-NEXT:  Loop Unroll: F[test] Loop %loop1
+; CHECK-NEXT:    Loop Size = 98
+; CHECK-NEXT:  PEELING loop %loop1 with iteration count 2!
+entry:
+  br label %loop1
+
+loop1:
+  %phi = phi i32 [ 1, %entry ], [ 0, %loop1.latch ]
+  br label %loop2
+
+loop2:
+  %phi3 = phi i64 [ 0, %loop1 ], [ %sext, %loop2.latch ]
+  br label %loop3
+
+loop3:
+  %phi5 = phi i64 [ %phi3, %loop2 ], [ %sext, %loop3.latch ]
+  %phi6 = phi i32 [ 1, %loop2 ], [ %add10, %loop3.latch ]
+  %trunc = trunc i64 %phi5 to i32
+  br i1 true, label %loop3.latch, label %exit
+
+loop3.latch:
+  %add = add i32 1, %phi
+  %sext = sext i32 %add to i64
+  %add10 = add i32 %phi6, 1
+  %icmp = icmp ugt i32 %add10, 2
+  br i1 %icmp, label %loop2.latch, label %loop3
+
+loop2.latch:
+  br i1 false, label %loop4.preheader, label %loop2
+
+loop4.preheader:
+  br label %loop4
+
+loop4:
+  br i1 false, label %loop1.latch, label %loop4
+
+loop1.latch:
+  br label %loop1
+
+exit:
+  %phi8 = phi i32 [ %trunc, %loop3 ]
+  ret void
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CHECK: {{.*}}