[CodeExtractor] Optionally keep code in original function.

[Meinersbur]

Summary

When extracting a region into a new function, optionally allow cloning basic blocks and instructions into the extracted function instead of moving them. The keeps the original code in the original function such that they can still be referenced -- and branched to -- in the original function.

The motivation is the use of CodeExtractor in the OpenMPIRBuilder. The implementation of createParallel first emits the parallel region into the lexical function, then uses CodeExtractor to outline that region into a new function. The problem here is that Clang's code generator will references some basic blocks for code inside as well as outside the region. This includes some special purpose block (EHResumeBlock, TerminateLandingPad, TerminateHandler, UnreachableBlock, ...) and cleanup/dtor code that is re-used from multiple scopes (see test case extract-block-cleanup.ll). Moving these blocks into a different function will result in malformed IR. The KeepOldBlocks option will instead clone the outlined code into a new function keeping the auxiliary code intact, relying on later DCE to remove code that indeed has become unreachable. Additionally, this code could also be uses as a fallback when threading/offloading is disabled via environment option.

Use of KeepOldBlocks by OpenMPIRBuilder is not part of this patch. For testing, we extend llvm-extract allowing the use of this option and thus making it more powerful.

[jdoerfert]

While this functionality is probably useful by itself, I am confused we need it in the OpenMP-IR-Builder.

In my mind that use case looked something like:
When createParallel starts we push new special blocks (EHResumeBlock, ...) onto the stack.
All cleanup uses the associated special blocks while body generation happens.
Once the region is done we revert back to the special blocks that were used before.

Why doesn't this work? I thought it already did.

[Meinersbur]

In D115216#3281303, @jdoerfert wrote:

In my mind that use case looked something like:
When createParallel starts we push new special blocks (EHResumeBlock, ...) onto the stack.
All cleanup uses the associated special blocks while body generation happens.
Once the region is done we revert back to the special blocks that were used before.

Why doesn't this work? I thought it already did.

Unfortunately is it a lot more complicated. There is no limited list of blocks that might be shared between outlined and outer function, the special blocks in CodeGenFunction (EHResumeBlock, TerminateLandingPad, TerminateHandler, UnreachableBlock) are just the most obvious ones. For instance, cleanup code is emitted only once and looks like this:

bb1:
  store i32 1, %c
  br label %common_cleanup

bb2:
  store i32 2, %c
  br label %common_cleanup

common_cleanup:
  [cleanup code]
  switch i32 %c, [
    [i32 1, label %bb1_continue],
    [i32 2, label %bb2_continue]
  }

bb1_continue:
  [...]

bb2_continue:
  [...]

And for every destruction a new switch case is added. Similar things happen with exceptions e.g. landing pads. What's worst, I have no exhaustive list of circumstances when sharing of code may happen and there is also the risk that the switched-out special block are stored in some other data structure and re-used after having left the region. I am not saying these cannot be fixed by switching-out some fields during the outlined region (there's also TerminateFunclets, ExceptionSlot, EHSelectorSlot, SEHCodeSlotStack, SEHInfo, ObjCEHValueStack, CurFuncIsThunk, ...), up to the entire CodeGenFunction but I vastly prefer a straightforward solution over making Clang's internals even more complicated. It's also one less concern for potential other front-ends than Clang/Flang.

That being said, a clean exception handling would indeed create a new handling context; currently exception handling inside the outlined region will also branch to the cleanup code in the outer function using the dtor pattern above. Actually triggering the exception at runtime is not well-formed in OpenMP, but the compiler should not crash crash because the cleanup code has been moved away. That is, copying the code instead of moving it will already allow us to compile C++ code without -fno-exceptions.

[Meinersbur]

• Rebase

[Meinersbur]

ping

[kiranchandramohan]

Thanks @Meinersbur for this patch. Apologies for the long wait. Just starting to go over it.

The summary is quite helpful to understand the need for this extension to the extractor from the Clang side. I guess this is currently not required for Flang/MLIR? Or is it needed to handle the cancellation constructs?

I think the KeepFunctions and KeepBlocks flags are useful to ensure that there are no changes to existing users. But I would also like to check with you whether keeping the functions and blocks are necessary for correctness when there are branches to the blocks other than the first one. If so would it be better to support it without flags? Also, do we have asserts presently to ensure that incorrect code is not generated in the presence of these branches when the flags are not specified?

llvm/include/llvm/Transforms/Utils/Cloning.h

117

Nit: Not about this patch.
DebugInfoFinder is missing from these comments.

120

Nit: There seems to be a style here for mentioning the parameter as the nth (7th here).

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

87

Nit: Unlsess -> Unless

[Meinersbur]

• Address review

[Meinersbur]

In D115216#3441487, @kiranchandramohan wrote:

The summary is quite helpful to understand the need for this extension to the extractor from the Clang side. I guess this is currently not required for Flang/MLIR? Or is it needed to handle the cancellation constructs?

As far as I know it it currently not necessary the OpenMP dialect lowing. Neither does MLIR currently support exceptions and even fir-dev does not support final procedures yet. However, support basic block sharing for code size reduction or other purposes might be required in the future.

But I would also like to check with you whether keeping the functions and blocks are necessary for correctness when there are branches to the blocks other than the first one. If so would it be better to support it without flags? Also, do we have asserts presently to ensure that incorrect code is not generated in the presence of these branches when the flags are not specified?

This patch only introduces using KeepOldBlocks for llvm-extract. I did not want to change the default behaviour for the tool, which is throwing away the parent function. Hence, if there are any branches to moved block are left, those branch instructions are thrown away with its parent function. The "create a call to extracted function" functionality offered by CodeExtractor is unused. Using --bb-keep-functions without --bb-keep-blocks may result in invalid IR if the --bb argument does not specify a proper region. This will be caught by the IR verifier in assert builds at CodeExtractor.cpp:1445. Are you suggesting to add the verifier for non-assert builds as well?

OpenMPIRBuilder will always use KeepOldBlocks mode.

llvm/include/llvm/Transforms/Utils/Cloning.h

117

CloneBasicBlock with InstSelect parameter is used in this patch.

The parameter DIFinder was added in D33655 and unrelated to this patch.

[Meinersbur]

Introduce --replace-with call to replace --bb-keep-functions and --bb-keep-blocks.
Makes the interface easier and --bb-keep-functions alone can result in invalid IR.

llvm/include/llvm/Transforms/IPO.h

 Pass *createSingleLoopExtractorPass();
 
 /// createBlockExtractorPass - This pass extracts all the specified blocks
 /// from the functions in the module.
 ///
 ModulePass *createBlockExtractorPass();
 ModulePass *
 createBlockExtractorPass(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
-                         bool EraseFunctions);
+                         bool EraseFunctions, bool KeepOldBlocks = false);
 ModulePass *
 createBlockExtractorPass(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
                              &GroupsOfBlocksToExtract,
-                         bool EraseFunctions);
+                         bool EraseFunctions, bool KeepOldBlocks = false);
 
 /// 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/include/llvm/Transforms/Utils/Cloning.h

 /// names, specify it as the optional third parameter.
 ///
 /// If you would like the basic block to be auto-inserted into the end of a
 /// function, you can specify it as the optional fourth parameter.
 ///
 /// If you would like to collect additional information about the cloned
 /// function, you can specify a ClonedCodeInfo object with the optional fifth
 /// parameter.
