[LoopNest][LoopFlatten] Change LoopFlattenPass to LoopNest pass

[eopXD]

Summary

This patch changes LoopFlattenPass from FunctionPass to LoopNestPass.

Utilize LoopNest and let function 'Flatten' generate information from it.

[eopXD]

Apply clang-format.

[eopXD]

Utilize LoopNest in Flatten, this shall eliminate repeating calls to FlattenLoopPair

[eopXD]

Apply clang-tidy.

[eopXD]

Reset update on files.

[Whitney]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

71

Please avoid unnecessary changes.

729–730

We can keep the LPM manager as FunctionPass, as after this change, a loop nest can be transform repeatedly the number of loops time.

[eopXD]

address comments.

[Whitney]

llvm/lib/Passes/PassRegistry.def

425

Please undo this no newline change.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

92

Please avoid unnecessary changes.
Same goes to other changes.

[eopXD]

revert the changes caused by clang-format.

[eopXD]

more recoveries from clang-format.

[eopXD]

Added new lines at end of file,
also removed the redundant header I added.

[dmgreen]

D99149 ended up needing to use the LPMUpdater as the loops may be removed. Does the same thing need to be done here? As loop flattening does inherently remove loops.

[eopXD]

D99149 ended up needing to use the LPMUpdater as the loops may be removed. Does the same thing need to be done here? As loop flattening does inherently remove loops.

Loop flattening aims to remove the inner loop of a nested loop,
but under LoopNestMode there is an assertion in LPMUpdater::markLoopAsDeleted that checks if the deleting loop the outer-most loop.
I think the assertion should not hold, the function shall allow markLoopAsDeleted for sub-loop in LoopNestMode,
and check that whether the mark-deleted-loop is a sub-loop, then call revisitCurrentLoop if so.

Thank you for spotting out the problem.

[eopXD]

added LPMUpdater for loop removal,
updated LPMUpdater::markLoopAsDeleted, allow sub-loop to be deleted in LoopNestMode.

[Whitney]

llvm/include/llvm/Transforms/Scalar/LoopPassManager.h

262 ↗

