[LoopExtract][NewPM] Port -loop-extract to NPM

-loop-extract-single is just -loop-extract on one loop.

-loop-extract depended on -break-crit-edges and -loop-simplify in the
legacy PM, but the NPM doesn't allow specifying pass dependencies like
that, so manually add those passes to the RUN lines where necessary.

Reviewed By: asbirlea

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

include/llvm/InitializePasses.h

@@ -235,7 +235,7 @@ void initializeLoopAccessLegacyAnalysisPass(PassRegistry&);
 void initializeLoopDataPrefetchLegacyPassPass(PassRegistry&);
 void initializeLoopDeletionLegacyPassPass(PassRegistry&);
 void initializeLoopDistributeLegacyPass(PassRegistry&);
-void initializeLoopExtractorPass(PassRegistry&);
+void initializeLoopExtractorLegacyPassPass(PassRegistry &);
 void initializeLoopGuardWideningLegacyPassPass(PassRegistry&);
 void initializeLoopFuseLegacyPass(PassRegistry&);
 void initializeLoopIdiomRecognizeLegacyPassPass(PassRegistry&);

include/llvm/Transforms/IPO/LoopExtractor.h

@@ -0,0 +1,32 @@
+//===- LoopExtractor.h - Extract each loop into a new function ------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A pass wrapper around the ExtractLoop() scalar transformation to extract each
+// 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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_IPO_LOOPEXTRACTOR_H
+#define LLVM_TRANSFORMS_IPO_LOOPEXTRACTOR_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+struct LoopExtractorPass : public PassInfoMixin<LoopExtractorPass> {
+  LoopExtractorPass(unsigned NumLoops = ~0) : NumLoops(NumLoops) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+
+private:
+  unsigned NumLoops;
+};
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_LOOPEXTRACTOR_H

lib/Passes/PassBuilder.cpp

@@ -98,6 +98,7 @@
 #include "llvm/Transforms/IPO/InferFunctionAttrs.h"
 #include "llvm/Transforms/IPO/Inliner.h"
 #include "llvm/Transforms/IPO/Internalize.h"
+#include "llvm/Transforms/IPO/LoopExtractor.h"
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
 #include "llvm/Transforms/IPO/MergeFunctions.h"
 #include "llvm/Transforms/IPO/OpenMPOpt.h"

lib/Passes/PassRegistry.def

@@ -67,6 +67,7 @@ MODULE_PASS("internalize", InternalizePass())
 MODULE_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 MODULE_PASS("ipsccp", IPSCCPPass())
 MODULE_PASS("print-ir-similarity", IRSimilarityAnalysisPrinterPass(dbgs()))
+MODULE_PASS("loop-extract", LoopExtractorPass())
 MODULE_PASS("lowertypetests", LowerTypeTestsPass())
 MODULE_PASS("metarenamer", MetaRenamerPass())
 MODULE_PASS("mergefunc", MergeFunctionsPass())
@@ -90,6 +91,7 @@ MODULE_PASS("rpo-function-attrs", ReversePostOrderFunctionAttrsPass())
 MODULE_PASS("sample-profile", SampleProfileLoaderPass())
 MODULE_PASS("scc-oz-module-inliner",
   buildInlinerPipeline(OptimizationLevel::Oz, ThinLTOPhase::None, DebugLogging))
+MODULE_PASS("loop-extract-single", LoopExtractorPass(1))
 MODULE_PASS("oz-module-optimizer",
   buildModuleOptimizationPipeline(OptimizationLevel::Oz, DebugLogging, /*LTOPreLink*/false))
 MODULE_PASS("strip", StripSymbolsPass())

lib/Transforms/IPO/IPO.cpp

