[LoopExtractor] Convert LoopExtractor from LoopPass to ModulePass

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.
Includes reverting rL82990 implications on the LoopExtractor.

Fixes PR3082 and PR8929.

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

lib/Transforms/IPO/LoopExtractor.cpp

@@ -15,7 +15,7 @@
 
 #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"
@@ -36,22 +36,30 @@ using namespace llvm;
 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) {
+      initializeLoopExtractorPass(*PassRegistry::getPassRegistry());
+    }
 
-    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>();
     }
   };
@@ -63,6 +71,7 @@ INITIALIZE_PASS_BEGIN(LoopExtractor, "loop-extract",
 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)
 
@@ -83,83 +92,131 @@ INITIALIZE_PASS(SingleLoopExtractor, "loop-extract-single",
 //
 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;
 
-  // 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.
-  bool ShouldExtractLoop = 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;
 
-  // Extract the loop if the entry block doesn't branch to the loop header.
-  Instruction *EntryTI =
-      L->getHeader()->getParent()->getEntryBlock().getTerminator();
-  if (!isa<BranchInst>(EntryTI) ||
-      !cast<BranchInst>(EntryTI)->isUnconditional() ||
-      EntryTI->getSuccessor(0) != L->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())) {
-        ShouldExtractLoop = true;
-        break;
-      }
+    Changed |= runOnFunction(F);
+    if (!NumLoops)
+      break;
+
+    // If this is the last function.
+    if (I == E)
+      break;
+
+    ++I;
   }
-
-  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) {
-    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);
-    }
-    ++NumExtracted;
-  }
-
   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.
+  if (std::next(LI.begin()) != LI.end())
+    return extractLoops(LI.begin(), LI.end(), LI, DT);
+
+  // Otherwise there is exactly one top-level loop.
+  Loop *TLL = *LI.begin();
+
+  // If the loop is in LoopSimplify form, then extract it only if this function
+  // is more than a minimal wrapper around the loop.
+  if (TLL->isLoopSimplifyForm()) {
+    bool ShouldExtractLoop = false;
+
+    // Extract the loop if the entry block doesn't branch to the loop header.
+    Instruction *EntryTI = F.getEntryBlock().getTerminator();
+    if (!isa<BranchInst>(EntryTI) ||
+        !cast<BranchInst>(EntryTI)->isUnconditional() ||
+        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;
+      TLL->getExitBlocks(ExitBlocks);
+      for (auto *ExitBlock : ExitBlocks)
+        if (!isa<ReturnInst>(ExitBlock->getTerminator())) {
+          ShouldExtractLoop = true;
+          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)
+      break;
+  }
+  return Changed;
+}
+
+bool LoopExtractor::extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT) {
+  assert(NumLoops != 0);
+  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)) {
+    LI.erase(L);
+    --NumLoops;
+    ++NumExtracted;
+    return true;
+  }
+  return false;
+}
+
 // createSingleLoopExtractorPass - This pass extracts one natural loop from the
 // program into a function if it can.  This is used by bugpoint.
 //

test/Feature/optnone-opt.ll

@@ -54,7 +54,6 @@ attributes #0 = { optnone noinline }
 
 ; 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'

test/Transforms/CodeExtractor/LoopExtractor.ll

@@ -0,0 +1,68 @@
+; RUN: opt < %s -loop-extract -S | FileCheck %s
+
+; This function has 2 simple loops and they should be extracted into 2 new functions.
+define void @test3() {
+; CHECK-LABEL: @test3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label %codeRepl1
+; CHECK:       codeRepl1:
+; CHECK-NEXT:    call void @test3.loop.0()
+; CHECK-NEXT:    br label %loop.0.loop.1_crit_edge
+; CHECK:       loop.0.loop.1_crit_edge:
+; CHECK-NEXT:    br label %codeRepl
+; CHECK:       codeRepl:
+; CHECK-NEXT:    call void @test3.loop.1()
+; CHECK-NEXT:    br label %exit
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+
+entry:
+  br label %loop.0
+
+loop.0:                                           ; preds = %loop.0, %entry
+  %index.0 = phi i32 [ 10, %entry ], [ %next.0, %loop.0 ]
+  tail call void @foo()
+  %next.0 = add nsw i32 %index.0, -1
+  %repeat.0 = icmp sgt i32 %index.0, 1
+  br i1 %repeat.0, label %loop.0, label %loop.1
+
+loop.1:                                           ; preds = %loop.0, %loop.1
+  %index.1 = phi i32 [ %next.1, %loop.1 ], [ 10, %loop.0 ]
+  tail call void @foo()
+  %next.1 = add nsw i32 %index.1, -1
+  %repeat.1 = icmp sgt i32 %index.1, 1
+  br i1 %repeat.1, label %loop.1, label %exit
+
+exit:                                             ; preds = %loop.1
+  ret void
+}
+
+declare void @foo()
+
+; CHECK-LABEL: define internal void @test3.loop.1()
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label %loop.1
+; CHECK:       exit.exitStub:
+; CHECK-NEXT:    ret void
+; CHECK:       loop.1:
+; CHECK-NEXT:    %index.1 = phi i32 [ %next.1, %loop.1.loop.1_crit_edge ], [ 10, %newFuncRoot ]
+; CHECK-NEXT:    tail call void @foo()
+; CHECK-NEXT:    %next.1 = add nsw i32 %index.1, -1
+; CHECK-NEXT:    %repeat.1 = icmp sgt i32 %index.1, 1
+; CHECK-NEXT:    br i1 %repeat.1, label %loop.1.loop.1_crit_edge, label %exit.exitStub
+; CHECK:       loop.1.loop.1_crit_edge:
+; CHECK-NEXT:    br label %loop.1
+
+; CHECK-LABEL: define internal void @test3.loop.0()
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label %loop.0
+; CHECK:       loop.0.loop.1_crit_edge.exitStub:
+; CHECK-NEXT:    ret void
+; CHECK:       loop.0:
+; CHECK-NEXT:    %index.0 = phi i32 [ 10, %newFuncRoot ], [ %next.0, %loop.0.loop.0_crit_edge ]
+; CHECK-NEXT:    tail call void @foo()
+; CHECK-NEXT:    %next.0 = add nsw i32 %index.0, -1
+; CHECK-NEXT:    %repeat.0 = icmp sgt i32 %index.0, 1
+; CHECK-NEXT:    br i1 %repeat.0, label %loop.0.loop.0_crit_edge, label %loop.0.loop.1_crit_edge.exitStub
+; CHECK:       loop.0.loop.0_crit_edge:
+; CHECK-NEXT:    br label %loop.0

