[BlockExtract][NewPM] Port -extract-blocks to NPM

[thakis]

llvm/lib/Passes/PassRegistry.def

53

nit: for most entries, the name matches the class name somewhat closely; here the words are flipped in name vs class name

[aeubanks]

llvm/lib/Passes/PassRegistry.def

53

I'm trying to match the legacy pass name for now, we can rename in the future after the NPM flip.

llvm/include/llvm/InitializePasses.h

 void initializeAssumptionCacheTrackerPass(PassRegistry&);
 void initializeAtomicExpandPass(PassRegistry&);
 void initializeAttributorLegacyPassPass(PassRegistry&);
 void initializeAttributorCGSCCLegacyPassPass(PassRegistry &);
 void initializeBasicBlockSectionsPass(PassRegistry &);
 void initializeBDCELegacyPassPass(PassRegistry&);
 void initializeBarrierNoopPass(PassRegistry&);
 void initializeBasicAAWrapperPassPass(PassRegistry&);
-void initializeBlockExtractorPass(PassRegistry &);
+void initializeBlockExtractorLegacyPassPass(PassRegistry &);
 void initializeBlockFrequencyInfoWrapperPassPass(PassRegistry&);
 void initializeBoundsCheckingLegacyPassPass(PassRegistry&);
 void initializeBranchFolderPassPass(PassRegistry&);
 void initializeBranchProbabilityInfoWrapperPassPass(PassRegistry&);
 void initializeBranchRelaxationPass(PassRegistry&);
 void initializeBreakCriticalEdgesPass(PassRegistry&);
 void initializeBreakFalseDepsPass(PassRegistry&);
 void initializeCanonicalizeAliasesLegacyPassPass(PassRegistry &);

llvm/include/llvm/Transforms/IPO/BlockExtractor.h

+//===- BlockExtractor.h - Extracts blocks into their own functions --------===//
+//
+// 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 pass extracts the specified basic blocks from the module into their
+// own functions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_IPO_BLOCKEXTRACTOR_H
+#define LLVM_TRANSFORMS_IPO_BLOCKEXTRACTOR_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+struct BlockExtractorPass : PassInfoMixin<BlockExtractorPass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_BLOCKEXTRACTOR_H

llvm/lib/Passes/PassBuilder.cpp

 #include "llvm/Transforms/Coroutines/CoroCleanup.h"
 #include "llvm/Transforms/Coroutines/CoroEarly.h"
 #include "llvm/Transforms/Coroutines/CoroElide.h"
 #include "llvm/Transforms/Coroutines/CoroSplit.h"
 #include "llvm/Transforms/HelloNew/HelloWorld.h"
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/IPO/Attributor.h"
+#include "llvm/Transforms/IPO/BlockExtractor.h"
 #include "llvm/Transforms/IPO/CalledValuePropagation.h"
 #include "llvm/Transforms/IPO/ConstantMerge.h"
 #include "llvm/Transforms/IPO/CrossDSOCFI.h"
 #include "llvm/Transforms/IPO/DeadArgumentElimination.h"
 #include "llvm/Transforms/IPO/ElimAvailExtern.h"
 #include "llvm/Transforms/IPO/ForceFunctionAttrs.h"
 #include "llvm/Transforms/IPO/FunctionAttrs.h"
 #include "llvm/Transforms/IPO/FunctionImport.h"

llvm/lib/Passes/PassRegistry.def

 MODULE_PASS("attributor", AttributorPass())
 MODULE_PASS("called-value-propagation", CalledValuePropagationPass())
 MODULE_PASS("canonicalize-aliases", CanonicalizeAliasesPass())
 MODULE_PASS("cg-profile", CGProfilePass())
 MODULE_PASS("constmerge", ConstantMergePass())
 MODULE_PASS("cross-dso-cfi", CrossDSOCFIPass())
 MODULE_PASS("deadargelim", DeadArgumentEliminationPass())
 MODULE_PASS("elim-avail-extern", EliminateAvailableExternallyPass())
+MODULE_PASS("extract-blocks", BlockExtractorPass())
 MODULE_PASS("forceattrs", ForceFunctionAttrsPass())
 MODULE_PASS("function-import", FunctionImportPass())
 MODULE_PASS("globaldce", GlobalDCEPass())
 MODULE_PASS("globalopt", GlobalOptPass())
 MODULE_PASS("globalsplit", GlobalSplitPass())
 MODULE_PASS("hotcoldsplit", HotColdSplittingPass())
 MODULE_PASS("hwasan", HWAddressSanitizerPass(false, false))
 MODULE_PASS("khwasan", HWAddressSanitizerPass(true, true))

