[IROutliner] Allowing PHINodes in Exit Blocks

[AndrewLitteken]

Summary

In addition to having multiple exit locations, there can be multiple blocks leading to the same exit location, which results in a potential phi node. If we find that multiple blocks within the region branch to the same block outside the region, resulting in a phi node, the code extractor pulls this phi node into the function and uses it as an output.

We make sure that this phi node is given an output slot, and that the two values are removed from the outputs if they are not used anywhere else outside of the region. Across the extracted regions, the phi nodes are combined into a single block for each potential output block, similar to the previous patch.

[AndrewLitteken]

Updating to handle a global in the split phi node, and correcting which successor gets reassigned when creating a new aggregate phi block.

[AndrewLitteken]

Updating to clean up the patch, splitting into several patches, and adding more tests.

[AndrewLitteken]

llvm/lib/Transforms/IPO/IROutliner.cpp

940

clang-format

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

56

Why this number?

272

757

run on?

788

avoid double negative with variable name?

798

this comment is a bit of a run-on

823

if you rename NotOnlyUsedInExitBlocks to OnlyUsedInExitBlocks and change the code accordingly, you can just return that variable.

848

in for loops like this, it's more idiomatic to write

for (unsigned I = 0, E = PN.getNumIncomingValues(); I < E; ++I)
   ...

855

Using empty() would be more explicit here.

857

Don't need else since you would have continued already

883

Add comments for what these types represent?

920

Can this use

if (any_of(...))
  return None;

?

933

Maybe pull PN->getNumIncomingValues into a variable to avoid recalculating it across the two loops?

937

actually, can this loop be merged into the previous one?

943

what happens if BBGVN doesn't have a value?

1137

comments?

1187

this looks like it could be

bool GVNMatches = all_of(...);

or even better:

if (all_of(...))
   return CurrPN;

1926

1990

What happens if OVGN doesn't have a value?

[AndrewLitteken]

llvm/lib/Transforms/IPO/IROutliner.cpp

56

It has to count backwards from -3 (or rather the unsigned version) since -2, and -1 are reserved as tombstone and empty values for unsigned integers in the DenseMap. I'll add it as a comment.

937

yes, I think so. But I think enough values would need to be passed into the lambda for this that I don't think it makes sense to put it in any_of if that's the case.

943

It shouldn't since the block of a Phi node containing should be the target of some branch in the region, and will have been marked in that case, but I will add an assertion to make sure it exists.

1990

Right now, it will throw an assertion in the Optional code that there isn't a value.

[AndrewLitteken]

Updating based on initial comments.

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

59

Can we assert somewhere that -2 and -1 are actually used by DenseMap?

774

avoid recalculating loop invariant

779

probably don't need the variable here

781

probably don't need the variable here either

808

drop braces

845

we calculate IncomingVals.size()s twice here

how about

unsigned NumIncoming = IncomingVals.size();

if (NumIncoming == 0)
  continue;

if (NumIncoming == 1) {
   ...
}

848

comments here?

897

928

comment?

1151

can you do this like

if (!AdjustArgNo) {
// The call instruction was already removed when we created a call to an outlined function. So, we have all the arguments already. Just grab the argument.
IVal = CI->getArgOperand(ArgNum);
continue;
}

if (Region.AggArgToConstant.count(ArgNum)) {
  // The argument is/isn't a constant...
  // do stuff
  continue;
}

// Other case

1197

use phis()?

1214

avoid recalculating loop invariant

1220

getFirstNonPHI can return nullptr. assert?

1925

comment?

1935

why is there an assert below and not here?

1981

code duplication with change above?

[AndrewLitteken]

llvm/lib/Transforms/IPO/IROutliner.cpp

59

I'm not sure where to do that? They are representing the tombstone value, and empty key value, so you can't use them as keys for other items.

1151

There is a shared action between all three of the cases. I can clean up what the checks are, but I don't think it makes sense to use this pattern to continue and then duplicate the canonical number finding below in each case. These cases are to make sure we are getting the canonical number for the correct value, and then we can use the same code below to find the canonical number.

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

59

The MachineOutliner does:

InstructionMapper() {
  // Make sure that the implementation of DenseMapInfo<unsigned> hasn't
  // changed.
  assert(DenseMapInfo<unsigned>::getEmptyKey() == (unsigned)-1 &&
         "DenseMapInfo<unsigned>'s empty key isn't -1!");
  assert(DenseMapInfo<unsigned>::getTombstoneKey() == (unsigned)-2 &&
         "DenseMapInfo<unsigned>'s tombstone key isn't -2!");
}

[AndrewLitteken]

Updating based on comments

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

774

don't recalculate loop invariant

also can this just be an any_of?

785

this comment could be more succinct, like "Check if the value is used by any instructions outside of the region."

786

this could probably also be an any_of

793

This sentence is pretty long and kind of confusing. Can you simplify it?

798

can you explain this a bit more?

813

this is kind of wordy, maybe rewrite it like

"Test whether \p ExitBlock contains any PhiNodes that should be considered outputs. PhiNodes should be outputs when...(blah)"

818

the comments here don't line up with the variable names in the function

this could be a more descriptive name as well, like "PotentialExitsFromRegion"

819

this could be named "BasicBlocksInRegion"

822

I think this variable would be better-named "OutputsReplacedByPhis"

824

maybe this could have a more descriptive name too, like "OutputsWithNonPhiUses" or "OutputsThatCouldNotBeReplacedWithPhis"

