[IROutliner][NFC] Add initial support for multiple exit paths

[AndrewLitteken]

Summary

The IROutliner could potentially support multiple basic blocks inside of the outlined region, which could result in multiple exit paths from different basic blocks, each with its own output stores necessary. This changes the data structures tracking returning basic blocks out of the functions, and the data structures tracking the output objects to be ready to support this. IT should not change any current behavior of outlining a single block.

[AndrewLitteken]

removing to be replaced with a separate patch.

llvm/include/llvm/Transforms/IPO/IROutliner.h

struct OutlinableGroup {
   FunctionType *OutlinedFunctionType = nullptr;
   /// The Function for the collective overall function.
   Function *OutlinedFunction = nullptr;
 
   /// Flag for whether we should not consider this group of OutlinableRegions
   /// for extraction.
   bool IgnoreGroup = false;
 
-  /// The return block for the overall function.
-  BasicBlock *EndBB = nullptr;
+  /// The return blocks for the overall function.
+  DenseMap<Value *, BasicBlock *> EndBBs;
+
+  /// The PHIBlocks with their corresponding return block based on the return
+  /// value as the key.
+  DenseMap<Value *, BasicBlock *> PHIBlocks;
 
   /// A set containing the different GVN store sets needed. Each array contains
   /// a sorted list of the different values that need to be stored into output
   /// registers.
   DenseSet<ArrayRef<unsigned>> OutputGVNCombinations;
 
   /// Flag for whether the \ref ArgumentTypes have been defined after the
   /// extraction of the first region.

llvm/lib/Transforms/IPO/IROutliner.cpp

 
 #include "llvm/Transforms/IPO/IROutliner.h"
 #include "llvm/Analysis/IRSimilarityIdentifier.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DIBuilder.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
 #include <map>
 #include <set>
Function *IROutliner::createFunction(Module &M, OutlinableGroup &Group,
 
   return Group.OutlinedFunction;
 }
 
