[llvm] [LoopUnroll] Rotate loop before unrolling inside of UnrollRuntimeLoopRemainder (PR #148243)

[Marek Sedláček via llvm-commits]

================
@@ -587,21 +536,116 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
       UnrollRuntimeEpilog.getNumOccurrences() ? UnrollRuntimeEpilog
                                               : isEpilogProfitable(L);
 
-  if (ULO.Runtime &&
-      !UnrollRuntimeLoopRemainder(L, ULO.Count, ULO.AllowExpensiveTripCount,
-                                  EpilogProfitability, ULO.UnrollRemainder,
-                                  ULO.ForgetAllSCEV, LI, SE, DT, AC, TTI,
-                                  PreserveLCSSA, ULO.SCEVExpansionBudget,
-                                  ULO.RuntimeUnrollMultiExit, RemainderLoop)) {
+  bool LoopRotated = false;
+  bool ReminderUnrolled = false;
+  if (ULO.Runtime) {
+    // Call unroll with disabled rotation, to see if it is possible without it.
+    ReminderUnrolled = UnrollRuntimeLoopRemainder(
+        L, ULO.Count, ULO.AllowExpensiveTripCount, EpilogProfitability,
+        ULO.UnrollRemainder, ULO.ForgetAllSCEV, LI, SE, DT, AC, TTI,
+        PreserveLCSSA, ULO.SCEVExpansionBudget, ULO.RuntimeUnrollMultiExit,
+        RemainderLoop);
+
+    // If unroll is not possible, then try with loop rotation.
+    if (!ReminderUnrolled) {
+      BasicBlock *OrigHeader = L->getHeader();
+      BranchInst *BI = dyn_cast<BranchInst>(OrigHeader->getTerminator());
+      if (BI && !BI->isUnconditional() &&
+          isa<SCEVCouldNotCompute>(SE->getExitCount(L, L->getLoopLatch())) &&
+          !isa<SCEVCouldNotCompute>(SE->getExitCount(L, OrigHeader))) {
+        LLVM_DEBUG(
+            dbgs() << "  Rotating loop to make the exit count computable.\n");
+        SimplifyQuery SQ{OrigHeader->getDataLayout()};
+        SQ.TLI = nullptr;
+        SQ.DT = DT;
+        SQ.AC = AC;
+        LoopRotated =
+            llvm::LoopRotation(L, LI, TTI, AC, DT, SE,
+                               /*MemorySSAUpdater*/ nullptr, SQ,
+                               /*RotationOnly*/ false, /*Threshold*/ 16,
+                               /*IsUtilMode*/ false, /*PrepareForLTO*/ false,
+                               [](Loop *, ScalarEvolution *) { return true; });
----------------
mark-sed wrote:

Now that I have tested it, I think having it this way is cleaner and nicer. When I tried to move the rotation inside of the `UnrollRuntimeLoopRemainder` I tried separating the legality checks into a function/lambda to call, but it relies no a lot of resources (Latch, Header, LatchBR, LatchExit, TripCountSC...) which are also used later on in the unrolling.
So it currently initializes some values, does a check, intializes other, does a check and so forth and then it unrolls. If I was to add rotation to this, then even with lambda that captures all, I will have to initialize all of them and them do the checks inside of the lambda, rotate and then re-initialize all of the values and run the check again. It becomes a mess and less efficient than just calling `UnrollRuntimeLoopRemainder` twice, since it currently can do early exit without initializing all the values.

https://github.com/llvm/llvm-project/pull/148243