(second option is kind of wordy, but it describes what's happening pretty well)

835

don't recalculate loop invariant

836

these could use more descriptive names

893

902

this could probably be called "BasicBlocksInRegion" or something

911

can this be named Cand or something?

914

since you don't use the index here, could this just be a loop like this?

for (Value *Incoming : PN->incoming_values()) {
   ...
}

917

comment is pretty wordy, split it up?

920

can you pull C.getGVN(V) into a variable, since you reuse it below?

921–922

pull Group.ArgumentTypes.size() into a variable and reuse it? it's calculated 4 times in this function.

927

is this guaranteed to have a value?

928

simpler?

930

can you add a comment explaining what you're going to do in this portion of the function here?

930

932

934

this comment has a lot of run-ons; can you split it up a little?

936

this comment is confusing. Is it two PhiNodes or one?

938

can you expand on the store sets a little more? I feel like it's kind of out of nowhere.

1066

simpler:

Find or create a BasicBlock in the outlined function containing PhiBlocks for \p RetVal.

1077

I think this could be called like "PhiBlockForRetVal"

1078

this could also use a more descriptive name

1102

• this can be a one-liner
• I think this should use cast?

1106

don't recalculate loop invariant

1118

simpler:

\returns The Value representing \p A at the call location for some outlined function.

1122

maybe this should be "the region corresponding to the outlined function" or something?

1123

or even simpler:

"True when \p A's location must be adjusted. This happens when... (stuff)."

I think this could also be named something like "HasAlreadyBeenReplacedWithOutlinedCall" or something?

1125

1158

this could use a better name (like above)

1161

formatting

1178

is this guaranteed to have a value?

1196

should this ever be null? can it be a reference?

1203

unfinished comment?

1206

this could use a more descriptive name or a comment

1222

PHINode *cloneImpl() const;

do you have to cast here? static_cast seems wrong too

1226

just call this IncomingVal?

1227

just call this IncomingBlock?

1231

call this FirstNonPhi?

1233

this could use a more descriptive name?

1916

this sentence is confusing

\returns The Value represented by a canonical number \p OutputCanon in \p Region.

I assume this is an output as well? Should this function be called "findOutputValueInRegion" or something?

[AndrewLitteken]

Updating based off of newest comments.

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

757

we can be even more concise here:

\returns true if \p V has any uses outside its region other than \p PN.

766

this is redundant, since you mention it in the first part of the comment.

774

Would it make sense to use a SmallVector here?

777

couple probably be a one-liner?

return (Idx != PhiLoc || V == PN.getIncomingValue(Idx) || !BlocksInRegion.contains(PN.getIncomingBlock(Idx)));

788

Return the any_of rather than using an if?

821

"it" is unclear IMO

also update comment for updated variable names?

1102

braces

1105

comment on the successors?

1207

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

1158

update comment to reflect variable name?

1916

1917

"one the first incoming value" doesn't parse for me

[AndrewLitteken]

Updating comment variables and few stylistic points.

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

283

This is being computed twice?

1094

1099

This comment would be more useful if it said something like "find the predecessors in order to (blah)"

1128

I think I'd prefer two functions here rather than one with a flag? It'd be clearer what's going on.

Something like

• getPassedArgumentAndAdjustArgumentLocation(Argument *A, OutlinableRegion &Region)
• getPassedArgumentInAlreadyOutlinedFunction(Argument *A, const OutlinableRegion &Region)

Also I think A can probably be const?

[AndrewLitteken]

Updating based on comments

[AndrewLitteken]

Updating diff with context

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

916

probably want a comment here explaining what's going on?

943

comment for this part?

1136

This can probably be a one-liner?

1160

Can probably eliminate the variable.

1266

Can probably just do this as a one-liner?

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

1085

1234

Missing /* ReplacedWithOutlinedCall = */ on bool param?

[AndrewLitteken]

Updating with one-liners and extra comments.

[paquette]

LGTM

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

struct OutlinableRegion {
   /// Marks whether we need to change the order of the arguments when mapping
   /// the old extracted function call to the new aggregate outlined function
   /// call.
   bool ChangedArgOrder = false;
 
   /// Marks whether this region ends in a branch, there is special handling
   /// required for the following basic blocks in this case.
   bool EndsInBranch = false;
+
+  BasicBlock* PHINodeAdjustBB = nullptr;
+
+  /// The PHIBlocks with their corresponding return block based on the return
+  /// value as the key.
+  DenseMap<Value *, BasicBlock *> PHIBlocks;
+
   /// Mapping of the argument number in the deduplicated function
   /// to a given constant, which is used when creating the arguments to the call
   /// to the newly created deduplicated function.  This is handled separately
   /// since the CodeExtractor does not recognize constants.
   DenseMap<unsigned, Constant *> AggArgToConstant;
 
   /// The global value numbers that are used as outputs for this section. Once
   /// extracted, each output will be stored to an output register.  This

llvm/lib/Transforms/IPO/IROutliner.cpp

 
 // This is a debug option to test small pieces of code to ensure that outlining
 // works correctly.
 static cl::opt<bool> NoCostModel(
     "ir-outlining-no-cost", cl::init(false), cl::ReallyHidden,
     cl::desc("Debug option to outline greedily, without restriction that "
              "calculated benefit outweighs cost"));
 
 /// Move the contents of \p SourceBB to before the last instruction of \p
 /// TargetBB.
 /// \param SourceBB - the BasicBlock to pull Instructions from.
 /// \param TargetBB - the BasicBlock to put Instruction into.
 static void moveBBContents(BasicBlock &SourceBB, BasicBlock &TargetBB) {
   BasicBlock::iterator BBCurr, BBEnd, BBNext;
   for (BBCurr = SourceBB.begin(), BBEnd = SourceBB.end(); BBCurr != BBEnd;
        BBCurr = BBNext) {
     BBNext = std::next(BBCurr);
     BBCurr->moveBefore(TargetBB, TargetBB.end());
   }
 }
for (Instruction &I : *BB) {
         break;
       }
     }
   }
 
   return Benefit;
 }
 
