[LoopUnrollAndJam] Change LoopUnrollAndJamPass to LoopNest pass

[uint256_t]

Summary

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

[uint256_t]

Followed clang-tidy's warning

[dmgreen]

This sounds good to me, so long as all the existing tests are still doing OK.

Does it work, in general, to have a loop pass that creates and destroys loops? Even if the outer loop is completely unrolled?

And does unroll and jam preserve memory-ssa correctly?

[uint256_t]

Does it work, in general, to have a loop pass that creates and destroys loops? Even if the outer loop is completely unrolled?

I should have added code like

if (Result == LoopUnrollResult::FullyUnrolled)
      LPM.markLoopAsDeleted(*L);

after calling tryToUnrollAndJamLoop.

And does unroll and jam preserve memory-ssa correctly?

I'm not sure, but according to https://reviews.llvm.org/D72230, I think

OptimizePM.addPass(createFunctionToLoopPassAdaptor(
        std::move(LPM), EnableMSSALoopDependency,
        /*UseBlockFrequencyInfo=*/true, DebugLogging));

should be simply like

OptimizePM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM));

(I may be misunderstanding the problem)

[uint256_t]

Updated the code

[Whitney]

Approve with minor comment.

llvm/lib/Passes/PassBuilder.cpp

1409

OptimizePM.addPass(createFunctionToLoopPassAdaptor(LoopUnrollAndJamPass(Level.getSpeedupLevel())));

[uint256_t]

Updated the code. Thanks.

[uint256_t]

The pre-merge checks failed, so I fixed it.

[uint256_t]

To avoid a failure by address sanitizer

[MaskRay]

This caused a layering violation.

LLVMScalarOpts (lib/Transforms/Scalar) depends on LLVMTransformUtils (lib/Transforms/Utils) however LLVMTransformUtils now a header dependency #include "llvm/Transforms/Scalar/LoopPassManager.h" on LLVMScalarOpts.

[uint256_t]

Thank you for reporting. I'll revert the commit and fix the problem.

[uint256_t]

Maybe fixed the problem.

[uint256_t]

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

261

I don't know if this change is ok, but I think CurrentL is always a top-level loop if it's loop nest mode.

[Whitney]

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

261

If CurrentL is always a top-level loop if it's loop nest mode, then why do we need to change from L.isOutermost() to CurrentL == &L?

[uint256_t]

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

261

I found that markLoopAsDeleted is sometimes called on loops that are already destroyed (e.g. by LI->erase(L)). Then L.isOutermost() is invalid since the destructor of L is called. (asan detected it)

[Whitney]

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

261

/// This runs the destructor of the loop object making it invalid to
/// reference afterward. The memory is retained so that the *pointer* to the
/// loop remains valid.
destroy()

So I agree L.isOutermost() should be changed.

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

52 ↗

