[PartialInliner] Inline vararg functions that forward varargs.

[fhahn]

Summary

This patch extends the partial inliner to support inlining parts of
vararg functions, if the vararg handling is done in the outlined part.

It adds a ForwardVarArgsTo argument to InlineFunction. If it is
non-null, all varargs passed to the inlined function will be added to
all calls to ForwardVarArgsTo.

The partial inliner takes care to only pass ForwardVarArgsTo if the
varargs handing is done in the outlined function. It checks that vastart
is not part of the function to be inlined.

test/Transforms/CodeExtractor/PartialInlineNoInline.ll (already part
of the repo) checks we do not do partial inlining if vastart is used in
a basic block that will be inlined.

[fhahn]

CC Tobias, as he asked me to look into that in D38585.

[grosser]

Hi Florian,

thanks for looking over this patch. Your new variant looks indeed a lot cleaner than my original patch. I like it.

Overall, I think your patch looks good. Just some minor nitpicks and maybe some more documentation for the new arguments you introduce.

Also, AFAIKS the outliner will (temporarily) generate invalid IR as the calls to the outlined function won't forward the vararg parameters. This is only "fixed" after inlining, when the va_start() and va_end() intrinsics again refer to the right function. While I don't feel we should "fix" this (e.g., by obtaining and forwarding the necessary parameters to the outlined function), we should certainly document this behavior.

Best,
Tobias

PS: Marking this as requesting changes to be informed about updates.

include/llvm/Transforms/Utils/Cloning.h

241

Maybe document ForwardVarArgsTo? How is it used? What is the use case, in which context is it valid to use, and how is code normally generated to require the use of this flag (outlining with AllowVarargs)?

include/llvm/Transforms/Utils/CodeExtractor.h

79

Document the new flag?

lib/Transforms/Utils/InlineFunction.cpp

1494

Document the new parameter?

1655

Unrelated!

1865

Start variable name with UpperCase?

[fhahn]

Thanks Tobias! I've updated the patch trying to address your comments.

[fhahn]

In D39607#917548, @grosser wrote:

Also, AFAIKS the outliner will (temporarily) generate invalid IR as the calls to the outlined function won't forward the vararg parameters. This is only "fixed" after inlining, when the va_start() and va_end() intrinsics again refer to the right function. While I don't feel we should "fix" this (e.g., by obtaining and forwarding the necessary parameters to the outlined function), we should certainly document this behavior.

Correct. I've added a comment to doFunctionOutlining, not sure if that is the right place. I think we would also have to ensure that vaend is also in the outlined function?

lib/Transforms/Utils/InlineFunction.cpp

1494

Hm it looks like the same function is documented here and in include/llvm/Transforms/Utils/Cloning.h. The first 2 paragraphs match, but include/llvm/Transforms/Utils/Cloning.h contains additional info. According to [1], the doxygen comments for public functions should only be at the declaration.

[1] https://llvm.org/docs/CodingStandards.html#comment-formatting

[grosser]

In D39607#917705, @fhahn wrote:

In D39607#917548, @grosser wrote:

Also, AFAIKS the outliner will (temporarily) generate invalid IR as the calls to the outlined function won't forward the vararg parameters. This is only "fixed" after inlining, when the va_start() and va_end() intrinsics again refer to the right function. While I don't feel we should "fix" this (e.g., by obtaining and forwarding the necessary parameters to the outlined function), we should certainly document this behavior.

Correct. I've added a comment to doFunctionOutlining, not sure if that is the right place.

That looks good to me.

I think we would also have to ensure that vaend is also in the outlined function?

Right. Both va_start and va_end need to be in the outlined function -- or more correct -- they cannot be in the part that is not outlined.

lib/Transforms/IPO/PartialInlining.cpp

835

As suggested by yourself, also check for va_end here.

[grosser]

Officially accept this revision. From my perspective this looks good.

Do you have some other partial inliner person who could give some additional feedback?

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

830

Is it possible to teach code extractor to handle this and return null instead of doing the cleanup here? CodeExtractor knows the set of blocks to be extracted or not.

[fhahn]

Added check for vaend and moved check to extractCodeRegion.

[fhahn]

In D39607#917790, @grosser wrote:

Officially accept this revision. From my perspective this looks good.

Do you have some other partial inliner person who could give some additional feedback?

I think it would be great if either @davidxl or @davide could sign-off on this.