+static Value *findOutputMapping(const DenseMap<Value *, Value *> OutputMappings,
+                                Value *Input) {
+  if (OutputMappings.find(Input) != OutputMappings.end())
+    return OutputMappings.find(Input)->second;
+  return Input;
+}
+
 /// Find whether \p Region matches the global value numbering to Constant
 /// mapping found so far.
 ///
 /// \param Region - The OutlinableRegion we are checking for constants
 /// \param GVNToConstant - The mapping of global value number to Constants.
 /// \param NotSame - The set of global value numbers that do not have the same
 /// constant in each region.
 /// \returns true if all Constants are the same in every use of a Constant in \p
 /// Region and false if not
 static bool
 collectRegionsConstants(OutlinableRegion &Region,
                         DenseMap<unsigned, Constant *> &GVNToConstant,
                         DenseSet<unsigned> &NotSame) {
   bool ConstantsTheSame = true;
findExtractedInputToOverallInputMapping(OutlinableRegion &Region,
   if (!Group.InputTypesSet) {
     Group.NumAggregateInputs = TypeIndex;
     Group.InputTypesSet = true;
   }
 
   Region.NumExtractedInputs = OriginalIndex;
 }
 
 /// Create a mapping of the output arguments for the \p Region to the output
 /// arguments of the overall outlined function.
 ///
 /// \param [in,out] Region - The region of code to be analyzed.
 /// \param [in] Outputs - The values found by the code extractor.
 static void
 findExtractedOutputToOverallOutputMapping(OutlinableRegion &Region,
                                           SetVector<Value *> &Outputs) {
   OutlinableGroup &Group = *Region.Parent;
   IRSimilarityCandidate &C = *Region.Candidate;
 
   std::vector<BasicBlock *> BE;
   DenseSet<BasicBlock *> BBSet;
   C.getBasicBlocks(BBSet, BE);
 
   // Find the exits to the region.
   SmallPtrSet<BasicBlock *, 1> Exits;
   for (BasicBlock *Block : BE) {
     for (BasicBlock *Succ : successors(Block)) {
       if (!BBSet.contains(Succ)) {
         Exits.insert(Succ);
       }
     }
   }
 
   // After determining which blocks exit to PHINodes, we add these PHINodes to
   // the set of outputs to be processed.  We also check the incoming values of
   // the PHINodes for whether they should no longer be considered outputs.
   DenseSet<Value *> PHIWrapped;
   for (BasicBlock *ExitBB : Exits) {
     for (PHINode &PN : ExitBB->phis()) {
       // Find all incoming values from the outlining region.
       SmallVector<unsigned, 2> IncomingVals;
       for (unsigned i = 0; i < PN.getNumIncomingValues(); ++i)
         if (BBSet.contains(PN.getIncomingBlock(i)))
           IncomingVals.push_back(i);
 
       // Do not process PHI if there is one (or fewer) predecessor from region.
       if (IncomingVals.size() <= 1)
         continue;
-      else {
-        Region.IgnoreRegion = true;
-        return;
+
+      Outputs.insert(&PN);
+
+      // Not all of the incoming values should be ignored as other inputs and
+      // outputs into the outlined region.  If they have other uses outside of
+      // the single PHINode we should not skip over it.
+      for (unsigned i : IncomingVals) {
+        Value *V = PN.getIncomingValue(i);
+        bool NotOnlyUsedInExitBlocks = true;
+        for (User *U : V->users()) {
+          Instruction *I = dyn_cast<Instruction>(U);
+          if (!I)
+            continue;
+          
+          // If the use of the item is inside the region, we skip it since use
+          // inside the region does tell us anything about the uses outside
+          // non exit blocks outside the region.
+          BasicBlock *Parent = I->getParent();
+          if (BBSet.contains(Parent))
+            continue;
+
+          // If it's not a PHINode then we definitely know the use matters.
+          if (!isa<PHINode>(I)) {
+            NotOnlyUsedInExitBlocks = false;
+            break;
+          }
+
+          // If we have a PHINode outside one of the exit locations, then it
+          // can be considered an outside use as well.
+          if (!Exits.contains(Parent)) {
+            NotOnlyUsedInExitBlocks = false;
+            break;
+          }
+        }
+        if (!NotOnlyUsedInExitBlocks)
+          continue;
+        PHIWrapped.insert(PN.getIncomingValue(i));
       }
     }
   }
 
   // This counts the argument number in the extracted function.
   unsigned OriginalIndex = Region.NumExtractedInputs;
 
   // This counts the argument number in the overall function.
   unsigned TypeIndex = Group.NumAggregateInputs;
   bool TypeFound;
   DenseSet<unsigned> AggArgsUsed;
 
   // Iterate over the output types and identify if there is an aggregate pointer
   // type whose base type matches the current output type. If there is, we mark
   // that we will use this output register for this value. If not we add another
   // type to the overall argument type list. We also store the GVNs used for
   // stores to identify which values will need to be moved into an special
   // block that holds the stores to the output registers.
+  SmallVector<unsigned, 1> GlobalValues;
+  DenseSet<unsigned> GVNStored;
   for (Value *Output : Outputs) {
     TypeFound = false;
     // We can do this since it is a result value, and will have a number
     // that is necessarily the same. BUT if in the future, the instructions
     // do not have to be in same order, but are functionally the same, we will
     // have to use a different scheme, as one-to-one correspondence is not
     // guaranteed.
-    unsigned GlobalValue = C.getGVN(Output).getValue();
+    GlobalValues.clear();
+
+    // Since values outside the region can be combined into PHINode when we
+    // have multiple exists, we check if we have a PHINode in the output, to
+    // mark that both of these outputs are getting caught.
+
+    // If we cannot find a GVN, this means that the input to the PHINode is not
+    // included in the region we are trying to analyze, meaning, that if it was
+    // outlined, we would be adding an extra input.  We ignore this case for
+    // now, and so ignore the region.
+
+    // TODO: Adapt to the extra input from the PHINode.
+    PHINode *PN = dyn_cast<PHINode>(Output);
+    if (PN && !BBSet.contains(PN->getParent())) {
+      for (unsigned i = 0; i < PN->getNumIncomingValues(); ++i) {
+        Value *V = PN->getIncomingValue(i);
+        if (!C.getGVN(V).hasValue()) {
+          Region.IgnoreRegion = true;
+          return;
+        }
+        if (isa<Constant>(V))
+          continue;
+        Instruction *I = dyn_cast<Instruction>(V);
+        if (I && !BBSet.contains(I->getParent()))
+          continue;
+        GlobalValues.push_back(C.getGVN(V).getValue());
+      }
+    } else
+      GlobalValues.push_back(C.getGVN(Output).getValue());
     unsigned ArgumentSize = Group.ArgumentTypes.size();
 
+    // If the output is combined in a PHINode, we make sure to skip over it.
+    if (PHIWrapped.contains(Output))
+      continue;
     for (unsigned Jdx = TypeIndex; Jdx < ArgumentSize; Jdx++) {
       if (Group.ArgumentTypes[Jdx] != PointerType::getUnqual(Output->getType()))
         continue;
 
       if (AggArgsUsed.contains(Jdx))
         continue;
 
       TypeFound = true;
       AggArgsUsed.insert(Jdx);
       Region.ExtractedArgToAgg.insert(std::make_pair(OriginalIndex, Jdx));
       Region.AggArgToExtracted.insert(std::make_pair(Jdx, OriginalIndex));
-      Region.GVNStores.push_back(GlobalValue);
+      for (unsigned GlobalValue : GlobalValues)
+        GVNStored.insert(GlobalValue);
       break;
     }
 
     // We were unable to find an unused type in the output type set that matches
     // the output, so we add a pointer type to the argument types of the overall
     // function to handle this output and create a mapping to it.
     if (!TypeFound) {
       Group.ArgumentTypes.push_back(PointerType::getUnqual(Output->getType()));
       AggArgsUsed.insert(Group.ArgumentTypes.size() - 1);
       Region.ExtractedArgToAgg.insert(
           std::make_pair(OriginalIndex, Group.ArgumentTypes.size() - 1));
       Region.AggArgToExtracted.insert(
           std::make_pair(Group.ArgumentTypes.size() - 1, OriginalIndex));
-      Region.GVNStores.push_back(GlobalValue);
+      for (unsigned GlobalValue : GlobalValues)
+        GVNStored.insert(GlobalValue);
     }
 
-    stable_sort(Region.GVNStores);
     OriginalIndex++;
     TypeIndex++;
   }
+  for (unsigned GVN : GVNStored)
+    Region.GVNStores.push_back(GVN);
+  stable_sort(Region.GVNStores);
 }
 
 void IROutliner::findAddInputsOutputs(Module &M, OutlinableRegion &Region,
                                       DenseSet<unsigned> &NotSame) {
   std::vector<unsigned> Inputs;
   SetVector<Value *> ArgInputs, Outputs;
 
   getCodeExtractorArguments(Region, Inputs, NotSame, OutputMappings, ArgInputs,
                             Outputs);
 
   if (Region.IgnoreRegion)
     return;
 
   // Map the inputs found by the CodeExtractor to the arguments found for
   // the overall function.
   findExtractedInputToOverallInputMapping(Region, Inputs, ArgInputs);
 
   // Map the outputs found by the CodeExtractor to the arguments found for
   // the overall function.
   findExtractedOutputToOverallOutputMapping(Region, Outputs);
 }
CallInst *replaceCalledFunction(Module &M, OutlinableRegion &Region) {
   // argument.
   if (Group.SwiftErrorArgument.hasValue())
     Call->addParamAttr(Group.SwiftErrorArgument.getValue(),
                        Attribute::SwiftError);
 
   return Call;
 }
 
 // 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.
 /// \param [in] FirstFunction - A flag to indicate whether we are using this
 /// function to define the overall outlined function for all the regions, or
 /// if we are operating on one of the following regions.
 static void
 replaceArgumentUses(OutlinableRegion &Region,
                     DenseMap<Value *, BasicBlock *> &OutputBBs,
+                    const DenseMap<Value *, Value *> &OutputMappings,
                     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);
     // The argument is an input, so we can simply replace it with the overall
     // argument value
     if (ArgIdx < Region.NumExtractedInputs) {
       LLVM_DEBUG(dbgs() << "Replacing uses of input " << *Arg << " in function "
                         << *Region.ExtractedFunction << " with " << *AggArg
                         << " in function " << *Group.OutlinedFunction << "\n");
       Arg->replaceAllUsesWith(AggArg);
       continue;
     }
 
     // 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);
     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, PHIVBBIt;
       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");
+
+      // If this is storing a PHINode, we must make sure it is included in the
+      // overall function.
+      StoreInst *SI = static_cast<StoreInst *>(NewI);
+      if (!isa<PHINode>(SI->getValueOperand()))
+        continue;
+      PHINode *PN = cast<PHINode>(SI->getValueOperand());
+      // If it has a value, it was not split by the code extractor, which
+      // is what we are looking for.
+      if (Region.Candidate->getGVN(PN).hasValue())
+        continue;
+
+      // Find if a PHIBlock exists for this return value already.  If it is
+      // the first time we are anaylzing this, we will not, so we record it.
+      PHIVBBIt = Group.PHIBlocks.find(RetVal);
+      bool Inserted = false;
+      if (FirstFunction) {
+        BasicBlock *PHIBlock = PN->getParent();
+        std::tie(PHIVBBIt, Inserted) =
+            Group.PHIBlocks.insert(std::make_pair(RetVal, PHIBlock));
+      }
+
+      Region.PHIBlocks.insert(std::make_pair(RetVal, PN->getParent()));
+
+      // If we did not find a block, we create one, and insert it into the
+      // overall function and record it.
+      if (PHIVBBIt == Group.PHIBlocks.end()) {
+        BasicBlock *PHIBlock;
+        PHIBlock = BasicBlock::Create(I->getContext(), "phi_block",
+                                      OutputBB->getParent());
+        std::tie(PHIVBBIt, Inserted) =
+            Group.PHIBlocks.insert(std::make_pair(RetVal, PHIBlock));
+
+        // We find the predecessors in the overall function, and make sure
+        // we now branch to this new block.
+        for (BasicBlock *Pred : predecessors(OutputBB)) {
+          Instruction *Term = Pred->getTerminator();
+          BranchInst *BI = static_cast<BranchInst *>(Term);
+          for (unsigned Succ = 0; Succ < BI->getNumSuccessors(); Succ++) {
+            if (BI->getSuccessor(Succ) != OutputBB)
+              continue;
+            BI->setSuccessor(Succ, PHIBlock);
+          }
+        }
+
+        BranchInst::Create(OutputBB, PHIBlock);
+      }
+      BasicBlock *OverallPhiBlock = PHIVBBIt->second;
+
+      // If this is the first function, we do not need to worry about mergiing
+      // this with any other block in the overall outlined function
+      if (FirstFunction)
+        continue;
+
+      // For our PHINode, we find the combined canonical numbering, and
+      // attempt to find a matching PHINode in the overall PHIBlock.  If we
+      // cannot, we copy the PHINode and move it into this new block.
+      DenseSet<unsigned> PNGVNs;
+      for (unsigned Idx = 0; Idx < PN->getNumIncomingValues(); Idx++) {
+        Value *IVal = PN->getIncomingValue(Idx);
+        if (Argument *A = dyn_cast<Argument>(IVal)) {
+          CallInst *CI = cast<CallInst>(Region.ExtractedFunction->user_back());
+          unsigned ArgNum = A->getArgNo();
+          IVal = CI->getArgOperand(ArgNum);
+        }
+        IVal = findOutputMapping(OutputMappings, IVal);
+        Optional<unsigned> GVN = Region.Candidate->getGVN(IVal);
+        assert(GVN.hasValue() && "No GVN for incoming value");
+        Optional<unsigned> CanonNum = Region.Candidate->getCanonicalNum(*GVN);
+        PNGVNs.insert(*CanonNum);
+      }
+  
+      bool Found = false;
+      OutlinableRegion *First = Group.Regions[0];
+      for (Instruction &Inst : *OverallPhiBlock) {
+        if (!isa<PHINode>(Inst))
+          continue;
+        bool GVNMatches = true;
+        PHINode *CurrPN = static_cast<PHINode *>(&Inst);
+        for (unsigned Idx = 0; Idx < CurrPN->getNumIncomingValues(); Idx++) {
+          Value *IVal = CurrPN->getIncomingValue(Idx);
+          OutlinableRegion *First = Group.Regions[0];
+          if (Argument *A = dyn_cast<Argument>(IVal)) {
+            CallInst *CI = cast<CallInst>(First->Call);
+            unsigned ArgNum = A->getArgNo();
+            IVal = CI->getArgOperand(ArgNum);
+          }
+          IVal = findOutputMapping(OutputMappings, IVal);
+          Optional<unsigned> GVN = First->Candidate->getGVN(IVal);
+          assert(GVN.hasValue() && "No GVN for incoming value");
+          Optional<unsigned> CanonNum = First->Candidate->getCanonicalNum(*GVN);
+          if (!PNGVNs.contains(*CanonNum)) {
+            GVNMatches = false;
+            break;
+          }
+        }
+        if (GVNMatches) {
+          PN->replaceAllUsesWith(CurrPN);
+          Found = true;
+          break;
+        }
+      }
+
+      if (!Found) {
+        PHINode *NewPN = static_cast<PHINode *>(PN->clone());
+        NewPN->insertBefore(&*OverallPhiBlock->begin());
+        for (unsigned Idx = 0; Idx < NewPN->getNumIncomingValues(); Idx++) {
+          Value *IVal = NewPN->getIncomingValue(Idx);
+          BasicBlock *IBlock = NewPN->getIncomingBlock(Idx);
+          if (Argument *A = dyn_cast<Argument>(IVal)) {
+            unsigned ArgNum = A->getArgNo();
+            BasicBlock *In = NewPN->getIncomingBlock(Idx);
+            Instruction *Inst = In->getFirstNonPHI();
+            Optional<unsigned> GVN = Region.Candidate->getGVN(Inst);
+            assert(GVN.hasValue() && "No GVN for incoming value");
+            Optional<unsigned> CanonNum =
+                Region.Candidate->getCanonicalNum(*GVN);
+            Optional<unsigned> FirstGVN =
+                First->Candidate->fromCanonicalNum(*CanonNum);
+            Optional<Value *> TempVal = First->Candidate->fromGVN(*FirstGVN);
+            Value *Val = Group.OutlinedFunction->getArg(ArgNum);
+            NewPN->setIncomingValue(Idx, Val);
+            NewPN->setIncomingBlock(Idx,
+                                    cast<Instruction>(*TempVal)->getParent());
+            continue;
+          }
+          
+          IVal = findOutputMapping(OutputMappings, IVal);
+          Optional<unsigned> GVN = Region.Candidate->getGVN(IVal);
+          assert(GVN.hasValue() && "No GVN for incoming value");
+          Optional<unsigned> CanonNum = Region.Candidate->getCanonicalNum(*GVN);
+          Optional<unsigned> FirstGVN =
+              First->Candidate->fromCanonicalNum(*CanonNum);
+          Optional<Value *> Val = First->Candidate->fromGVN(*FirstGVN);
+          NewPN->setIncomingValue(Idx, *Val);
+          if (!isa<Instruction>(*Val)) {
+            Instruction *FirstI = &IBlock->front();
+            Optional<unsigned> GVN2 = Region.Candidate->getGVN(FirstI);
+            assert(GVN2.hasValue() && "No GVN for incoming value");
+            Optional<unsigned> CanonNum2 = Region.Candidate->getCanonicalNum(*GVN2);
+            Optional<unsigned> FirstGVN2 =
+                First->Candidate->fromCanonicalNum(*CanonNum2);
+            Val = First->Candidate->fromGVN(*FirstGVN2);
+          }
+          NewPN->setIncomingBlock(Idx, cast<Instruction>(*Val)->getParent());
+        }
+        PN->replaceAllUsesWith(NewPN);
+      }
     }
 
     // If we added an edge for basic blocks without a predecessor, we remove it
     // here.
     if (EdgeAdded)
       DT.deleteEdge(&DominatingFunction->getEntryBlock(), BB);
     I->eraseFromParent();
 
static void alignOutputBlockWithAggFunc(
   // from the overall function, so that the store is storing the correct
   // value from the overall function.
   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);
+  // Add PHIBlocks to exclusionary zones since it is possible they may not be
+  // same across the extracted functions.
+  for (std::pair<Value *, BasicBlock *> &VBPair : Region.PHIBlocks)
+    ExcludeBBs.insert(VBPair.second);
+  for (std::pair<Value *, BasicBlock *> &VBPair : Region.Parent->PHIBlocks)
+    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!");
 
   unsigned NumInstructions = ExtractedFunctionInsts.size();
   for (unsigned Idx = 0; Idx < NumInstructions; Idx++) {
     Value *V = ExtractedFunctionInsts[Idx];
 
     if (OutputMappings.find(V) != OutputMappings.end())
       V = OutputMappings.find(V)->second;
     Optional<unsigned> GVN = Region.Candidate->getGVN(V);
 
     // If we have found one of the stored values for output, replace the value
     // with the corresponding one from the overall function.
     if (GVN.hasValue() && ValuesToFind.erase(GVN.getValue())) {
+      Value *NewV = OverallFunctionInsts[Idx];
+      PHINode *OverallPN = nullptr;
+      // If there is a PHINode following the instruction, it means it is not
+      // really an output, but the following PHINode is, we make sure to update
+      // the uses of the PHINode as well to ensure that our output blocks in the
+      // overall function are consistent.
+      for (User *U : NewV->users())
+        if (PHINode *PN = dyn_cast<PHINode>(U))
+          OverallPN = PN;
+
+      for (User *U : V->users()) {
+        if (PHINode *PN = dyn_cast<PHINode>(U)) {
+          if (!OverallPN) {
+            BasicBlock *Old = static_cast<Instruction *>(V)->getParent();
+            BasicBlock *New = static_cast<Instruction *>(NewV)->getParent();
+            PN->replaceIncomingBlockWith(Old, New);
+            continue;
+          }
+          PN->replaceAllUsesWith(OverallPN);
+        }
+      }
+
       V->replaceAllUsesWith(OverallFunctionInsts[Idx]);
       if (ValuesToFind.size() == 0)
         break;
     }
 
     if (ValuesToFind.size() == 0)
       break;
   }
 
+  CallInst *CI = cast<CallInst>(Region.ExtractedFunction->user_back());
+  for (Value *V : CI->args()) {
+    Optional<unsigned> GVN = Region.Candidate->getGVN(V);
+    if (GVN.hasValue())
+      ValuesToFind.erase(*GVN);
+  }
+
   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.
   bool AllRemoved = true;
   Value *RetValueForBB;
   BasicBlock *NewBB;
   SmallVector<Value *, 4> ToRemove;

 /// 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.
+/// \param [in] OutputMappings - Extracted functions to erase from module
+/// once outlining is complete.
 static void fillOverallFunction(
     Module &M, OutlinableGroup &CurrentGroup,
     std::vector<DenseMap<Value *, BasicBlock *>> &OutputStoreBBs,
-    std::vector<Function *> &FuncsToRemove) {
+    std::vector<Function *> &FuncsToRemove,
+    const DenseMap<Value *, Value *> &OutputMappings) {
   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");
   moveFunctionData(*CurrentOS->ExtractedFunction,
                    *CurrentGroup.OutlinedFunction, CurrentGroup.EndBBs);
BasicBlock *NewBB = BasicBlock::Create(
         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, NewBBs, true);
+  replaceArgumentUses(*CurrentOS, NewBBs, OutputMappings, true);
   replaceConstants(*CurrentOS);
 
   // 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 *>());
void IROutliner::deduplicateExtractedSections(
     Module &M, OutlinableGroup &CurrentGroup,
     std::vector<Function *> &FuncsToRemove, unsigned &OutlinedFunctionNum) {
   createFunction(M, CurrentGroup, OutlinedFunctionNum);
 
   std::vector<DenseMap<Value *, BasicBlock *>> OutputStoreBBs;
 
   OutlinableRegion *CurrentOS;
 
-  fillOverallFunction(M, CurrentGroup, OutputStoreBBs, FuncsToRemove);
+  fillOverallFunction(M, CurrentGroup, OutputStoreBBs, FuncsToRemove,
+                      OutputMappings);
 
   for (unsigned Idx = 1; Idx < CurrentGroup.Regions.size(); Idx++) {
     CurrentOS = CurrentGroup.Regions[Idx];
     AttributeFuncs::mergeAttributesForOutlining(*CurrentGroup.OutlinedFunction,
                                                *CurrentOS->ExtractedFunction);
 
     // 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(),
           Twine("output_block_") + Twine(static_cast<unsigned>(Idx)) +
               Twine("_") + Twine(static_cast<unsigned>(BBIdx++)),
           CurrentGroup.OutlinedFunction);
       NewBBs.insert(std::make_pair(VtoBB.first, NewBB));
     }
-    replaceArgumentUses(*CurrentOS, NewBBs);
+    replaceArgumentUses(*CurrentOS, NewBBs, OutputMappings);
     alignOutputBlockWithAggFunc(CurrentGroup, *CurrentOS, NewBBs,
                                 CurrentGroup.EndBBs, OutputMappings,
                                 OutputStoreBBs);
 
     CurrentOS->Call = replaceCalledFunction(M, *CurrentOS);
     FuncsToRemove.push_back(CurrentOS->ExtractedFunction);
   }
 
for (OutlinableRegion *Region : CurrentGroup.Regions) {
     RegionBenefit += Region->getBenefit(TTI);
     LLVM_DEBUG(dbgs() << "Adding: " << RegionBenefit
                       << " saved instructions to overfall benefit.\n");
   }
 
   return RegionBenefit;
 }
 
 InstructionCost
 IROutliner::findCostOutputReloads(OutlinableGroup &CurrentGroup) {
   InstructionCost OverallCost = 0;
   for (OutlinableRegion *Region : CurrentGroup.Regions) {
     TargetTransformInfo &TTI = getTTI(*Region->StartBB->getParent());
 
     // Each output incurs a load after the call, so we add that to the cost.
     for (unsigned OutputGVN : Region->GVNStores) {
       Optional<Value *> OV = Region->Candidate->fromGVN(OutputGVN);
       assert(OV.hasValue() && "Could not find value for GVN?");
       Value *V = OV.getValue();
       InstructionCost LoadCost =
           TTI.getMemoryOpCost(Instruction::Load, V->getType(), Align(1), 0,
                               TargetTransformInfo::TCK_CodeSize);
 
       LLVM_DEBUG(dbgs() << "Adding: " << LoadCost
                         << " instructions to cost for output of type "
                         << *V->getType() << "\n");
       OverallCost += LoadCost;
     }
   }
 
   return OverallCost;
 }
 
 /// Find the extra instructions needed to handle any output values for the
 /// region.
 ///