(On Diff #347666)

Is this change still needed?

[Whitney]

LGTM, thanks!

[hvdijk]

I'm seeing test failures now that this has been committed. These tests pass again after reverting this. (Tested on 21653600 directly and re-tested on 12f53e53, and tested the revert on top of the latter.) However, I have not found the same errors in the CI results, https://lab.llvm.org/buildbot/#/changes/22300 shows these tests as passing even in configurations that enable assertions. Is it possible that something in here, or something relied upon by something in here, is not fully deterministic and produces different results in different build configurations? I will try different build configurations to see if I can find one where the tests pass, and compare just what the two different builds do, but if you can spot a potential problem straightaway and suggest something to try that would be very helpful.

LLVM :: Transforms/LoopUnrollAndJam/dependencies.ll
LLVM :: Transforms/LoopUnrollAndJam/dependencies_multidims.ll
LLVM :: Transforms/LoopUnrollAndJam/disable.ll
LLVM :: Transforms/LoopUnrollAndJam/pragma-explicit.ll
LLVM :: Transforms/LoopUnrollAndJam/unroll-and-jam.ll

Of these, the unroll-and-jam failure looks the most interesting. The others have wrong output, but this one hard-fails with an assertion failure for me:

FAIL: LLVM :: Transforms/LoopUnrollAndJam/unroll-and-jam.ll (1 of 1)
******************** TEST 'LLVM :: Transforms/LoopUnrollAndJam/unroll-and-jam.ll' FAILED ********************
Script:
--
: 'RUN: at line 2';   /home/harald/llvm-project/build/bin/opt -basic-aa -tbaa -loop-unroll-and-jam -allow-unroll-and-jam -unroll-and-jam-count=4 -unroll-remainder < /home/harald/llvm-project/llvm/test/Transforms/LoopUnrollAndJam/unroll-and-jam.ll -S | /home/harald/llvm-project/build/bin/FileCheck /home/harald/llvm-project/llvm/test/Transforms/LoopUnrollAndJam/unroll-and-jam.ll
: 'RUN: at line 3';   /home/harald/llvm-project/build/bin/opt -aa-pipeline=tbaa,basic-aa -passes='loop-unroll-and-jam' -allow-unroll-and-jam -unroll-and-jam-count=4 -unroll-remainder < /home/harald/llvm-project/llvm/test/Transforms/LoopUnrollAndJam/unroll-and-jam.ll -S | /home/harald/llvm-project/build/bin/FileCheck /home/harald/llvm-project/llvm/test/Transforms/LoopUnrollAndJam/unroll-and-jam.ll
--
Exit Code: 2

Command Output (stderr):
--
opt: /home/harald/llvm-project/llvm/lib/Transforms/Utils/LoopSimplify.cpp:731: bool llvm::simplifyLoop(llvm::Loop *, llvm::DominatorTree *, llvm::LoopInfo *, llvm::ScalarEvolution *, llvm::AssumptionCache *, llvm::MemorySSAUpdater *, bool): Assertion `L->isRecursivelyLCSSAForm(*DT, *LI) && "Requested to preserve LCSSA, but it's already broken."' failed.
PLEASE submit a bug report to https://bugs.llvm.org/ and include the crash backtrace.
Stack dump:
0.	Program arguments: /home/harald/llvm-project/build/bin/opt -basic-aa -tbaa -loop-unroll-and-jam -allow-unroll-and-jam -unroll-and-jam-count=4 -unroll-remainder -S
1.	Running pass 'Function Pass Manager' on module '<stdin>'.
2.	Running pass 'Loop Pass Manager' on function '@test8'
3.	Running pass 'Unroll and Jam loops' on basic block '%for.outer'
 #0 0x0000000001b4abf3 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/harald/llvm-project/build/bin/opt+0x1b4abf3)
 #1 0x0000000001b4898e llvm::sys::RunSignalHandlers() (/home/harald/llvm-project/build/bin/opt+0x1b4898e)
 #2 0x0000000001b4b0ba SignalHandler(int) Signals.cpp:0:0
 #3 0x00007ffa3a712d00 __restore_rt sigaction.c:0:0
 #4 0x00007ffa3a269d4d raise (/lib64/libc.so.6+0x37d4d)
 #5 0x00007ffa3a254526 abort (/lib64/libc.so.6+0x22526)
 #6 0x00007ffa3a25441f _nl_load_domain.cold loadmsgcat.c:0:0
 #7 0x00007ffa3a262ac2 (/lib64/libc.so.6+0x30ac2)
 #8 0x0000000001bf538f (/home/harald/llvm-project/build/bin/opt+0x1bf538f)
 #9 0x0000000001bffba3 llvm::UnrollLoop(llvm::Loop*, llvm::UnrollLoopOptions, llvm::LoopInfo*, llvm::ScalarEvolution*, llvm::DominatorTree*, llvm::AssumptionCache*, llvm::TargetTransformInfo const*, llvm::OptimizationRemarkEmitter*, bool, llvm::Loop**) (/home/harald/llvm-project/build/bin/opt+0x1bffba3)
#10 0x0000000001c0ed22 llvm::UnrollRuntimeLoopRemainder(llvm::Loop*, unsigned int, bool, bool, bool, bool, llvm::LoopInfo*, llvm::ScalarEvolution*, llvm::DominatorTree*, llvm::AssumptionCache*, llvm::TargetTransformInfo const*, bool, llvm::Loop**) (/home/harald/llvm-project/build/bin/opt+0x1c0ed22)
#11 0x0000000001c01c65 llvm::UnrollAndJamLoop(llvm::Loop*, unsigned int, unsigned int, unsigned int, bool, llvm::LoopInfo*, llvm::ScalarEvolution*, llvm::DominatorTree*, llvm::AssumptionCache*, llvm::TargetTransformInfo const*, llvm::OptimizationRemarkEmitter*, llvm::Loop**) (/home/harald/llvm-project/build/bin/opt+0x1c01c65)
#12 0x000000000194f77a tryToUnrollAndJamLoop(llvm::Loop*, llvm::DominatorTree&, llvm::LoopInfo*, llvm::ScalarEvolution&, llvm::TargetTransformInfo const&, llvm::AssumptionCache&, llvm::DependenceInfo&, llvm::OptimizationRemarkEmitter&, int) LoopUnrollAndJamPass.cpp:0:0
#13 0x000000000194e6e7 (anonymous namespace)::LoopUnrollAndJam::runOnLoop(llvm::Loop*, llvm::LPPassManager&) LoopUnrollAndJamPass.cpp:0:0
#14 0x0000000000b6a8c5 llvm::LPPassManager::runOnFunction(llvm::Function&) (/home/harald/llvm-project/build/bin/opt+0xb6a8c5)
#15 0x0000000001362558 llvm::FPPassManager::runOnFunction(llvm::Function&) (/home/harald/llvm-project/build/bin/opt+0x1362558)
#16 0x0000000001368eb1 llvm::FPPassManager::runOnModule(llvm::Module&) (/home/harald/llvm-project/build/bin/opt+0x1368eb1)
#17 0x0000000001362c16 llvm::legacy::PassManagerImpl::run(llvm::Module&) (/home/harald/llvm-project/build/bin/opt+0x1362c16)
#18 0x000000000072c858 main (/home/harald/llvm-project/build/bin/opt+0x72c858)
#19 0x00007ffa3a2557dd __libc_start_main (/lib64/libc.so.6+0x237dd)
#20 0x00000000007158aa _start (/home/harald/llvm-project/build/bin/opt+0x7158aa)
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /home/harald/llvm-project/build/bin/FileCheck /home/harald/llvm-project/llvm/test/Transforms/LoopUnrollAndJam/unroll-and-jam.ll

[hvdijk]

Ah, problem found, I think: this breaks when ENABLE_EXPERIMENTAL_NEW_PASS_MANAGER=OFF, which I never explicitly set but had inherited as I was using a CMakeCache.txt generated when that was the default. As far as I know this is, for now, a supported configuration. The new pass manager was enabled by default by D95380 and one of the major differences listed there is "LCSSA and LoopSimplify are run before all loop passes"; it seems like this change may be relying on that. Should this perhaps be reverted until the problem can be fixed?

[uint256_t]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

@hvdijk Could you try reversing the running order of LoopSimplify and LCSSAWrapperPass, and check if the problem still reproduces?

[hvdijk]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

Tried now, and I see the same results if I reverse those two lines.

[uint256_t]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

Thank you for reporting! I didn't care about the legacy pass manager. Maybe I'll make a new patch to fix the problem, so I think this need not be reverted.
Do you mind if unroll-and-jam fails on legacy pass manager now?

[hvdijk]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

It's not my call to make, but according to the description of D95380, users are supposed to be able to switch to the legacy pass manager if they encounter problems with the new pass manager. It may be acceptable for the two pass managers to produce slightly different results, but if we now get crashes with the legacy pass manager, users may not have the option to switch if needed. Unless the legacy pass manager is removed entirely (which should at least have some discussion on the mailing list), I think at least the crash will need fixing.

[uint256_t]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

I reverted the commit to fix the problem. thanks.

[hvdijk]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

Thanks. Happy to re-test if you later have a new version, if you want.

[uint256_t]

Code updated. No crash on legacy pass manager in my environment.

[uint256_t]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

In my environment there's no more crash on legacy pass manager. Could you please re-test yourself?

[hvdijk]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

Nice, I see no crash, and the tests I mentioned all pass even with the old pass manager. One other test fails (Transforms/SimpleLoopUnswitch/implicit-null-checks.ll) but it's not a crash, just different test output. I'll check deeper later to make sure there's no real issue, but I think it's probably fine.

[uint256_t]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

The failure on 'Transforms/SimpleLoopUnswitch/implicit-null-checks.ll' seems to reproduce even if reverting this patch. I'm not sure if I had better fix it.

[hvdijk]

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

514

Oh, if you are already seeing it fail even without re-committing this, then sure, no need to fix it as part of this. :) I am not entirely sure whether you should get this re-approved before re-pushing, but it looks good to me.

[uint256_t]

Just a small change.

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

public:
   /// Note that this loop must either be the current loop or a subloop of the
   /// current loop. This routine must be called prior to removing the loop from
   /// the loop nest.
   ///
   /// If this is called for the current loop, in addition to clearing any
   /// state, this routine will mark that the current loop should be skipped by
   /// the rest of the pass management infrastructure.
   void markLoopAsDeleted(Loop &L, llvm::StringRef Name) {
-    assert((!LoopNestMode || L.isOutermost()) &&
+    assert((!LoopNestMode || CurrentL == &L) &&
            "L should be a top-level loop in loop-nest mode.");
     LAM.clear(L, Name);
     assert((&L == CurrentL || CurrentL->contains(&L)) &&
            "Cannot delete a loop outside of the "
            "subloop tree currently being processed.");
     if (&L == CurrentL)
       SkipCurrentLoop = true;
   }

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

 //===- LoopUnrollAndJamPass.h -----------------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_SCALAR_LOOPUNROLLANDJAMPASS_H
 #define LLVM_TRANSFORMS_SCALAR_LOOPUNROLLANDJAMPASS_H
 
 #include "llvm/IR/PassManager.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
 
 namespace llvm {
 class Function;
 
 /// A simple loop rotation transformation.
 class LoopUnrollAndJamPass : public PassInfoMixin<LoopUnrollAndJamPass> {
   const int OptLevel;
 
 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
 
 #endif // LLVM_TRANSFORMS_SCALAR_LOOPUNROLLANDJAMPASS_H

llvm/lib/Passes/PassBuilder.cpp

if (IsLTO) {
     // again. Unroll small loops to hide loop backedge latency and saturate any
     // parallel execution resources of an out-of-order processor. We also then
     // need to clean up redundancies and loop invariant code.
     // FIXME: It would be really good to use a loop-integrated instruction
     // combiner for cleanup here so that the unrolling and LICM can be pipelined
     // 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)));
     FPM.addPass(WarnMissedTransformationsPass());
   }
 
   if (!IsLTO) {
     // Eliminate loads by forwarding stores from the previous iteration to loads
if (!IsLTO) {
     // Unroll small loops to hide loop backedge latency and saturate any
     // parallel execution resources of an out-of-order processor. We also then
     // need to clean up redundancies and loop invariant code.
     // FIXME: It would be really good to use a loop-integrated instruction
     // combiner for cleanup here so that the unrolling and LICM can be pipelined
     // 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)));
     FPM.addPass(WarnMissedTransformationsPass());
     FPM.addPass(InstCombinePass());
     FPM.addPass(
         RequireAnalysisPass<OptimizationRemarkEmitterAnalysis, Function>());
PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
   // llvm.loop.distribute=true or when -enable-loop-distribute is specified.
   OptimizePM.addPass(LoopDistributePass());
 
   // Populates the VFABI attribute with the scalar-to-vector mappings
   // from the TargetLibraryInfo.
   OptimizePM.addPass(InjectTLIMappings());
 
   addVectorPasses(Level, OptimizePM, /* IsLTO */ false);
 
   // Split out cold code. Splitting is done late to avoid hiding context from
   // other optimizations and inadvertently regressing performance. The tradeoff
   // is that this has a higher code size cost than splitting early.
   if (EnableHotColdSplit && !LTOPreLink)
     MPM.addPass(HotColdSplittingPass());
 
   // Search the code for similar regions of code. If enough similar regions can
   // be found where extracting the regions into their own function will decrease

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("lower-constant-intrinsics", LowerConstantIntrinsicsPass())
 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("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("loop-interchange", LoopInterchangePass())
 LOOP_PASS("loop-rotate", LoopRotatePass())
 LOOP_PASS("no-op-loop", NoOpLoopPass())
 LOOP_PASS("print", PrintLoopPass(dbgs()))
 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()))
 LOOP_PASS("print<iv-users>", IVUsersPrinterPass(dbgs()))
 LOOP_PASS("print<loopnest>", LoopNestPrinterPass(dbgs()))
 LOOP_PASS("print<loop-cache-cost>", LoopCachePrinterPass(dbgs()))
 LOOP_PASS("loop-predication", LoopPredicationPass())
 LOOP_PASS("guard-widening", GuardWideningPass())

llvm/lib/Transforms/Scalar/LoopUnrollAndJamPass.cpp

 #include "llvm/ADT/PriorityWorklist.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #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"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/PassRegistry.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils.h"
+#include "llvm/Transforms/Utils/LCSSA.h"
 #include "llvm/Transforms/Utils/LoopPeel.h"
 #include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/UnrollLoop.h"
 #include <cassert>
 #include <cstdint>
 #include <vector>
 
tryToUnrollAndJamLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
   // If loop has an unroll count pragma or unrolled by explicitly set count
   // mark loop as unrolled to prevent unrolling beyond that requested.
   if (UnrollResult != LoopUnrollResult::FullyUnrolled && IsCountSetExplicitly)
     L->setLoopAlreadyUnrolled();
 
   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,
-                                  int OptLevel) {
+                                  OptimizationRemarkEmitter &ORE, int OptLevel,
+                                  LPMUpdater &U) {
   bool DidSomething = false;
+  ArrayRef<Loop *> Loops = LN.getLoops();
+  Loop *OutmostLoop = &LN.getOutermostLoop();
 
-  // The loop unroll and jam pass requires loops to be in simplified form, and
-  // 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
+  // Add the loop nests in the reverse order of LN. See method
   // declaration.
   SmallPriorityWorklist<Loop *, 4> Worklist;
-  appendLoopsToWorklist(LI, Worklist);
+  appendLoopsToWorklist(Loops, Worklist);
   while (!Worklist.empty()) {
     Loop *L = Worklist.pop_back_val();
+    std::string LoopName = std::string(L->getName());
     LoopUnrollResult Result =
         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);
   }
 
   return DidSomething;
 }
 
 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 {
-    if (skipFunction(F))
+  bool runOnLoop(Loop *L, LPPassManager &LPM) override {
+    if (skipLoop(L))
       return false;
 
-    auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-    LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
-    ScalarEvolution &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
-    const TargetTransformInfo &TTI =
-        getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
-    auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+    auto *F = L->getHeader()->getParent();
+    auto &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
+    auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     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);
+
+    if (Result == LoopUnrollResult::FullyUnrolled)
+      LPM.markLoopAsDeleted(*L);
+
+    return Result != LoopUnrollResult::Unmodified;
   }
 
   /// This transformation requires natural loop information & requires that
   /// loop preheaders be inserted into the CFG...
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DependenceAnalysisWrapperPass>();
     AU.addRequired<OptimizationRemarkEmitterWrapperPass>();
+    getLoopAnalysisUsage(AU);
   }
 };
 
 } // end anonymous namespace
 
 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)
 INITIALIZE_PASS_DEPENDENCY(DependenceAnalysisWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(OptimizationRemarkEmitterWrapperPass)
 INITIALIZE_PASS_END(LoopUnrollAndJam, "loop-unroll-and-jam",
                     "Unroll and Jam loops", false, false)
 
 Pass *llvm::createLoopUnrollAndJamPass(int OptLevel) {
   return new LoopUnrollAndJam(OptLevel);
 }
 
-PreservedAnalyses LoopUnrollAndJamPass::run(Function &F,
-                                            FunctionAnalysisManager &AM) {
-  ScalarEvolution &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
-  LoopInfo &LI = AM.getResult<LoopAnalysis>(F);
-  TargetTransformInfo &TTI = AM.getResult<TargetIRAnalysis>(F);
-  AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
-  DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
-  DependenceInfo &DI = AM.getResult<DependenceAnalysis>(F);
-  OptimizationRemarkEmitter &ORE =
-      AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
+PreservedAnalyses LoopUnrollAndJamPass::run(LoopNest &LN,
+                                            LoopAnalysisManager &AM,
+                                            LoopStandardAnalysisResults &AR,
+                                            LPMUpdater &U) {
+  Function &F = *LN.getParent();
+
+  DependenceInfo DI(&F, &AR.AA, &AR.SE, &AR.LI);
+  OptimizationRemarkEmitter ORE(&F);
 
-  if (!tryToUnrollAndJamLoop(F, DT, LI, SE, TTI, AC, DI, ORE, OptLevel))
+  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/test/Transforms/LoopUnrollAndJam/innerloop.ll

 ; 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
 ; outer of 3 nested loops. The (new) inner loops need to be added to LI.
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 
 define i32 @test() {