Index: include/llvm/Transforms/Utils/CodeExtractor.h =================================================================== --- include/llvm/Transforms/Utils/CodeExtractor.h +++ include/llvm/Transforms/Utils/CodeExtractor.h @@ -124,7 +124,7 @@ private: void severSplitPHINodes(BasicBlock *&Header); void splitReturnBlocks(); - + BasicBlock *splitExitBlock(BasicBlock *Exit); Function *constructFunction(const ValueSet &inputs, const ValueSet &outputs, BasicBlock *header, Index: lib/Transforms/IPO/PartialInlining.cpp =================================================================== --- lib/Transforms/IPO/PartialInlining.cpp +++ lib/Transforms/IPO/PartialInlining.cpp @@ -97,33 +97,6 @@ // use the inliner functionality when we're done hacking. F->replaceAllUsesWith(DuplicateFunction); - // Special hackery is needed with PHI nodes that have inputs from more than - // one extracted block. For simplicity, just split the PHIs into a two-level - // sequence of PHIs, some of which will go in the extracted region, and some - // of which will go outside. - BasicBlock *PreReturn = NewReturnBlock; - NewReturnBlock = NewReturnBlock->splitBasicBlock( - NewReturnBlock->getFirstNonPHI()->getIterator()); - BasicBlock::iterator I = PreReturn->begin(); - Instruction *Ins = &NewReturnBlock->front(); - while (I != PreReturn->end()) { - PHINode *OldPhi = dyn_cast(I); - if (!OldPhi) - break; - - PHINode *RetPhi = PHINode::Create(OldPhi->getType(), 2, "", Ins); - OldPhi->replaceAllUsesWith(RetPhi); - Ins = NewReturnBlock->getFirstNonPHI(); - - RetPhi->addIncoming(&*I, PreReturn); - RetPhi->addIncoming(OldPhi->getIncomingValueForBlock(NewEntryBlock), - NewEntryBlock); - OldPhi->removeIncomingValue(NewEntryBlock); - - ++I; - } - NewEntryBlock->getTerminator()->replaceUsesOfWith(PreReturn, NewReturnBlock); - // Gather up the blocks that we're going to extract. std::vector ToExtract; ToExtract.push_back(NewNonReturnBlock); Index: lib/Transforms/Utils/CodeExtractor.cpp =================================================================== --- lib/Transforms/Utils/CodeExtractor.cpp +++ lib/Transforms/Utils/CodeExtractor.cpp @@ -41,6 +41,7 @@ #include "llvm/Transforms/Utils/BasicBlockUtils.h" #include #include +#include using namespace llvm; #define DEBUG_TYPE "code-extractor" @@ -169,6 +170,21 @@ return false; } +/// movePredsTo - Move the predecessors of \p Block +/// that are either within or outside of the provided \p Blocks depending +/// on the \p IfWithinSet to \p New. +static void movePredsTo(BasicBlock *Block, BasicBlock *New, + const SetVector Blocks, + bool IfWithinSet) { + // Move the predecessors of Block to New. + SmallPtrSet Preds; + for (BasicBlock *Pred : predecessors(Block)) + if (Blocks.count(Pred) == IfWithinSet) + Preds.insert(Pred); + for (BasicBlock *Pred : Preds) + Pred->getTerminator()->replaceUsesOfWith(Block, New); +} + void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs) const { for (BasicBlock *BB : Blocks) { @@ -289,6 +305,49 @@ } } +/// splitExitBlock - Split an exit block that has a phi with more than 1 +/// incoming +/// block that is within the extracted region, returns the new exit block. +/// +BasicBlock *CodeExtractor::splitExitBlock(BasicBlock *Exit) { + // Special hackery is needed with PHI nodes that have inputs from more than + // one extracted block. For simplicity, just split the PHIs into a two-level + // sequence of PHIs, some of which will go in the extracted region, and some + // of which will go outside. + BasicBlock *ExtractedReturn = Exit; + Exit = SplitBlock(Exit, Exit->getFirstNonPHI(), DT); + + // Move the branches of the non extracted blocks to the new exit block. + movePredsTo(ExtractedReturn, Exit, Blocks, /*IfWithinSet*/ false); + + // Create new phi nodes in the exit block. + Instruction *InsPoint = &Exit->front(); + for (Instruction &I : *ExtractedReturn) { + PHINode *Phi = dyn_cast(&I); + if (!Phi) + break; + + // Create a new phi and redirect the extracted phi to it. + PHINode *NewPhi = + PHINode::Create(Phi->getType(), 0, "split.exit", InsPoint); + Phi->replaceAllUsesWith(NewPhi); + NewPhi->addIncoming(Phi, ExtractedReturn); + + // Move non extracted incoming blocks to the new phi. + for (unsigned i = 0, e = Phi->getNumIncomingValues(); i < e; ++i) { + BasicBlock *IncomingBlock = Phi->getIncomingBlock(i); + if (Blocks.count(IncomingBlock)) + continue; + NewPhi->addIncoming(Phi->getIncomingValue(i), IncomingBlock); + Phi->removeIncomingValue(i, /*DeletePHIIfEmpty*/ false); + --i; + --e; + } + } + Blocks.insert(ExtractedReturn); + return Exit; +} + /// constructFunction - make a function based on inputs and outputs, as follows: /// f(in0, ..., inN, out0, ..., outN) /// @@ -787,13 +846,10 @@ "newFuncRoot"); newFuncRoot->getInstList().push_back(BranchInst::Create(header)); - // Find inputs to, outputs from the code region. - findInputsOutputs(inputs, outputs); - // Calculate the exit blocks for the extracted region and the total exit // weights for each of those blocks. DenseMap ExitWeights; - SmallPtrSet ExitBlocks; + std::unordered_map ExitBlocks; for (BasicBlock *Block : Blocks) { for (succ_iterator SI = succ_begin(Block), SE = succ_end(Block); SI != SE; ++SI) { @@ -803,12 +859,33 @@ BlockFrequency &BF = ExitWeights[*SI]; BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, *SI); } - ExitBlocks.insert(*SI); + ++ExitBlocks[*SI]; } } } NumExitBlocks = ExitBlocks.size(); + // Split the phi nodes of each of the exit blocks if there is more than 1 + // incoming block from the extracted region. + for (auto I = ExitBlocks.begin(), E = ExitBlocks.end(); I != E;) { + BasicBlock *Exit = I->first; + if (I->second < 2 || !isa(&Exit->front())) { + ++I; + continue; + } + + // Split the block and swap the values with the old exit block before + // removing it. + BasicBlock *SplitBlock = splitExitBlock(Exit); + ExitWeights[SplitBlock] = ExitWeights[Exit]; + ExitBlocks.emplace(SplitBlock, 0); + ExitWeights.erase(Exit); + ExitBlocks.erase(I++); + } + + // Find inputs to, outputs from the code region. + findInputsOutputs(inputs, outputs); + // Construct new function based on inputs/outputs & add allocas for all defs. Function *newFunction = constructFunction(inputs, outputs, header, newFuncRoot, Index: test/Transforms/CodeExtractor/MultipleExitBranchProb.ll =================================================================== --- test/Transforms/CodeExtractor/MultipleExitBranchProb.ll +++ test/Transforms/CodeExtractor/MultipleExitBranchProb.ll @@ -31,4 +31,4 @@ !2 = !{!"branch_weights", i32 5, i32 5} !3 = !{!"branch_weights", i32 4, i32 1} -; CHECK: [[COUNT1]] = !{!"branch_weights", i32 8, i32 31} +; CHECK: [[COUNT1]] = !{!"branch_weights", i32 {{(31, i32 8)|(8, i32 31)}}} Index: test/Transforms/CodeExtractor/SplitExitBlockPhiNodes.ll =================================================================== --- /dev/null +++ test/Transforms/CodeExtractor/SplitExitBlockPhiNodes.ll @@ -0,0 +1,27 @@ +; RUN: opt < %s -partial-inliner -S | FileCheck %s +; This testcase checks CodeExtractors ability to split +; an exit block that contains a phi node with more than +; 1 entry from an extracted block. + +define internal i32 @inlinedFunc(i1 %cond) { +entry: + br i1 %cond, label %if.then, label %return +if.then: + br i1 %cond, label %internal.if.then, label %internal.if.else +internal.if.then: + br label %return +internal.if.else: + br label %return +return: ; preds = %entry + %return.val = phi i32 [ 0, %entry ], [ 1, %internal.if.then ], [ 2, %internal.if.else ] + ret i32 %return.val +} + +define internal i32 @dummyCaller(i1 %cond) { +entry: + %val = call i32 @inlinedFunc(i1 %cond) + ret i32 %val +} + +; CHECK-LABEL: @inlinedFunc.1_if.then +; CHECK: phi i32