-BasicBlock *CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
-                            const Twine &NameSuffix = "", Function *F = nullptr,
-                            ClonedCodeInfo *CodeInfo = nullptr,
-                            DebugInfoFinder *DIFinder = nullptr);
+///
+/// If you would like to clone only a subset of instructions in the basic block,
+/// you can specify a callback returning true only for those instructions that
+/// are to be cloned.
+BasicBlock *
+CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
+                DebugInfoFinder *DIFinder = nullptr,
+                function_ref<bool(const Instruction *)> InstSelect = {});
 
 /// Return a copy of the specified function and add it to that
 /// function's module.  Also, any references specified in the VMap are changed
 /// to refer to their mapped value instead of the original one.  If any of the
 /// arguments to the function are in the VMap, the arguments are deleted from
 /// the resultant function.  The VMap is updated to include mappings from all of
 /// the instructions and basicblocks in the function from their old to new
 /// values.  The final argument captures information about the cloned code if

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

class CodeExtractor {
     const bool AggregateArgs;
     BlockFrequencyInfo *BFI;
     BranchProbabilityInfo *BPI;
     AssumptionCache *AC;
 
     // If true, varargs functions can be extracted.
     bool AllowVarArgs;
 
+    /// If true, copies the code into the extracted function instead of moving
+    /// it.
+    bool KeepOldBlocks;
+
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
     unsigned NumExitBlocks = std::numeric_limits<unsigned>::max();
     Type *RetTy;
 
     /// Lists of blocks that are branched from the code region to be extracted.
     /// Each block is contained at most once. Its order defines the return value
     /// of the extracted function, when leaving the extracted function via the
public:
     /// Given a sequence of basic blocks where the first block in the sequence
     /// dominates the rest, prepare a code extractor object for pulling this
     /// sequence out into its new function. When a DominatorTree is also given,
     /// extra checking and transformations are enabled. If AllowVarArgs is true,
     /// vararg functions can be extracted. This is safe, if all vararg handling
     /// code is extracted, including vastart. If AllowAlloca is true, then
     /// extraction of blocks containing alloca instructions would be possible,
     /// however code extractor won't validate whether extraction is legal.
+    ///
+    /// If KeepOldBlocks is true, the original instances of the extracted region
+    /// remains in the original function so they can still be branched to from
+    /// non-extracted blocks. However, only branches to the first block will
+    /// call the extracted function.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr,
-                  AssumptionCache *AC = nullptr,
-                  bool AllowVarArgs = false, bool AllowAlloca = false,
-                  std::string Suffix = "");
+                  AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
+                  bool AllowAlloca = false, std::string Suffix = "",
+                  bool KeepOldBlocks = false);
 
     /// Create a code extractor for a loop body.
     ///
     /// Behaves just like the generic code sequence constructor, but uses the
     /// block sequence of the loop.
     CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
                   BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr,

llvm/lib/Transforms/IPO/BlockExtractor.cpp

 
 static cl::opt<bool>
     BlockExtractorEraseFuncs("extract-blocks-erase-funcs",
                              cl::desc("Erase the existing functions"),
                              cl::Hidden);
 namespace {
 class BlockExtractor {
 public:
-  BlockExtractor(bool EraseFunctions) : EraseFunctions(EraseFunctions) {}
+  BlockExtractor(bool EraseFunctions, bool KeepOldBlocks = false)
+      : EraseFunctions(EraseFunctions), KeepOldBlocks(KeepOldBlocks) {}
   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;
+  bool KeepOldBlocks;
   /// 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;
   BlockExtractorLegacyPass(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
-                           bool EraseFunctions)
-      : ModulePass(ID), BE(EraseFunctions) {
+                           bool EraseFunctions, bool KeepOldBlocks)
+      : ModulePass(ID), BE(EraseFunctions, KeepOldBlocks) {
     // 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);
     }
     BE.init(MassagedGroupsOfBlocks);
   }
 
   BlockExtractorLegacyPass(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
                                &GroupsOfBlocksToExtract,
-                           bool EraseFunctions)
-      : ModulePass(ID), BE(EraseFunctions) {
+                           bool EraseFunctions, bool KeepOldBlocks)
+      : ModulePass(ID), BE(EraseFunctions, KeepOldBlocks) {
     BE.init(GroupsOfBlocksToExtract);
   }
 
   BlockExtractorLegacyPass()
-      : BlockExtractorLegacyPass(SmallVector<BasicBlock *, 0>(), false) {}
+      : BlockExtractorLegacyPass(SmallVector<BasicBlock *, 0>(), false, false) {
+  }
 };
 
 } // end anonymous namespace
 
 char BlockExtractorLegacyPass::ID = 0;
 INITIALIZE_PASS(BlockExtractorLegacyPass, "extract-blocks",
                 "Extract basic blocks from module", false, false)
 
 ModulePass *llvm::createBlockExtractorPass() {
   return new BlockExtractorLegacyPass();
 }
 ModulePass *llvm::createBlockExtractorPass(
-    const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
-  return new BlockExtractorLegacyPass(BlocksToExtract, EraseFunctions);
+    const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions,
+    bool KeepOldBlocks) {
+  return new BlockExtractorLegacyPass(BlocksToExtract, EraseFunctions,
+                                      KeepOldBlocks);
 }
 ModulePass *llvm::createBlockExtractorPass(
     const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
         &GroupsOfBlocksToExtract,
-    bool EraseFunctions) {
-  return new BlockExtractorLegacyPass(GroupsOfBlocksToExtract, EraseFunctions);
+    bool EraseFunctions, bool KeepOldBlocks) {
+  return new BlockExtractorLegacyPass(GroupsOfBlocksToExtract, EraseFunctions,
+                                      KeepOldBlocks);
 }
 
 /// 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.
for (BasicBlock *BB : BBs) {
                         << "\n");
       BlocksToExtractVec.push_back(BB);
       if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
         BlocksToExtractVec.push_back(II->getUnwindDest());
       ++NumExtracted;
       Changed = true;
     }
     CodeExtractorAnalysisCache CEAC(*BBs[0]->getParent());
-    Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion(CEAC);
+    Function *F =
+        CodeExtractor(BlocksToExtractVec, /* DT */ nullptr,
+                      /* AggregateArgs*/ false, /* BFI */ nullptr,
+                      /* BPI */ nullptr, /* AC */ nullptr,
+                      /* AllowVarArgs */ false, /* AllowAlloca */ false,
+                      /* Suffix */ "", KeepOldBlocks)
+            .extractCodeRegion(CEAC);
     if (F)
       LLVM_DEBUG(dbgs() << "Extracted group '" << (*BBs.begin())->getName()
                         << "' in: " << F->getName() << '\n');
     else
       LLVM_DEBUG(dbgs() << "Failed to extract for group '"
                         << (*BBs.begin())->getName() << "'\n");
   }
 

llvm/lib/Transforms/Utils/CloneFunction.cpp

 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <map>
 using namespace llvm;
 
 #define DEBUG_TYPE "clone-function"
 
 /// See comments in Cloning.h.
-BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
-                                  const Twine &NameSuffix, Function *F,
-                                  ClonedCodeInfo *CodeInfo,
-                                  DebugInfoFinder *DIFinder) {
+BasicBlock *
+llvm::CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
+                      ClonedCodeInfo *CodeInfo, DebugInfoFinder *DIFinder,
+                      function_ref<bool(const Instruction *)> InstSelect) {
   BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "", F);
   if (BB->hasName())
     NewBB->setName(BB->getName() + NameSuffix);
 
   bool hasCalls = false, hasDynamicAllocas = false;
   Module *TheModule = F ? F->getParent() : nullptr;
 
   // Loop over all instructions, and copy them over.
   for (const Instruction &I : *BB) {
+    if (InstSelect && !InstSelect(&I))
+      continue;
+
     if (DIFinder && TheModule)
       DIFinder->processInstruction(*TheModule, I);
 
     Instruction *NewInst = I.clone();
     if (I.hasName())
       NewInst->setName(I.getName() + NameSuffix);
     NewBB->getInstList().push_back(NewInst);
     VMap[&I] = NewInst; // Add instruction map to value.

llvm/lib/Transforms/Utils/CodeExtractor.cpp

 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/SSAUpdater.h"
 #include <cassert>
 #include <cstdint>
 #include <iterator>
 #include <map>
 #include <set>
 #include <utility>
 #include <vector>
 
static bool isBlockValidForExtraction(const BasicBlock &BB,
   }
 
   return true;
 }
 
 /// Build a set of blocks to extract if the input blocks are viable.
 static SetVector<BasicBlock *>
 buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