(On Diff #347340)

Can you please explain why should revisitCurrentLoop?

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

734

There can be more than one loop nest in a function. You may need a loop like for (auto *InnerLoop : LI->getLoopsInPreorder()) { here.

[eopXD]

@Whitney : Can you please explain why should revisitCurrentLoop?

I think I have misused revisitCurrentLoop here, the loop don't need to be revisited again,
because Flatten will be called for go through the possible inner/outer loop pairs.

@dmgreen
In LoopNestMode, sub-loops are not in the work list of LPMUpdater, so inner-loops can't be marked by deleted.
I think wrong for removing the assertion, since LAM.clear(L, Name); won't find the sub-loop, and markLoopAsDelete shall not be called for the inner-loop.
May you correct me if I am wrong.

[eopXD]

1. revert change for LPMUpdater::markLoopAsDeleted, remove the call also.
2. iterate through Loops and record top-level loops for legacy LoopFlattenPass, a seperate for-loop is called because LoopFlatten may affect the range.

[eopXD]

remove redundant LPMUpdater parameter.

[dmgreen]

OK. So because we only delete the inner loop, not the current one, the loop nest pass remains valid? That sounds Ok then.

I think the changes from UnrollAndJam were found with a sanitizer build, so it might be worth testing it with one of those (or letting the buildbots do it post-commit, and keeping an eye one them to make sure they remain OK).

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

682

We can remove the LPMUpdater again if it's not needed/used?

740

I think this can just be for (Loop *L : *LI), then it will only include the outermost loops.

[eopXD]

address comments.

[eopXD]

We can remove the LPMUpdater again if it's not needed/used?

I think for a LoopNestPass (or a LoopPass), the LPMUpdater is necessary parameter.
As LoopNestPass is added by createFunctionToLoopPassAdaptor< LoopPassManager >(),
it needs to takes LPMUpdater & even though it is not used.

[Whitney]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

724

Why do we need these 3 new dependencies?

[eopXD]

@Whitney : do we need these 3 new dependencies?

I've deleted the parts for loop simplification and form LCSSA inside function Flatten,
and thought this would lead to the legacy pass depending on the passes I added.

I've just commented out the 3 passes and the tests still passed when I added -enable-new-pm=0 (using the legacy pass).
I think this means the 3 dependencies are unnecessary then, I will remove my addition.

[eopXD]

remove redundant pass dependencies added.

[eopXD]

remove redundant header added.

[Whitney]

LGTM.

[eopXD]

Reverting my commit to investigate for build failure.

[eopXD]

Resolve merge conflict correctly. (previously caused build fail)

[eopXD]

Rebase code to main.

[eopXD]

Remove declaration of LoopUnrollAndJam as a FunctionPass.

Previously I resolved the merge conflict incorrecly.
LoopUnrollAndJam is now a LoopPass.

llvm/include/llvm/Transforms/Scalar.h

 //===-- Scalar.h - Scalar Transformations -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines prototypes for accessor functions that expose passes
 // in the Scalar transformations library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_SCALAR_H
 #define LLVM_TRANSFORMS_SCALAR_H
 
 #include "llvm/Transforms/Utils/SimplifyCFGOptions.h"
+#include "llvm/Analysis/LoopPass.h"
 #include <functional>
 
 namespace llvm {
 
 class Function;
 class FunctionPass;
 class ModulePass;
 class Pass;

 // cache-friendly memory access patterns.
 //
 Pass *createLoopInterchangePass();
 
 //===----------------------------------------------------------------------===//
 //
 // LoopFlatten - This pass flattens nested loops into a single loop.
 //
-FunctionPass *createLoopFlattenPass();
+LoopPass *createLoopFlattenPass();
 
 //===----------------------------------------------------------------------===//
 //
 // LoopStrengthReduce - This pass is strength reduces GEP instructions that use
 // a loop's canonical induction variable as one of their indices.
 //
 Pass *createLoopStrengthReducePass();
 

llvm/include/llvm/Transforms/Scalar/LoopFlatten.h

 #include "llvm/Transforms/Scalar/LoopPassManager.h"
 
 namespace llvm {
 
 class LoopFlattenPass : public PassInfoMixin<LoopFlattenPass> {
 public:
   LoopFlattenPass() = default;
 
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  PreservedAnalyses run(LoopNest &LN, LoopAnalysisManager &LAM,
+                        LoopStandardAnalysisResults &AR, LPMUpdater &U);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_LOOPFLATTEN_H

llvm/lib/Passes/PassBuilder.cpp

PassBuilder::buildO1FunctionSimplificationPipeline(OptimizationLevel Level,
   FPM.addPass(
       RequireAnalysisPass<OptimizationRemarkEmitterAnalysis, Function>());
   FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM1),
                                               EnableMSSALoopDependency,
                                               /*UseBlockFrequencyInfo=*/true));
   FPM.addPass(SimplifyCFGPass());
   FPM.addPass(InstCombinePass());
   if (EnableLoopFlatten)
-    FPM.addPass(LoopFlattenPass());
+    FPM.addPass(createFunctionToLoopPassAdaptor(LoopFlattenPass()));
   // The loop passes in LPM2 (LoopFullUnrollPass) do not preserve MemorySSA.
   // *All* loop passes must preserve it, in order to be able to use it.
   FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM2),
                                               /*UseMemorySSA=*/false,
                                               /*UseBlockFrequencyInfo=*/false));
 
   // Delete small array after loop unroll.
   FPM.addPass(SROA());
PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
   FPM.addPass(
       RequireAnalysisPass<OptimizationRemarkEmitterAnalysis, Function>());
   FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM1),
                                               EnableMSSALoopDependency,
                                               /*UseBlockFrequencyInfo=*/true));
   FPM.addPass(SimplifyCFGPass());
   FPM.addPass(InstCombinePass());
   if (EnableLoopFlatten)
-    FPM.addPass(LoopFlattenPass());
+    FPM.addPass(createFunctionToLoopPassAdaptor(LoopFlattenPass()));
   // The loop passes in LPM2 (LoopIdiomRecognizePass, IndVarSimplifyPass,
   // LoopDeletionPass and LoopFullUnrollPass) do not preserve MemorySSA.
   // *All* loop passes must preserve it, in order to be able to use it.
   FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM2),
                                               /*UseMemorySSA=*/false,
                                               /*UseBlockFrequencyInfo=*/false));
 
   // Delete small array after loop unroll.
PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
   MainFPM.addPass(MemCpyOptPass());
 
   // Nuke dead stores.
   MainFPM.addPass(DSEPass());
   MainFPM.addPass(MergedLoadStoreMotionPass());
 
   // More loops are countable; try to optimize them.
   if (EnableLoopFlatten && Level.getSpeedupLevel() > 1)
-    MainFPM.addPass(LoopFlattenPass());
+    MainFPM.addPass(createFunctionToLoopPassAdaptor(LoopFlattenPass()));
 
   if (EnableConstraintElimination)
     MainFPM.addPass(ConstraintEliminationPass());
 
   LoopPassManager LPM;
   LPM.addPass(IndVarSimplifyPass());
   LPM.addPass(LoopDeletionPass());
   // FIXME: Add loop interchange.

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("lower-matrix-intrinsics", LowerMatrixIntrinsicsPass())
 FUNCTION_PASS("lower-matrix-intrinsics-minimal", LowerMatrixIntrinsicsPass(true))
 FUNCTION_PASS("lower-widenable-condition", LowerWidenableConditionPass())
 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())
 FUNCTION_PASS("mem2reg", PromotePass())
 FUNCTION_PASS("memcpyopt", MemCpyOptPass())
 FUNCTION_PASS("mergeicmps", MergeICmpsPass())
 FUNCTION_PASS("mergereturn", UnifyFunctionExitNodesPass())
 FUNCTION_PASS("nary-reassociate", NaryReassociatePass())
 FUNCTION_PASS("newgvn", NewGVNPass())

 LOOP_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 LOOP_PASS("licm", LICMPass())
 LOOP_PASS("loop-idiom", LoopIdiomRecognizePass())
 LOOP_PASS("loop-instsimplify", LoopInstSimplifyPass())
 LOOP_PASS("loop-interchange", LoopInterchangePass())
 LOOP_PASS("loop-rotate", LoopRotatePass())
 LOOP_PASS("no-op-loop", NoOpLoopPass())
 LOOP_PASS("print", PrintLoopPass(dbgs()))
+LOOP_PASS("loop-flatten", LoopFlattenPass())
 LOOP_PASS("loop-deletion", LoopDeletionPass())
 LOOP_PASS("loop-simplifycfg", LoopSimplifyCFGPass())
 LOOP_PASS("loop-reduce", LoopStrengthReducePass())
 LOOP_PASS("indvars", IndVarSimplifyPass())
 LOOP_PASS("loop-unroll-full", LoopFullUnrollPass())
 LOOP_PASS("print-access-info", LoopAccessInfoPrinterPass(dbgs()))
 LOOP_PASS("print<ddg>", DDGAnalysisPrinterPass(dbgs()))
 LOOP_PASS("print<iv-users>", IVUsersPrinterPass(dbgs()))

 #ifndef LOOP_PASS_WITH_PARAMS
 #define LOOP_PASS_WITH_PARAMS(NAME, CREATE_PASS, PARSER)
 #endif
 LOOP_PASS_WITH_PARAMS("unswitch",
                            [](bool NonTrivial) {
                              return SimpleLoopUnswitchPass(NonTrivial);
                            },
                            parseLoopUnswitchOptions)
 #undef LOOP_PASS_WITH_PARAMS

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

 // div/mod to reconstruct the original values, so this wouldn't be profitable.
 //
 // We also need to prove that N*M will not overflow.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/LoopFlatten.h"
 #include "llvm/Analysis/AssumptionCache.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"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"

 
 static cl::opt<bool>
     AssumeNoOverflow("loop-flatten-assume-no-overflow", cl::Hidden,
                      cl::init(false),
                      cl::desc("Assume that the product of the two iteration "
                               "limits will never overflow"));
 
 static cl::opt<bool>