@@ -39,7 +39,7 @@ void llvm::initializeIPO(PassRegistry &Registry) {
   initializeSimpleInlinerPass(Registry);
   initializeInferFunctionAttrsLegacyPassPass(Registry);
   initializeInternalizeLegacyPassPass(Registry);
-  initializeLoopExtractorPass(Registry);
+  initializeLoopExtractorLegacyPassPass(Registry);
   initializeBlockExtractorPass(Registry);
   initializeSingleLoopExtractorPass(Registry);
   initializeLowerTypeTestsPass(Registry);

lib/Transforms/IPO/LoopExtractor.cpp

@@ -13,12 +13,14 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Transforms/IPO/LoopExtractor.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
@@ -36,51 +38,71 @@ using namespace llvm;
 STATISTIC(NumExtracted, "Number of loops extracted");
 
 namespace {
-  struct LoopExtractor : public ModulePass {
-    static char ID; // Pass identification, replacement for typeid
+struct LoopExtractorLegacyPass : public ModulePass {
+  static char ID; // Pass identification, replacement for typeid
 
-    // The number of natural loops to extract from the program into functions.
-    unsigned NumLoops;
+  unsigned NumLoops;
 
-    explicit LoopExtractor(unsigned numLoops = ~0)
-        : ModulePass(ID), NumLoops(numLoops) {
-      initializeLoopExtractorPass(*PassRegistry::getPassRegistry());
-    }
+  explicit LoopExtractorLegacyPass(unsigned NumLoops = ~0)
+      : ModulePass(ID), NumLoops(NumLoops) {
+    initializeLoopExtractorLegacyPassPass(*PassRegistry::getPassRegistry());
+  }
 
-    bool runOnModule(Module &M) override;
-    bool runOnFunction(Function &F);
+  bool runOnModule(Module &M) override;
 
-    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.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<LoopInfoWrapperPass>();
+    AU.addPreserved<LoopInfoWrapperPass>();
+    AU.addRequiredID(LoopSimplifyID);
+    AU.addUsedIfAvailable<AssumptionCacheTracker>();
+  }
+};
 
-    void getAnalysisUsage(AnalysisUsage &AU) const override {
-      AU.addRequiredID(BreakCriticalEdgesID);
-      AU.addRequired<DominatorTreeWrapperPass>();
-      AU.addRequired<LoopInfoWrapperPass>();
-      AU.addPreserved<LoopInfoWrapperPass>();
-      AU.addRequiredID(LoopSimplifyID);
-      AU.addUsedIfAvailable<AssumptionCacheTracker>();
-    }
-  };
-}
+struct LoopExtractor {
+  explicit LoopExtractor(
+      unsigned NumLoops,
+      function_ref<DominatorTree &(Function &)> LookupDomTree,
+      function_ref<LoopInfo &(Function &)> LookupLoopInfo,
+      function_ref<AssumptionCache *(Function &)> LookupAssumptionCache)
+      : NumLoops(NumLoops), LookupDomTree(LookupDomTree),
+        LookupLoopInfo(LookupLoopInfo),
+        LookupAssumptionCache(LookupAssumptionCache) {}
+  bool runOnModule(Module &M);
 
-char LoopExtractor::ID = 0;
-INITIALIZE_PASS_BEGIN(LoopExtractor, "loop-extract",
+private:
+  // The number of natural loops to extract from the program into functions.
+  unsigned NumLoops;
+
+  function_ref<DominatorTree &(Function &)> LookupDomTree;
+  function_ref<LoopInfo &(Function &)> LookupLoopInfo;
+  function_ref<AssumptionCache *(Function &)> LookupAssumptionCache;
+
+  bool runOnFunction(Function &F);
+
+  bool extractLoops(Loop::iterator From, Loop::iterator To, LoopInfo &LI,
+                    DominatorTree &DT);
+  bool extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT);
+};
+} // namespace
+
+char LoopExtractorLegacyPass::ID = 0;
+INITIALIZE_PASS_BEGIN(LoopExtractorLegacyPass, "loop-extract",
                       "Extract loops into new functions", false, false)
 INITIALIZE_PASS_DEPENDENCY(BreakCriticalEdges)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
-INITIALIZE_PASS_END(LoopExtractor, "loop-extract",
+INITIALIZE_PASS_END(LoopExtractorLegacyPass, "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) {}
-  };
+struct SingleLoopExtractor : public LoopExtractorLegacyPass {
+  static char ID; // Pass identification, replacement for typeid
+  SingleLoopExtractor() : LoopExtractorLegacyPass(1) {}
+};
 } // End anonymous namespace
 
 char SingleLoopExtractor::ID = 0;