lib/Transforms/IPO/PartialInlining.cpp

830

Yes I moved the code to the beginning of extractCodeRegion, before anything is modified. I could not find a place were CodeExtractor already iterates over the set of blocks not the be extracted, so it is iterating over the basic blocks of the original function and excludes blocks to be extracted.

[grosser]

In D39607#919138, @fhahn wrote:

In D39607#917790, @grosser wrote:

Officially accept this revision. From my perspective this looks good.

Do you have some other partial inliner person who could give some additional feedback?

I think it would be great if either @davidxl or @davide could sign-off on this.

Yes, this would be good, indeed!

I also like the suggestion from @davidxl to move the check into the extractor.

[davidxl]

We do life-time marker and static alloca shrink wrapping optimization for outlining. Do you see opportunity for vaarg shrink-wrapping too? If yes, it may be considered as a follow up.

lib/Transforms/IPO/PartialInlining.cpp

153

This looks like stale comment.

lib/Transforms/Utils/CodeExtractor.cpp

82

Is there a need for this parameter?

119

Having the check here may prevent other independent legality check (in the future) from taking effect. If it better to move it down.

121–126

if (AllowVarArgs)

continue;

else

return false;

981

nit: use std::any_of

[davide]

Really nice :) Some comments inline to begin with.

include/llvm/Transforms/Utils/CodeExtractor.h

91–95

This should presumably have a comment explaining what the additional bool controls.

lib/Transforms/Utils/CodeExtractor.cpp

179–181

Mind to add a comment next to the false with the member variable name :) ?

lib/Transforms/Utils/InlineFunction.cpp

1829–1830

Are you switching this because of iterator invalidation issues?
In this case, maybe it would be cleaner to collect the set of instructions to be removed in a set, and then remove them at the end of this loop?

[fhahn]

Thank you very much for having a look. Updated to address the comments.

[fhahn]

lib/Transforms/IPO/PartialInlining.cpp

153

It still applies, the outliner will generate a call to outlined vararg function, but no varargs will be passed. When we do the inlining, the varargs from the original call site will be forwarded to the outlined function. Maybe that comment needs re-phrasing to make it more clear? Or maybe the comment should be at a different place?

lib/Transforms/Utils/CodeExtractor.cpp

82

Unfortunately isBlockValidForExtraction is a static member function, so I think we need this one.

[davidxl]

lgtm

[davide]

LGTM too. Thanks Florian.

include/llvm/Transforms/Utils/Cloning.h

 /// instruction 'call B' is inlined, and 'B' calls 'C', then the call to 'C' now
 /// exists in the instruction stream.  Similarly this will inline a recursive
 /// function by one level.
 ///
 /// Note that while this routine is allowed to cleanup and optimize the
 /// *inlined* code to minimize the actual inserted code, it must not delete
 /// code in the caller as users of this routine may have pointers to
 /// instructions in the caller that need to remain stable.
+///
+/// If ForwardVarArgsTo is passed, inlining a function with varargs is allowed
+/// and all varargs at the callsite will be passed to any calls to
+/// ForwardVarArgsTo. The caller of InlineFunction has to make sure any varargs
+/// are only used by ForwardVarArgsTo.
 bool InlineFunction(CallInst *C, InlineFunctionInfo &IFI,
                     AAResults *CalleeAAR = nullptr, bool InsertLifetime = true);
 bool InlineFunction(InvokeInst *II, InlineFunctionInfo &IFI,
                     AAResults *CalleeAAR = nullptr, bool InsertLifetime = true);
 bool InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
-                    AAResults *CalleeAAR = nullptr, bool InsertLifetime = true);
+                    AAResults *CalleeAAR = nullptr, bool InsertLifetime = true,
+                    Function *ForwardVarArgsTo = nullptr);
 
 /// \brief Clones a loop \p OrigLoop.  Returns the loop and the blocks in \p
 /// Blocks.
 ///
 /// Updates LoopInfo and DominatorTree assuming the loop is dominated by block
 /// \p LoopDomBB.  Insert the new blocks before block specified in \p Before.
 /// Note: Only innermost loops are supported.
 Loop *cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,

include/llvm/Transforms/Utils/CodeExtractor.h

class CodeExtractor {
     using ValueSet = SetVector<Value *>;
 