if (NumOutputBranches == 0)
     NumOutputBranches++;
 
   CurrentGroup.BranchesToOutside = NumOutputBranches;
 
   for (const ArrayRef<unsigned> &OutputUse :
        CurrentGroup.OutputGVNCombinations) {
     for (unsigned GVN : OutputUse) {
       Optional<Value *> OV = Candidate.fromGVN(GVN);
       assert(OV.hasValue() && "Could not find value for GVN?");
       Value *V = OV.getValue();
       InstructionCost StoreCost =
           TTI.getMemoryOpCost(Instruction::Load, V->getType(), Align(1), 0,
                               TargetTransformInfo::TCK_CodeSize);
 
       // An instruction cost is added for each store set that needs to occur for
       // various output combinations inside the function, plus a branch to
       // return to the exit block.
       LLVM_DEBUG(dbgs() << "Adding: " << StoreCost
                         << " instructions to cost for output of type "
                         << *V->getType() << "\n");
       OutputCost += StoreCost * NumOutputBranches;
     }
 
     InstructionCost BranchCost =
         TTI.getCFInstrCost(Instruction::Br, TargetTransformInfo::TCK_CodeSize);
     LLVM_DEBUG(dbgs() << "Adding " << BranchCost << " to the current cost for"

llvm/test/Transforms/IROutliner/outlining-branches-phi-nodes.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs
+; RUN: opt -S -verify -ir-sim-branches -iroutliner -ir-outlining-no-cost < %s | FileCheck %s
+
+; Here we have multiple exits, but the different sources, same outputs are
+; needed, this checks that they are compressed, and moved into the appropriate
+; output blocks.
+
+define void @outline_outputs1() #0 {
+entry:
+  %output = alloca i32, align 4
+  %result = alloca i32, align 4
+  %output2 = alloca i32, align 4
+  %result2 = alloca i32, align 4
+  %a = alloca i32, align 4
+  %b = alloca i32, align 4
+  br label %block_2
+block_1:
+  %a2 = alloca i32, align 4
+  %b2 = alloca i32, align 4
+  br label %block_2
+block_2:
+  %a2val = load i32, i32* %a
+  %b2val = load i32, i32* %b
+  %add2 = add i32 2, %a2val
+  %mul2 = mul i32 2, %b2val
+  br label %block_5
+block_3:
+  %aval = load i32, i32* %a
+  %bval = load i32, i32* %b
+  %add = add i32 2, %aval
+  %mul = mul i32 2, %bval
+  br label %block_4
+block_4:
+  store i32 %add, i32* %output, align 4
+  store i32 %mul, i32* %result, align 4
+  br label %block_6
+block_5:
+  store i32 %add2, i32* %output, align 4
+  store i32 %mul2, i32* %result, align 4
+  br label %block_6
+block_6:
+  %diff = phi i32 [%aval, %block_4], [%a2val, %block_5]
+  ret void
+}
+
+define void @outline_outputs2() #0 {
+entry:
+  %output = alloca i32, align 4
+  %result = alloca i32, align 4
+  %output2 = alloca i32, align 4
+  %result2 = alloca i32, align 4
+  %a = alloca i32, align 4
+  %b = alloca i32, align 4
+  br label %block_2
+block_1:
+  %a2 = alloca i32, align 4
+  %b2 = alloca i32, align 4
+  br label %block_2
+block_2:
+  %a2val = load i32, i32* %a
+  %b2val = load i32, i32* %b
+  %add2 = add i32 2, %a2val
+  %mul2 = mul i32 2, %b2val
+  br label %block_5
+block_3:
+  %aval = load i32, i32* %a
+  %bval = load i32, i32* %b
+  %add = add i32 2, %aval
+  %mul = mul i32 2, %bval
+  br label %block_4
+block_4:
+  store i32 %add, i32* %output, align 4
+  store i32 %mul, i32* %result, align 4
+  br label %block_6
+block_5:
+  store i32 %add2, i32* %output, align 4
+  store i32 %mul2, i32* %result, align 4
+  br label %block_6
+block_6:
+  %diff = phi i32 [%aval, %block_4], [%a2val, %block_5]
+  ret void
+}
+
+; CHECK-LABEL: @outline_outputs1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DIFF_CE_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[OUTPUT:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[RESULT:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[OUTPUT2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[RESULT2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    br label [[BLOCK_2:%.*]]
+; CHECK:       block_1:
+; CHECK-NEXT:    [[A2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    br label [[BLOCK_2]]
+; CHECK:       block_2:
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[DIFF_CE_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32* [[A]], i32* [[B]], i32* [[OUTPUT]], i32* [[RESULT]], i32* [[DIFF_CE_LOC]])
+; CHECK-NEXT:    [[DIFF_CE_RELOAD:%.*]] = load i32, i32* [[DIFF_CE_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    br label [[BLOCK_6:%.*]]
+; CHECK:       block_6:
+; CHECK-NEXT:    [[DIFF:%.*]] = phi i32 [ [[DIFF_CE_RELOAD]], [[BLOCK_2]] ]
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @outline_outputs2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DIFF_CE_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[OUTPUT:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[RESULT:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[OUTPUT2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[RESULT2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    br label [[BLOCK_2:%.*]]
+; CHECK:       block_1:
+; CHECK-NEXT:    [[A2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    br label [[BLOCK_2]]
+; CHECK:       block_2:
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[DIFF_CE_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32* [[A]], i32* [[B]], i32* [[OUTPUT]], i32* [[RESULT]], i32* [[DIFF_CE_LOC]])
+; CHECK-NEXT:    [[DIFF_CE_RELOAD:%.*]] = load i32, i32* [[DIFF_CE_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    br label [[BLOCK_6:%.*]]
+; CHECK:       block_6:
+; CHECK-NEXT:    [[DIFF:%.*]] = phi i32 [ [[DIFF_CE_RELOAD]], [[BLOCK_2]] ]
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK:  define internal void @outlined_ir_func_0(
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label [[BLOCK_2_TO_OUTLINE:%.*]]
+; CHECK:       block_6.exitStub:
+; CHECK-NEXT:    store i32 [[DIFF_CE:%.*]], i32* [[TMP4:%.*]], align 4
+; CHECK-NEXT:    ret void
+; CHECK:       block_2_to_outline:
+; CHECK-NEXT:    [[A2VAL:%.*]] = load i32, i32* [[TMP0:%.*]], align 4
+; CHECK-NEXT:    [[B2VAL:%.*]] = load i32, i32* [[TMP1:%.*]], align 4
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 2, [[A2VAL]]
+; CHECK-NEXT:    [[MUL2:%.*]] = mul i32 2, [[B2VAL]]
+; CHECK-NEXT:    br label [[BLOCK_5:%.*]]
+; CHECK:       block_3:
+; CHECK-NEXT:    [[AVAL:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    [[BVAL:%.*]] = load i32, i32* [[TMP1]], align 4
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 2, [[AVAL]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i32 2, [[BVAL]]
+; CHECK-NEXT:    br label [[BLOCK_4:%.*]]
+; CHECK:       block_4:
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[TMP2:%.*]], align 4
+; CHECK-NEXT:    store i32 [[MUL]], i32* [[TMP3:%.*]], align 4
+; CHECK-NEXT:    br label [[BLOCK_6_SPLIT:%.*]]
+; CHECK:       block_5:
+; CHECK-NEXT:    store i32 [[ADD2]], i32* [[TMP2]], align 4
+; CHECK-NEXT:    store i32 [[MUL2]], i32* [[TMP3]], align 4
+; CHECK-NEXT:    br label [[BLOCK_6_SPLIT]]
+; CHECK:       block_6.split:
+; CHECK-NEXT:    [[DIFF_CE]] = phi i32 [ [[AVAL]], [[BLOCK_4]] ], [ [[A2VAL]], [[BLOCK_5]] ]
+; CHECK-NEXT:    br label [[BLOCK_6_EXITSTUB:%.*]]
+;

llvm/test/Transforms/IROutliner/outlining-multiple-exits-diff-outputs.ll

 ; CHECK-NEXT:    [[DIFF:%.*]] = sub i32 [[A2VAL_RELOAD]], [[B2VAL_RELOAD]]
 ; CHECK-NEXT:    ret void
 ; CHECK:       block_7:
 ; CHECK-NEXT:    [[QUOT:%.*]] = udiv i32 [[ADD_RELOAD]], [[MUL_RELOAD]]
 ; CHECK-NEXT:    ret void
 ;
 ;
 ; CHECK: define internal i1 @outlined_ir_func_0(
-; CHECK:  newFuncRoot:
+; CHECK-NEXT:  newFuncRoot:
 ; CHECK-NEXT:    br label [[BLOCK_2_TO_OUTLINE:%.*]]
 ; CHECK:       block_6.exitStub:
-; CHECK-NEXT:    switch i32 [[TMP8:%.*]], label [[FINAL_BLOCK_1:%.*]] [
-; CHECK-NEXT:    i32 0, label [[OUTPUT_BLOCK_0_1:%.*]]
-; CHECK-NEXT:    i32 1, label [[OUTPUT_BLOCK_1_1:%.*]]
-; CHECK-NEXT:    ]
-; CHECK:       block_7.exitStub:
-; CHECK-NEXT:    switch i32 [[TMP8]], label [[FINAL_BLOCK_0:%.*]] [
+; CHECK-NEXT:    switch i32 [[TMP8:%.*]], label [[FINAL_BLOCK_0:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[OUTPUT_BLOCK_0_0:%.*]]
 ; CHECK-NEXT:    i32 1, label [[OUTPUT_BLOCK_1_0:%.*]]
 ; CHECK-NEXT:    ]
+; CHECK:       block_7.exitStub:
+; CHECK-NEXT:    switch i32 [[TMP8]], label [[FINAL_BLOCK_1:%.*]] [
+; CHECK-NEXT:    i32 0, label [[OUTPUT_BLOCK_0_1:%.*]]
+; CHECK-NEXT:    i32 1, label [[OUTPUT_BLOCK_1_1:%.*]]
+; CHECK-NEXT:    ]
 ; CHECK:       block_2_to_outline:
 ; CHECK-NEXT:    [[A2VAL:%.*]] = load i32, i32* [[TMP0:%.*]], align 4
 ; CHECK-NEXT:    [[B2VAL:%.*]] = load i32, i32* [[TMP1:%.*]], align 4
 ; CHECK-NEXT:    [[ADD2:%.*]] = add i32 2, [[A2VAL]]
 ; CHECK-NEXT:    [[MUL2:%.*]] = mul i32 2, [[B2VAL]]
 ; CHECK-NEXT:    br label [[BLOCK_5:%.*]]
 ; CHECK:       block_3:
 ; CHECK-NEXT:    [[AVAL:%.*]] = load i32, i32* [[TMP0]], align 4
 ; CHECK-NEXT:    [[BVAL:%.*]] = load i32, i32* [[TMP1]], align 4
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 2, [[AVAL]]
 ; CHECK-NEXT:    [[MUL:%.*]] = mul i32 2, [[BVAL]]
 ; CHECK-NEXT:    br label [[BLOCK_4:%.*]]
 ; CHECK:       block_4:
 ; CHECK-NEXT:    store i32 [[ADD]], i32* [[TMP2:%.*]], align 4
 ; CHECK-NEXT:    store i32 [[MUL]], i32* [[TMP3:%.*]], align 4
 ; CHECK-NEXT:    br label [[BLOCK_6_EXITSTUB:%.*]]
 ; CHECK:       block_5:
 ; CHECK-NEXT:    store i32 [[ADD2]], i32* [[TMP2]], align 4
 ; CHECK-NEXT:    store i32 [[MUL2]], i32* [[TMP3]], align 4
 ; CHECK-NEXT:    br label [[BLOCK_7_EXITSTUB:%.*]]
 ; CHECK:       output_block_0_0:
-; CHECK-NEXT:    store i32 [[A2VAL]], i32* [[TMP4:%.*]], align 4
-; CHECK-NEXT:    store i32 [[B2VAL]], i32* [[TMP5:%.*]], align 4
-; CHECK-NEXT:    br label [[FINAL_BLOCK_0]]
-; CHECK:       output_block_0_1:
 ; CHECK-NEXT:    store i32 [[AVAL]], i32* [[TMP6:%.*]], align 4
 ; CHECK-NEXT:    store i32 [[BVAL]], i32* [[TMP7:%.*]], align 4
+; CHECK-NEXT:    br label [[FINAL_BLOCK_0]]
+; CHECK:       output_block_0_1:
+; CHECK-NEXT:    store i32 [[A2VAL]], i32* [[TMP4:%.*]], align 4
+; CHECK-NEXT:    store i32 [[B2VAL]], i32* [[TMP5:%.*]], align 4
 ; CHECK-NEXT:    br label [[FINAL_BLOCK_1]]
 ; CHECK:       output_block_1_0:
-; CHECK-NEXT:    store i32 [[A2VAL]], i32* [[TMP4]], align 4
-; CHECK-NEXT:    store i32 [[B2VAL]], i32* [[TMP5]], align 4
-; CHECK-NEXT:    br label [[FINAL_BLOCK_0]]
-; CHECK:       output_block_1_1:
 ; CHECK-NEXT:    store i32 [[ADD]], i32* [[TMP6]], align 4
 ; CHECK-NEXT:    store i32 [[MUL]], i32* [[TMP7]], align 4
+; CHECK-NEXT:    br label [[FINAL_BLOCK_0]]
+; CHECK:       output_block_1_1:
+; CHECK-NEXT:    store i32 [[A2VAL]], i32* [[TMP4]], align 4
+; CHECK-NEXT:    store i32 [[B2VAL]], i32* [[TMP5]], align 4
 ; CHECK-NEXT:    br label [[FINAL_BLOCK_1]]
 ; CHECK:       final_block_0:
-; CHECK-NEXT:    ret i1 false
-; CHECK:       final_block_1:
 ; CHECK-NEXT:    ret i1 true
+; CHECK:       final_block_1:
+; CHECK-NEXT:    ret i1 false
 ;