@@ -90,12 +112,30 @@ INITIALIZE_PASS(SingleLoopExtractor, "loop-extract-single",
 // createLoopExtractorPass - This pass extracts all natural loops from the
 // program into a function if it can.
 //
-Pass *llvm::createLoopExtractorPass() { return new LoopExtractor(); }
+Pass *llvm::createLoopExtractorPass() { return new LoopExtractorLegacyPass(); }
 
-bool LoopExtractor::runOnModule(Module &M) {
+bool LoopExtractorLegacyPass::runOnModule(Module &M) {
   if (skipModule(M))
     return false;
 
+  bool Changed = false;
+  auto LookupDomTree = [this](Function &F) -> DominatorTree & {
+    return this->getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
+  };
+  auto LookupLoopInfo = [this, &Changed](Function &F) -> LoopInfo & {
+    return this->getAnalysis<LoopInfoWrapperPass>(F, &Changed).getLoopInfo();
+  };
+  auto LookupACT = [this](Function &F) -> AssumptionCache * {
+    if (auto *ACT = this->getAnalysisIfAvailable<AssumptionCacheTracker>())
+      return ACT->lookupAssumptionCache(F);
+    return nullptr;
+  };
+  return LoopExtractor(NumLoops, LookupDomTree, LookupLoopInfo, LookupACT)
+             .runOnModule(M) ||
+         Changed;
+}
+
+bool LoopExtractor::runOnModule(Module &M) {
   if (M.empty())
     return false;
 
@@ -132,13 +172,13 @@ bool LoopExtractor::runOnFunction(Function &F) {
     return false;
 
   bool Changed = false;
-  LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>(F, &Changed).getLoopInfo();
+  LoopInfo &LI = LookupLoopInfo(F);
 
   // If there are no loops in the function.
   if (LI.empty())
     return Changed;
 
-  DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
+  DominatorTree &DT = LookupDomTree(F);
 
   // If there is more than one top-level loop in this function, extract all of
   // the loops.
@@ -203,10 +243,8 @@ bool LoopExtractor::extractLoops(Loop::iterator From, Loop::iterator To,
 
 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);
+  AssumptionCache *AC = LookupAssumptionCache(Func);
   CodeExtractorAnalysisCache CEAC(Func);
   CodeExtractor Extractor(DT, *L, false, nullptr, nullptr, AC);
   if (Extractor.extractCodeRegion(CEAC)) {
@@ -224,3 +262,24 @@ bool LoopExtractor::extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT) {
 Pass *llvm::createSingleLoopExtractorPass() {
   return new SingleLoopExtractor();
 }
+
+PreservedAnalyses LoopExtractorPass::run(Module &M, ModuleAnalysisManager &AM) {
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  auto LookupDomTree = [&FAM](Function &F) -> DominatorTree & {
+    return FAM.getResult<DominatorTreeAnalysis>(F);
+  };
+  auto LookupLoopInfo = [&FAM](Function &F) -> LoopInfo & {
+    return FAM.getResult<LoopAnalysis>(F);
+  };
+  auto LookupAssumptionCache = [&FAM](Function &F) -> AssumptionCache * {
+    return FAM.getCachedResult<AssumptionAnalysis>(F);
+  };
+  if (!LoopExtractor(NumLoops, LookupDomTree, LookupLoopInfo,
+                     LookupAssumptionCache)
+           .runOnModule(M))
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  PA.preserve<LoopAnalysis>();
+  return PA;
+}

test/Transforms/CodeExtractor/2004-03-13-LoopExtractorCrash.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -disable-output
+; RUN: opt < %s -loop-simplify -loop-extract -disable-output
 
 define void @solve() {
 entry:

test/Transforms/CodeExtractor/2004-03-14-DominanceProblem.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -disable-output
+; RUN: opt < %s -loop-simplify -loop-extract -disable-output
 ; This testcase is failing the loop extractor because not all exit blocks 
 ; are dominated by all of the live-outs.
 

test/Transforms/CodeExtractor/2004-03-14-NoSwitchSupport.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract-single -disable-output
+; RUN: opt < %s -loop-simplify -loop-extract-single -disable-output
 
 define void @ab() {
 entry:

test/Transforms/CodeExtractor/2004-03-17-MissedLiveIns.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -disable-output
+; RUN: opt < %s -loop-simplify -loop-extract -disable-output
 
 define void @sendMTFValues() {
 entry:

test/Transforms/CodeExtractor/2004-03-17-UpdatePHIsOutsideRegion.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -disable-output
+; RUN: opt < %s -loop-simplify -loop-extract -disable-output
 
 define void @maketree() {
 entry:

test/Transforms/CodeExtractor/2004-03-18-InvokeHandling.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -disable-output
+; RUN: opt < %s -loop-simplify -loop-extract -disable-output
 
 declare i32 @_IO_getc()
 

test/Transforms/CodeExtractor/BlockAddressReference.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -S | FileCheck %s
+; RUN: opt < %s -loop-simplify -loop-extract -S | FileCheck %s
 
 @label = common local_unnamed_addr global i8* null
 

test/Transforms/CodeExtractor/BlockAddressSelfReference.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -S | FileCheck %s
+; RUN: opt < %s -loop-simplify -loop-extract -S | FileCheck %s
 
 @choum.addr = internal unnamed_addr constant [3 x i8*] [i8* blockaddress(@choum, %bb10), i8* blockaddress(@choum, %bb14), i8* blockaddress(@choum, %bb18)]
 

test/Transforms/CodeExtractor/LoopExtractor.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -S | FileCheck %s
+; RUN: opt < %s -break-crit-edges -loop-simplify -loop-extract -S | FileCheck %s
 
 ; This function has 2 simple loops and they should be extracted into 2 new functions.
 define void @test3() {

test/Transforms/CodeExtractor/LoopExtractor_alloca.ll

@@ -1,4 +1,4 @@
-; RUN: opt -debugify -loop-extract -S < %s | FileCheck %s
+; RUN: opt -debugify -loop-simplify -loop-extract -S < %s | FileCheck %s
 
 ; This tests 2 cases:
 ; 1. loop1 should be extracted into a function, without extracting %v1 alloca.

test/Transforms/CodeExtractor/LoopExtractor_crash.ll

@@ -1,5 +1,5 @@
-; RUN: opt < %s -inline -loop-extract -S | FileCheck %s
-; RUN: opt < %s -argpromotion -loop-extract -S | FileCheck %s
+; RUN: opt < %s -inline -loop-simplify -loop-extract -S | FileCheck %s
+; RUN: opt < %s -argpromotion -loop-simplify -loop-extract -S | FileCheck %s
 
 ; This test used to trigger an assert (PR8929).
 

test/Transforms/CodeExtractor/LoopExtractor_min_wrapper.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -loop-extract -S | FileCheck %s
+; RUN: opt < %s -break-crit-edges -loop-simplify -loop-extract -S | FileCheck %s
 
 ; This function is just a minimal wrapper around a loop and should not be extracted.
 define void @test() {