[LoopExtractor] Convert LoopExtractor from LoopPass to ModulePass

[ekatz]

Summary

The LoopExtractor created new functions (by definition), which violates the restrictions of a LoopPass.
The correct implementation of this pass should be as a ModulePass.
Also revert r82990 implications on the LoopExtractor.

Fixes PR3082 and PR8929.

[vsk]

Thanks, this looks reasonable to me. Please wait for another +1, as I'm not terribly familiar with the loop pass machinery.

[ekatz]

Found the revision that changed the pass from FunctionPass to LoopPass - r82990.
This change actually removed (for no apparent reason) the handling of sub-loops, in case the function is a minimal wrapper around a singular top-level loop, added in r12403.
I re-added this code.

In addition, I extracted some of the logics into separate functions, for better readability.

[asbirlea]

Some comments, but it looks reasonable overall.

llvm/lib/Transforms/IPO/LoopExtractor.cpp

43

Could you add a comment on how NumLoops is used?

106

Unused after the latest change?

131

Nit: rename to runOnFunction for consistency? I don't mean having an override method.
To give a similar example, function passes may have a static or class specific runOnBasicBlock without being a basic block pass.

142–143

for (auto *ExitBlock : ExitBlocks)

230

for (auto *ExitBlock : ExitBlocks)

[ekatz]

Applied requested changes.

[ekatz]

Removed redundant handling of EH pads at exit blocks, as it is already handled in CodeExtractor.

[asbirlea]

Since I don't really understand the CodeExtractor, I'm relying on the author or another reviewer to confirm the correctness here.
This patch by itself looks good, but please wait to see if there are additional comments for the other reviewers before checking in.

llvm/lib/Transforms/IPO/LoopExtractor.cpp

61–62

Does this preserve LoopInfo? If yes, addPreserved, if unsure best to leave as is.

Similar question on the DomTree.
This is an item where I don't honestly understand if there are potential correctness implications in extractLoops if the DT is not preserved while iterating.

[ekatz]

llvm/lib/Transforms/IPO/LoopExtractor.cpp

61–62

You are correct. I will add the preservation of LoopInfo, as it is.
The DomTree is not.

[ekatz]

Notify that the LoopInfo is preserved.

[ekatz]

This comment was removed by ekatz.

[ekatz]

This comment was removed by ekatz.

[ekatz]

This comment was removed by ekatz.

[ekatz]

Fix test.

[fhahn]

Would it be possible to also add some of the test cases for the linked bugs?

llvm/lib/Transforms/IPO/LoopExtractor.cpp

144

It might be slightly better to use size (from STLExtra.h), as it will fail to compile, if the iterators are not random access iterators. Or just check std::next(LI.begin()) != std::end()

149

Do we actually need this special case when it's a module pass? IIUC, it was only needed with the LoopPass to avoid extracting loops we already extracted into new functions earlier right?

181

I am not sure I see why we would try to extract the sub loops here? The original loopPass only tried to extract top-level loops, unless I missed something.

213

The original NumLoops was on a per-function basis I guess, but now it's per-module, right? I think it would be better to keep it on a per-function basis, i.e. NumLoops == the max number of loops to extract per function.

llvm/test/Feature/optnone-opt.ll

57

The pass is still run right, just somewhere else? Could you check this in the right place?

[ekatz]

llvm/lib/Transforms/IPO/LoopExtractor.cpp

144

I agree. size (from STLExtras.h) is better.
Regarding the std::next solution, I thought it is a bit less clear, but it is actually even a better solution for the general case. So, I'll change it to that.

149

This particular line was added in SVN commit r86175 (Git commit rGa83ac2d9e7b990df6978598d9dec860b5ef13a9a). Personally, I do not see a good reason not to extract loops in a non simplify-form, but the change (that added it) does not explain much, so I decided to keep it (maybe there was a related bug?).
If you are positive that it is not needed, then I'll remove it.

181

As described in the "review summary", this is a kind of revert of SVN commit r82990 (Git commit rG9a7320c7110ee7f5450f8c46abe361de769886c0). I don't know why the extraction of the sub-loops was removed in that change, other than, probably, the change to LoopPass. Seems like an unintentional change/bug.

213

