[UnrollAndJam] Do not preserve loop nests if a loop was fully unrolled. (#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

2 files changed, 46 insertions, 5 deletions

diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp
index f1d1f3b..4fe74c7 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 0000000..f8affdb
--- /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
+}
+