[llvm-extract] Support extracting basic blocks

Summary:
Currently, there is no way to extract a basic block from a function easily. This patch
extends llvm-extract to extract the specified basic block(s).

Reviewers: loladiro, rafael, bogner

Reviewed By: bogner

Subscribers: hintonda, mgorny, qcolombet, llvm-commits

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

llvm-svn: 323266

Author: Volkan Keles

llvm/include/llvm/InitializePasses.h

@@ -73,7 +73,7 @@ void initializeAtomicExpandPass(PassRegistry&);
 void initializeBDCELegacyPassPass(PassRegistry&);
 void initializeBarrierNoopPass(PassRegistry&);
 void initializeBasicAAWrapperPassPass(PassRegistry&);
-void initializeBlockExtractorPassPass(PassRegistry&);
+void initializeBlockExtractorPass(PassRegistry &);
 void initializeBlockFrequencyInfoWrapperPassPass(PassRegistry&);
 void initializeBoundsCheckingLegacyPassPass(PassRegistry&);
 void initializeBranchFolderPassPass(PassRegistry&);

llvm/include/llvm/Transforms/IPO.h

@@ -179,10 +179,13 @@ Pass *createLoopExtractorPass();
 ///
 Pass *createSingleLoopExtractorPass();
 
-/// createBlockExtractorPass - This pass extracts all blocks (except those
-/// specified in the argument list) from the functions in the module.
+/// createBlockExtractorPass - This pass extracts all the specified blocks
+/// from the functions in the module.
 ///
 ModulePass *createBlockExtractorPass();
+ModulePass *
+createBlockExtractorPass(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
+                         bool EraseFunctions);
 
 /// createStripDeadPrototypesPass - This pass removes any function declarations
 /// (prototypes) that are not used.

llvm/lib/Transforms/IPO/BlockExtractor.cpp

@@ -0,0 +1,174 @@
+//===- BlockExtractor.cpp - Extracts blocks into their own functions ------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass extracts the specified basic blocks from the module into their
+// own functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.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<BasicBlock *, 16> Blocks;
+  bool EraseFunctions;
+  SmallVector<std::pair<std::string, std::string>, 32> BlocksByName;
+
+public:
+  static char ID;
+  BlockExtractor(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
+                 bool EraseFunctions)
+      : ModulePass(ID), Blocks(BlocksToExtract.begin(), BlocksToExtract.end()),
+        EraseFunctions(EraseFunctions) {
+    if (!BlockExtractorFile.empty())
+      loadFile();
+  }
+  BlockExtractor() : BlockExtractor(SmallVector<BasicBlock *, 0>(), false) {}
+  bool runOnModule(Module &M) override;
+
+private:
+  void loadFile();
+  void splitLandingPadPreds(Function &F);
+};
+} // end anonymous namespace
+
+char BlockExtractor::ID = 0;
+INITIALIZE_PASS(BlockExtractor, "extract-blocks",
+                "Extract basic blocks from module", false, false)
+
+ModulePass *llvm::createBlockExtractorPass() { return new BlockExtractor(); }
+ModulePass *llvm::createBlockExtractorPass(
+    const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
+  return new BlockExtractor(BlocksToExtract, EraseFunctions);
+}
+
+/// Gets all of the blocks specified in the input file.
+void BlockExtractor::loadFile() {
+  auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
+  if (std::error_code EC = ErrOrBuf.getError())
+    report_fatal_error("BlockExtractor couldn't load the file.");
+  // Read the file.
+  auto &Buf = *ErrOrBuf;
+  SmallVector<StringRef, 16> Lines;
+  Buf->getBuffer().split(Lines, '\n', /*MaxSplit=*/-1,
+                         /*KeepEmpty=*/false);
+  for (const auto &Line : Lines) {
+    auto FBPair = Line.split(' ');
+    BlocksByName.push_back({FBPair.first, FBPair.second});
+  }
+}
+
+/// Extracts the landing pads to make sure all of them have only one
+/// predecessor.
+void BlockExtractor::splitLandingPadPreds(Function &F) {
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      if (!isa<InvokeInst>(&I))
+        continue;
+      InvokeInst *II = cast<InvokeInst>(&I);
+      BasicBlock *Parent = II->getParent();
+      BasicBlock *LPad = II->getUnwindDest();
+
+      // Look through the landing pad's predecessors. If one of them ends in an
+      // 'invoke', then we want to split the landing pad.
+      bool Split = false;
+      for (auto PredBB : predecessors(LPad)) {
+        if (PredBB->isLandingPad() && PredBB != Parent &&
+            isa<InvokeInst>(Parent->getTerminator())) {
+          Split = true;
+          break;
+        }
+      }
+
+      if (!Split)
+        continue;
+
+      SmallVector<BasicBlock *, 2> NewBBs;
+      SplitLandingPadPredecessors(LPad, Parent, ".1", ".2", NewBBs);
+    }
+  }
+}
+
+bool BlockExtractor::runOnModule(Module &M) {
+
+  bool Changed = false;
+
+  // Get all the functions.
+  SmallVector<Function *, 4> Functions;
+  for (Function &F : M) {
+    splitLandingPadPreds(F);
+    Functions.push_back(&F);
+  }
+
+  // Get all the blocks specified in the input file.
+  for (const auto &BInfo : BlocksByName) {
+    Function *F = M.getFunction(BInfo.first);
+    if (!F)
+      report_fatal_error("Invalid function name specified in the input file");
+    auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
+      return BB.getName().equals(BInfo.second);
+    });
+    if (Res == F->end())
+      report_fatal_error("Invalid block name specified in the input file");
+    Blocks.push_back(&*Res);
+  }
+
+  // Extract basic blocks.
+  for (BasicBlock *BB : Blocks) {
+    // Check if the module contains BB.
+    if (BB->getParent()->getParent() != &M)
+      report_fatal_error("Invalid basic block");
+    DEBUG(dbgs() << "BlockExtractor: Extracting " << BB->getParent()->getName()
+                 << ":" << BB->getName() << "\n");
+    SmallVector<BasicBlock *, 2> BlocksToExtractVec;
+    BlocksToExtractVec.push_back(BB);
+    if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
+      BlocksToExtractVec.push_back(II->getUnwindDest());
+    CodeExtractor(BlocksToExtractVec).extractCodeRegion();
+    ++NumExtracted;
+    Changed = true;
+  }
+
+  // Erase the functions.
+  if (EraseFunctions || BlockExtractorEraseFuncs) {
+    for (Function *F : Functions) {
+      DEBUG(dbgs() << "BlockExtractor: Deleting " << F->getName() << "\n");
+      F->eraseFromParent();
+    }
+    // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
+    for (Function &F : M)
+      F.setLinkage(GlobalValue::ExternalLinkage);
+    Changed = true;
+  }
+
+  return Changed;
+}