     // Various bits of state computed on construction.
     DominatorTree *const DT;
     const bool AggregateArgs;
     BlockFrequencyInfo *BFI;
     BranchProbabilityInfo *BPI;
 
+    // If true, varargs functions can be extracted.
+    bool AllowVarArgs;
+
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
     unsigned NumExitBlocks = std::numeric_limits<unsigned>::max();
     Type *RetTy;
 
   public:
     /// \brief Create a code extractor for a sequence of blocks.
     ///
     /// 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.
+    /// 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.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
-                  BranchProbabilityInfo *BPI = nullptr);
+                  BranchProbabilityInfo *BPI = nullptr,
+                  bool AllowVarArgs = false);
 
     /// \brief 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);
 
     /// \brief Check to see if a block is valid for extraction.
     ///
-    /// Blocks containing EHPads, allocas, invokes, or vastarts are not valid.
-    static bool isBlockValidForExtraction(const BasicBlock &BB);
+    /// Blocks containing EHPads, allocas and invokes are not valid. If
+    /// AllowVarArgs is true, blocks with vastart can be extracted. This is
+    /// safe, if all vararg handling code is extracted, including vastart.
+    static bool isBlockValidForExtraction(const BasicBlock &BB,
+                                          bool AllowVarArgs);
 
     /// \brief Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
     /// returns false.
     Function *extractCodeRegion();
 
     /// \brief Test whether this code extractor is eligible.

lib/Transforms/IPO/PartialInlining.cpp

struct FunctionCloner {
     FunctionCloner(Function *F, FunctionOutliningInfo *OI);
     ~FunctionCloner();
 
     // Prepare for function outlining: making sure there is only
     // one incoming edge from the extracted/outlined region to
     // the return block.
     void NormalizeReturnBlock();
 
-    // Do function outlining:
+    // Do function outlining.
+    // NOTE: For vararg functions that do the vararg handling in the outlined
+    //       function, we temporarily generate IR that does not properly
+    //       forward varargs to the outlined function. Calling InlineFunction
+    //       will update calls to the outlined functions to properly forward
+    //       the varargs.
     Function *doFunctionOutlining();
 
     Function *OrigFunc = nullptr;
     Function *ClonedFunc = nullptr;
     Function *OutlinedFunc = nullptr;
     BasicBlock *OutliningCallBB = nullptr;
     // ClonedFunc is inlined in one of its callers after function
     // outlining.
Function *PartialInlinerImpl::FunctionCloner::doFunctionOutlining() {
 
   // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
   LoopInfo LI(DT);
   BranchProbabilityInfo BPI(*ClonedFunc, LI);
   ClonedFuncBFI.reset(new BlockFrequencyInfo(*ClonedFunc, BPI, LI));
 
   // Extract the body of the if.
   OutlinedFunc = CodeExtractor(ToExtract, &DT, /*AggregateArgs*/ false,
-                               ClonedFuncBFI.get(), &BPI)
+                               ClonedFuncBFI.get(), &BPI,
+                               /* AllowVarargs */ true)
                      .extractCodeRegion();
 
   if (OutlinedFunc) {
     OutliningCallBB = PartialInlinerImpl::getOneCallSiteTo(OutlinedFunc)
         .getInstruction()
         ->getParent();
     assert(OutliningCallBB->getParent() == ClonedFunc);
   }
 