-    WidenIV("loop-flatten-widen-iv", cl::Hidden,
-            cl::init(true),
+    WidenIV("loop-flatten-widen-iv", cl::Hidden, cl::init(true),
             cl::desc("Widen the loop induction variables, if possible, so "
                      "overflow checks won't reject flattening"));
 
 struct FlattenInfo {
   Loop *OuterLoop = nullptr;
   Loop *InnerLoop = nullptr;
   PHINode *InnerInductionPHI = nullptr;
   PHINode *OuterInductionPHI = nullptr;
   Value *InnerLimit = nullptr;
   Value *OuterLimit = nullptr;
   BinaryOperator *InnerIncrement = nullptr;
   BinaryOperator *OuterIncrement = nullptr;
   BranchInst *InnerBranch = nullptr;
   BranchInst *OuterBranch = nullptr;
   SmallPtrSet<Value *, 4> LinearIVUses;
   SmallPtrSet<PHINode *, 4> InnerPHIsToTransform;
 
   // Whether this holds the flatten info before or after widening.
   bool Widened = false;
 
-  FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL) {};
+  FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL){};
 };
 
 // Finds the induction variable, increment and limit for a simple loop that we
 // can flatten.
 static bool findLoopComponents(
     Loop *L, SmallPtrSetImpl<Instruction *> &IterationInstructions,
     PHINode *&InductionPHI, Value *&Limit, BinaryOperator *&Increment,
     BranchInst *&BackBranch, ScalarEvolution *SE) {
static bool checkIVUsers(FlattenInfo &FI) {
   //
   //   (OuterPHI * InnerLimit) + InnerPHI
   //
   // Any uses of the induction variables not matching that pattern would
   // require a div/mod to reconstruct in the flattened loop, so the
   // transformation wouldn't be profitable.
 
   Value *InnerLimit = FI.InnerLimit;
-  if (FI.Widened &&
-      (isa<SExtInst>(InnerLimit) || isa<ZExtInst>(InnerLimit)))
+  if (FI.Widened && (isa<SExtInst>(InnerLimit) || isa<ZExtInst>(InnerLimit)))
     InnerLimit = cast<Instruction>(InnerLimit)->getOperand(0);
 
   // Check that all uses of the inner loop's induction variable match the
   // expected pattern, recording the uses of the outer IV.
   SmallPtrSet<Value *, 4> ValidOuterPHIUses;
   for (User *U : FI.InnerInductionPHI->users()) {
     if (U == FI.InnerIncrement)
       continue;
 
-    // After widening the IVs, a trunc instruction might have been introduced, so
-    // look through truncs.
+    // After widening the IVs, a trunc instruction might have been introduced,
+    // so look through truncs.
     if (isa<TruncInst>(U)) {
       if (!U->hasOneUse())
         return false;
       U = *U->user_begin();
     }
 
     LLVM_DEBUG(dbgs() << "Found use of inner induction variable: "; U->dump());
 
     Value *MatchedMul;
     Value *MatchedItCount;
     bool IsAdd = match(U, m_c_Add(m_Specific(FI.InnerInductionPHI),
                                   m_Value(MatchedMul))) &&
                  match(MatchedMul, m_c_Mul(m_Specific(FI.OuterInductionPHI),
                                            m_Value(MatchedItCount)));
 
     // Matches the same pattern as above, except it also looks for truncs
     // on the phi, which can be the result of widening the induction variables.
-    bool IsAddTrunc = match(U, m_c_Add(m_Trunc(m_Specific(FI.InnerInductionPHI)),
-                                       m_Value(MatchedMul))) &&
-                      match(MatchedMul,
-                            m_c_Mul(m_Trunc(m_Specific(FI.OuterInductionPHI)),
+    bool IsAddTrunc =
+        match(U, m_c_Add(m_Trunc(m_Specific(FI.InnerInductionPHI)),
+        match(MatchedMul, m_c_Mul(m_Trunc(m_Specific(FI.OuterInductionPHI)),
                                   m_Value(MatchedItCount)));
 
     if ((IsAdd || IsAddTrunc) && MatchedItCount == InnerLimit) {
       LLVM_DEBUG(dbgs() << "Use is optimisable\n");
       ValidOuterPHIUses.insert(MatchedMul);
       FI.LinearIVUses.insert(U);
     } else {
       LLVM_DEBUG(dbgs() << "Did not match expected pattern, bailing\n");
       return false;
     }
   }
 
   // Check that there are no uses of the outer IV other than the ones found
   // as part of the pattern above.
   for (User *U : FI.OuterInductionPHI->users()) {
     if (U == FI.OuterIncrement)
       continue;
 
-    auto IsValidOuterPHIUses = [&] (User *U) -> bool {
-      LLVM_DEBUG(dbgs() << "Found use of outer induction variable: "; U->dump());
+    auto IsValidOuterPHIUses = [&](User *U) -> bool {
+      LLVM_DEBUG(dbgs() << "Found use of outer induction variable: ";
+                 U->dump());
       if (!ValidOuterPHIUses.count(U)) {
         LLVM_DEBUG(dbgs() << "Did not match expected pattern, bailing\n");
         return false;
       }
       LLVM_DEBUG(dbgs() << "Use is optimisable\n");
       return true;
     };
 
     if (auto *V = dyn_cast<TruncInst>(U)) {
       for (auto *K : V->users()) {
         if (!IsValidOuterPHIUses(K))
           return false;
       }
       continue;
     }
 
     if (!IsValidOuterPHIUses(U))
       return false;
   }
 
   LLVM_DEBUG(dbgs() << "checkIVUsers: OK\n";
              dbgs() << "Found " << FI.LinearIVUses.size()
                     << " value(s) that can be replaced:\n";
-             for (Value *V : FI.LinearIVUses) {
+             for (Value *V
+                  : FI.LinearIVUses) {
                dbgs() << "  ";
                V->dump();
              });
   return true;
 }
 
 // Return an OverflowResult dependant on if overflow of the multiplication of
 // InnerLimit and OuterLimit can be assumed not to happen.
static OverflowResult checkOverflow(FlattenInfo &FI, DominatorTree *DT,
 
   return OverflowResult::MayOverflow;
 }
 
 static bool CanFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                                ScalarEvolution *SE, AssumptionCache *AC,
                                const TargetTransformInfo *TTI) {
   SmallPtrSet<Instruction *, 8> IterationInstructions;
-  if (!findLoopComponents(FI.InnerLoop, IterationInstructions, FI.InnerInductionPHI,
-                          FI.InnerLimit, FI.InnerIncrement, FI.InnerBranch, SE))
-    return false;
-  if (!findLoopComponents(FI.OuterLoop, IterationInstructions, FI.OuterInductionPHI,
-                          FI.OuterLimit, FI.OuterIncrement, FI.OuterBranch, SE))
+  if (!findLoopComponents(FI.InnerLoop, IterationInstructions,
+                          FI.InnerInductionPHI, FI.InnerLimit,
+                          FI.InnerIncrement, FI.InnerBranch, SE))
+    return false;
+  if (!findLoopComponents(FI.OuterLoop, IterationInstructions,
+                          FI.OuterInductionPHI, FI.OuterLimit,
+                          FI.OuterIncrement, FI.OuterBranch, SE))
     return false;
 
   // Both of the loop limit values must be invariant in the outer loop
   // (non-instructions are all inherently invariant).
   if (!FI.OuterLoop->isLoopInvariant(FI.InnerLimit)) {
     LLVM_DEBUG(dbgs() << "inner loop limit not invariant\n");
     return false;
   }

 
 static bool DoFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                               ScalarEvolution *SE, AssumptionCache *AC,
                               const TargetTransformInfo *TTI) {
   Function *F = FI.OuterLoop->getHeader()->getParent();
   LLVM_DEBUG(dbgs() << "Checks all passed, doing the transformation\n");
   {
     using namespace ore;
-    OptimizationRemark Remark(DEBUG_TYPE, "Flattened", FI.InnerLoop->getStartLoc(),
+    OptimizationRemark Remark(DEBUG_TYPE, "Flattened",
+                              FI.InnerLoop->getStartLoc(),
                               FI.InnerLoop->getHeader());
     OptimizationRemarkEmitter ORE(F);
     Remark << "Flattened into outer loop";
     ORE.emit(Remark);
   }
 
-  Value *NewTripCount =
-      BinaryOperator::CreateMul(FI.InnerLimit, FI.OuterLimit, "flatten.tripcount",
-                                FI.OuterLoop->getLoopPreheader()->getTerminator());
+  Value *NewTripCount = BinaryOperator::CreateMul(
+      FI.InnerLimit, FI.OuterLimit, "flatten.tripcount",
   LLVM_DEBUG(dbgs() << "Created new trip count in preheader: ";
              NewTripCount->dump());
 
   // Fix up PHI nodes that take values from the inner loop back-edge, which
   // we are about to remove.
   FI.InnerInductionPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
 
   // The old Phi will be optimised away later, but for now we can't leave
static bool DoFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   // variables with the one new one.
   IRBuilder<> Builder(FI.OuterInductionPHI->getParent()->getTerminator());
   for (Value *V : FI.LinearIVUses) {
     Value *OuterValue = FI.OuterInductionPHI;
     if (FI.Widened)
       OuterValue = Builder.CreateTrunc(FI.OuterInductionPHI, V->getType(),
                                        "flatten.trunciv");
 
-    LLVM_DEBUG(dbgs() << "Replacing: "; V->dump();
-               dbgs() << "with:      "; OuterValue->dump());
+    LLVM_DEBUG(dbgs() << "Replacing: "; V->dump(); dbgs() << "with:      ";
+               OuterValue->dump());
     V->replaceAllUsesWith(OuterValue);
   }
 
   // Tell LoopInfo, SCEV and the pass manager that the inner loop has been
   // deleted, and any information that have about the outer loop invalidated.
   SE->forgetLoop(FI.OuterLoop);
   SE->forgetLoop(FI.InnerLoop);
   LI->erase(FI.InnerLoop);
static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   unsigned MaxLegalSize = DL.getLargestLegalIntTypeSizeInBits();
   auto *MaxLegalType = DL.getLargestLegalIntType(M->getContext());
 
   // If both induction types are less than the maximum legal integer width,
   // promote both to the widest type available so we know calculating
   // (OuterLimit * InnerLimit) as the new trip count is safe.
   if (InnerType != OuterType ||
       InnerType->getScalarSizeInBits() >= MaxLegalSize ||
-      MaxLegalType->getScalarSizeInBits() < InnerType->getScalarSizeInBits() * 2) {
+      MaxLegalType->getScalarSizeInBits() <
+          InnerType->getScalarSizeInBits() * 2) {
     LLVM_DEBUG(dbgs() << "Can't widen the IV\n");
     return false;
   }
 
   SCEVExpander Rewriter(*SE, DL, "loopflatten");
   SmallVector<WideIVInfo, 2> WideIVs;
   SmallVector<WeakTrackingVH, 4> DeadInsts;
-  WideIVs.push_back( {FI.InnerInductionPHI, MaxLegalType, false });
-  WideIVs.push_back( {FI.OuterInductionPHI, MaxLegalType, false });
+  WideIVs.push_back({FI.InnerInductionPHI, MaxLegalType, false});
+  WideIVs.push_back({FI.OuterInductionPHI, MaxLegalType, false});
   unsigned ElimExt = 0;
   unsigned Widened = 0;
 
   for (const auto &WideIV : WideIVs) {
-    PHINode *WidePhi = createWideIV(WideIV, LI, SE, Rewriter, DT, DeadInsts,
-                                    ElimExt, Widened, true /* HasGuards */,
-                                    true /* UsePostIncrementRanges */);
+    PHINode *WidePhi =
+        createWideIV(WideIV, LI, SE, Rewriter, DT, DeadInsts, ElimExt, Widened,
+                     true /* HasGuards */, true /* UsePostIncrementRanges */);
     if (!WidePhi)
       return false;
     LLVM_DEBUG(dbgs() << "Created wide phi: "; WidePhi->dump());
     LLVM_DEBUG(dbgs() << "Deleting old phi: "; WideIV.NarrowIV->dump());
     RecursivelyDeleteDeadPHINode(WideIV.NarrowIV);
   }
   // After widening, rediscover all the loop components.
   assert(Widened && "Widened IV expected");
if (OR == OverflowResult::AlwaysOverflowsHigh ||
     LLVM_DEBUG(dbgs() << "Multiply might overflow, not flattening\n");
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "Multiply cannot overflow, modifying loop in-place\n");
   return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
 }
 
-bool Flatten(DominatorTree *DT, LoopInfo *LI, ScalarEvolution *SE,
+bool Flatten(LoopNest &LN, DominatorTree *DT, LoopInfo *LI, ScalarEvolution *SE,
              AssumptionCache *AC, TargetTransformInfo *TTI) {
   bool Changed = false;
-  for (auto *InnerLoop : LI->getLoopsInPreorder()) {
+  for (Loop *InnerLoop : LN.getLoops()) {
     auto *OuterLoop = InnerLoop->getParentLoop();
     if (!OuterLoop)
       continue;
     FlattenInfo FI(OuterLoop, InnerLoop);
     Changed |= FlattenLoopPair(FI, DT, LI, SE, AC, TTI);
   }
   return Changed;
 }
 
-PreservedAnalyses LoopFlattenPass::run(Function &F,
-                                       FunctionAnalysisManager &AM) {
-  auto *DT = &AM.getResult<DominatorTreeAnalysis>(F);
-  auto *LI = &AM.getResult<LoopAnalysis>(F);
-  auto *SE = &AM.getResult<ScalarEvolutionAnalysis>(F);
-  auto *AC = &AM.getResult<AssumptionAnalysis>(F);
-  auto *TTI = &AM.getResult<TargetIRAnalysis>(F);
+PreservedAnalyses LoopFlattenPass::run(LoopNest &LN, LoopAnalysisManager &LAM,
+                                       LoopStandardAnalysisResults &AR,
+                                       LPMUpdater &U) {
 
   bool Changed = false;
 
   // The loop flattening pass requires loops to be
   // in simplified form, and also needs LCSSA. Running
   // this pass will simplify all loops that contain inner loops,
   // regardless of whether anything ends up being flattened.
-  for (const auto &L : *LI) {
-    if (L->isInnermost())
-      continue;
-    Changed |=
-        simplifyLoop(L, DT, LI, SE, AC, nullptr, false /* PreserveLCSSA */);
-    Changed |= formLCSSARecursively(*L, *DT, LI, SE);
-  }
-
-  Changed |= Flatten(DT, LI, SE, AC, TTI);
+  Changed |= Flatten(LN, &AR.DT, &AR.LI, &AR.SE, &AR.AC, &AR.TTI);
 
   if (!Changed)
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
 }
 
 namespace {
-class LoopFlattenLegacyPass : public FunctionPass {
+class LoopFlattenLegacyPass : public LoopPass {
 public:
   static char ID; // Pass ID, replacement for typeid
-  LoopFlattenLegacyPass() : FunctionPass(ID) {
+  LoopFlattenLegacyPass() : LoopPass(ID) {
     initializeLoopFlattenLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   // Possibly flatten loop L into its child.
-  bool runOnFunction(Function &F) override;
+  bool runOnLoop(Loop *L, LPPassManager &LPM) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     getLoopAnalysisUsage(AU);
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addPreserved<TargetTransformInfoWrapperPass>();
     AU.addRequired<AssumptionCacheTracker>();
     AU.addPreserved<AssumptionCacheTracker>();
   }
 };
 } // namespace
 
 char LoopFlattenLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopFlattenLegacyPass, "loop-flatten", "Flattens loops",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_END(LoopFlattenLegacyPass, "loop-flatten", "Flattens loops",
                     false, false)
 
-FunctionPass *llvm::createLoopFlattenPass() { return new LoopFlattenLegacyPass(); }
+LoopPass *llvm::createLoopFlattenPass() { return new LoopFlattenLegacyPass(); }
 
-bool LoopFlattenLegacyPass::runOnFunction(Function &F) {
+bool LoopFlattenLegacyPass::runOnLoop(Loop *L, LPPassManager &LPM) {
+  Function &F = *L->getHeader()->getParent();
   ScalarEvolution *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
+  auto LN = LoopNest::getLoopNest(*L, *SE);
   LoopInfo *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   DominatorTree *DT = DTWP ? &DTWP->getDomTree() : nullptr;
   auto &TTIP = getAnalysis<TargetTransformInfoWrapperPass>();
   auto *TTI = &TTIP.getTTI(F);
   auto *AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
-  return Flatten(DT, LI, SE, AC, TTI);
+  return Flatten(*LN, DT, LI, SE, AC, TTI);
 }