-                        bool AllowVarArgs, bool AllowAlloca) {
+                        bool AllowVarArgs, bool AllowAlloca,
+                        bool KeepOldBlocks) {
   assert(!BBs.empty() && "The set of blocks to extract must be non-empty");
   SetVector<BasicBlock *> Result;
 
   // Loop over the blocks, adding them to our set-vector, and aborting with an
   // empty set if we encounter invalid blocks.
   for (BasicBlock *BB : BBs) {
     // If this block is dead, don't process it.
     if (DT && !DT->isReachableFromEntry(BB))
if (BB == Result.front()) {
       if (BB->isEHPad()) {
         LLVM_DEBUG(dbgs() << "The first block cannot be an unwind block\n");
         return {};
       }
       continue;
     }
 
     // All blocks other than the first must not have predecessors outside of
-    // the subgraph which is being extracted.
+    // the subgraph which is being extracted. KeepOldBlocks relaxes this
+    // requirement.
+    if (!KeepOldBlocks) {
       for (auto *PBB : predecessors(BB))
         if (!Result.count(PBB)) {
-        LLVM_DEBUG(dbgs() << "No blocks in this region may have entries from "
-                             "outside the region except for the first block!\n"
-                          << "Problematic source BB: " << BB->getName() << "\n"
-                          << "Problematic destination BB: " << PBB->getName()
-                          << "\n");
+          LLVM_DEBUG(dbgs()
+                     << "No blocks in this region may have entries from "
           return {};
         }
     }
+  }
 
   return Result;
 }
 
 CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
                              bool AggregateArgs, BlockFrequencyInfo *BFI,
                              BranchProbabilityInfo *BPI, AssumptionCache *AC,
                              bool AllowVarArgs, bool AllowAlloca,
-                             std::string Suffix)
+                             std::string Suffix, bool KeepOldBlocks)
     : DT(DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
       BPI(BPI), AC(AC), AllowVarArgs(AllowVarArgs),
-      Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs, AllowAlloca)),
+      KeepOldBlocks(KeepOldBlocks),
+      Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs, AllowAlloca,
+                                     KeepOldBlocks)),
       Suffix(Suffix) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs,
                              BlockFrequencyInfo *BFI,
                              BranchProbabilityInfo *BPI, AssumptionCache *AC,
                              std::string Suffix)
     : DT(&DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
-      BPI(BPI), AC(AC), AllowVarArgs(false),
+      BPI(BPI), AC(AC), AllowVarArgs(false), KeepOldBlocks(false),
       Blocks(buildExtractionBlockSet(L.getBlocks(), &DT,
                                      /* AllowVarArgs */ false,
-                                     /* AllowAlloca */ false)),
+                                     /* AllowAlloca */ false,
+                                     /* KeepOldBlocks */ false)),
       Suffix(Suffix) {}
 
 /// definedInRegion - Return true if the specified value is defined in the
 /// extracted region.
 static bool definedInRegion(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (Blocks.count(I->getParent()))
       return true;

 }
 
 void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
                                       const ValueSet &SinkCands) const {
   for (BasicBlock *BB : Blocks) {
     // If a used value is defined outside the region, it's an input.  If an
     // instruction is used outside the region, it's an output.
     for (Instruction &II : *BB) {
+      // Ignore assumptions if not been removed yet.
+      if (isa<AssumeInst>(II))
+        continue;
+
       for (auto &OI : II.operands()) {
         Value *V = OI;
         if (!SinkCands.count(V) && definedInCaller(Blocks, V))
           Inputs.insert(V);
       }
 
       for (User *U : II.users())
         if (!definedInRegion(Blocks, U)) {
CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
   Function *oldFunction = header->getParent();
 
   normalizeCFGForExtraction(header);
 
+  if (!KeepOldBlocks) {
     // Transforms/HotColdSplit/stale-assume-in-original-func.ll
     // Remove @llvm.assume calls that will be moved to the new function from the
     // old function's assumption cache.
     for (BasicBlock *Block : Blocks) {
       for (Instruction &I : llvm::make_early_inc_range(*Block)) {
         if (auto *AI = dyn_cast<AssumeInst>(&I)) {
           if (AC)
             AC->unregisterAssumption(AI);
           AI->eraseFromParent();
         }
       }
     }
+  }
 
   ValueSet SinkingCands, HoistingCands;
   BasicBlock *CommonExit = nullptr;
   findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
   assert(HoistingCands.empty() || CommonExit);
 
   // analysis, after ret splitting (for values returned)
   // Find inputs to, outputs from the code region.
for (BasicBlock *Succ : SwitchCases) {
         BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, Succ);
       }
     }
   }
 
   // Determine position for the replacement code. Do so before header is moved
   // to the new function.
   BasicBlock *ReplIP = header;
+  if (!KeepOldBlocks) {
     while (ReplIP && Blocks.count(ReplIP))
       ReplIP = ReplIP->getNextNode();
+  }
 
   // Construct new function based on inputs/outputs & add allocas for all defs.
   std::string SuffixToUse =
       Suffix.empty()
           ? (header->getName().empty() ? "extracted" : header->getName().str())
           : Suffix;
   Function *newFunction = constructFunctionDeclaration(
       inputs, outputs, EntryFreq, oldFunction->getName() + "." + SuffixToUse);
void CodeExtractor::normalizeCFGForExtraction(BasicBlock *&Header) {
   severSplitPHINodesOfEntry(Header);
 
   // If a PHI in an exit block has multiple invoming values from the outlined
   // region, create a new PHI for those values within the region such that only
   // PHI itself becomes an output value, not each of its incoming values
   // individually.
   recomputeExitBlocks();
   severSplitPHINodesOfExits();
+
+  // If the option was given, ensure there are no PHI nodes at all in the exit
+  // nodes themselves.
+  if (KeepOldBlocks) {
+    for (BasicBlock *Block : Blocks) {
+      for (BasicBlock *Succ : make_early_inc_range(successors(Block))) {
+        if (Blocks.count(Succ))
+          continue;
+
+        if (!Succ->getSinglePredecessor())
+          Succ = SplitEdge(Block, Succ, DT);
+
+        // Ensure no PHI node in exit block (still possible with single
+        // predecessor, e.g. LCSSA)
+        while (auto *P = dyn_cast<PHINode>(&Succ->front())) {
+          assert(P->getNumIncomingValues() == 1);
+          P->replaceAllUsesWith(P->getIncomingValue(0));
+          P->eraseFromParent();
+        }
+      }
+    }
+
+    // Exit nodes may have changed by SplitEdge.
+    recomputeExitBlocks();
+  }
 }
 
 void CodeExtractor::recomputeExitBlocks() {
   SwitchCases.clear();
 
   SmallPtrSet<BasicBlock *, 1> ExitBlocks;
   for (BasicBlock *Block : Blocks) {
     for (BasicBlock *Succ : successors(Block)) {
void CodeExtractor::emitFunctionBody(
   Function *oldFunction = header->getParent();
   LLVMContext &Context = oldFunction->getContext();
 
   // The new function needs a root node because other nodes can branch to the
   // head of the region, but the entry node of a function cannot have preds.
   BasicBlock *newFuncRoot =
       BasicBlock::Create(Context, "newFuncRoot", newFunction);
 
+  // The map of values from the original function to the corresponding values in
+  // the extracted function; only used with KeepOldBlocks.
+  ValueToValueMapTy VMap;
+
+  // Additional instructions not in a extracted block whose operands need to be
+  // remapped.
+  SmallVector<Instruction *> AdditionalRemap;
+
+  // Copy or move (depending on KeepOldBlocks) an instruction to the new
+  // function.
+  auto MoveOrCopyInst = [this, newFuncRoot, &VMap,
+                         &AdditionalRemap](Instruction *I) -> Instruction * {
+    BasicBlock::iterator IP = newFuncRoot->getFirstInsertionPt();
+    if (!KeepOldBlocks) {
+      I->moveBefore(*newFuncRoot, IP);
+      return I;
+    }
+
+    Instruction *ClonedI = I->clone();
+    ClonedI->setName(I->getName());
+    newFuncRoot->getInstList().insert(IP, ClonedI);
+    AdditionalRemap.push_back(ClonedI);
+    VMap[I] = ClonedI;
+    return ClonedI;
+  };
+
   // Now sink all instructions which only have non-phi uses inside the region.
   // Group the allocas at the start of the block, so that any bitcast uses of
   // the allocas are well-defined.
   for (auto *II : SinkingCands) {
     if (!isa<AllocaInst>(II)) {
-      cast<Instruction>(II)->moveBefore(*newFuncRoot,
-                                        newFuncRoot->getFirstInsertionPt());
+      MoveOrCopyInst(cast<Instruction>(II));
     }
   }
   for (auto *II : SinkingCands) {
     if (auto *AI = dyn_cast<AllocaInst>(II)) {
-      AI->moveBefore(*newFuncRoot, newFuncRoot->getFirstInsertionPt());
+      MoveOrCopyInst(AI);
     }
   }
 
   // Create an iterator to name all of the arguments we inserted.
   Function::arg_iterator AI = newFunction->arg_begin();
 
   // Rewrite all users of the inputs in the extracted region to use the
   // arguments (or appropriate addressing into struct) instead.
if (AggregateArgs) {
       RewriteVal = new LoadInst(StructArgTy->getElementType(i), GEP,
                                 "loadgep_" + inputs[i]->getName(), TI);
     } else
       RewriteVal = &*AI++;
 
     NewValues.push_back(RewriteVal);
   }
 
+  if (KeepOldBlocks) {
+    // Copy blocks and instrutions to newFunction.
+    for (BasicBlock *Block : Blocks) {
+      BasicBlock *CBB = CloneBasicBlock(
+          Block, VMap, {}, newFunction, /* CodeInfo */ nullptr,
+          /* DIFinder */ nullptr,
+          [](const Instruction *I) -> bool { return !isa<AssumeInst>(I); });
+
+      // Add basic block mapping.
+      VMap[Block] = CBB;
+
+      // It is only legal to clone a function if a block address within that
+      // function is never referenced outside of the function.  Given that, we
+      // want to map block addresses from the old function to block addresses in
+      // the clone. (This is different from the generic ValueMapper
+      // implementation, which generates an invalid blockaddress when
+      // cloning a function.)
+      if (Block->hasAddressTaken()) {
+        Constant *OldBBAddr = BlockAddress::get(oldFunction, Block);
+        VMap[OldBBAddr] = BlockAddress::get(newFunction, CBB);
+      }
+
+      // Non-header block may have branches from outside the region. These
+      // continue to branch to the original blocks, hence remove their PHI
+      // entries.
+      if (Block != header)
+        for (auto &&P : CBB->phis()) {
+          unsigned NumIncoming = P.getNumIncomingValues();
+          for (int Idx = NumIncoming - 1; Idx >= 0; --Idx) {
+            BasicBlock *IncomingBlock = P.getIncomingBlock(Idx);
+            if (Blocks.count(IncomingBlock))
+              continue;
+            P.removeIncomingValue(Idx, /*DeletePHIIfEmpty=*/false);
+          }
+        }
+    }
+
+    for (auto P : enumerate(inputs))
+      VMap[P.value()] = NewValues[P.index()];
+
+  } else {
     moveCodeToFunction(newFunction);
 
     for (unsigned i = 0, e = inputs.size(); i != e; ++i) {
       Value *RewriteVal = NewValues[i];
 
       std::vector<User *> Users(inputs[i]->user_begin(), inputs[i]->user_end());
       for (User *use : Users)
         if (Instruction *inst = dyn_cast<Instruction>(use))
           if (Blocks.count(inst->getParent()))
             inst->replaceUsesOfWith(inputs[i], RewriteVal);
     }
+  }
 
   // Create stubs for the original exit blocks.
   std::map<BasicBlock *, BasicBlock *> ExitBlockMap;
   for (auto P : enumerate(SwitchCases)) {
     BasicBlock *OldTarget = P.value();
     size_t SuccNum = P.index();
 
     BasicBlock *&NewTarget = ExitBlockMap[OldTarget];
     assert(!NewTarget && "Switch cases muast be unique");
 
     // If we don't already have an exit stub for this non-extracted
     // destination, create one now!
     NewTarget = BasicBlock::Create(Context, OldTarget->getName() + ".exitStub",
                                    newFunction);
+    if (KeepOldBlocks)
+      VMap[OldTarget] = NewTarget;
 
     Value *brVal = nullptr;
     assert(NumExitBlocks < 0xffff && "too many exit blocks for switch");
     switch (NumExitBlocks) {
     case 0:
     case 1:
       break; // No value needed.
     case 2:  // Conditional branch, return a bool
for (BasicBlock *Block : Blocks) {
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) {
       if (Blocks.count(TI->getSuccessor(i)))
         continue;
       BasicBlock *OldTarget = TI->getSuccessor(i);
       // add a new basic block which returns the appropriate value
       BasicBlock *NewTarget = ExitBlockMap[OldTarget];
       assert(NewTarget && "Unknown target block!");
 
+      if (!KeepOldBlocks) {
         // rewrite the original branch instruction with this new target
         TI->setSuccessor(i, NewTarget);
+      } else {
+        VMap[OldTarget] = NewTarget;
+      }
     }
   }
 
+  // Update values references to point to the new function.
+  if (KeepOldBlocks) {
+    for (BasicBlock *Pred : predecessors(header)) {
+      if (VMap.count(Pred))
+        continue;
+      VMap[Pred] = newFuncRoot;
+    }
+
+    for (Instruction *II : AdditionalRemap)
+      RemapInstruction(II, VMap, RF_NoModuleLevelChanges);
+
+    // Loop over all of the instructions in the new function, fixing up operand
+    // references as we go. This uses VMap to do all the hard work.
+    for (BasicBlock *Block : Blocks) {
+      WeakTrackingVH NewBlock = VMap.lookup(Block);
+      if (!NewBlock)
+        continue;
+
+      // Loop over all instructions, fixing each one as we find it...
+      for (Instruction &II : cast<BasicBlock>(*NewBlock))
+        RemapInstruction(&II, VMap, RF_NoModuleLevelChanges);
+    }
+  } else {
     // Loop over all of the PHI nodes in the header and exit blocks, and change
     // any references to the old incoming edge to be the new incoming edge.
     for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {
       PHINode *PN = cast<PHINode>(I);
       for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
         if (!Blocks.count(PN->getIncomingBlock(i)))
           PN->setIncomingBlock(i, newFuncRoot);
     }
+  }
+
+  BasicBlock *NewHeader =
+      KeepOldBlocks ? cast<BasicBlock>(VMap.lookup(header)) : header;
+  assert(NewHeader && "Header must have been cloned/moved to newFunction");
 
   // Connect newFunction entry block to new header.
-  BranchInst *BranchI2 = BranchInst::Create(header, newFuncRoot);
+  BranchInst *BranchI2 = BranchInst::Create(NewHeader, newFuncRoot);
   applyFirstDebugLoc(oldFunction, Blocks.getArrayRef(), BranchI2);
 
   // Store the arguments right after the definition of output value.
   // This should be proceeded after creating exit stubs to be ensure that invoke
   // result restore will be placed in the outlined function.
   Function::arg_iterator OAI = newFunction->arg_begin() + inputs.size();
   for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
     auto *OutI = dyn_cast<Instruction>(outputs[i]);
     if (!OutI)
       continue;
 
+    if (KeepOldBlocks)
+      OutI = cast<Instruction>(VMap.lookup(OutI));
+
     // Find proper insertion point.
     BasicBlock::iterator InsertPt;
     // In case OutI is an invoke, we insert the store at the beginning in the
     // 'normal destination' BB. Otherwise we insert the store right after OutI.
     if (auto *InvokeI = dyn_cast<InvokeInst>(OutI))
       InsertPt = InvokeI->getNormalDest()->getFirstInsertionPt();
     else if (auto *Phi = dyn_cast<PHINode>(OutI))
       InsertPt = Phi->getParent()->getFirstInsertionPt();
void CodeExtractor::insertReplacerCall(
   std::vector<User *> Users(header->user_begin(), header->user_end());
   for (auto &U : Users)
     // The BasicBlock which contains the branch is not in the region
     // modify the branch target to a new block
     if (Instruction *I = dyn_cast<Instruction>(U))
       if (I->isTerminator() && I->getFunction() == oldFunction)
         I->replaceUsesOfWith(header, codeReplacer);
 
+  if (KeepOldBlocks) {
+    // Change references to output values after the call to use either the value
+    // written by the extracted function or the original value if we skipped the
+    // call. Use SSAUpdater to propagate the new PHI since the CFG has changed.
+
+    SSAUpdater SSA;
+    for (auto P : enumerate(outputs)) {
+      size_t OutIdx = P.index();
+      Instruction *OldVal = cast<Instruction>(P.value());
+      Value *NewVal = Reloads[OutIdx];
+
+      SSA.Initialize(OldVal->getType(),
+                     (OldVal->getName() + ".merge_with_extracted").str());
+      SSA.AddAvailableValue(codeReplacer, NewVal);
+
+      // Could help SSAUpdater by determining in advance which output values are
+      // available in which exit blocks (from DT).
+      SSA.AddAvailableValue(OldVal->getParent(), OldVal);
+
+      for (Use &U : make_early_inc_range(OldVal->uses())) {
+        auto *User = dyn_cast<Instruction>(U.getUser());
+        if (!User)
+          continue;
+        BasicBlock *EffectiveUser = User->getParent();
+        if (auto *PHI = dyn_cast<PHINode>(User))
+          EffectiveUser = PHI->getIncomingBlock(U);
+
+        if (EffectiveUser == codeReplacer || Blocks.count(EffectiveUser))
+          continue;
+
+        SSA.RewriteUseAfterInsertions(U);
+      }
+    }
+  } else {
     // When moving the code region it is sufficient to replace all uses to the
     // extracted function values. Since the original definition's block
     // dominated its use, it will also be dominated by codeReplacer's switch
     // which joined multiple exit blocks.
+
     for (BasicBlock *ExitBB : SwitchCases)
       for (PHINode &PN : ExitBB->phis()) {
         Value *IncomingCodeReplacerVal = nullptr;
         for (unsigned i = 0, e = PN.getNumIncomingValues(); i != e; ++i) {
           // Ignore incoming values from outside of the extracted region.
           if (!Blocks.count(PN.getIncomingBlock(i)))
             continue;
 
           // Ensure that there is only one incoming value from codeReplacer.
           if (!IncomingCodeReplacerVal) {
             PN.setIncomingBlock(i, codeReplacer);
             IncomingCodeReplacerVal = PN.getIncomingValue(i);
           } else
             assert(IncomingCodeReplacerVal == PN.getIncomingValue(i) &&
                    "PHI has two incompatbile incoming values from codeRepl");
         }
       }
 
     for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
       Value *load = Reloads[i];
-    std::vector<User *> Users(outputs[i]->user_begin(), outputs[i]->user_end());
+      std::vector<User *> Users(outputs[i]->user_begin(),
+                                outputs[i]->user_end());
       for (unsigned u = 0, e = Users.size(); u != e; ++u) {
         Instruction *inst = cast<Instruction>(Users[u]);
         if (inst->getParent()->getParent() == oldFunction)
           inst->replaceUsesOfWith(outputs[i], load);
       }
     }
+  }
 
   // Update the branch weights for the exit block.
   if (BFI && NumExitBlocks > 1)
     calculateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI);
 }
 
 bool CodeExtractor::verifyAssumptionCache(const Function &OldFunc,
                                           const Function &NewFunc,

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

+; RUN: llvm-extract -S -bb "foo:region_start;extractonly;cleanup;fallback;region_end" %s --bb-keep-functions --bb-keep-blocks | FileCheck %s
+
+
+; CHECK-LABEL: define void @foo(i32* %arg, i1 %c) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %c, label %codeRepl, label %outsideonly
+; CHECK-EMPTY:
+; CHECK-NEXT:  outsideonly:
+; CHECK-NEXT:    store i32 0, i32* %arg, align 4
+; CHECK-NEXT:    br label %cleanup
+; CHECK-EMPTY:
+; CHECK-NEXT:  codeRepl:
+; CHECK-NEXT:    %targetBlock = call i1 @foo.region_start(i32* %arg)
+; CHECK-NEXT:    br i1 %targetBlock, label %cleanup.return_crit_edge, label %region_end.split
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_start:
+; CHECK-NEXT:    br label %extractonly
+; CHECK-EMPTY:
+; CHECK-NEXT:  extractonly:
+; CHECK-NEXT:    store i32 1, i32* %arg, align 4
+; CHECK-NEXT:    br label %cleanup
+; CHECK-EMPTY:
+; CHECK-NEXT:  cleanup:
+; CHECK-NEXT:    %dest = phi i8 [ 0, %outsideonly ], [ 1, %extractonly ]
+; CHECK-NEXT:    switch i8 %dest, label %fallback [
+; CHECK-NEXT:      i8 0, label %cleanup.return_crit_edge
+; CHECK-NEXT:      i8 1, label %region_end
+; CHECK-NEXT:    ]
+; CHECK-EMPTY:
+; CHECK-NEXT:  cleanup.return_crit_edge:
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  fallback:
+; CHECK-NEXT:    unreachable
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_end:
+; CHECK-NEXT:    br label %region_end.split
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_end.split:
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  outsidecont:
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  return:
+; CHECK-NEXT:    ret void
+; CHECK-NEXT:  }
+
+
+; CHECK-LABEL: define internal i1 @foo.region_start(i32* %arg) {
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label %region_start
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_start:
+; CHECK-NEXT:    br label %extractonly
+; CHECK-EMPTY:
+; CHECK-NEXT:  extractonly:
+; CHECK-NEXT:    store i32 1, i32* %arg, align 4
+; CHECK-NEXT:    br label %cleanup
+; CHECK-EMPTY:
+; CHECK-NEXT:  cleanup:
+; CHECK-NEXT:    %dest = phi i8 [ 1, %extractonly ]
+; CHECK-NEXT:    switch i8 %dest, label %fallback [
+; CHECK-NEXT:      i8 0, label %cleanup.return_crit_edge.exitStub
+; CHECK-NEXT:      i8 1, label %region_end
+; CHECK-NEXT:    ]
+; CHECK-EMPTY:
+; CHECK-NEXT:  fallback:
+; CHECK-NEXT:    unreachable
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_end:
+; CHECK-NEXT:    br label %region_end.split.exitStub
+; CHECK-EMPTY:
+; CHECK-NEXT:  cleanup.return_crit_edge.exitStub:
+; CHECK-NEXT:    ret i1 true
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_end.split.exitStub:
+; CHECK-NEXT:    ret i1 false
+; CHECK-NEXT:  }
+
+
+
+define void @foo(i32* %arg, i1 %c) {
+entry:
+  br i1 %c, label %region_start, label %outsideonly
+
+outsideonly:
+  store i32 0, i32* %arg, align 4
+  br label %cleanup
+
+region_start:
+  br label %extractonly
+
+extractonly:
+  store i32 1, i32* %arg, align 4
+  br label %cleanup
+
+cleanup:
+  %dest = phi i8 [0, %outsideonly], [1, %extractonly]
+  switch i8 %dest, label %fallback [
+    i8 0, label %return
+    i8 1, label %region_end
+  ]
+
+fallback:
+  unreachable
+
+region_end:
+  br label %return
+
+outsidecont:
+  br label %return
+
+return:
+  ret void
+}

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

+; RUN: llvm-extract -S -bb "func:region_start;exiting0;exiting1" %s --bb-keep-functions --bb-keep-blocks | FileCheck %s
+
+
+; CHECK-LABEL: define void @func(i32* %arg, i1 %c0, i1 %c1, i1 %c2, i8 %dest) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %B.ce.loc = alloca i32, align 4
+; CHECK-NEXT:    %c.loc = alloca i32, align 4
+; CHECK-NEXT:    %b.loc = alloca i32, align 4
+; CHECK-NEXT:    %a.loc = alloca i32, align 4
+; CHECK-NEXT:    br i1 %c0, label %codeRepl, label %exit
+; CHECK-EMPTY:
+; CHECK-NEXT:  codeRepl:                                        
+; CHECK-NEXT:    %lt.cast = bitcast i32* %a.loc to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast)
+; CHECK-NEXT:    %lt.cast1 = bitcast i32* %b.loc to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast1)
+; CHECK-NEXT:    %lt.cast2 = bitcast i32* %c.loc to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast2)
+; CHECK-NEXT:    %lt.cast3 = bitcast i32* %B.ce.loc to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast3)
+; CHECK-NEXT:    %targetBlock = call i16 @func.region_start(i1 %c1, i1 %c2, i8 %dest, i32* %a.loc, i32* %b.loc, i32* %c.loc, i32* %B.ce.loc)
+; CHECK-NEXT:    %a.reload = load i32, i32* %a.loc, align 4
+; CHECK-NEXT:    %b.reload = load i32, i32* %b.loc, align 4
+; CHECK-NEXT:    %c.reload = load i32, i32* %c.loc, align 4
+; CHECK-NEXT:    %B.ce.reload = load i32, i32* %B.ce.loc, align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* %lt.cast)
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* %lt.cast1)
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* %lt.cast2)
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* %lt.cast3)
+; CHECK-NEXT:    switch i16 %targetBlock, label %exit0 [
+; CHECK-NEXT:      i16 0, label %exiting0.exit_crit_edge
+; CHECK-NEXT:      i16 1, label %fallback
+; CHECK-NEXT:      i16 2, label %exit1
+; CHECK-NEXT:      i16 3, label %exit2
+; CHECK-NEXT:    ]
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_start:                                
+; CHECK-NEXT:    %a = add i32 42, 1
+; CHECK-NEXT:    br i1 %c1, label %exiting0, label %exiting1
+; CHECK-EMPTY:
+; CHECK-NEXT:  exiting0:                                     
+; CHECK-NEXT:    %b = add i32 42, 2
+; CHECK-NEXT:    br i1 %c2, label %exiting0.exit_crit_edge, label %exit0.split
+; CHECK-EMPTY:
+; CHECK-NEXT:  exiting0.exit_crit_edge:                         
+; CHECK-NEXT:    %b.merge_with_extracted7 = phi i32 [ %b.reload, %codeRepl ], [ %b, %exiting0 ]
+; CHECK-NEXT:    br label %exit
+; CHECK-EMPTY:
+; CHECK-NEXT:  exiting1:                                        
+; CHECK-NEXT:    %c = add i32 42, 3
+; CHECK-NEXT:    switch i8 %dest, label %fallback [
+; CHECK-NEXT:      i8 0, label %exit0.split
+; CHECK-NEXT:      i8 1, label %exit1
+; CHECK-NEXT:      i8 2, label %exit2
+; CHECK-NEXT:      i8 3, label %exit0.split
+; CHECK-NEXT:    ]
+; CHECK-EMPTY:
+; CHECK-NEXT:  fallback:                                        
+; CHECK-NEXT:    unreachable
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit:                                          
+; CHECK-NEXT:    %A = phi i32 [ 42, %entry ], [ %b.merge_with_extracted7, %exiting0.exit_crit_edge ]
+; CHECK-NEXT:    store i32 %A, i32* %arg, align 4
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit0.split:                                     
+; CHECK-NEXT:    %b.merge_with_extracted6 = phi i32 [ %b, %exiting0 ], [ undef, %exiting1 ], [ undef, %exiting1 ]
+; CHECK-NEXT:    %B.ce = phi i32 [ %b, %exiting0 ], [ %a, %exiting1 ], [ %a, %exiting1 ]
+; CHECK-NEXT:    br label %exit0
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit0:                                         
+; CHECK-NEXT:    %B.ce.merge_with_extracted = phi i32 [ %B.ce.reload, %codeRepl ], [ %B.ce, %exit0.split ]
+; CHECK-NEXT:    %b.merge_with_extracted = phi i32 [ %b.reload, %codeRepl ], [ %b.merge_with_extracted6, %exit0.split ]
+; CHECK-NEXT:    %a.merge_with_extracted5 = phi i32 [ %a.reload, %codeRepl ], [ %a, %exit0.split ]
+; CHECK-NEXT:    store i32 %a.merge_with_extracted5, i32* %arg, align 4
+; CHECK-NEXT:    store i32 %B.ce.merge_with_extracted, i32* %arg, align 4
+; CHECK-NEXT:    br label %after
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit1:                                         
+; CHECK-NEXT:    %c.merge_with_extracted8 = phi i32 [ %c.reload, %codeRepl ], [ %c, %exiting1 ]
+; CHECK-NEXT:    %a.merge_with_extracted4 = phi i32 [ %a.reload, %codeRepl ], [ %a, %exiting1 ]
+; CHECK-NEXT:    br label %after
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit2:                                         
+; CHECK-NEXT:    %c.merge_with_extracted = phi i32 [ %c.reload, %codeRepl ], [ %c, %exiting1 ]
+; CHECK-NEXT:    store i32 %c.merge_with_extracted, i32* %arg, align 4
+; CHECK-NEXT:    store i32 %c.merge_with_extracted, i32* %arg, align 4
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  after:                                         
+; CHECK-NEXT:    %a.merge_with_extracted = phi i32 [ %a.merge_with_extracted5, %exit0 ], [ %a.merge_with_extracted4, %exit1 ]
+; CHECK-NEXT:    %D = phi i32 [ %b.merge_with_extracted, %exit0 ], [ %c.merge_with_extracted8, %exit1 ]
+; CHECK-NEXT:    store i32 %a.merge_with_extracted, i32* %arg, align 4
+; CHECK-NEXT:    store i32 %D, i32* %arg, align 4
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  return:                                         
+; CHECK-NEXT:    ret void
+; CHECK-NEXT:  }
+
+
+; CHECK-LABEL: define internal i16 @func.region_start(i1 %c1, i1 %c2, i8 %dest, i32* %a.out, i32* %b.out, i32* %c.out, i32* %B.ce.out) {
+; CHECK-NEXT:  newFuncRoot:                                    
+; CHECK-NEXT:    br label %region_start
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_start:                                   
+; CHECK-NEXT:    %a = add i32 42, 1
+; CHECK-NEXT:    store i32 %a, i32* %a.out, align 4
+; CHECK-NEXT:    br i1 %c1, label %exiting0, label %exiting1
+; CHECK-EMPTY:
+; CHECK-NEXT:  exiting0:                                     
+; CHECK-NEXT:    %b = add i32 42, 2
+; CHECK-NEXT:    store i32 %b, i32* %b.out, align 4
+; CHECK-NEXT:    br i1 %c2, label %exiting0.exit_crit_edge.exitStub, label %exit0.split
+; CHECK-EMPTY:
+; CHECK-NEXT:  exiting1:                                     
+; CHECK-NEXT:    %c = add i32 42, 3
+; CHECK-NEXT:    store i32 %c, i32* %c.out, align 4
+; CHECK-NEXT:    switch i8 %dest, label %fallback.exitStub [
+; CHECK-NEXT:      i8 0, label %exit0.split
+; CHECK-NEXT:      i8 1, label %exit1.exitStub
+; CHECK-NEXT:      i8 2, label %exit2.exitStub
+; CHECK-NEXT:      i8 3, label %exit0.split
+; CHECK-NEXT:    ]
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit0.split:                                    
+; CHECK-NEXT:    %B.ce = phi i32 [ %b, %exiting0 ], [ %a, %exiting1 ], [ %a, %exiting1 ]
+; CHECK-NEXT:    store i32 %B.ce, i32* %B.ce.out, align 4
+; CHECK-NEXT:    br label %exit0.exitStub
+; CHECK-EMPTY:
+; CHECK-NEXT:  exiting0.exit_crit_edge.exitStub:                
+; CHECK-NEXT:    ret i16 0
+; CHECK-EMPTY:
+; CHECK-NEXT:  fallback.exitStub:                             
+; CHECK-NEXT:    ret i16 1
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit1.exitStub:                                
+; CHECK-NEXT:    ret i16 2
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit2.exitStub:                               
+; CHECK-NEXT:    ret i16 3
+; CHECK-EMPTY:
+; CHECK-NEXT:  exit0.exitStub:                                
+; CHECK-NEXT:    ret i16 4
+; CHECK-NEXT:  }
+
+
+define void @func(i32* %arg, i1 %c0, i1 %c1, i1 %c2, i8 %dest) {
+entry:
+  br i1 %c0, label %region_start, label %exit
+
+region_start:
+  %a = add i32 42, 1
+  br i1 %c1, label %exiting0, label %exiting1
+
+exiting0:
+  %b = add i32 42, 2
+  br i1 %c2, label %exit, label %exit0
+
+exiting1:
+  %c = add i32 42, 3
+  switch i8 %dest, label %fallback [
+    i8 0, label %exit0
+    i8 1, label %exit1
+    i8 2, label %exit2
+    i8 3, label %exit0
+  ]
+
+fallback:
+  unreachable
+
+exit:
+  %A = phi i32 [ 42, %entry ], [ %b, %exiting0 ]
+  store i32 %A, i32* %arg
+  br label %return
+
+exit0:
+  %B = phi i32 [ %b, %exiting0 ], [ %a, %exiting1 ] , [ %a, %exiting1 ]
+  store i32 %a, i32* %arg
+  store i32 %B, i32* %arg
+  br label %after
+
+exit1:
+  br label %after
+
+exit2:
+  %C = phi i32 [ %c, %exiting1 ]
+  store i32 %c, i32* %arg
+  store i32 %C, i32* %arg
+  br label %return
+
+after:
+  %D = phi i32 [ %b, %exit0 ], [ %c, %exit1 ]
+  store i32 %a, i32* %arg
+  store i32 %D, i32* %arg
+  br label %return
+
+return:
+  ret void
+}

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