It used to be per program, as it is now; as the previous LoopPass was created only once per program (as well as the new ModulePass). The result remains the same.

llvm/test/Feature/optnone-opt.ll

57

It is still run, but it does not have an indication that the pass is skipped. Not sure why it was designed that way (probably wasn't), but skipped ModulePasses (via a call to skipModule) do not print any indication, while Function/Loop/Region-passes do.

[ekatz]

In D69069#1768576, @fhahn wrote:

Would it be possible to also add some of the test cases for the linked bugs?

This patch fixes a crash. Maybe just add a .ll file that shouldn't crash (with the fix)?!

[ekatz]

@fhahn, what do you say?

[fhahn]

In D69069#1812947, @ekatz wrote:

In D69069#1768576, @fhahn wrote:

Would it be possible to also add some of the test cases for the linked bugs?

This patch fixes a crash. Maybe just add a .ll file that shouldn't crash (with the fix)?!

that would be ideal.

llvm/lib/Transforms/IPO/LoopExtractor.cpp

149

I think the position of my original comment got messed up. I was referring to the distinction between functions with a single loop and functions with multiple loops. I *think* this was only needed with the LoopPass to avoid extracting loops we already extracted into new functions earlier which might result in an infinite loop.

I.e. just have the following starting at line 140.

DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
return extractLoops(LI.begin(), LI.end(), LI, DT);

The test coverage isn't great, but it seems to pass all tests ;)

213

Right.

[ekatz]

llvm/lib/Transforms/IPO/LoopExtractor.cpp

149

Not sure the position got messed up - it is probably my fault. I do not sleep much these days, with my 1 year old teething or whatever he decides the reason to wake up every hour at night...

Anyway, this code is from SVN commit rL12403 (Git commit rG2f155d87). It is from the time that the pass was a FunctionPass. It tries to distinguish between a function that has more than a single loop with a simple function wrapper around a loop (which can be created by other means than the extractor).

[ekatz]

llvm/lib/Transforms/IPO/LoopExtractor.cpp

149

I'll add a test for that particular case.

[fhahn]

This LGTM, pending the test cases.

llvm/lib/Transforms/IPO/LoopExtractor.cpp

149

Anyway, this code is from SVN commit rL12403 (Git commit rG2f155d87). It is from the time that the pass was a FunctionPass. It tries to distinguish between a function that has more than a single loop with a simple function wrapper around a loop (which can be created by other means than the extractor).

Right, but it cannot happen while the LoopExtractor is running, with it being a module pass now, so we cannot get stuck in a loop where we extract a loop into a wrapper function and then try to extract the loop from the newly created function and get stuck.

Of course it still might be slightly more preferable to avoid extracting loops from trivial wrapper functions, but IMO the extra complexity does not really warrant the smallish benefit on a few probably not very common cases (also the checks are quite brittle and I guess not hit very often in practice (unless running loop-extract multiple times...))

But I am not feeling very strongly about that point :)

[ekatz]

Added test cases for the bugs, as well as for testing a simple extraction and a no extraction in the edge case where the function is just a minimal wrapper around a loop.

[ekatz]

llvm/lib/Transforms/IPO/LoopExtractor.cpp

149

The new test LoopExtractor_min_wrapper.ll demonstrates the use case.

[ekatz]

@fhahn , do I have a green light?

[fhahn]

LGTM, with comments addressed before committing. Thanks!

llvm/lib/Transforms/IPO/LoopExtractor.cpp

151

nit: the comment is a bit confusing, as it checks if the loop is in simplify-form. Maybe say something like: If the loop is in LoopSimplify form, check if the function is a minimal container around the loop.

llvm/test/Transforms/CodeExtractor/LoopExtractor.ll

3 ↗