   return OutlinedFunc;
 }
 
 PartialInlinerImpl::FunctionCloner::~FunctionCloner() {
   // Ditch the duplicate, since we're done with it, and rewrite all remaining
   // users (function pointers, etc.) back to the original function.
   ClonedFunc->replaceAllUsesWith(OrigFunc);
   ClonedFunc->eraseFromParent();
   if (!IsFunctionInlined) {
     // Remove the function that is speculatively created if there is no
     // reference.
     if (OutlinedFunc)
for (User *User : Users) {
 
     // Construct remark before doing the inlining, as after successful inlining
     // the callsite is removed.
     OptimizationRemark OR(DEBUG_TYPE, "PartiallyInlined", CS.getInstruction());
     OR << ore::NV("Callee", Cloner.OrigFunc) << " partially inlined into "
        << ore::NV("Caller", CS.getCaller());
 
     InlineFunctionInfo IFI(nullptr, GetAssumptionCache, PSI);
-    if (!InlineFunction(CS, IFI))
+    if (!InlineFunction(CS, IFI, nullptr, true, Cloner.OutlinedFunc))
       continue;
 
     ORE.emit(OR);
 
     // Now update the entry count:
     if (CalleeEntryCountV && CallSiteToProfCountMap.count(User)) {
       uint64_t CallSiteCount = CallSiteToProfCountMap[User];
       CalleeEntryCountV -= std::min(CalleeEntryCountV, CallSiteCount);

lib/Transforms/Utils/CodeExtractor.cpp

 // for functions produced by the code extractor. This is useful when converting
 // extracted functions to pthread-based code, as only one argument (void*) can
 // be passed in to pthread_create().
 static cl::opt<bool>
 AggregateArgsOpt("aggregate-extracted-args", cl::Hidden,
                  cl::desc("Aggregate arguments to code-extracted functions"));
 
 /// \brief Test whether a block is valid for extraction.
-bool CodeExtractor::isBlockValidForExtraction(const BasicBlock &BB) {
+bool CodeExtractor::isBlockValidForExtraction(const BasicBlock &BB,
+                                              bool AllowVarArgs) {
   // Landing pads must be in the function where they were inserted for cleanup.
   if (BB.isEHPad())
     return false;
   // taking the address of a basic block moved to another function is illegal
   if (BB.hasAddressTaken())
     return false;
 
   // don't hoist code that uses another basicblock address, as it's likely to
if (isa<Instruction>(Curr) && cast<Instruction>(Curr)->getParent() != &BB)
       continue;
 
     for (auto const &U : Curr->operands()) {
       if (auto *UU = dyn_cast<User>(U))
         ToVisit.push_back(UU);
     }
   }
 
-  // Don't hoist code containing allocas, invokes, or vastarts.
+  // Don't hoist code containing allocas or invokes. If explicitly requested,
+  // allow vastart.
   for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
     if (isa<AllocaInst>(I) || isa<InvokeInst>(I))
       return false;
     if (const CallInst *CI = dyn_cast<CallInst>(I))
       if (const Function *F = CI->getCalledFunction())
-        if (F->getIntrinsicID() == Intrinsic::vastart)
+        if (F->getIntrinsicID() == Intrinsic::vastart) {
+          if (AllowVarArgs)
+            continue;
+          else
             return false;
         }
+  }
 
   return true;
 }
 
 /// \brief Build a set of blocks to extract if the input blocks are viable.
 static SetVector<BasicBlock *>
-buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs, DominatorTree *DT) {
+buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
+                        bool AllowVarArgs) {
   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))
       continue;
 
     if (!Result.insert(BB))
       llvm_unreachable("Repeated basic blocks in extraction input");
-    if (!CodeExtractor::isBlockValidForExtraction(*BB)) {
+    if (!CodeExtractor::isBlockValidForExtraction(*BB, AllowVarArgs)) {
       Result.clear();
       return Result;
     }
   }
 
 #ifndef NDEBUG
   for (SetVector<BasicBlock *>::iterator I = std::next(Result.begin()),
                                          E = Result.end();
        I != E; ++I)
     for (pred_iterator PI = pred_begin(*I), PE = pred_end(*I);
          PI != PE; ++PI)
       assert(Result.count(*PI) &&
              "No blocks in this region may have entries from outside the region"
              " except for the first block!");
 #endif
 
   return Result;
 }
 
 CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
                              bool AggregateArgs, BlockFrequencyInfo *BFI,
-                             BranchProbabilityInfo *BPI)
+                             BranchProbabilityInfo *BPI, bool AllowVarArgs)
     : DT(DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
-      BPI(BPI), Blocks(buildExtractionBlockSet(BBs, DT)) {}
+      BPI(BPI), AllowVarArgs(AllowVarArgs),
+      Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs)) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs,
                              BlockFrequencyInfo *BFI,
                              BranchProbabilityInfo *BPI)
     : DT(&DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
-      BPI(BPI), Blocks(buildExtractionBlockSet(L.getBlocks(), &DT)) {}
+      BPI(BPI), Blocks(buildExtractionBlockSet(L.getBlocks(), &DT,
+                                               /* AllowVarArgs */ false)) {}
 
 /// 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;
   return false;
 }
Function *CodeExtractor::constructFunction(const ValueSet &inputs,
 
   StructType *StructTy;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
     StructTy = StructType::get(M->getContext(), paramTy);
     paramTy.clear();
     paramTy.push_back(PointerType::getUnqual(StructTy));
   }
   FunctionType *funcType =
