[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 88686be177f4cef24c115af081881e14b4cfd7e2 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 emits the required LCSSA
Phi instructions.
---
 llvm/lib/Transforms/Utils/SimplifyIndVar.cpp  | 13 +++++
 .../gh-issue77118-broken-lcssa-form.ll        | 58 +++++++++++++++++++
 2 files changed, 71 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 66bba1ca2f1d7c..499dde75f920ac 100644
--- a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
@@ -26,6 +26,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
+#include "llvm/Transforms/Utils/LoopUtils.h"
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
@@ -643,10 +644,22 @@ bool SimplifyIndvar::replaceIVUserWithLoopInvariant(Instruction *I) {
   }
 
   auto *Invariant = Rewriter.expandCodeFor(S, I->getType(), IP);
+  bool NeedToEmitLCSSAPhis = false;
+  if (!LI->replacementPreservesLCSSAForm(I, Invariant))
+    NeedToEmitLCSSAPhis = true;
 
   I->replaceAllUsesWith(Invariant);
   LLVM_DEBUG(dbgs() << "INDVARS: Replace IV user: " << *I
                     << " with loop invariant: " << *S << '\n');
+
+  if (NeedToEmitLCSSAPhis) {
+    SmallVector<Instruction *, 1> NeedsLCSSAPhis;
+    NeedsLCSSAPhis.push_back(dyn_cast<Instruction>(Invariant));
+    formLCSSAForInstructions(NeedsLCSSAPhis, *DT, *LI, SE);
+    LLVM_DEBUG(dbgs() << "INDVARS: replacement of IV user: " << *I
+                      << " with loop invariant: " << *S
+                      << " breaks LCSSA form, inserting LCSSA Phis" << '\n');
+  }
   ++NumFoldedUser;
   Changed = true;
   DeadInsts.emplace_back(I);
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..3381aa5e261ebe
--- /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 = 95
+; 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: {{.*}}