llvm/lib/Transforms/IPO/CMakeLists.txt

@@ -2,6 +2,7 @@ add_llvm_library(LLVMipo
   AlwaysInliner.cpp
   ArgumentPromotion.cpp
   BarrierNoopPass.cpp
+  BlockExtractor.cpp
   CalledValuePropagation.cpp
   ConstantMerge.cpp
   CrossDSOCFI.cpp

llvm/lib/Transforms/IPO/IPO.cpp

@@ -40,7 +40,7 @@ void llvm::initializeIPO(PassRegistry &Registry) {
   initializeInferFunctionAttrsLegacyPassPass(Registry);
   initializeInternalizeLegacyPassPass(Registry);
   initializeLoopExtractorPass(Registry);
-  initializeBlockExtractorPassPass(Registry);
+  initializeBlockExtractorPass(Registry);
   initializeSingleLoopExtractorPass(Registry);
   initializeLowerTypeTestsPass(Registry);
   initializeMergeFunctionsPass(Registry);

llvm/lib/Transforms/IPO/LoopExtractor.cpp

@@ -158,155 +158,3 @@ bool LoopExtractor::runOnLoop(Loop *L, LPPassManager &LPM) {
 Pass *llvm::createSingleLoopExtractorPass() {
   return new SingleLoopExtractor();
 }
-
-
-// BlockFile - A file which contains a list of blocks that should not be
-// extracted.
-static cl::opt<std::string>
-BlockFile("extract-blocks-file", cl::value_desc("filename"),
-          cl::desc("A file containing list of basic blocks to not extract"),
-          cl::Hidden);
-
-namespace {
-  /// BlockExtractorPass - This pass is used by bugpoint to extract all blocks
-  /// from the module into their own functions except for those specified by the
-  /// BlocksToNotExtract list.
-  class BlockExtractorPass : public ModulePass {
-    void LoadFile(const char *Filename);
-    void SplitLandingPadPreds(Function *F);
-
-    std::vector<BasicBlock*> BlocksToNotExtract;
-    std::vector<std::pair<std::string, std::string> > BlocksToNotExtractByName;
-  public:
-    static char ID; // Pass identification, replacement for typeid
-    BlockExtractorPass() : ModulePass(ID) {
-      if (!BlockFile.empty())
-        LoadFile(BlockFile.c_str());
-    }
-
-    bool runOnModule(Module &M) override;
-  };
-}
-
-char BlockExtractorPass::ID = 0;
-INITIALIZE_PASS(BlockExtractorPass, "extract-blocks",
-                "Extract Basic Blocks From Module (for bugpoint use)",
-                false, false)
-
-// createBlockExtractorPass - This pass extracts all blocks (except those
-// specified in the argument list) from the functions in the module.
-//
-ModulePass *llvm::createBlockExtractorPass() {
-  return new BlockExtractorPass();
-}
-
-void BlockExtractorPass::LoadFile(const char *Filename) {
-  // Load the BlockFile...
-  std::ifstream In(Filename);
-  if (!In.good()) {
-    errs() << "WARNING: BlockExtractor couldn't load file '" << Filename
-           << "'!\n";
-    return;
-  }
-  while (In) {
-    std::string FunctionName, BlockName;
-    In >> FunctionName;
-    In >> BlockName;
-    if (!BlockName.empty())
-      BlocksToNotExtractByName.push_back(
-          std::make_pair(FunctionName, BlockName));
-  }
-}
-
-/// SplitLandingPadPreds - The landing pad needs to be extracted with the invoke
-/// instruction. The critical edge breaker will refuse to break critical edges
-/// to a landing pad. So do them here. After this method runs, all landing pads
-/// should have only one predecessor.
-void BlockExtractorPass::SplitLandingPadPreds(Function *F) {
-  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
-    InvokeInst *II = dyn_cast<InvokeInst>(I);
-    if (!II) continue;
-    BasicBlock *Parent = II->getParent();
-    BasicBlock *LPad = II->getUnwindDest();
-
-    // Look through the landing pad's predecessors. If one of them ends in an
-    // 'invoke', then we want to split the landing pad.
-    bool Split = false;
-    for (pred_iterator
-           PI = pred_begin(LPad), PE = pred_end(LPad); PI != PE; ++PI) {
-      BasicBlock *BB = *PI;
-      if (BB->isLandingPad() && BB != Parent &&
-          isa<InvokeInst>(Parent->getTerminator())) {
-        Split = true;
-        break;
-      }
-    }
-
-    if (!Split) continue;
-
-    SmallVector<BasicBlock*, 2> NewBBs;
-    SplitLandingPadPredecessors(LPad, Parent, ".1", ".2", NewBBs);
-  }
-}
-
-bool BlockExtractorPass::runOnModule(Module &M) {
-  if (skipModule(M))
-    return false;
-
-  std::set<BasicBlock*> TranslatedBlocksToNotExtract;
-  for (unsigned i = 0, e = BlocksToNotExtract.size(); i != e; ++i) {
-    BasicBlock *BB = BlocksToNotExtract[i];
-    Function *F = BB->getParent();
-
-    // Map the corresponding function in this module.
-    Function *MF = M.getFunction(F->getName());
-    assert(MF->getFunctionType() == F->getFunctionType() && "Wrong function?");
-
-    // Figure out which index the basic block is in its function.
-    Function::iterator BBI = MF->begin();
-    std::advance(BBI, std::distance(F->begin(), Function::iterator(BB)));
-    TranslatedBlocksToNotExtract.insert(&*BBI);
-  }
-
-  while (!BlocksToNotExtractByName.empty()) {
-    // There's no way to find BBs by name without looking at every BB inside
-    // every Function. Fortunately, this is always empty except when used by
-    // bugpoint in which case correctness is more important than performance.
-
-    std::string &FuncName  = BlocksToNotExtractByName.back().first;
-    std::string &BlockName = BlocksToNotExtractByName.back().second;
-
-    for (Function &F : M) {
-      if (F.getName() != FuncName) continue;
-
-      for (BasicBlock &BB : F) {
-        if (BB.getName() != BlockName) continue;
-
-        TranslatedBlocksToNotExtract.insert(&BB);
-      }
-    }
-
-    BlocksToNotExtractByName.pop_back();
-  }
-
-  // Now that we know which blocks to not extract, figure out which ones we WANT
-  // to extract.
-  std::vector<BasicBlock*> BlocksToExtract;
-  for (Function &F : M) {
-    SplitLandingPadPreds(&F);
-    for (BasicBlock &BB : F)
-      if (!TranslatedBlocksToNotExtract.count(&BB))
-        BlocksToExtract.push_back(&BB);
-  }
-
-  for (BasicBlock *BlockToExtract : BlocksToExtract) {
-    SmallVector<BasicBlock*, 2> BlocksToExtractVec;
-    BlocksToExtractVec.push_back(BlockToExtract);
-    if (const InvokeInst *II =
-            dyn_cast<InvokeInst>(BlockToExtract->getTerminator()))
-      BlocksToExtractVec.push_back(II->getUnwindDest());
-    CodeExtractor(BlocksToExtractVec).extractCodeRegion();
-  }
-
-  return !BlocksToExtract.empty();
-}