-                  FunctionType::get(RetTy, paramTy, false);
+                  FunctionType::get(RetTy, paramTy,
+                                    AllowVarArgs && oldFunction->isVarArg());
 
   // Create the new function
   Function *newFunction = Function::Create(funcType,
                                            GlobalValue::InternalLinkage,
                                            oldFunction->getName() + "_" +
                                            header->getName(), M);
   // If the old function is no-throw, so is the new one.
   if (oldFunction->doesNotThrow())
TI->setMetadata(
       LLVMContext::MD_prof,
       MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
 }
 
 Function *CodeExtractor::extractCodeRegion() {
   if (!isEligible())
     return nullptr;
 
-  ValueSet inputs, outputs, SinkingCands, HoistingCands;
-  BasicBlock *CommonExit = nullptr;
-
   // 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();
+
+  // For functions with varargs, check that varargs handling is only done in the
+  // outlined function, i.e vastart and vaend are only used in outlined blocks.
+  if (AllowVarArgs && oldFunction->getFunctionType()->isVarArg()) {
+    auto containsVarArgIntrinsic = [](Instruction &I) {
+      if (const CallInst *CI = dyn_cast<CallInst>(&I))
+        if (const Function *F = CI->getCalledFunction())
+          return F->getIntrinsicID() == Intrinsic::vastart ||
+                 F->getIntrinsicID() == Intrinsic::vaend;
+      return false;
+    };
+
+    for (auto &BB : *oldFunction) {
+      if (Blocks.count(&BB))
+        continue;
+      if (llvm::any_of(BB, containsVarArgIntrinsic))
+        return nullptr;
+    }
+  }
+  ValueSet inputs, outputs, SinkingCands, HoistingCands;
+  BasicBlock *CommonExit = nullptr;
 
   // Calculate the entry frequency of the new function before we change the root
   //   block.
   BlockFrequency EntryFreq;
   if (BFI) {
     assert(BPI && "Both BPI and BFI are required to preserve profile info");
     for (BasicBlock *Pred : predecessors(header)) {
       if (Blocks.count(Pred))
         continue;
       EntryFreq +=
           BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, header);
     }
   }
 
   // If we have to split PHI nodes or the entry block, do so now.
   severSplitPHINodes(header);
 
   // If we have any return instructions in the region, split those blocks so
   // that the return is not in the region.
   splitReturnBlocks();
 
-  Function *oldFunction = header->getParent();
-
   // This takes place of the original loop
   BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(), 
                                                 "codeRepl", oldFunction,
                                                 header);
 
   // 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(header->getContext(), 

lib/Transforms/Utils/InlineFunction.cpp

 /// caller. This returns false if it is not possible to inline this call.
 /// The program is still in a well defined state if this occurs though.
 ///
 /// Note that this only does one level of inlining.  For example, if the
 /// instruction 'call B' is inlined, and 'B' calls 'C', then the call to 'C' now
 /// exists in the instruction stream.  Similarly this will inline a recursive
 /// function by one level.
 bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
