[llvm] [UnrollAndJam] Do not preserve loop nests if a loop was fully unrolled. (PR #133510)

[David Green via llvm-commits]

https://github.com/davemgreen created https://github.com/llvm/llvm-project/pull/133510

If UnJ completely unrolls a loop and removes it entirely, the loop remains in the current loop nest. If the loop nest gets reused the loops will no longer be valid. As there is no way to remove a loop from a LoopNest, this patch removes the preserve of the LoopNestAnalysis so that it will be regenerated.

Fixes #124518

>From a514943b220771ce800824d8f54eef794e5ea20b Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Fri, 28 Mar 2025 19:32:53 +0000
Subject: [PATCH] [UnrollAndJam] Do not preserve loop nests if a loop was fully
 unrolled.

If UnJ completely unrolls a loop and removes it entirely, the loop remains in
the current loop nest. If the loop nest gets reused the loops will no longer be
valid. As there is no way to remove a loop from a LoopNest, this patch removes
the preserve of the LoopNestAnalysis so that it will be regenerated.

Fixes #124518
---
 .../Scalar/LoopUnrollAndJamPass.cpp           | 15 +++++---
 .../LoopUnrollAndJam/delete_middle_loop.ll    | 36 +++++++++++++++++++
 2 files changed, 46 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopUnrollAndJam/delete_middle_loop.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
index f1d1f3bc1e307..4fe74c7c3bbcd 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
@@ -425,7 +425,7 @@ static bool tryToUnrollAndJamLoop(LoopNest &LN, DominatorTree &DT, LoopInfo &LI,
                                   const TargetTransformInfo &TTI,
                                   AssumptionCache &AC, DependenceInfo &DI,
                                   OptimizationRemarkEmitter &ORE, int OptLevel,
-                                  LPMUpdater &U) {
+                                  LPMUpdater &U, bool &AnyLoopRemoved) {
   bool DidSomething = false;
   ArrayRef<Loop *> Loops = LN.getLoops();
   Loop *OutmostLoop = &LN.getOutermostLoop();
@@ -441,8 +441,11 @@ static bool tryToUnrollAndJamLoop(LoopNest &LN, DominatorTree &DT, LoopInfo &LI,
         tryToUnrollAndJamLoop(L, DT, &LI, SE, TTI, AC, DI, ORE, OptLevel);
     if (Result != LoopUnrollResult::Unmodified)
       DidSomething = true;
-    if (L == OutmostLoop && Result == LoopUnrollResult::FullyUnrolled)
-      U.markLoopAsDeleted(*L, LoopName);
+    if (Result == LoopUnrollResult::FullyUnrolled) {
+      if (L == OutmostLoop)
+        U.markLoopAsDeleted(*L, LoopName);
+      AnyLoopRemoved = true;
+    }
   }
 
   return DidSomething;
@@ -457,11 +460,13 @@ PreservedAnalyses LoopUnrollAndJamPass::run(LoopNest &LN,
   DependenceInfo DI(&F, &AR.AA, &AR.SE, &AR.LI);
   OptimizationRemarkEmitter ORE(&F);
 
+  bool AnyLoopRemoved = false;
   if (!tryToUnrollAndJamLoop(LN, AR.DT, AR.LI, AR.SE, AR.TTI, AR.AC, DI, ORE,
-                             OptLevel, U))
+                             OptLevel, U, AnyLoopRemoved))
     return PreservedAnalyses::all();
 
   auto PA = getLoopPassPreservedAnalyses();
-  PA.preserve<LoopNestAnalysis>();
+  if (!AnyLoopRemoved)
+    PA.preserve<LoopNestAnalysis>();
   return PA;
 }
diff --git a/llvm/test/Transforms/LoopUnrollAndJam/delete_middle_loop.ll b/llvm/test/Transforms/LoopUnrollAndJam/delete_middle_loop.ll
new file mode 100644
index 0000000000000..f8affdb821903
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnrollAndJam/delete_middle_loop.ll
@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes="loop(invalidate<all>,loop-unroll-and-jam,loop-unroll-and-jam)" -allow-unroll-and-jam -unroll-and-jam-count=4 < %s -S | FileCheck %s
+
+; This test completely unrolls the middle loop out of a 3-deep loop nest.
+
+define i16 @test_it() {
+; CHECK-LABEL: define i16 @test_it() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[FOR_COND:.*]]
+; CHECK:       [[FOR_COND_LOOPEXIT:.*]]:
+; CHECK-NEXT:    br label %[[FOR_COND]]
+; CHECK:       [[FOR_COND]]:
+; CHECK-NEXT:    br label %[[DO_BODY2:.*]]
+; CHECK:       [[DO_BODY2]]:
+; CHECK-NEXT:    br label %[[WHILE_COND3:.*]]
+; CHECK:       [[WHILE_COND3]]:
+; CHECK-NEXT:    br i1 true, label %[[DO_COND:.*]], label %[[WHILE_COND3]]
+; CHECK:       [[DO_COND]]:
+; CHECK-NEXT:    br label %[[FOR_COND_LOOPEXIT]]
+;
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %do.cond, %entry
+  br label %do.body2
+
+do.body2:                                         ; preds = %do.cond, %for.cond
+  br label %while.cond3
+
+while.cond3:                                      ; preds = %while.cond3, %do.body2
+  br i1 true, label %do.cond, label %while.cond3
+
+do.cond:                                          ; preds = %while.cond3
+  br i1 true, label %for.cond, label %do.body2
+}
+