(On Diff #239663)

I think it would be good to check the full IR after extraction.

llvm/test/Transforms/CodeExtractor/LoopExtractor_crash.ll

4 ↗

(On Diff #239663)

There are no arguments to promote and no functions to inline? Is -inline/-argpromotion needed? Also, please add some rudimentary checks. (you could use llvm/tools/update_test_checks.py)

llvm/test/Transforms/CodeExtractor/LoopExtractor_infinite.ll

3 ↗

(On Diff #239663)

Is running -mergereturn required? Or is it enough to have the output of merge return as test case? Also it would be good to have some rudimentary checks.

[ekatz]

llvm/test/Transforms/CodeExtractor/LoopExtractor_crash.ll

4 ↗

(On Diff #239663)

inline and argpromotion shouldn't do anything but certainly not crash. This crash is observed in those cases.

llvm/test/Transforms/CodeExtractor/LoopExtractor_infinite.ll

3 ↗

(On Diff #239663)

opt must be run with mergereturn and loop-extract to observe the infinite loop.

[ekatz]

llvm/lib/Transforms/IPO/LoopExtractor.cpp

151

Done.

llvm/test/Transforms/CodeExtractor/LoopExtractor.ll

3 ↗

(On Diff #239663)

Added full IR check.

llvm/test/Transforms/CodeExtractor/LoopExtractor_crash.ll

4 ↗

(On Diff #239663)

I don't think there is any reason to test the resulting IR. This test is only testing that there is no crash.

[fhahn]

llvm/test/Transforms/CodeExtractor/LoopExtractor_crash.ll

4 ↗

(On Diff #239663)

Sure, but it would still be good to check that the pass does what we expect (not crashing is a relatively low bar). I guess this is a case where the function is just a minimal container? If that's the case, that's worth explicitly checking (and probably best to also add a short comment & a reference to the bug).

also without check, the pass could also delete the whole function in order not to crash ;)

[ekatz]

Added IR checks of all the tests.

[ekatz]

llvm/test/Transforms/CodeExtractor/LoopExtractor_crash.ll

4 ↗

(On Diff #239663)

I agree. Done.

[ekatz]

@fhahn does it look right?

[fhahn]

Looks good, thanks!

[fhahn]

Is it possible that this broke the sanitizer bot? http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/38687/steps/check-llvm%20asan/logs/stdio

[ekatz]

In D69069#1872630, @fhahn wrote:

Is it possible that this broke the sanitizer bot? http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/38687/steps/check-llvm%20asan/logs/stdio

Seems like it... :/
I wonder if it is because of the duplication of the dependencies in:

INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)

as LoopSimplify already depends on LoopInfoWrapperPass which depends on DominatorTreeWrapperPass.

I'll try to remove the last 2.

[ekatz]

Should be fixed in rGd8a2ea9fd5c7

[fhahn]

In D69069#1872969, @ekatz wrote:

Should be fixed in rGd8a2ea9fd5c7

looks like it, thanks!

llvm/lib/Transforms/IPO/LoopExtractor.cpp

 // top-level loop into its own new function. If the loop is the ONLY loop in a
 // given function, it is not touched. This is a pass most useful for debugging
 // via bugpoint.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
-#include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 #include <fstream>
 #include <set>
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-extract"
 
 STATISTIC(NumExtracted, "Number of loops extracted");
 
 namespace {
-  struct LoopExtractor : public LoopPass {
+  struct LoopExtractor : public ModulePass {
     static char ID; // Pass identification, replacement for typeid
+
+    // The number of natural loops to extract from the program into functions.
     unsigned NumLoops;
 
     explicit LoopExtractor(unsigned numLoops = ~0)
-      : LoopPass(ID), NumLoops(numLoops) {
-        initializeLoopExtractorPass(*PassRegistry::getPassRegistry());
-      }
+        : ModulePass(ID), NumLoops(numLoops) {
 
-    bool runOnLoop(Loop *L, LPPassManager &) override;
+    bool runOnModule(Module &M) override;
+    bool runOnFunction(Function &F);
+
+    bool extractLoops(Loop::iterator From, Loop::iterator To, LoopInfo &LI,
+                      DominatorTree &DT);
+    bool extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT);
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.addRequiredID(BreakCriticalEdgesID);
       AU.addRequiredID(LoopSimplifyID);
       AU.addRequired<DominatorTreeWrapperPass>();
       AU.addRequired<LoopInfoWrapperPass>();
+      AU.addPreserved<LoopInfoWrapperPass>();
       AU.addUsedIfAvailable<AssumptionCacheTracker>();
     }
   };
 }
 
 char LoopExtractor::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopExtractor, "loop-extract",
                       "Extract loops into new functions", false, false)
 INITIALIZE_PASS_DEPENDENCY(BreakCriticalEdges)
 INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_END(LoopExtractor, "loop-extract",
                     "Extract loops into new functions", false, false)
 
 namespace {
   /// SingleLoopExtractor - For bugpoint.
   struct SingleLoopExtractor : public LoopExtractor {
     static char ID; // Pass identification, replacement for typeid
     SingleLoopExtractor() : LoopExtractor(1) {}
   };
 } // End anonymous namespace
 
 char SingleLoopExtractor::ID = 0;
 INITIALIZE_PASS(SingleLoopExtractor, "loop-extract-single",
                 "Extract at most one loop into a new function", false, false)
 
 // createLoopExtractorPass - This pass extracts all natural loops from the
 // program into a function if it can.
 //
 Pass *llvm::createLoopExtractorPass() { return new LoopExtractor(); }
 
-bool LoopExtractor::runOnLoop(Loop *L, LPPassManager &LPM) {
-  if (skipLoop(L))
+bool LoopExtractor::runOnModule(Module &M) {
+  if (skipModule(M))
     return false;
 
-  // Only visit top-level loops.
-  if (L->getParentLoop())
+  if (M.empty())
     return false;
 
-  // If LoopSimplify form is not available, stay out of trouble.
-  if (!L->isLoopSimplifyForm())
+  if (!NumLoops)
     return false;
 
-  DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-  LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   bool Changed = false;
 
+  // The end of the function list may change (new functions will be added at the
+  // end), so we run from the first to the current last.
+  auto I = M.begin(), E = --M.end();
+  while (true) {
+    Function &F = *I;
+
+    Changed |= runOnFunction(F);
+    if (!NumLoops)
+      break;
+
+    // If this is the last function.
+    if (I == E)
+      break;
+
+    ++I;
+  }
+  return Changed;
+}
+
+bool LoopExtractor::runOnFunction(Function &F) {
+  // Do not modify `optnone` functions.
+  if (F.hasOptNone())
+    return false;
+
+  if (F.empty())
+    return false;
+
+  LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();
+
+  // If there are no loops in the function.
+  if (LI.empty())
+    return false;
+
+  DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
+
   // If there is more than one top-level loop in this function, extract all of
-  // the loops. Otherwise there is exactly one top-level loop; in this case if
-  // this function is more than a minimal wrapper around the loop, extract
-  // the loop.
+  // the loops.
+  if (std::distance(LI.begin(), LI.end()) > 1)
+    return extractLoops(LI.begin(), LI.end(), LI, DT);
+
+  // Otherwise there is exactly one top-level loop; in this case if this
+  // function is more than a minimal wrapper around the loop, extract the loop.
+  Loop *TLL = *LI.begin();
+
+  // If LoopSimplify form is not available, stay out of trouble.
+  if (TLL->isLoopSimplifyForm()) {
     bool ShouldExtractLoop = false;
 
     // Extract the loop if the entry block doesn't branch to the loop header.
-  Instruction *EntryTI =
-      L->getHeader()->getParent()->getEntryBlock().getTerminator();
+    Instruction *EntryTI = F.getEntryBlock().getTerminator();
     if (!isa<BranchInst>(EntryTI) ||
         !cast<BranchInst>(EntryTI)->isUnconditional() ||
-      EntryTI->getSuccessor(0) != L->getHeader()) {
+        EntryTI->getSuccessor(0) != TLL->getHeader()) {
       ShouldExtractLoop = true;
     } else {
       // Check to see if any exits from the loop are more than just return
       // blocks.
-    SmallVector<BasicBlock*, 8> ExitBlocks;
-    L->getExitBlocks(ExitBlocks);
-    for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i)
-      if (!isa<ReturnInst>(ExitBlocks[i]->getTerminator())) {
+      SmallVector<BasicBlock *, 8> ExitBlocks;
+      TLL->getExitBlocks(ExitBlocks);
+      for (auto *ExitBlock : ExitBlocks)
+        if (!isa<ReturnInst>(ExitBlock->getTerminator())) {
           ShouldExtractLoop = true;
           break;
         }
     }
 
-  if (ShouldExtractLoop) {
-    // We must omit EH pads. EH pads must accompany the invoke
-    // instruction. But this would result in a loop in the extracted
-    // function. An infinite cycle occurs when it tries to extract that loop as
-    // well.
-    SmallVector<BasicBlock*, 8> ExitBlocks;
-    L->getExitBlocks(ExitBlocks);
-    for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i)
-      if (ExitBlocks[i]->isEHPad()) {
-        ShouldExtractLoop = false;
-        break;
-      }
-  }
+    if (ShouldExtractLoop)
+      return extractLoop(TLL, LI, DT);
+  }
+
+  // Okay, this function is a minimal container around the specified loop.
+  // If we extract the loop, we will continue to just keep extracting it
+  // infinitely... so don't extract it. However, if the loop contains any
+  // sub-loops, extract them.
+  return extractLoops(TLL->begin(), TLL->end(), LI, DT);
+}
+
+bool LoopExtractor::extractLoops(Loop::iterator From, Loop::iterator To,
+                                 LoopInfo &LI, DominatorTree &DT) {
+  bool Changed = false;
+  SmallVector<Loop *, 8> Loops;
+
+  // Save the list of loops, as it may change.
+  Loops.assign(From, To);
+  for (Loop *L : Loops) {
+    // If LoopSimplify form is not available, stay out of trouble.
+    if (L->isLoopSimplifyForm())
+      continue;
+
+    Changed |= extractLoop(L, LI, DT);
+    if (!NumLoops)
+  return Changed;
 
-  if (ShouldExtractLoop) {
-    if (NumLoops == 0) return Changed;
-    --NumLoops;
-    AssumptionCache *AC = nullptr;
-    Function &Func = *L->getHeader()->getParent();
-    if (auto *ACT = getAnalysisIfAvailable<AssumptionCacheTracker>())
-      AC = ACT->lookupAssumptionCache(Func);
-    CodeExtractorAnalysisCache CEAC(Func);
-    CodeExtractor Extractor(DT, *L, false, nullptr, nullptr, AC);
-    if (Extractor.extractCodeRegion(CEAC) != nullptr) {
-      Changed = true;
-      // After extraction, the loop is replaced by a function call, so
-      // we shouldn't try to run any more loop passes on it.
-      LPM.markLoopAsDeleted(*L);
-      LI.erase(L);
-    }
+bool LoopExtractor::extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT) {
+  assert(NumLoops != 0);
+  if (Extractor.extractCodeRegion(CEAC)) {
+    --NumLoops;
     ++NumExtracted;
+    return true;
   }
-
-  return Changed;
+  return false;
 }
 
 // createSingleLoopExtractorPass - This pass extracts one natural loop from the
 // program into a function if it can.  This is used by bugpoint.
 //
 Pass *llvm::createSingleLoopExtractorPass() {
   return new SingleLoopExtractor();
 }

llvm/test/Feature/optnone-opt.ll

 ; OPT-O2O3-DAG: Skipping pass 'SLP Vectorizer'
 
 ; Additional IR passes that opt doesn't turn on by default.
 ; OPT-MORE-DAG: Skipping pass 'Dead Code Elimination'
 ; OPT-MORE-DAG: Skipping pass 'Dead Instruction Elimination'
 
 ; Loop IR passes that opt doesn't turn on by default.
 ; OPT-LOOP-DAG: Skipping pass 'Delete dead loops'
-; OPT-LOOP-DAG: Skipping pass 'Extract loops into new functions'
 ; OPT-LOOP-DAG: Skipping pass 'Induction Variable Simplification'
 ; OPT-LOOP-DAG: Skipping pass 'Loop Invariant Code Motion'
 ; OPT-LOOP-DAG: Skipping pass 'Loop Strength Reduction'
 ; OPT-LOOP-DAG: Skipping pass 'Recognize loop idioms'
 ; OPT-LOOP-DAG: Skipping pass 'Reroll loops'
 ; OPT-LOOP-DAG: Skipping pass 'Rotate Loops'
 ; OPT-LOOP-DAG: Skipping pass 'Simplify instructions in loops'
 ; OPT-LOOP-DAG: Skipping pass 'Unroll loops'
 ; OPT-LOOP-DAG: Skipping pass 'Unswitch loops'