+; RUN: llvm-extract -S -bb "foo:region_start" %s --bb-keep-functions --bb-keep-blocks | FileCheck %s
+
+
+
+; CHECK-LABEL: define void @foo() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %a = alloca i32, align 4
+; CHECK-NEXT:    %b = alloca i32, align 4
+; CHECK-NEXT:    br label %codeRepl
+; CHECK-EMPTY:
+; CHECK-NEXT:  codeRepl:
+; CHECK-NEXT:    call void @foo.region_start(i32* %b)
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_start:
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i32(i64 4, i32* nonnull %a)
+; CHECK-NEXT:    store i32 43, i32* %a, align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i32(i64 4, i32* nonnull %a)
+; CHECK-NEXT:    store i32 44, i32* %b, align 4
+; CHECK-NEXT:    br label %return
+; CHECK-EMPTY:
+; CHECK-NEXT:  return:
+; CHECK-NEXT:    ret void
+; CHECK-NEXT:  }
+
+
+; CHECK-LABEL: define internal void @foo.region_start(i32* %b) {
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    %a = alloca i32, align 4
+; CHECK-NEXT:    br label %region_start
+; CHECK-EMPTY:
+; CHECK-NEXT:  region_start:
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i32(i64 4, i32* nonnull %a)
+; CHECK-NEXT:    store i32 43, i32* %a, align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i32(i64 4, i32* nonnull %a)
+; CHECK-NEXT:    store i32 44, i32* %b, align 4
+; CHECK-NEXT:    br label %return.exitStub
+; CHECK-EMPTY:
+; CHECK-NEXT:  return.exitStub:
+; CHECK-NEXT:    ret void
+; CHECK-NEXT:  }
+
+
+
+
+
+
+declare void @llvm.lifetime.start.p0i32(i64, i32* nocapture)
+declare void @llvm.lifetime.end.p0i32(i64, i32* nocapture)
+
+
+define void @foo() {
+entry:
+  %a = alloca i32, align 4
+  %b = alloca i32, align 4
+  br label %region_start
+
+region_start:
+  call void @llvm.lifetime.start.p0i32(i64 4, i32* nonnull %a)
+  store i32 43, i32* %a
+  call void @llvm.lifetime.end.p0i32(i64 4, i32* nonnull %a)
+  store i32 44, i32* %b
+  br label %return
+
+return:
+  ret void
+}

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