test/Transforms/CodeExtractor/LoopExtractor_crash.ll

@@ -0,0 +1,46 @@
+; RUN: opt < %s -inline -loop-extract -S | FileCheck %s
+; RUN: opt < %s -argpromotion -loop-extract -S | FileCheck %s
+
+; This test used to trigger an assert (PR8929).
+
+define void @test() {
+; CHECK-LABEL: define void @test()
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label %codeRepl
+; CHECK:       codeRepl:
+; CHECK-NEXT:    call void @test.loopentry()
+; CHECK-NEXT:    br label %loopexit
+; CHECK:       loopexit:
+; CHECK-NEXT:    br label %exit
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+
+entry:
+  br label %loopentry
+
+loopentry:                                        ; preds = %loopbody, %entry
+  br i1 undef, label %loopbody, label %loopexit
+
+loopbody:                                         ; preds = %codeRepl1
+  call void @foo()
+  br label %loopentry
+
+loopexit:                                         ; preds = %codeRepl
+  br label %exit
+
+exit:                                             ; preds = %loopexit
+  ret void
+}
+
+declare void @foo()
+
+; CHECK-LABEL: define internal void @test.loopentry()
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label %loopentry
+; CHECK:       loopexit.exitStub:
+; CHECK-NEXT:    ret void
+; CHECK:       loopentry:
+; CHECK-NEXT:    br i1 false, label %loopbody, label %loopexit.exitStub
+; CHECK:       loopbody:
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    br label %loopentry

test/Transforms/CodeExtractor/LoopExtractor_infinite.ll

@@ -0,0 +1,47 @@
+; RUN: opt < %s -mergereturn -loop-extract -S | FileCheck %s
+
+; This test used to enter an infinite loop, until out of memory (PR3082).
+
+define void @test() {
+; CHECK-LABEL: define void @test()
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label %codeRepl
+; CHECK:       codeRepl:
+; CHECK-NEXT:    %targetBlock = call i1 @test.loopentry()
+; CHECK-NEXT:    br i1 %targetBlock, label %exit.1, label %exit.0
+; CHECK:       exit.0:
+; CHECK-NEXT:    br label %UnifiedReturnBlock
+; CHECK:       exit.1:
+; CHECK-NEXT:    br label %UnifiedReturnBlock
+; CHECK:       UnifiedReturnBlock:
+; CHECK-NEXT:    ret void
+
+entry:
+  br label %loopentry
+
+loopentry:                                        ; preds = %loopexit, %entry
+  br i1 undef, label %exit.1, label %loopexit
+
+loopexit:                                         ; preds = %loopentry
+  br i1 undef, label %loopentry, label %exit.0
+
+exit.0:                                           ; preds = %loopexit
+  ret void
+
+exit.1:                                           ; preds = %loopentry
+  ret void
+}
+
+; CHECK-LABEL: define internal i1 @test.loopentry()
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label %loopentry
+; CHECK:       exit.1.exitStub:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       exit.0.exitStub:
+; CHECK-NEXT:    ret i1 false
+; CHECK:       loopentry:
+; CHECK-NEXT:    br i1 true, label %exit.1.exitStub, label %loopexit
+; CHECK:       loopexit:
+; CHECK-NEXT:    br i1 false, label %loopexit.loopentry_crit_edge, label %exit.0.exitStub
+; CHECK:       loopexit.loopentry_crit_edge:
+; CHECK-NEXT:    br label %loopentry

test/Transforms/CodeExtractor/LoopExtractor_min_wrapper.ll

@@ -0,0 +1,35 @@
+; RUN: opt < %s -loop-extract -S | FileCheck %s
+
+; This function is just a minimal wrapper around a loop and should not be extracted.
+define void @test() {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label %loop
+; CHECK:       loop:
+; CHECK-NEXT:    %index = phi i32 [ 0, %entry ], [ %next, %loop.loop_crit_edge ]
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    %next = add nsw i32 %index, -1
+; CHECK-NEXT:    %repeat = icmp sgt i32 %index, 1
+; CHECK-NEXT:    br i1 %repeat, label %loop.loop_crit_edge, label %exit
+; CHECK:       loop.loop_crit_edge:
+; CHECK-NEXT:    br label %loop
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+
+entry:
+  br label %loop
+
+loop:                                             ; preds = %loop, %entry
+  %index = phi i32 [ 0, %entry ], [ %next, %loop ]
+  call void @foo()
+  %next = add nsw i32 %index, -1
+  %repeat = icmp sgt i32 %index, 1
+  br i1 %repeat, label %loop, label %exit
+
+exit:                                             ; preds = %loop
+  ret void
+}
+
+declare void @foo()
+
+; CHECK-NOT: define