[llvm-extract] Support extracting basic blocks

[volkan]

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).

[bogner]

include/llvm/Transforms/IPO.h

182–188 ↗

(On Diff #128351)

The doc comment is either confusing or wrong here, I think. Could you please update it?

lib/Transforms/IPO/BlockExtractor.cpp

76–77 ↗

(On Diff #128351)

Should this use report_fatal_error or something instead of llvm_unreachable?

134–135 ↗

(On Diff #128351)

These should either be fatal errors (see above comment) or asserts. "if (x) unreachable" is a strange pattern.

136–138 ↗

(On Diff #128351)

llvm::find_if would save a few characters, if you prefer.

155 ↗

(On Diff #128351)

I guess CodeExtractor takes care of naming the extracted function, but could you double check that it does the right thing for name collisions? IE, if you extract bb1 from foo, but there's already a function called foo_bb1, what happens?

test/tools/llvm-extract/extract-block.ll

3 ↗

(On Diff #128351)

Worth checking that it got the right block? Could check for "%tmp5 = ", etc.

test/tools/llvm-extract/extract-invalid-block.ll

3 ↗

(On Diff #128351)

typo, s/contains/contain/ (pretty sure this test doesn't pass right now)

tools/bugpoint/ExtractFunction.cpp

410 ↗

(On Diff #128351)

Why this change?

tools/llvm-extract/llvm-extract.cpp

241–244 ↗

(On Diff #128351)

Not sure it's worth it, but I wonder if it would make sense to allow something like "-bb func1:bb2 -bb func2:bb3" to support extracting blocks from multiple functions.

[volkan]

lib/Transforms/IPO/BlockExtractor.cpp

155 ↗

(On Diff #128351)

Yes, it renames the function as foo_bb1.1 in that case.

tools/bugpoint/ExtractFunction.cpp

410 ↗

(On Diff #128351)

I think I renamed the option and reverted, but I forgot to remove this change. I will remove it.

[volkan]

Updated based on the feedback.

[bogner]

tools/llvm-extract/llvm-extract.cpp

321–322 ↗

(On Diff #131120)

On interesting side effect of the "-bb function:block" approach is that we'll probably do something kind of strange if someone calls this like "llvm-extract -func f0 -bb f0:bb1" - will the result be functions f0 and f0_bb1, or will f0 be deleted and we only get f0_bb1? I suspect the latter. This is probably fine though - it seems unlikely that someone would really want to do that.

[volkan]

tools/llvm-extract/llvm-extract.cpp

321–322 ↗

(On Diff #131120)

Yes, it's going to be the latter. BlockExtractor erases the existing functions.

[bogner]

Looks good. If people don't like the behaviour of mixing -func and -bb we can revisit later.

llvm/trunk/include/llvm/InitializePasses.h

 void initializeAlignmentFromAssumptionsPass(PassRegistry&);
 void initializeAlwaysInlinerLegacyPassPass(PassRegistry&);
 void initializeArgPromotionPass(PassRegistry&);
 void initializeAssumptionCacheTrackerPass(PassRegistry&);
 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&);
 void initializeBranchProbabilityInfoWrapperPassPass(PassRegistry&);
 void initializeBranchRelaxationPass(PassRegistry&);
 void initializeBreakCriticalEdgesPass(PassRegistry&);
 void initializeBreakFalseDepsPass(PassRegistry&);
 void initializeCallSiteSplittingLegacyPassPass(PassRegistry&);

llvm/trunk/include/llvm/Transforms/IPO.h

 ///
 Pass *createLoopExtractorPass();
 
 /// createSingleLoopExtractorPass - This pass extracts one natural loop from the
 /// program into a function if it can.  This is used by bugpoint.
 ///
 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.
 ModulePass *createStripDeadPrototypesPass();
 
 //===----------------------------------------------------------------------===//
 /// createReversePostOrderFunctionAttrsPass - This pass walks SCCs of the call
 /// graph in RPO to deduce and propagate function attributes. Currently it

llvm/trunk/lib/Transforms/IPO/BlockExtractor.cpp

+//===- 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/trunk/lib/Transforms/IPO/CMakeLists.txt

 add_llvm_library(LLVMipo
   AlwaysInliner.cpp
   ArgumentPromotion.cpp
   BarrierNoopPass.cpp
+  BlockExtractor.cpp
   CalledValuePropagation.cpp
   ConstantMerge.cpp
   CrossDSOCFI.cpp
   DeadArgumentElimination.cpp
   ElimAvailExtern.cpp
   ExtractGV.cpp
   ForceFunctionAttrs.cpp
   FunctionAttrs.cpp

llvm/trunk/lib/Transforms/IPO/IPO.cpp

void llvm::initializeIPO(PassRegistry &Registry) {
   initializeGlobalOptLegacyPassPass(Registry);
   initializeGlobalSplitPass(Registry);
   initializeIPCPPass(Registry);
   initializeAlwaysInlinerLegacyPassPass(Registry);
   initializeSimpleInlinerPass(Registry);
   initializeInferFunctionAttrsLegacyPassPass(Registry);
   initializeInternalizeLegacyPassPass(Registry);
   initializeLoopExtractorPass(Registry);
-  initializeBlockExtractorPassPass(Registry);
+  initializeBlockExtractorPass(Registry);
   initializeSingleLoopExtractorPass(Registry);
   initializeLowerTypeTestsPass(Registry);
   initializeMergeFunctionsPass(Registry);
   initializePartialInlinerLegacyPassPass(Registry);
   initializePostOrderFunctionAttrsLegacyPassPass(Registry);
   initializeReversePostOrderFunctionAttrsLegacyPassPass(Registry);
   initializePruneEHPass(Registry);
   initializeStripDeadPrototypesLegacyPassPass(Registry);

llvm/trunk/lib/Transforms/IPO/LoopExtractor.cpp

 }
 
 // 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();
 }
-
-
-// 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();
-}

llvm/trunk/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
+
+; CHECK-NO-ERASE: @foo(
+; CHECK-NO-ERASE: @foo_bb9(
+; CHECK-NO-ERASE: @foo_bb20(
+; CHECK-ERASE-NOT: @foo(
+; CHECK-ERASE: @foo_bb9(
+; CHECK-ERASE: @foo_bb20(
+define i32 @foo(i32 %arg, i32 %arg1) {
+bb:
+  %tmp5 = icmp sgt i32 %arg, 0
+  %tmp8 = icmp sgt i32 %arg1, 0
+  %or.cond = and i1 %tmp5, %tmp8
+  br i1 %or.cond, label %bb9, label %bb14
+
+bb9:                                              ; preds = %bb
+  %tmp12 = shl i32 %arg1, 2
+  %tmp13 = add nsw i32 %tmp12, %arg
+  br label %bb30
+
+bb14:                                             ; preds = %bb
+  %0 = and i32 %arg1, %arg
+  %1 = icmp slt i32 %0, 0
+  br i1 %1, label %bb20, label %bb26
+
+bb20:                                             ; preds = %bb14
+  %tmp22 = mul nsw i32 %arg, 3
+  %tmp24 = sdiv i32 %arg1, 6
+  %tmp25 = add nsw i32 %tmp24, %tmp22
+  br label %bb30
+
+bb26:                                             ; preds = %bb14
+  %tmp29 = sub nsw i32 %arg, %arg1
+  br label %bb30
+
+bb30:                                             ; preds = %bb26, %bb20, %bb9
+  %tmp.0 = phi i32 [ %tmp13, %bb9 ], [ %tmp25, %bb20 ], [ %tmp29, %bb26 ]
+  ret i32 %tmp.0
+}
+

llvm/trunk/test/Transforms/BlockExtractor/invalid-block.ll

+; RUN: echo 'bar invalidbb' > %t
+; RUN: not opt -S -extract-blocks -extract-blocks-file=%t %s 2>&1 | FileCheck %s
+
+; CHECK: Invalid block
+define void @bar() {
+bb:
+  ret void
+}
+

llvm/trunk/test/Transforms/BlockExtractor/invalid-function.ll

+; RUN: echo 'foo bb' > %t
+; RUN: not opt -S -extract-blocks -extract-blocks-file=%t %s 2>&1 | FileCheck %s
+
+; CHECK: Invalid function
+define void @bar() {
+bb:
+  ret void
+}
+

llvm/trunk/test/tools/llvm-extract/extract-block.ll

+; RUN: llvm-extract -S -bb foo:bb4 %s | FileCheck %s
+
+; CHECK: @foo_bb4
+; CHECK: %tmp5
+define i32 @foo(i32 %arg) {
+bb:
+  %tmp = alloca i32, align 4
+  %tmp1 = alloca i32, align 4
+  store i32 %arg, i32* %tmp1, align 4
+  %tmp2 = load i32, i32* %tmp1, align 4
+  %tmp3 = icmp sgt i32 %tmp2, 0
+  br i1 %tmp3, label %bb4, label %bb7
+
+bb4:                                              ; preds = %bb
+  %tmp5 = load i32, i32* %tmp1, align 4
+  %tmp6 = add nsw i32 %tmp5, 1
+  store i32 %tmp6, i32* %tmp1, align 4
+  store i32 %tmp6, i32* %tmp, align 4
+  br label %bb8
+
+bb7:                                              ; preds = %bb
+  store i32 0, i32* %tmp, align 4
+  br label %bb8
+
+bb8:                                              ; preds = %bb7, %bb4
+  %tmp9 = load i32, i32* %tmp, align 4
+  ret i32 %tmp9
+}
+

llvm/trunk/test/tools/llvm-extract/extract-invalid-block.ll

+; RUN: not llvm-extract -S -bb foo:invalidbb %s 2>&1 | FileCheck %s
+
+; CHECK: function foo doesn't contain a basic block named 'invalidbb'!
+define i32 @foo(i32 %arg) {
+bb:
+  %tmp = alloca i32, align 4
+  %tmp1 = alloca i32, align 4
+  store i32 %arg, i32* %tmp1, align 4
+  %tmp2 = load i32, i32* %tmp1, align 4
+  %tmp3 = icmp sgt i32 %tmp2, 0
+  br i1 %tmp3, label %bb4, label %bb7
+
+bb4:                                              ; preds = %bb
+  %tmp5 = load i32, i32* %tmp1, align 4
+  %tmp6 = add nsw i32 %tmp5, 1
+  store i32 %tmp6, i32* %tmp1, align 4
+  store i32 %tmp6, i32* %tmp, align 4
+  br label %bb8
+
+bb7:                                              ; preds = %bb
+  store i32 0, i32* %tmp, align 4
+  br label %bb8
+
+bb8:                                              ; preds = %bb7, %bb4
+  %tmp9 = load i32, i32* %tmp, align 4
+  ret i32 %tmp9
+}
+

llvm/trunk/test/tools/llvm-extract/extract-multiple-blocks.ll

+; RUN: llvm-extract -S -bb foo:bb4 -bb foo:bb7 %s | FileCheck %s
+
+; CHECK: @foo_bb4
+; CHECK: @foo_bb7
+define i32 @foo(i32 %arg) {
+bb:
+  %tmp = alloca i32, align 4
+  %tmp1 = alloca i32, align 4
+  store i32 %arg, i32* %tmp1, align 4
+  %tmp2 = load i32, i32* %tmp1, align 4
+  %tmp3 = icmp sgt i32 %tmp2, 0
+  br i1 %tmp3, label %bb4, label %bb7
+
+bb4:                                              ; preds = %bb
+  %tmp5 = load i32, i32* %tmp1, align 4
+  %tmp6 = add nsw i32 %tmp5, 1
+  store i32 %tmp6, i32* %tmp1, align 4
+  store i32 %tmp6, i32* %tmp, align 4
+  br label %bb8
+
+bb7:                                              ; preds = %bb
+  store i32 0, i32* %tmp, align 4
+  br label %bb8
+
+bb8:                                              ; preds = %bb7, %bb4
+  %tmp9 = load i32, i32* %tmp, align 4
+  ret i32 %tmp9
+}
+

llvm/trunk/tools/bugpoint/ExtractFunction.cpp

if (!Temp) {
     outs() << "*** Basic Block extraction failed!\n";
     errs() << "Error creating temporary file: " << toString(Temp.takeError())
            << "\n";
     EmitProgressBitcode(M, "basicblockextractfail", true);
     return nullptr;
   }
   DiscardTemp Discard{*Temp};
 
+  // Extract all of the blocks except the ones in BBs.
+  SmallVector<BasicBlock *, 32> BlocksToExtract;
+  for (Function &F : *M)
+    for (BasicBlock &BB : F)
+      // Check if this block is going to be extracted.
+      if (std::find(BBs.begin(), BBs.end(), &BB) == BBs.end())
+        BlocksToExtract.push_back(&BB);
+
   raw_fd_ostream OS(Temp->FD, /*shouldClose*/ false);
-  for (std::vector<BasicBlock *>::const_iterator I = BBs.begin(), E = BBs.end();
-       I != E; ++I) {
-    BasicBlock *BB = *I;
+  for (BasicBlock *BB : BBs) {
     // If the BB doesn't have a name, give it one so we have something to key
     // off of.
     if (!BB->hasName())
       BB->setName("tmpbb");
     OS << BB->getParent()->getName() << " " << BB->getName() << "\n";
   }
   OS.flush();
   if (OS.has_error()) {

llvm/trunk/tools/llvm-extract/llvm-extract.cpp

 
 // ExtractRegExpFuncs - The functions, matched via regular expression, to
 // extract from the module.
 static cl::list<std::string>
 ExtractRegExpFuncs("rfunc", cl::desc("Specify function(s) to extract using a "
                                      "regular expression"),
                    cl::ZeroOrMore, cl::value_desc("rfunction"));
 
+// ExtractBlocks - The blocks to extract from the module.
+static cl::list<std::string>
+    ExtractBlocks("bb",
+                  cl::desc("Specify <function, basic block> pairs to extract"),
+                  cl::ZeroOrMore, cl::value_desc("function:bb"));
+
 // ExtractAlias - The alias to extract from the module.
 static cl::list<std::string>
 ExtractAliases("alias", cl::desc("Specify alias to extract"),
                cl::ZeroOrMore, cl::value_desc("alias"));
 
 
 // ExtractRegExpAliases - The aliases, matched via regular expression, to
 // extract from the module.
for (size_t i = 0, e = ExtractRegExpFuncs.size(); i != e; ++i) {
     }
     if (!match) {
       errs() << argv[0] << ": program doesn't contain global named '"
              << ExtractRegExpFuncs[i] << "'!\n";
       return 1;
     }
   }
 
+  // Figure out which BasicBlocks we should extract.
+  SmallVector<BasicBlock *, 4> BBs;
+  for (StringRef StrPair : ExtractBlocks) {
+    auto BBInfo = StrPair.split(':');
+    // Get the function.
+    Function *F = M->getFunction(BBInfo.first);
+    if (!F) {
+      errs() << argv[0] << ": program doesn't contain a function named '"
+             << BBInfo.first << "'!\n";
+      return 1;
+    }
+    // Do not materialize this function.
+    GVs.insert(F);
+    // Get the basic block.
+    auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
+      return BB.getName().equals(BBInfo.second);
+    });
+    if (Res == F->end()) {
+      errs() << argv[0] << ": function " << F->getName()
+             << " doesn't contain a basic block named '" << BBInfo.second
+             << "'!\n";
+      return 1;
+    }
+    BBs.push_back(&*Res);
+  }
+
   // Use *argv instead of argv[0] to work around a wrong GCC warning.
   ExitOnError ExitOnErr(std::string(*argv) + ": error reading input: ");
 
   if (Recursive) {
     std::vector<llvm::Function *> Workqueue;
     for (GlobalValue *GV : GVs) {
       if (auto *F = dyn_cast<Function>(GV)) {
         Workqueue.push_back(F);
if (!DeleteFn) {
     Extract.run(*M);
 
     // Now that we have all the GVs we want, mark the module as fully
     // materialized.
     // FIXME: should the GVExtractionPass handle this?
     ExitOnErr(M->materializeAll());
   }
 
+  // Extract the specified basic blocks from the module and erase the existing
+  // functions.
+  if (!ExtractBlocks.empty()) {
+    legacy::PassManager PM;
+    PM.add(createBlockExtractorPass(BBs, true));
+    PM.run(*M);
+  }
+
   // In addition to deleting all other functions, we also want to spiff it
   // up a little bit.  Do this now.
   legacy::PassManager Passes;
 
   if (!DeleteFn)
     Passes.add(createGlobalDCEPass());           // Delete unreachable globals
   Passes.add(createStripDeadDebugInfoPass());    // Remove dead debug info
   Passes.add(createStripDeadPrototypesPass());   // Remove dead func decls