[LoopUnrollAndJam] Change LoopUnrollAndJamPass to LoopNest pass

This patch changes LoopUnrollAndJamPass from FunctionPass to LoopNest pass.
The next patch will utilize LoopNest to effectively handle loop nests.

Reviewed By: Whitney

Differential Revision: https://reviews.llvm.org/D99149

llvm/include/llvm/Transforms/Scalar/LoopUnrollAndJamPass.h

@@ -10,6 +10,7 @@
 #define LLVM_TRANSFORMS_SCALAR_LOOPUNROLLANDJAMPASS_H
 
 #include "llvm/IR/PassManager.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
 
 namespace llvm {
 class Function;
@@ -20,7 +21,8 @@ class LoopUnrollAndJamPass : public PassInfoMixin<LoopUnrollAndJamPass> {
 
 public:
   explicit LoopUnrollAndJamPass(int OptLevel = 2) : OptLevel(OptLevel) {}
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  PreservedAnalyses run(LoopNest &L, LoopAnalysisManager &AM,
+                        LoopStandardAnalysisResults &AR, LPMUpdater &U);
 };
 
 } // end namespace llvm

llvm/include/llvm/Transforms/Utils/UnrollLoop.h

@@ -17,6 +17,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
 
 namespace llvm {
 
@@ -97,7 +98,7 @@ LoopUnrollResult UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
                                   LoopInfo *LI, ScalarEvolution *SE,
                                   DominatorTree *DT, AssumptionCache *AC,
                                   const TargetTransformInfo *TTI,
-                                  OptimizationRemarkEmitter *ORE,
+                                  OptimizationRemarkEmitter *ORE, LPMUpdater *U,
                                   Loop **EpilogueLoop = nullptr);
 
 bool isSafeToUnrollAndJam(Loop *L, ScalarEvolution &SE, DominatorTree &DT,

llvm/lib/Passes/PassBuilder.cpp

@@ -1207,7 +1207,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
     // across the loop nests.
     // We do UnrollAndJam in a separate LPM to ensure it happens before unroll
     if (EnableUnrollAndJam && PTO.LoopUnrolling)
-      FPM.addPass(LoopUnrollAndJamPass(Level.getSpeedupLevel()));
+      FPM.addPass(createFunctionToLoopPassAdaptor(
+          LoopUnrollAndJamPass(Level.getSpeedupLevel())));
     FPM.addPass(LoopUnrollPass(LoopUnrollOptions(
         Level.getSpeedupLevel(), /*OnlyWhenForced=*/!PTO.LoopUnrolling,
         PTO.ForgetAllSCEVInLoopUnroll)));
@@ -1290,7 +1291,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
     // across the loop nests.
     // We do UnrollAndJam in a separate LPM to ensure it happens before unroll
     if (EnableUnrollAndJam && PTO.LoopUnrolling) {
-      FPM.addPass(LoopUnrollAndJamPass(Level.getSpeedupLevel()));
+      FPM.addPass(createFunctionToLoopPassAdaptor(
+          LoopUnrollAndJamPass(Level.getSpeedupLevel())));
     }
     FPM.addPass(LoopUnrollPass(LoopUnrollOptions(
         Level.getSpeedupLevel(), /*OnlyWhenForced=*/!PTO.LoopUnrolling,

llvm/lib/Passes/PassRegistry.def

@@ -247,7 +247,6 @@ FUNCTION_PASS("guard-widening", GuardWideningPass())
 FUNCTION_PASS("load-store-vectorizer", LoadStoreVectorizerPass())
 FUNCTION_PASS("loop-simplify", LoopSimplifyPass())
 FUNCTION_PASS("loop-sink", LoopSinkPass())
-FUNCTION_PASS("loop-unroll-and-jam", LoopUnrollAndJamPass())
 FUNCTION_PASS("loop-flatten", LoopFlattenPass())
 FUNCTION_PASS("lowerinvoke", LowerInvokePass())
 FUNCTION_PASS("lowerswitch", LowerSwitchPass())
@@ -399,6 +398,7 @@ LOOP_PASS("loop-deletion", LoopDeletionPass())
 LOOP_PASS("loop-simplifycfg", LoopSimplifyCFGPass())
 LOOP_PASS("loop-reduce", LoopStrengthReducePass())
 LOOP_PASS("indvars", IndVarSimplifyPass())
+LOOP_PASS("loop-unroll-and-jam", LoopUnrollAndJamPass())
 LOOP_PASS("loop-unroll-full", LoopFullUnrollPass())
 LOOP_PASS("print-access-info", LoopAccessInfoPrinterPass(dbgs()))
 LOOP_PASS("print<ddg>", DDGAnalysisPrinterPass(dbgs()))

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

@@ -22,6 +22,7 @@
 #include "llvm/Analysis/DependenceAnalysis.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
@@ -278,11 +279,10 @@ static bool computeUnrollAndJamCount(
   return false;
 }
 
-static LoopUnrollResult
+static LoopUnrollResult tryToUnrollAndJamLoop(
-tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
+    Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
-                      ScalarEvolution &SE, const TargetTransformInfo &TTI,
+    const TargetTransformInfo &TTI, AssumptionCache &AC, DependenceInfo &DI,
-                      AssumptionCache &AC, DependenceInfo &DI,
+    OptimizationRemarkEmitter &ORE, int OptLevel, LPMUpdater *U) {
-                      OptimizationRemarkEmitter &ORE, int OptLevel) {
   TargetTransformInfo::UnrollingPreferences UP =
       gatherUnrollingPreferences(L, SE, TTI, nullptr, nullptr, OptLevel, None,
                                  None, None, None, None, None);
@@ -385,7 +385,7 @@ tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
   Loop *EpilogueOuterLoop = nullptr;
   LoopUnrollResult UnrollResult = UnrollAndJamLoop(
       L, UP.Count, OuterTripCount, OuterTripMultiple, UP.UnrollRemainder, LI,
-      &SE, &DT, &AC, &TTI, &ORE, &EpilogueOuterLoop);
+      &SE, &DT, &AC, &TTI, &ORE, U, &EpilogueOuterLoop);
 
   // Assign new loop attributes.
   if (EpilogueOuterLoop) {
@@ -424,33 +424,23 @@ tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
   return UnrollResult;
 }
 
-static bool tryToUnrollAndJamLoop(Function &F, DominatorTree &DT, LoopInfo &LI,
+static bool tryToUnrollAndJamLoop(LoopNest &LN, DominatorTree &DT, LoopInfo &LI,
                                   ScalarEvolution &SE,
                                   const TargetTransformInfo &TTI,
                                   AssumptionCache &AC, DependenceInfo &DI,
-                                  OptimizationRemarkEmitter &ORE,
+                                  OptimizationRemarkEmitter &ORE, int OptLevel,
-                                  int OptLevel) {
+                                  LPMUpdater &U) {
   bool DidSomething = false;
+  ArrayRef<Loop *> Loops = LN.getLoops();
 
-  // The loop unroll and jam pass requires loops to be in simplified form, and
+  // Add the loop nests in the reverse order of LN. See method
-  // also needs LCSSA. Since simplification may add new inner loops, it has to
-  // run before the legality and profitability checks. This means running the
-  // loop unroll and jam pass will simplify all loops, regardless of whether
-  // anything end up being unroll and jammed.
-  for (auto &L : LI) {
-    DidSomething |=
-        simplifyLoop(L, &DT, &LI, &SE, &AC, nullptr, false /* PreserveLCSSA */);
-    DidSomething |= formLCSSARecursively(*L, DT, &LI, &SE);
-  }
-
-  // Add the loop nests in the reverse order of LoopInfo. See method
   // declaration.
   SmallPriorityWorklist<Loop *, 4> Worklist;
-  appendLoopsToWorklist(LI, Worklist);
+  appendLoopsToWorklist(Loops, Worklist);
   while (!Worklist.empty()) {
     Loop *L = Worklist.pop_back_val();
     LoopUnrollResult Result =
-        tryToUnrollAndJamLoop(L, DT, &LI, SE, TTI, AC, DI, ORE, OptLevel);
+        tryToUnrollAndJamLoop(L, DT, &LI, SE, TTI, AC, DI, ORE, OptLevel, &U);
     if (Result != LoopUnrollResult::Unmodified)
       DidSomething = true;
   }
@@ -460,29 +450,35 @@ static bool tryToUnrollAndJamLoop(Function &F, DominatorTree &DT, LoopInfo &LI,
 
 namespace {
 
-class LoopUnrollAndJam : public FunctionPass {
+class LoopUnrollAndJam : public LoopPass {
 public:
   static char ID; // Pass ID, replacement for typeid
   unsigned OptLevel;
 
-  LoopUnrollAndJam(int OptLevel = 2) : FunctionPass(ID), OptLevel(OptLevel) {
+  LoopUnrollAndJam(int OptLevel = 2) : LoopPass(ID), OptLevel(OptLevel) {
     initializeLoopUnrollAndJamPass(*PassRegistry::getPassRegistry());
   }
 
-  bool runOnFunction(Function &F) override {
+  bool runOnLoop(Loop *L, LPPassManager &LPM) override {
-    if (skipFunction(F))
+    if (skipLoop(L))
       return false;
 
-    auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+    auto *F = L->getHeader()->getParent();
-    LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+    auto &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
-    ScalarEvolution &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
+    auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
-    const TargetTransformInfo &TTI =
-        getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
-    auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
     auto &DI = getAnalysis<DependenceAnalysisWrapperPass>().getDI();
+    auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+    auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(*F);
     auto &ORE = getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE();
+    auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(*F);
 
-    return tryToUnrollAndJamLoop(F, DT, LI, SE, TTI, AC, DI, ORE, OptLevel);
+    LoopUnrollResult Result = tryToUnrollAndJamLoop(L, DT, LI, SE, TTI, AC, DI,
+                                                    ORE, OptLevel, nullptr);
+
+    if (Result == LoopUnrollResult::FullyUnrolled)
+      LPM.markLoopAsDeleted(*L);
+
+    return Result != LoopUnrollResult::Unmodified;
   }
 
   /// This transformation requires natural loop information & requires that
@@ -505,7 +501,10 @@ char LoopUnrollAndJam::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopUnrollAndJam, "loop-unroll-and-jam",
                       "Unroll and Jam loops", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(LoopPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
+INITIALIZE_PASS_DEPENDENCY(LCSSAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
@@ -518,19 +517,20 @@ Pass *llvm::createLoopUnrollAndJamPass(int OptLevel) {
   return new LoopUnrollAndJam(OptLevel);
 }
 
-PreservedAnalyses LoopUnrollAndJamPass::run(Function &F,
+PreservedAnalyses LoopUnrollAndJamPass::run(LoopNest &LN,
-                                            FunctionAnalysisManager &AM) {
+                                            LoopAnalysisManager &AM,
-  ScalarEvolution &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
+                                            LoopStandardAnalysisResults &AR,
-  LoopInfo &LI = AM.getResult<LoopAnalysis>(F);
+                                            LPMUpdater &U) {
-  TargetTransformInfo &TTI = AM.getResult<TargetIRAnalysis>(F);
+  Function &F = *LN.getParent();
-  AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
-  DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
-  DependenceInfo &DI = AM.getResult<DependenceAnalysis>(F);
-  OptimizationRemarkEmitter &ORE =
-      AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
 
-  if (!tryToUnrollAndJamLoop(F, DT, LI, SE, TTI, AC, DI, ORE, OptLevel))
+  DependenceInfo DI(&F, &AR.AA, &AR.SE, &AR.LI);
+  OptimizationRemarkEmitter ORE(&F);
+
+  if (!tryToUnrollAndJamLoop(LN, AR.DT, AR.LI, AR.SE, AR.TTI, AR.AC, DI, ORE,
+                             OptLevel, U))
     return PreservedAnalyses::all();
 
-  return getLoopPassPreservedAnalyses();
+  auto PA = getLoopPassPreservedAnalyses();
+  PA.preserve<LoopNestAnalysis>();
+  return PA;
 }

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

@@ -49,6 +49,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/GenericDomTree.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
@@ -221,12 +222,11 @@ static void moveHeaderPhiOperandsToForeBlocks(BasicBlock *Header,
   If EpilogueLoop is non-null, it receives the epilogue loop (if it was
   necessary to create one and not fully unrolled).
 */
-LoopUnrollResult
+LoopUnrollResult llvm::UnrollAndJamLoop(
-llvm::UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
+    Loop *L, unsigned Count, unsigned TripCount, unsigned TripMultiple,
-                       unsigned TripMultiple, bool UnrollRemainder,
+    bool UnrollRemainder, LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT,
-                       LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT,
+    AssumptionCache *AC, const TargetTransformInfo *TTI,
-                       AssumptionCache *AC, const TargetTransformInfo *TTI,
+    OptimizationRemarkEmitter *ORE, LPMUpdater *U, Loop **EpilogueLoop) {
-                       OptimizationRemarkEmitter *ORE, Loop **EpilogueLoop) {
 
   // When we enter here we should have already checked that it is safe
   BasicBlock *Header = L->getHeader();
@@ -605,8 +605,11 @@ llvm::UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
   ++NumUnrolledAndJammed;
 
   // Update LoopInfo if the loop is completely removed.
-  if (CompletelyUnroll)
+  if (CompletelyUnroll) {
+    if (U)
+      U->markLoopAsDeleted(*L, std::string(L->getName()));
     LI->erase(L);
+  }
 
 #ifndef NDEBUG
   // We shouldn't have done anything to break loop simplify form or LCSSA.

llvm/test/Transforms/LoopUnrollAndJam/innerloop.ll

@@ -1,5 +1,5 @@
 ; RUN: opt -loop-unroll-and-jam -allow-unroll-and-jam -verify-loop-info < %s -S | FileCheck %s
-; RUN: opt -passes='loop-unroll-and-jam,verify<loops>' -allow-unroll-and-jam < %s -S | FileCheck %s
+; RUN: opt -passes='loop(loop-unroll-and-jam),verify<loops>' -allow-unroll-and-jam < %s -S | FileCheck %s
 
 ; Check that the newly created loops to not fail to be added to LI
 ; This test deliberately disables UnJ on the middle loop, performing it instead on the