llvm/lib/Transforms/IPO/BlockExtractor.cpp

 //===- BlockExtractor.cpp - Extracts blocks into their own functions ------===//
 //
 // 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 pass extracts the specified basic blocks from the module into their
 // own functions.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Transforms/IPO/BlockExtractor.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.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"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "block-extractor"
 
 STATISTIC(NumExtracted, "Number of basic blocks extracted");
 
 static cl::opt<std::string> BlockExtractorFile(
     "extract-blocks-file", cl::value_desc("filename"),
     cl::desc("A file containing list of basic blocks to extract"), cl::Hidden);
 
 cl::opt<bool> BlockExtractorEraseFuncs("extract-blocks-erase-funcs",
                                        cl::desc("Erase the existing functions"),
                                        cl::Hidden);
 namespace {
-class BlockExtractor : public ModulePass {
-  SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupsOfBlocks;
-  bool EraseFunctions;
-  /// Map a function name to groups of blocks.
-  SmallVector<std::pair<std::string, SmallVector<std::string, 4>>, 4>
-      BlocksByName;
-
+class BlockExtractor {
+public:
+  BlockExtractor(bool EraseFunctions) : EraseFunctions(EraseFunctions) {}
+  bool runOnModule(Module &M);
   void init(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
                 &GroupsOfBlocksToExtract) {
     for (const SmallVectorImpl<BasicBlock *> &GroupOfBlocks :
          GroupsOfBlocksToExtract) {
       SmallVector<BasicBlock *, 16> NewGroup;
       NewGroup.append(GroupOfBlocks.begin(), GroupOfBlocks.end());
       GroupsOfBlocks.emplace_back(NewGroup);
     }
     if (!BlockExtractorFile.empty())
       loadFile();
   }
 
+private:
+  SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupsOfBlocks;
+  bool EraseFunctions;
+  /// Map a function name to groups of blocks.
+  SmallVector<std::pair<std::string, SmallVector<std::string, 4>>, 4>
+      BlocksByName;
+
+  void loadFile();
+  void splitLandingPadPreds(Function &F);
+};
+
+class BlockExtractorLegacyPass : public ModulePass {
+  BlockExtractor BE;
+  bool runOnModule(Module &M) override;
+
 public:
   static char ID;
-  BlockExtractor(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
+  BlockExtractorLegacyPass(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
                            bool EraseFunctions)
-      : ModulePass(ID), EraseFunctions(EraseFunctions) {
+      : ModulePass(ID), BE(EraseFunctions) {
     // We want one group per element of the input list.
     SmallVector<SmallVector<BasicBlock *, 16>, 4> MassagedGroupsOfBlocks;
     for (BasicBlock *BB : BlocksToExtract) {
       SmallVector<BasicBlock *, 16> NewGroup;
       NewGroup.push_back(BB);
       MassagedGroupsOfBlocks.push_back(NewGroup);
     }
-    init(MassagedGroupsOfBlocks);
+    BE.init(MassagedGroupsOfBlocks);
   }
 
-  BlockExtractor(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
+  BlockExtractorLegacyPass(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
                                &GroupsOfBlocksToExtract,
                            bool EraseFunctions)
-      : ModulePass(ID), EraseFunctions(EraseFunctions) {
-    init(GroupsOfBlocksToExtract);
+      : ModulePass(ID), BE(EraseFunctions) {
+    BE.init(GroupsOfBlocksToExtract);
   }
 
-  BlockExtractor() : BlockExtractor(SmallVector<BasicBlock *, 0>(), false) {}
-  bool runOnModule(Module &M) override;
-
-private:
-  void loadFile();
-  void splitLandingPadPreds(Function &F);
+  BlockExtractorLegacyPass()
+      : BlockExtractorLegacyPass(SmallVector<BasicBlock *, 0>(), false) {}
 };
+
 } // end anonymous namespace
 
-char BlockExtractor::ID = 0;
-INITIALIZE_PASS(BlockExtractor, "extract-blocks",
+char BlockExtractorLegacyPass::ID = 0;
+INITIALIZE_PASS(BlockExtractorLegacyPass, "extract-blocks",
                 "Extract basic blocks from module", false, false)
 
-ModulePass *llvm::createBlockExtractorPass() { return new BlockExtractor(); }
+ModulePass *llvm::createBlockExtractorPass() {
+  return new BlockExtractorLegacyPass();
+}
 ModulePass *llvm::createBlockExtractorPass(
     const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
-  return new BlockExtractor(BlocksToExtract, EraseFunctions);
+  return new BlockExtractorLegacyPass(BlocksToExtract, EraseFunctions);
 }
 ModulePass *llvm::createBlockExtractorPass(
     const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
         &GroupsOfBlocksToExtract,
     bool EraseFunctions) {
-  return new BlockExtractor(GroupsOfBlocksToExtract, EraseFunctions);
+  return new BlockExtractorLegacyPass(GroupsOfBlocksToExtract, EraseFunctions);
 }
 
 /// Gets all of the blocks specified in the input file.
 void BlockExtractor::loadFile() {
   auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
   if (ErrOrBuf.getError())
     report_fatal_error("BlockExtractor couldn't load the file.");
   // Read the file.
if (EraseFunctions || BlockExtractorEraseFuncs) {
     // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
     for (Function &F : M)
       F.setLinkage(GlobalValue::ExternalLinkage);
     Changed = true;
   }
 
   return Changed;
 }
+
+bool BlockExtractorLegacyPass::runOnModule(Module &M) {
+  return BE.runOnModule(M);
+}
+
+PreservedAnalyses BlockExtractorPass::run(Module &M,
+                                          ModuleAnalysisManager &AM) {
+  BlockExtractor BE(false);
+  BE.init(SmallVector<SmallVector<BasicBlock *, 16>, 0>());
+  return BE.runOnModule(M) ? PreservedAnalyses::none()
+                           : PreservedAnalyses::all();
+}

llvm/lib/Transforms/IPO/IPO.cpp

void llvm::initializeIPO(PassRegistry &Registry) {
   initializeGlobalOptLegacyPassPass(Registry);
   initializeGlobalSplitPass(Registry);
   initializeHotColdSplittingLegacyPassPass(Registry);
   initializeAlwaysInlinerLegacyPassPass(Registry);
   initializeSimpleInlinerPass(Registry);
   initializeInferFunctionAttrsLegacyPassPass(Registry);
   initializeInternalizeLegacyPassPass(Registry);
   initializeLoopExtractorLegacyPassPass(Registry);
-  initializeBlockExtractorPass(Registry);
+  initializeBlockExtractorLegacyPassPass(Registry);
   initializeSingleLoopExtractorPass(Registry);
   initializeLowerTypeTestsPass(Registry);
   initializeMergeFunctionsLegacyPassPass(Registry);
   initializePartialInlinerLegacyPassPass(Registry);
   initializeAttributorLegacyPassPass(Registry);
   initializeAttributorCGSCCLegacyPassPass(Registry);
   initializePostOrderFunctionAttrsLegacyPassPass(Registry);
   initializeReversePostOrderFunctionAttrsLegacyPassPass(Registry);

llvm/test/Transforms/BlockExtractor/extract-blocks.ll

 ; RUN: echo 'foo bb9' > %t
 ; RUN: echo 'foo bb20' >> %t
 ; RUN: opt -S -extract-blocks -extract-blocks-file=%t %s | FileCheck %s --check-prefix=CHECK-NO-ERASE
 ; RUN: opt -S -extract-blocks -extract-blocks-file=%t -extract-blocks-erase-funcs %s | FileCheck %s --check-prefix=CHECK-ERASE
+; RUN: opt -S -passes=extract-blocks -extract-blocks-file=%t %s | FileCheck %s --check-prefix=CHECK-NO-ERASE
+; RUN: opt -S -passes=extract-blocks -extract-blocks-file=%t -extract-blocks-erase-funcs %s | FileCheck %s --check-prefix=CHECK-ERASE
 
 ; CHECK-NO-ERASE: @foo(
 ; CHECK-NO-ERASE: @foo.bb9(
 ; CHECK-NO-ERASE: @foo.bb20(
 ; CHECK-ERASE: declare i32 @foo(
 ; CHECK-ERASE: @foo.bb9(
 ; CHECK-ERASE: @foo.bb20(
 define i32 @foo(i32 %arg, i32 %arg1) {