-                          AAResults *CalleeAAR, bool InsertLifetime) {
+                          AAResults *CalleeAAR, bool InsertLifetime,
+                          Function *ForwardVarArgsTo) {
   Instruction *TheCall = CS.getInstruction();
   assert(TheCall->getParent() && TheCall->getFunction()
          && "Instruction not in function!");
 
   // If IFI has any state in it, zap it before we fill it in.
   IFI.reset();
 
   Function *CalledFunc = CS.getCalledFunction();
   if (!CalledFunc ||              // Can't inline external function or indirect
-      CalledFunc->isDeclaration() || // call, or call to a vararg function!
-      CalledFunc->getFunctionType()->isVarArg()) return false;
+      CalledFunc->isDeclaration() ||
+      (!ForwardVarArgsTo && CalledFunc->isVarArg())) // call, or call to a vararg function!
+      return false;
 
   // The inliner does not know how to inline through calls with operand bundles
   // in general ...
   if (CS.hasOperandBundles()) {
     for (int i = 0, e = CS.getNumOperandBundles(); i != e; ++i) {
       uint32_t Tag = CS.getOperandBundleAt(i).getTagID();
       // ... but it knows how to inline through "deopt" operand bundles ...
       if (Tag == LLVMContext::OB_deopt)
bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
 
   { // Scope to destroy VMap after cloning.
     ValueToValueMapTy VMap;
     // Keep a list of pair (dst, src) to emit byval initializations.
     SmallVector<std::pair<Value*, Value*>, 4> ByValInit;
 
     auto &DL = Caller->getParent()->getDataLayout();
 
-    assert(CalledFunc->arg_size() == CS.arg_size() &&
-           "No varargs calls can be inlined!");
+    assert((CalledFunc->arg_size() == CS.arg_size() || ForwardVarArgsTo) &&
+           "Varargs calls can only be inlined if the Varargs are forwarded!");
 
     // Calculate the vector of arguments to pass into the function cloner, which
     // matches up the formal to the actual argument values.
     CallSite::arg_iterator AI = CS.arg_begin();
     unsigned ArgNo = 0;
     for (Function::arg_iterator I = CalledFunc->arg_begin(),
          E = CalledFunc->arg_end(); I != E; ++I, ++AI, ++ArgNo) {
       Value *ActualArg = *AI;
 
       // When byval arguments actually inlined, we need to make the copy implied
       // by them explicit.  However, we don't do this if the callee is readonly
       // or readnone, because the copy would be unneeded: the callee doesn't
       // modify the struct.
       if (CS.isByValArgument(ArgNo)) {
         ActualArg = HandleByValArgument(ActualArg, TheCall, CalledFunc, IFI,
                                         CalledFunc->getParamAlignment(ArgNo));
         if (ActualArg != *AI)
           ByValInit.push_back(std::make_pair(ActualArg, (Value*) *AI));
       }
 
       VMap[&*I] = ActualArg;
     }
 
     // Add alignment assumptions if necessary. We do this before the inlined
     // instructions are actually cloned into the caller so that we can easily
     // check what will be known at the start of the inlined code.
     AddAlignmentAssumptions(CS, IFI);
 
     // We want the inliner to prune the code as it copies.  We would LOVE to
     // have no dead or constant instructions leftover after inlining occurs
     // (which can happen, e.g., because an argument was constant), but we'll be
for (BasicBlock::iterator I = FirstNewBlock->begin(),
           InsertPoint, FirstNewBlock->getInstList(), AI->getIterator(), I);
     }
     // Move any dbg.declares describing the allocas into the entry basic block.
     DIBuilder DIB(*Caller->getParent());
     for (auto &AI : IFI.StaticAllocas)
       replaceDbgDeclareForAlloca(AI, AI, DIB, /*Deref=*/false);
   }
 
+  SmallVector<Value*,4> VarArgsToForward;
+  for (unsigned i = CalledFunc->getFunctionType()->getNumParams();
+       i < CS.getNumArgOperands(); i++)
+    VarArgsToForward.push_back(CS.getArgOperand(i));
+
   bool InlinedMustTailCalls = false, InlinedDeoptimizeCalls = false;
   if (InlinedFunctionInfo.ContainsCalls) {
     CallInst::TailCallKind CallSiteTailKind = CallInst::TCK_None;
     if (CallInst *CI = dyn_cast<CallInst>(TheCall))
       CallSiteTailKind = CI->getTailCallKind();
 
     for (Function::iterator BB = FirstNewBlock, E = Caller->end(); BB != E;
          ++BB) {
-      for (Instruction &I : *BB) {
+      for (auto II = BB->begin(); II != BB->end();) {
+        Instruction &I = *II++;
         CallInst *CI = dyn_cast<CallInst>(&I);
         if (!CI)
           continue;
 
         if (Function *F = CI->getCalledFunction())
           InlinedDeoptimizeCalls |=
               F->getIntrinsicID() == Intrinsic::experimental_deoptimize;
 
for (Function::iterator BB = FirstNewBlock, E = Caller->end(); BB != E;
         ChildTCK = std::min(CallSiteTailKind, ChildTCK);
         CI->setTailCallKind(ChildTCK);
         InlinedMustTailCalls |= CI->isMustTailCall();
 
         // Calls inlined through a 'nounwind' call site should be marked
         // 'nounwind'.
         if (MarkNoUnwind)
           CI->setDoesNotThrow();
+
+        if (ForwardVarArgsTo && CI->getCalledFunction() == ForwardVarArgsTo) {
+          SmallVector<Value*, 6> Params(CI->arg_operands());
+          Params.append(VarArgsToForward.begin(), VarArgsToForward.end());
+          CallInst *Call = CallInst::Create(CI->getCalledFunction(), Params, "", CI);
+          CI->replaceAllUsesWith(Call);
+          CI->eraseFromParent();
+        }
       }
     }
   }
 
   // Leave lifetime markers for the static alloca's, scoping them to the
   // function we just inlined.
   if (InsertLifetime && !IFI.StaticAllocas.empty()) {
     IRBuilder<> builder(&FirstNewBlock->front());

test/Transforms/CodeExtractor/PartialInlineVarArg.ll

+; RUN: opt < %s -partial-inliner -S -skip-partial-inlining-cost-analysis | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -S -skip-partial-inlining-cost-analysis | FileCheck %s
+
+@stat = external global i32, align 4
+
+define i32 @vararg(i32 %count, ...) {
+entry:
+  %vargs = alloca i8*, align 8
+  %stat1 = load i32, i32* @stat, align 4
+  %cmp = icmp slt i32 %stat1, 0
+  br i1 %cmp, label %bb2, label %bb1
+
+bb1:                                              ; preds = %entry
+  %vg1 = add nsw i32 %stat1, 1
+  store i32 %vg1, i32* @stat, align 4
+  %vargs1 = bitcast i8** %vargs to i8*
+  call void @llvm.va_start(i8* %vargs1)
+  %va1 = va_arg i8** %vargs, i32
+  call void @foo(i32 %count, i32 %va1) #2
+  call void @llvm.va_end(i8* %vargs1)
+  br label %bb2
+
+bb2:                                              ; preds = %bb1, %entry
+  %res = phi i32 [ 1, %bb1 ], [ 0, %entry ]
+  ret i32 %res
+}
+
+declare void @foo(i32, i32)
+declare void @llvm.va_start(i8*)
+declare void @llvm.va_end(i8*)
+
+define i32 @caller1(i32 %arg) {
+bb:
+  %tmp = tail call i32 (i32, ...) @vararg(i32 %arg)
+  ret i32 %tmp
+}
+; CHECK-LABEL: @caller1
+; CHECK: codeRepl.i:
+; CHECK-NEXT:  call void (i32, i8**, i32, ...) @vararg.2_bb1(i32 %stat1.i, i8** %vargs.i, i32 %arg)
+
+define i32 @caller2(i32 %arg, float %arg2) {
+bb:
+  %tmp = tail call i32 (i32, ...) @vararg(i32 %arg, i32 10, float %arg2)
+  ret i32 %tmp
+}
+
+; CHECK-LABEL: @caller2
+; CHECK: codeRepl.i:
+; CHECK-NEXT:  call void (i32, i8**, i32, ...) @vararg.2_bb1(i32 %stat1.i, i8** %vargs.i, i32 %arg, i32 10, float %arg2)
+
+; Test case to check that we do not extract a vararg function, if va_end is in
+; a block that is not outlined.
+define i32 @vararg_not_legal(i32 %count, ...) {
+entry:
+  %vargs = alloca i8*, align 8
+  %vargs0 = bitcast i8** %vargs to i8*
+  %stat1 = load i32, i32* @stat, align 4
+  %cmp = icmp slt i32 %stat1, 0
+  br i1 %cmp, label %bb2, label %bb1
+
+bb1:                                              ; preds = %entry
+  %vg1 = add nsw i32 %stat1, 1
+  store i32 %vg1, i32* @stat, align 4
+  %vargs1 = bitcast i8** %vargs to i8*
+  call void @llvm.va_start(i8* %vargs1)
+  %va1 = va_arg i8** %vargs, i32
+  call void @foo(i32 %count, i32 %va1)
+  br label %bb2
+
+bb2:                                              ; preds = %bb1, %entry
+  %res = phi i32 [ 1, %bb1 ], [ 0, %entry ]
+  %ptr = phi i8* [ %vargs1, %bb1 ], [ %vargs0, %entry]
+  call void @llvm.va_end(i8* %ptr)
+  ret i32 %res
+}
+
+; CHECK-LABEL: @caller3
+; CHECK: tail call i32 (i32, ...) @vararg_not_legal(i32 %arg, i32 %arg)
+define i32 @caller3(i32 %arg) {
+bb:
+  %res = tail call i32 (i32, ...) @vararg_not_legal(i32 %arg, i32 %arg)
+  ret i32 %res
+}