 /// Move each BasicBlock in \p Old to \p New.
 ///
 /// \param [in] Old - The function to move the basic blocks from.
 /// \param [in] New - The function to move the basic blocks to.
+/// \param [out] NewEnds - The return blocks of the new overall function.
 /// \returns the first return block for the function in New.
-static BasicBlock *moveFunctionData(Function &Old, Function &New) {
+static void moveFunctionData(Function &Old, Function &New,
+                             DenseMap<Value *, BasicBlock *> &NewEnds) {
   Function::iterator CurrBB, NextBB, FinalBB;
-  BasicBlock *NewEnd = nullptr;
   std::vector<Instruction *> DebugInsts;
   for (CurrBB = Old.begin(), FinalBB = Old.end(); CurrBB != FinalBB;
        CurrBB = NextBB) {
     NextBB = std::next(CurrBB);
     CurrBB->removeFromParent();
     CurrBB->insertInto(&New);
     Instruction *I = CurrBB->getTerminator();
-    if (isa<ReturnInst>(I))
-      NewEnd = &(*CurrBB);
+    if (ReturnInst *RI = dyn_cast<ReturnInst>(I)) {
+      Value *RetVal = RI->getReturnValue();
+      NewEnds.insert(std::make_pair(RetVal, &(*CurrBB)));
+    }
 
     for (Instruction &Val : *CurrBB) {
       // We must handle the scoping of called functions differently than
       // other outlined instructions.
       if (!isa<CallInst>(&Val)) {
         // Remove the debug information for outlined functions.
         Val.setDebugLoc(DebugLoc());
         continue;
for (Instruction &Val : *CurrBB) {
         Val.setDebugLoc(DI);
       }
     }
 
     for (Instruction *I : DebugInsts)
       I->eraseFromParent();
   }
 
-  assert(NewEnd && "No return instruction for new function?");
-  return NewEnd;
+  assert(NewEnds.size() > 0 && "No return instruction for new function?");
 }
 
 /// Find the the constants that will need to be lifted into arguments
 /// as they are not the same in each instance of the region.
 ///
 /// \param [in] C - The IRSimilarityCandidate containing the region we are
 /// analyzing.
 /// \param [in] NotSame - The set of global value numbers that do not have a

 
 // Within an extracted function, replace the argument uses of the extracted
 // region with the arguments of the function for an OutlinableGroup.
 //
 /// \param [in] Region - The region of extracted code to be changed.
 /// \param [in,out] OutputBB - The BasicBlock for the output stores for this
 /// region.
 static void replaceArgumentUses(OutlinableRegion &Region,
-                                BasicBlock *OutputBB) {
+                                DenseMap<Value *, BasicBlock *> &OutputBBs,
+                                bool FirstFunction=false) {
   OutlinableGroup &Group = *Region.Parent;
   assert(Region.ExtractedFunction && "Region has no extracted function?");
 
+  Function *DominatingFunction = Region.ExtractedFunction;
+  if (FirstFunction)
+    DominatingFunction = Group.OutlinedFunction;
+  DominatorTree DT(*DominatingFunction);
+
   for (unsigned ArgIdx = 0; ArgIdx < Region.ExtractedFunction->arg_size();
        ArgIdx++) {
     assert(Region.ExtractedArgToAgg.find(ArgIdx) !=
                Region.ExtractedArgToAgg.end() &&
            "No mapping from extracted to outlined?");
     unsigned AggArgIdx = Region.ExtractedArgToAgg.find(ArgIdx)->second;
     Argument *AggArg = Group.OutlinedFunction->getArg(AggArgIdx);
     Argument *Arg = Region.ExtractedFunction->getArg(ArgIdx);
for (unsigned ArgIdx = 0; ArgIdx < Region.ExtractedFunction->arg_size();
     // If we are replacing an output, we place the store value in its own
     // block inside the overall function before replacing the use of the output
     // in the function.
     assert(Arg->hasOneUse() && "Output argument can only have one use");
     User *InstAsUser = Arg->user_back();
     assert(InstAsUser && "User is nullptr!");
 
     Instruction *I = cast<Instruction>(InstAsUser);
-    I->setDebugLoc(DebugLoc());
+    BasicBlock *BB = I->getParent();
+    SmallVector<BasicBlock *, 4> Descendants;
+    DT.getDescendants(BB, Descendants);
+    bool EdgeAdded = false;
+    if (Descendants.size() == 0) {
+      EdgeAdded = true;
+      DT.insertEdge(&DominatingFunction->getEntryBlock(), BB);
+      DT.getDescendants(BB, Descendants);
+    }
+
+    // Iterate over the following blocks, looking for return instructions,
+    // if we find one, find the corresponding output block for the return value
+    // and move our store instruction there.
+    for (BasicBlock *DescendBB : Descendants) {
+      Instruction *Term = DescendBB->getTerminator();
+      if (!isa<ReturnInst>(Term))
+        continue;
+      ReturnInst *RI = cast<ReturnInst>(Term);
+      Value *RetVal = RI->getReturnValue();
+      DenseMap<Value *, BasicBlock *>::iterator VBBIt = OutputBBs.find(RetVal);
+      assert(VBBIt != OutputBBs.end() && "Could not find output value!");
+
+      Instruction *NewI = I->clone();
+      NewI->setDebugLoc(DebugLoc());
+      BasicBlock *OutputBB = VBBIt->second;
+      OutputBB->getInstList().push_back(NewI);
     LLVM_DEBUG(dbgs() << "Move store for instruction " << *I << " to "
                       << *OutputBB << "\n");
 
-    I->moveBefore(*OutputBB, OutputBB->end());
+    }
+
+    // If we added an edge for basic blocks without a predecessor, we remove it
+    // here.
+    if (EdgeAdded)
+      DT.deleteEdge(&DominatingFunction->getEntryBlock(), BB);
+    I->eraseFromParent();
 
     LLVM_DEBUG(dbgs() << "Replacing uses of output " << *Arg << " in function "
                       << *Region.ExtractedFunction << " with " << *AggArg
                       << " in function " << *Group.OutlinedFunction << "\n");
     Arg->replaceAllUsesWith(AggArg);
   }
 }
 

 }
 
 /// It is possible that there is a basic block that already performs the same
 /// stores. This returns a duplicate block, if it exists
 ///
 /// \param OutputBB [in] the block we are looking for a duplicate of.
 /// \param OutputStoreBBs [in] The existing output blocks.
 /// \returns an optional value with the number output block if there is a match.
-Optional<unsigned>
-findDuplicateOutputBlock(BasicBlock *OutputBB,
-                         ArrayRef<BasicBlock *> OutputStoreBBs) {
+Optional<unsigned> findDuplicateOutputBlock(
+    DenseMap<Value *, BasicBlock *> &OutputBBs,
+    std::vector<DenseMap<Value *, BasicBlock *>> &OutputStoreBBs) {
 
   bool WrongInst = false;
   bool WrongSize = false;
   unsigned MatchingNum = 0;
-  for (BasicBlock *CompBB : OutputStoreBBs) {
+  // We compare the new set output blocks to the other sets of outut blocks.
+  // If they are the same number, and have identical instructions, they are
+  // considered to be the same.
+  for (DenseMap<Value *, BasicBlock *> &CompBBs : OutputStoreBBs) {
     WrongInst = false;
+    WrongSize = false;
+    for (std::pair<Value *, BasicBlock *> &VToB : CompBBs) {
+      DenseMap<Value *, BasicBlock *>::iterator OutputBBIt =
+          OutputBBs.find(VToB.first);
+      if (OutputBBIt == OutputBBs.end()) {
+        WrongSize = true;
+        break;
+      }
+
+      BasicBlock *CompBB = VToB.second;
+      BasicBlock *OutputBB = OutputBBIt->second;
     if (CompBB->size() - 1 != OutputBB->size()) {
       WrongSize = true;
-      MatchingNum++;
-      continue;
+        break;
     }
 
     WrongSize = false;
     BasicBlock::iterator NIt = OutputBB->begin();
     for (Instruction &I : *CompBB) {
       if (isa<BranchInst>(&I))
         continue;
 
       if (!I.isIdenticalTo(&(*NIt))) {
         WrongInst = true;
         break;
       }
 
       NIt++;
     }
+    }
+
     if (!WrongInst && !WrongSize)
       return MatchingNum;
 
     MatchingNum++;
   }
 
   return None;
 }
 
 /// For the outlined section, move needed the StoreInsts for the output
 /// registers into their own block. Then, determine if there is a duplicate
 /// output block already created.
 ///
 /// \param [in] OG - The OutlinableGroup of regions to be outlined.
 /// \param [in] Region - The OutlinableRegion that is being analyzed.
 /// \param [in,out] OutputBB - the block that stores for this region will be
 /// placed in.
 /// \param [in] EndBB - the final block of the extracted function.
 /// \param [in] OutputMappings - OutputMappings the mapping of values that have
 /// been replaced by a new output value.
 /// \param [in,out] OutputStoreBBs - The existing output blocks.
-static void
-alignOutputBlockWithAggFunc(OutlinableGroup &OG, OutlinableRegion &Region,
-                            BasicBlock *OutputBB, BasicBlock *EndBB,
+static void alignOutputBlockWithAggFunc(
+    OutlinableGroup &OG, OutlinableRegion &Region,
+    DenseMap<Value *, BasicBlock *> OutputBBs,
+    DenseMap<Value *, BasicBlock *> EndBBs,
                             const DenseMap<Value *, Value *> &OutputMappings,
-                            std::vector<BasicBlock *> &OutputStoreBBs) {
+    std::vector<DenseMap<Value *, BasicBlock *>> &OutputStoreBBs) {
   DenseSet<unsigned> ValuesToFind(Region.GVNStores.begin(),
                                   Region.GVNStores.end());
 
   // We iterate over the instructions in the extracted function, and find the
   // global value number of the instructions.  If we find a value that should
   // be contained in a store, we replace the uses of the value with the value
   // from the overall function, so that the store is storing the correct
   // value from the overall function.
-  DenseSet<BasicBlock *> ExcludeBBs(OutputStoreBBs.begin(),
-                                    OutputStoreBBs.end());
-  ExcludeBBs.insert(OutputBB);
+  DenseSet<BasicBlock *> ExcludeBBs;
+  for (DenseMap<Value *, BasicBlock *> &BBMap : OutputStoreBBs)
+    for (std::pair<Value *, BasicBlock *> &VBPair : BBMap)
+      ExcludeBBs.insert(VBPair.second);
+  for (std::pair<Value *, BasicBlock *> &VBPair : OutputBBs)
+    ExcludeBBs.insert(VBPair.second);
   std::vector<Instruction *> ExtractedFunctionInsts =
       collectRelevantInstructions(*(Region.ExtractedFunction), ExcludeBBs);
   std::vector<Instruction *> OverallFunctionInsts =
       collectRelevantInstructions(*OG.OutlinedFunction, ExcludeBBs);
 
   assert(ExtractedFunctionInsts.size() == OverallFunctionInsts.size() &&
          "Number of relevant instructions not equal!");
 
for (unsigned Idx = 0; Idx < NumInstructions; Idx++) {
     if (ValuesToFind.size() == 0)
       break;
   }
 
   assert(ValuesToFind.size() == 0 && "Not all store values were handled!");
 
   // If the size of the block is 0, then there are no stores, and we do not
   // need to save this block.
-  if (OutputBB->size() == 0) {
+  bool AllRemoved = true;
+  Value *RetValueForBB;
+  BasicBlock *NewBB;
+  SmallVector<Value *, 4> ToRemove;
+  for (std::pair<Value *, BasicBlock *> &VtoBB : OutputBBs) {
+    RetValueForBB = VtoBB.first;
+    NewBB = VtoBB.second;
+  
+    if (NewBB->size() == 0) {
+      NewBB->eraseFromParent();
+      ToRemove.push_back(RetValueForBB);
+      continue;
+    }
+    
+    AllRemoved = false;
+  }
+
+  for (Value *V : ToRemove)
+    OutputBBs.erase(V);
+
+  if (AllRemoved) {
     Region.OutputBlockNum = -1;
-    OutputBB->eraseFromParent();
     return;
   }
 
-  // Determine is there is a duplicate block.
+  // Determine is there is a duplicate set of blocks.
   Optional<unsigned> MatchingBB =
-      findDuplicateOutputBlock(OutputBB, OutputStoreBBs);
+      findDuplicateOutputBlock(OutputBBs, OutputStoreBBs);
 
-  // If there is, we remove the new output block.  If it does not,
-  // we add it to our list of output blocks.
+  // If there is, we remove the new output blocks.  If it does not,
+  // we add it to our list of sets of output blocks.
   if (MatchingBB.hasValue()) {
     LLVM_DEBUG(dbgs() << "Set output block for region in function"
                       << Region.ExtractedFunction << " to "
                       << MatchingBB.getValue());
 
     Region.OutputBlockNum = MatchingBB.getValue();
-    OutputBB->eraseFromParent();
+    for (std::pair<Value *, BasicBlock *> &VtoBB : OutputBBs)
+      VtoBB.second->eraseFromParent();
     return;
   }
 
   Region.OutputBlockNum = OutputStoreBBs.size();
 
+  OutputStoreBBs.push_back(DenseMap<Value *, BasicBlock *>());
+  for (std::pair<Value *, BasicBlock *> &VtoBB : OutputBBs) {
+    RetValueForBB = VtoBB.first;
+    NewBB = VtoBB.second;
+    DenseMap<Value *, BasicBlock *>::iterator VBBIt =
+        EndBBs.find(RetValueForBB);
   LLVM_DEBUG(dbgs() << "Create output block for region in"
                     << Region.ExtractedFunction << " to "
-                    << *OutputBB);
-  OutputStoreBBs.push_back(OutputBB);
-  BranchInst::Create(EndBB, OutputBB);
+                      << *NewBB);
+    BranchInst::Create(VBBIt->second, NewBB);
+    OutputStoreBBs.back().insert(std::make_pair(RetValueForBB, NewBB));
+  }
 }
 
 /// Create the switch statement for outlined function to differentiate between
 /// all the output blocks.
 ///
 /// For the outlined section, determine if an outlined block already exists that
 /// matches the needed stores for the extracted section.
 /// \param [in] M - The module we are outlining from.
 /// \param [in] OG - The group of regions to be outlined.
 /// \param [in] EndBB - The final block of the extracted function.
 /// \param [in,out] OutputStoreBBs - The existing output blocks.
-void createSwitchStatement(Module &M, OutlinableGroup &OG, BasicBlock *EndBB,
-                           ArrayRef<BasicBlock *> OutputStoreBBs) {
+void createSwitchStatement(
+    Module &M, OutlinableGroup &OG, DenseMap<Value *, BasicBlock *> &EndBBs,
+    ArrayRef<DenseMap<Value *, BasicBlock *>> OutputStoreBBs) {
   // We only need the switch statement if there is more than one store
   // combination.
   if (OG.OutputGVNCombinations.size() > 1) {
     Function *AggFunc = OG.OutlinedFunction;
-    // Create a final block
-    BasicBlock *ReturnBlock =
-        BasicBlock::Create(M.getContext(), "final_block", AggFunc);
+    // Create a final block for each different return block.
+    unsigned FinalBlockIdx = 0;
+    for (std::pair<Value *, BasicBlock *> &OutputBlock : EndBBs) {
+      BasicBlock *EndBB = OutputBlock.second;
+      BasicBlock *ReturnBlock = BasicBlock::Create(
+          M.getContext(),
+          "final_block_" + Twine(static_cast<unsigned>(FinalBlockIdx++)),
+          AggFunc);
     Instruction *Term = EndBB->getTerminator();
     Term->moveBefore(*ReturnBlock, ReturnBlock->end());
-    // Put the switch statement in the old end basic block for the function with
-    // a fall through to the new return block
+      // Put the switch statement in the old end basic block for the function
+      // with a fall through to the new return block.
     LLVM_DEBUG(dbgs() << "Create switch statement in " << *AggFunc << " for "
                       << OutputStoreBBs.size() << "\n");
     SwitchInst *SwitchI =
         SwitchInst::Create(AggFunc->getArg(AggFunc->arg_size() - 1),
                            ReturnBlock, OutputStoreBBs.size(), EndBB);
 
     unsigned Idx = 0;
-    for (BasicBlock *BB : OutputStoreBBs) {
-      SwitchI->addCase(ConstantInt::get(Type::getInt32Ty(M.getContext()), Idx),
-                       BB);
+      for (DenseMap<Value *, BasicBlock *> OutputStoreBB : OutputStoreBBs) {
+        DenseMap<Value *, BasicBlock *>::iterator OSBBIt =
+            OutputStoreBB.find(OutputBlock.first);
+
+        if (OSBBIt == OutputStoreBB.end())
+          continue;
+        BasicBlock *BB = OSBBIt->second;
+        SwitchI->addCase(
+            ConstantInt::get(Type::getInt32Ty(M.getContext()), Idx), BB);
       Term = BB->getTerminator();
       Term->setSuccessor(0, ReturnBlock);
       Idx++;
       }
+    }
     return;
   }
 
-  // If there needs to be stores, move them from the output block to the end
-  // block to save on branching instructions.
+  // If there needs to be stores, move them from the output blocks to their
+  // corresponding ending block.
   if (OutputStoreBBs.size() == 1) {
     LLVM_DEBUG(dbgs() << "Move store instructions to the end block in "
                       << *OG.OutlinedFunction << "\n");
-    BasicBlock *OutputBlock = OutputStoreBBs[0];
-    Instruction *Term = OutputBlock->getTerminator();
+    DenseMap<Value *, BasicBlock *> OutputBlocks = OutputStoreBBs[0];
+    for (std::pair<Value *, BasicBlock *> &VBPair : OutputBlocks) {
+      DenseMap<Value *, BasicBlock *>::iterator EndBBIt = EndBBs.find(VBPair.first);
+      assert(EndBBIt != EndBBs.end() && "Could not find end block");
+      BasicBlock *EndBB = EndBBIt->second;
+      BasicBlock *OutputBB = VBPair.second;
+      Instruction *Term = OutputBB->getTerminator();
     Term->eraseFromParent();
     Term = EndBB->getTerminator();
-    moveBBContents(*OutputBlock, *EndBB);
+      moveBBContents(*OutputBB, *EndBB);
     Term->moveBefore(*EndBB, EndBB->end());
-    OutputBlock->eraseFromParent();
+      OutputBB->eraseFromParent();
+    }
   }
 }
 
 /// Fill the new function that will serve as the replacement function for all of
 /// the extracted regions of a certain structure from the first region in the
 /// list of regions.  Replace this first region's extracted function with the
 /// new overall function.
 ///
 /// \param [in] M - The module we are outlining from.
 /// \param [in] CurrentGroup - The group of regions to be outlined.
 /// \param [in,out] OutputStoreBBs - The output blocks for each different
 /// set of stores needed for the different functions.
 /// \param [in,out] FuncsToRemove - Extracted functions to erase from module
 /// once outlining is complete.
-static void fillOverallFunction(Module &M, OutlinableGroup &CurrentGroup,
-                                std::vector<BasicBlock *> &OutputStoreBBs,
-                                std::vector<Function *> &FuncsToRemove) {
+static void fillOverallFunction(
+    Module &M, OutlinableGroup &CurrentGroup,
+    std::vector<DenseMap<Value *, BasicBlock *>> &OutputStoreBBs,
   OutlinableRegion *CurrentOS = CurrentGroup.Regions[0];
 
   // Move first extracted function's instructions into new function.
   LLVM_DEBUG(dbgs() << "Move instructions from "
                     << *CurrentOS->ExtractedFunction << " to instruction "
                     << *CurrentGroup.OutlinedFunction << "\n");
 
-  CurrentGroup.EndBB = moveFunctionData(*CurrentOS->ExtractedFunction,
-                                        *CurrentGroup.OutlinedFunction);
+  moveFunctionData(*CurrentOS->ExtractedFunction,
+                   *CurrentGroup.OutlinedFunction, CurrentGroup.EndBBs);
 
   // Transfer the attributes from the function to the new function.
   for (Attribute A :
        CurrentOS->ExtractedFunction->getAttributes().getFnAttributes())
     CurrentGroup.OutlinedFunction->addFnAttr(A);
 
-  // Create an output block for the first extracted function.
+  // Create a new set of output blocks for the first extracted function.
+  DenseMap<Value *, BasicBlock *> NewBBs;
+  unsigned Idx = 0;
+  for (std::pair<Value *, BasicBlock *> &VtoBB : CurrentGroup.EndBBs) {
   BasicBlock *NewBB = BasicBlock::Create(
-      M.getContext(), Twine("output_block_") + Twine(static_cast<unsigned>(0)),
+        M.getContext(),
+        Twine("output_block_") + Twine(static_cast<unsigned>(0)) + Twine("_") +
+            Twine(static_cast<unsigned>(Idx++)),
       CurrentGroup.OutlinedFunction);
+    NewBBs.insert(std::make_pair(VtoBB.first, NewBB));
+  }
   CurrentOS->OutputBlockNum = 0;
 
-  replaceArgumentUses(*CurrentOS, NewBB);
+  replaceArgumentUses(*CurrentOS, NewBBs, true);
   replaceConstants(*CurrentOS);
 
-  // If the new basic block has no new stores, we can erase it from the module.
-  // It it does, we create a branch instruction to the last basic block from the
-  // new one.
+  // If a new basic block has no new stores, we can erase it from the module.
+  // If it does, we create a branch instruction to the basic block to the return
+  // block for the function.
+  BasicBlock *NewBB;
+  Value *RetValueForBB;
+  OutputStoreBBs.push_back(DenseMap<Value *, BasicBlock *>());
+  bool AllRemoved = true;
+  SmallVector<Value *, 4> ToRemove;
+  for (std::pair<Value *, BasicBlock *> &VtoBB : NewBBs) {
+    RetValueForBB = VtoBB.first;
+    NewBB = VtoBB.second;
+  
   if (NewBB->size() == 0) {
     CurrentOS->OutputBlockNum = -1;
     NewBB->eraseFromParent();
-  } else {
-    BranchInst::Create(CurrentGroup.EndBB, NewBB);
-    OutputStoreBBs.push_back(NewBB);
+      ToRemove.push_back(RetValueForBB);
+      continue;
+  }
+
+    AllRemoved = false;
+    DenseMap<Value *, BasicBlock *>::iterator VBBIt =
+        CurrentGroup.EndBBs.find(RetValueForBB);
+    BasicBlock *EndBB = VBBIt->second;
+    BranchInst::Create(EndBB, NewBB);
+    OutputStoreBBs.back().insert(std::make_pair(RetValueForBB, NewBB));
   }
 
+  if (AllRemoved)
+    CurrentOS->OutputBlockNum = -1;
+
   // Replace the call to the extracted function with the outlined function.
   CurrentOS->Call = replaceCalledFunction(M, *CurrentOS);
 
   // We only delete the extracted functions at the end since we may need to
   // reference instructions contained in them for mapping purposes.
   FuncsToRemove.push_back(CurrentOS->ExtractedFunction);
 }
 
 void IROutliner::deduplicateExtractedSections(
     Module &M, OutlinableGroup &CurrentGroup,
     std::vector<Function *> &FuncsToRemove, unsigned &OutlinedFunctionNum) {
   createFunction(M, CurrentGroup, OutlinedFunctionNum);
 
-  std::vector<BasicBlock *> OutputStoreBBs;
+  std::vector<DenseMap<Value *, BasicBlock *>> OutputStoreBBs;
 
   OutlinableRegion *CurrentOS;
 
   fillOverallFunction(M, CurrentGroup, OutputStoreBBs, FuncsToRemove);
 
   for (unsigned Idx = 1; Idx < CurrentGroup.Regions.size(); Idx++) {
     CurrentOS = CurrentGroup.Regions[Idx];
     AttributeFuncs::mergeAttributesForOutlining(*CurrentGroup.OutlinedFunction,
                                                *CurrentOS->ExtractedFunction);
 
-    // Create a new BasicBlock to hold the needed store instructions.
+    // Create a set of BasicBlocks, one for each return block, to hold the
+    // needed store instructions.
+    DenseMap<Value *, BasicBlock *> NewBBs;
+    unsigned BBIdx = 0;
+    for (std::pair<Value *, BasicBlock *> &VtoBB : CurrentGroup.EndBBs) {
     BasicBlock *NewBB = BasicBlock::Create(
-        M.getContext(), "output_block_" + std::to_string(Idx),
+          M.getContext(),
+          Twine("output_block_") + Twine(static_cast<unsigned>(Idx)) +
+              Twine("_") + Twine(static_cast<unsigned>(BBIdx++)),
         CurrentGroup.OutlinedFunction);
-    replaceArgumentUses(*CurrentOS, NewBB);
-
-    alignOutputBlockWithAggFunc(CurrentGroup, *CurrentOS, NewBB,
-                                CurrentGroup.EndBB, OutputMappings,
+      NewBBs.insert(std::make_pair(VtoBB.first, NewBB));
+    }
+    replaceArgumentUses(*CurrentOS, NewBBs);
+    alignOutputBlockWithAggFunc(CurrentGroup, *CurrentOS, NewBBs,
+                                CurrentGroup.EndBBs, OutputMappings,
                                 OutputStoreBBs);
 
     CurrentOS->Call = replaceCalledFunction(M, *CurrentOS);
     FuncsToRemove.push_back(CurrentOS->ExtractedFunction);
   }
 
   // Create a switch statement to handle the different output schemes.
-  createSwitchStatement(M, CurrentGroup, CurrentGroup.EndBB, OutputStoreBBs);
+  createSwitchStatement(M, CurrentGroup, CurrentGroup.EndBBs, OutputStoreBBs);
 
   OutlinedFunctionNum++;
 }
 
 bool IROutliner::compatibleWithOutlined(OutlinableRegion &Region) {
   IRSimilarityCandidate *IRSC = Region.Candidate;
   unsigned StartIdx = IRSC->getStartIdx();
   unsigned EndIdx = IRSC->getEndIdx();