[UnrollAndJam] Do not preserve loop nests if a loop was fully unrolled.

[davemgreen]

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

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: David Green (davemgreen)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/133510.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp (+10-5)
• (added) llvm/test/Transforms/LoopUnrollAndJam/delete_middle_loop.ll (+36)

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
+}
+

[aeubanks]

if the PA.preserve<LoopNestAnalysis>(); is correct in cases where we don't remove the loop, then lgtm

[maekawatoshiki]

This looks nice to me.