-; RUN: llvm-extract -S -bb foo:bb4 %s | FileCheck %s
+; RUN: llvm-extract -S -bb foo:bb4 %s                                      | FileCheck %s --check-prefixes=CHECK,KILL
+; RUN: llvm-extract -S -bb foo:bb4 %s --bb-keep-functions --bb-keep-blocks | FileCheck %s --check-prefixes=CHECK,KEEP
+
 
 ; CHECK: declare void @bar()
 define void @bar() {
 bb:
   ret void
 }
 
 ; CHECK-NOT: @unused()
 define void @unused() {
 bb:
   ret void
 }
 
-; CHECK: @foo.bb4
+; KEEP-LABEL: define i32 @foo(i32 %arg) {
+; KEEP:       call void @foo.bb4
+
+; KILL-LABEL: define dso_local void @foo.bb4(
+; KEEP-LABEL: define internal void @foo.bb4(
 ; CHECK: call void @bar()
 ; 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

llvm/test/tools/llvm-extract/extract-blocks-with-groups.ll

-; RUN: llvm-extract -bb 'foo:if;then;else' -bb 'bar:bb14;bb20' -S %s  | FileCheck %s
+; RUN: llvm-extract -bb 'foo:if;then;else' -bb 'bar:bb14;bb20' -S %s                                      | FileCheck %s --check-prefixes=CHECK,KILL
+; RUN: llvm-extract -bb 'foo:if;then;else' -bb 'bar:bb14;bb20' -S %s --bb-keep-functions --bb-keep-blocks | FileCheck %s --check-prefixes=CHECK,KEEP
 ; Extract two groups of basic blocks in two different functions.
 
 
+; KEEP-LABEL: define i32 @foo(i32 %arg, i32 %arg1) {
+; KEEP:         call void @foo.if.split(
+
+; KEEP-LABEL: define i32 @bar(i32 %arg, i32 %arg1) {
+; KEEP:         %targetBlock = call i1 @bar.bb14(
+
+
 ; The first extracted function is the region composed by the
 ; blocks if, then, and else from foo.
-; CHECK: define dso_local void @foo.if.split(i32 %arg1, i32 %arg, i32* %tmp.0.ce.out) {
+; KILL-LABEL: define dso_local void @foo.if.split(i32 %arg1, i32 %arg, i32* %tmp.0.ce.out) {
+; KEEP-LABEL: define internal void @foo.if.split(i32 %arg1, i32 %arg, i32* %tmp.0.ce.out) {
 ; CHECK: newFuncRoot:
 ; CHECK:   br label %if.split
 ;
 ; CHECK: then:                                             ; preds = %if.split
 ; CHECK:   %tmp12 = shl i32 %arg1, 2
 ; CHECK:   %tmp13 = add nsw i32 %tmp12, %arg
 ; CHECK:   br label %end.split
 ;
 ; CHECK: else:                                             ; preds = %if.split
 ; CHECK:   %tmp22 = mul nsw i32 %arg, 3
 ; CHECK:   %tmp24 = sdiv i32 %arg1, 6
 ; CHECK:   %tmp25 = add nsw i32 %tmp24, %tmp22
 ; CHECK:   br label %end.split
 ;
 ; CHECK: if.split:                                         ; preds = %newFuncRoot
 ; CHECK:   %tmp5 = icmp sgt i32 %arg, 0
 ; CHECK:   %tmp8 = icmp sgt i32 %arg1, 0
 ; CHECK:   %or.cond = and i1 %tmp5, %tmp8
 ; CHECK:   br i1 %or.cond, label %then, label %else
 ;
-; CHECK: end.split:                                        ; preds = %then, %else
+; CHECK: end.split:
 ; CHECK:   %tmp.0.ce = phi i32 [ %tmp13, %then ], [ %tmp25, %else ]
 ; CHECK:   store i32 %tmp.0.ce, i32* %tmp.0.ce.out
 ; CHECK:   br label %end.exitStub
 ;
 ; CHECK: end.exitStub:                                     ; preds = %end.split
 ; CHECK:   ret void
 ; CHECK: }
 
 ; The second extracted function is the region composed by the blocks
 ; bb14 and bb20 from bar.
-; CHECK: define dso_local i1 @bar.bb14(i32 %arg1, i32 %arg, i32* %tmp25.out) {
+; KILL-LABEL: define dso_local i1 @bar.bb14(i32 %arg1, i32 %arg, i32* %tmp25.out) {
+; KEEP-LABEL: define internal i1 @bar.bb14(i32 %arg1, i32 %arg, i32* %tmp25.out) {
 ; CHECK: newFuncRoot:
 ; CHECK:   br label %bb14
 ;
 ; CHECK: bb14:                                             ; preds = %newFuncRoot
 ; CHECK:   %tmp0 = and i32 %arg1, %arg
 ; CHECK:   %tmp1 = icmp slt i32 %tmp0, 0
 ; CHECK:   br i1 %tmp1, label %bb20, label %bb26.exitStub
 ;
 ; CHECK: bb20:                                             ; preds = %bb14
 ; CHECK:   %tmp22 = mul nsw i32 %arg, 3
 ; CHECK:   %tmp24 = sdiv i32 %arg1, 6
 ; CHECK:   %tmp25 = add nsw i32 %tmp24, %tmp22
 ; CHECK:   store i32 %tmp25, i32* %tmp25.out
-; CHECK:   br label %bb30.exitStub
+; KILL:    br label %bb30.exitStub
+; KEEP:    br label %bb20.split.exitStub
 ;
 ; CHECK: bb26.exitStub:                                    ; preds = %bb14
 ; CHECK:   ret i1 true
 ;
-; CHECK: bb30.exitStub:                                    ; preds = %bb20
+; KILL:  bb30.exitStub:                                    ; preds = %bb20
+; KEEP:  bb20.split.exitStub:
 ; CHECK:   ret i1 false
 ; CHECK: }
 
 define i32 @foo(i32 %arg, i32 %arg1) {
 if:
   %tmp5 = icmp sgt i32 %arg, 0
   %tmp8 = icmp sgt i32 %arg1, 0
   %or.cond = and i1 %tmp5, %tmp8

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

 
 // ExtractBlocks - The blocks to extract from the module.
 static cl::list<std::string> ExtractBlocks(
     "bb",
     cl::desc(
         "Specify <function, basic block1[;basic block2...]> pairs to extract.\n"
         "Each pair will create a function.\n"
         "If multiple basic blocks are specified in one pair,\n"
-        "the first block in the sequence should dominate the rest.\n"
+        "the first block in the sequence should dominate the rest (Unlsess "
+        "using --bb-keep-blocks).\n"
         "eg:\n"
         "  --bb=f:bb1;bb2 will extract one function with both bb1 and bb2;\n"
         "  --bb=f:bb1 --bb=f:bb2 will extract two functions, one with bb1, one "
         "with bb2."),
     cl::ZeroOrMore, cl::value_desc("function:bb1[;bb2...]"),
     cl::cat(ExtractCat));
 
+static cl::opt<bool> KeepFunctions(
+    "bb-keep-functions",
+    cl::desc(
+        "When extracting blocks from functions, keep the original functions; "
+        "extracted code is replaced by function call to new function"),
+    cl::cat(ExtractCat));
+
+static cl::opt<bool> KeepBlocks(
+    "bb-keep-blocks",
+    cl::desc("Keep extracted blocks in original function after outlining. This "
+             "permits branches to any selected basic block from outside the "
+             "selection and overlapping code regions, but only branches to the "
+             "first in the group will call the extracted function."),
+    cl::cat(ExtractCat));
+
 // 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"),
                    cl::cat(ExtractCat));
 
 // ExtractRegExpAliases - The aliases, matched via regular expression, to
 // extract from the module.
for (auto &P : BBMap) {
           return 1;
         }
         BBs.push_back(&*Res);
       }
       GroupOfBBs.push_back(BBs);
     }
 
     legacy::PassManager PM;
-    PM.add(createBlockExtractorPass(GroupOfBBs, true));
+    PM.add(createBlockExtractorPass(GroupOfBBs, !KeepFunctions, KeepBlocks));
     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)