[CodeExtractor] Enable partial aggregate arguments

[ggeorgakoudis]

Summary

Enable CodeExtractor to construct output functions that partially aggregate inputs/outputs in their argument list. A use case is the OMPIRBuilder to create outlined functions for parallel regions that aggregate in a struct the payload variables for the region while passing as scalars thread and bound identifiers.

[vsk]

Sorry it's taken me so long to get to this.

partially aggregate inputs/outputs in their argument list

Could you explain what this means, and what the pros/cons might be compared to any alternatives? It'd also help to see a test case.

[ggeorgakoudis]

Add test

[ggeorgakoudis]

In D96854#2594840, @vsk wrote:

Sorry it's taken me so long to get to this.

partially aggregate inputs/outputs in their argument list

Could you explain what this means, and what the pros/cons might be compared to any alternatives? It'd also help to see a test case.

Hi @vsk,

No problem! Let me make the use case of the OpenMP IR builder concrete, I'll do some simplifications that do not affect the point. Currently, the OMPIRBuilder uses CodeExtractor to outline an OpenMP callback as:

void omp.outlined(int global_tid, int bound_tid, int* arg0, int* arg1, ..., int* argn)

where global_tid, bound_tid are OpenMP runtime filled values passed to the outlined functions and arg0, arg1, ...,argn are inputs/outputs found by CodeExtractor. To implement parallel execution calling the outlined call function, OMPIRBuilder emits the call to the OpenMP runtime fork_join function, which use ellipsis to pass the variadic number of parameters to the OpenMP runtime:

__kmpc_fork_call(int argc, omp.outlined, ..,)

so the ellipsis contains the arguments to the outlined function (arg0, arg1, ..., argn). The OpenMP runtime library fills the values for the preceding arguments global_tid, bound_tid when calling omp.outlined and forwards the rest of the arguments through a cumbersome dispatch function that unwraps the variadic arguments and uses a switch-case to call the function pointer of the callback as in:

switch(argc) {
  case 1: fp_to_omp.outlined(global_tid, bound_tid, vararg[0]); return;
  case 2: fp_to_omp.outlined(global_tid, bound_tid, vararg[0], vararg[1]); return;
  ...
}

We would like to remove this ellipsis interface because it creates various problems: there is a hardcoded limit on the number of arguments that the runtime forwards (limited by the switch-case style unwrapping), it has been the source of ABI bugs, and makes hard to analyze and optimize OpenMP code in LLVM. For this we would like to aggregate the input/output arguments to the outlined function but leave the runtime-filled arguments unaggregated:

void omp.outlined(int global_tid, int bound_tid, struct structArg)

This patch enables to exclude arguments from the aggregate by extending extractCodeRegion in CodeExtractor with a parameter of which arguments to exclude (assuming AggregateArgs has been set when creating the CodeExtractor). In this specific use case for OpenMP, the exclude arguments are global_tid and bound_tid.

I have added a unit test that tests this functionality. When we complete the change in OMPIRBuilder to use partial aggregation from CodeExtractor we will add also IR tests that will test this functionality too.

[ggeorgakoudis]

Ping!

[vsk]

Thanks for explaining. I'd suggest making ExcludedAggArgs part of a CodeExtractor instances internal state: e.g. the client may call CE.addArgExludedFromAggregate(Value *) some number of times before CE.extractCodeRegion(). This way, the client doesn't need to maintain a SetVector, and the rest of the interface isn't polluted with an option that's specific to the AggregateArg case.

[ggeorgakoudis]

In D96854#2645458, @vsk wrote:

Thanks for explaining. I'd suggest making ExcludedAggArgs part of a CodeExtractor instances internal state: e.g. the client may call CE.addArgExludedFromAggregate(Value *) some number of times before CE.extractCodeRegion(). This way, the client doesn't need to maintain a SetVector, and the rest of the interface isn't polluted with an option that's specific to the AggregateArg case.

Hi @vsk, thank you for your feedback. I think the proposed interface is more flexible and easier to support than making exclusions internal to CodeExtractor and having repeated calls through addArgExludedFromAggregate to set them. My line of thinking is that (1) it is more flexible to let the client manage the state and request extraction through extractCodeRegion with it, for example it makes possible to extract the same region with different exclusions without creating another CodeExtractor instance, (2) there is only one, backwards behavior compatible change to the external interface of CodeExtractor, that is for extractCodeRegion, so changes in other methods (constructFunction, emitCallAndSwitchStatement) are internal and changes do not pollute this external API. Please let me know what you think.

[vsk]

for example it makes possible to extract the same region with different exclusions without creating another CodeExtractor instance,

As extractCodeRegion mutates the original function, I assumed it was not possible to reuse a CodeExtractor instance in this way. Is there an in-tree example of CE instance reuse I can take a look at?

so changes in other methods (constructFunction, emitCallAndSwitchStatement) are internal and changes do not pollute this external API

I don't quite follow. Wouldn't introducing a separate API for adding arg exclusions would also be backwards compatible?

[ggeorgakoudis]

In D96854#2695426, @vsk wrote:

for example it makes possible to extract the same region with different exclusions without creating another CodeExtractor instance,

As extractCodeRegion mutates the original function, I assumed it was not possible to reuse a CodeExtractor instance in this way. Is there an in-tree example of CE instance reuse I can take a look at?

Unfortunately, there isn't an example so far. The CodeExtractor instance can be re-used because the analysis of CE stays the same. The exclusion from aggregates affects only the extraction of the outlined function when calling extractCodeRegion.

so changes in other methods (constructFunction, emitCallAndSwitchStatement) are internal and changes do not pollute this external API

I don't quite follow. Wouldn't introducing a separate API for adding arg exclusions would also be backwards compatible?

It will be backwards compatible but IMO it unnecessarily binds the CE instance to a specific way of argument creation for the outlined function. In the same sense, the AggregateArgs that is part of the CE constructor could be only an argument to extractCodeRegion since CE analysis is orthogonal to it.

[ggeorgakoudis]

@vsk I'll re-structure the implementation to use internal state and an addArgExludedFromAggregate interface. I don't have/can't think right now a use case for the flexibility I'm proposing.

[vsk]

In D96854#2707304, @ggeorgakoudis wrote:

In D96854#2695426, @vsk wrote:

for example it makes possible to extract the same region with different exclusions without creating another CodeExtractor instance,

As extractCodeRegion mutates the original function, I assumed it was not possible to reuse a CodeExtractor instance in this way. Is there an in-tree example of CE instance reuse I can take a look at?

Unfortunately, there isn't an example so far. The CodeExtractor instance can be re-used because the analysis of CE stays the same. The exclusion from aggregates affects only the extraction of the outlined function when calling extractCodeRegion.

Oh, I think I see what you meant - was it that the CE analysis cache can be re-used? The motivation there was eliminating quadratic compile-time for repeated outlining (D68616); with that change some cached analysis of the caller became reusable, but not the extractor instance itself.

so changes in other methods (constructFunction, emitCallAndSwitchStatement) are internal and changes do not pollute this external API

I don't quite follow. Wouldn't introducing a separate API for adding arg exclusions would also be backwards compatible?

It will be backwards compatible but IMO it unnecessarily binds the CE instance to a specific way of argument creation for the outlined function.
[snip]

I believe this would be true of any API we pick today.

[ggeorgakoudis]

In D96854#2709726, @vsk wrote:

Oh, I think I see what you meant - was it that the CE analysis cache can be re-used? The motivation there was eliminating quadratic compile-time for repeated outlining (D68616); with that change some cached analysis of the caller became reusable, but not the extractor instance itself.

Yes, that's what I meant! The CE analysis is reused and the extractCodeRegion should only define the way to extract the analyzed region to an outlined function.

I believe this would be true of any API we pick today.

I think it doesn't have to be this way, if we move the AggregateArgs flag to be an argument to extractCodeRegion along with the exclusions. For the use case in the OpenMP IR builder either interface will work. Since you have a more holistic use of the clients of CodeExtractor, it makes sense that you make the call. Shall I go forth with the addArgExludedFromAggregate interface?

[ggeorgakoudis]

Ping @vsk! Waiting for your decision :)

[vsk]

@ggeorgakoudis apologies again for the delay here.

My opinion hasn't changed, I still feel it'd make for a simpler interface to introduce a "CE.addArgExcludedFromAggregate(Value *)" API (though I admit that that API name isn't ideal).

[ggeorgakoudis]

Update interface for comments

[ggeorgakoudis]

Ping @vsk

[vsk]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

833

I notice there's a paramTy = ScalarParamTy: can we delete paramTy entirely? Having two copies of the same thing creates a risk of the copies diverging, adding complexity.

861

Suggest - assert(StructValues.empty() || AggregateArgs && "StructValues updated for arg structs only")

[ggeorgakoudis]

Update for comments

[ggeorgakoudis]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

833

In paramTy we concatenate the scalar parameter types and the aggregate type. The comment on diverging copies makes sense. I'll change paramTy to ParamTy, remove ScalarParamTy, and directly push non-aggregate arguments and concat the aggregate type.

861

Makes sense, will add the check

[ggeorgakoudis]

Ping @vsk :)

[vsk]

Looks good to me, thanks.
(And apologies for more delay: I'm not working on llvm much these days, Jessica or Jon (both cc'd) may be able to provide better turnaround time going forward.)

[jdoerfert]

Was accepted before, now mark it as such.

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

public:
     static bool verifyAssumptionCache(const Function &OldFunc,
                                       const Function &NewFunc,
                                       AssumptionCache *AC);
 
     /// Test whether this code extractor is eligible.
     ///
     /// Based on the blocks used when constructing the code extractor,
     /// determine whether it is eligible for extraction.
-    /// 
+    ///
     /// Checks that varargs handling (with vastart and vaend) is only done in
     /// the outlined blocks.
     bool isEligible() const;
 
     /// Compute the set of input values and output values for the code.
     ///
     /// These can be used either when performing the extraction or to evaluate
     /// the expected size of a call to the extracted function. Note that this
public:
     ///
     /// CommonExitBlock is block outside the outline region. It is the common
     /// successor of blocks inside the region. If there exists a single block
     /// inside the region that is the predecessor of CommonExitBlock, that block
     /// will be returned. Otherwise CommonExitBlock will be split and the
     /// original block will be added to the outline region.
     BasicBlock *findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock);
 
+    /// Exclude a value from aggregate argument passing when extracting a code
+    /// region, passing it instead as a scalar.
+    void excludeArgFromAggregate(Value *Arg);
+
   private:
     struct LifetimeMarkerInfo {
       bool SinkLifeStart = false;
       bool HoistLifeEnd = false;
       Instruction *LifeStart = nullptr;
       Instruction *LifeEnd = nullptr;
     };
 
+    ValueSet ExcludeArgsFromAggregate;
+
     LifetimeMarkerInfo
     getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
                        Instruction *Addr, BasicBlock *ExitBlock) const;
 
     void severSplitPHINodesOfEntry(BasicBlock *&Header);
     void severSplitPHINodesOfExits(const SmallPtrSetImpl<BasicBlock *> &Exits);
     void splitReturnBlocks();
 

llvm/lib/Transforms/Utils/CodeExtractor.cpp

Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   switch (NumExitBlocks) {
   case 0:
   case 1: RetTy = Type::getVoidTy(header->getContext()); break;
   case 2: RetTy = Type::getInt1Ty(header->getContext()); break;
   default: RetTy = Type::getInt16Ty(header->getContext()); break;
   }
 
   std::vector<Type *> paramTy;
+  std::vector<Type *> ScalarParamTy;
+  std::vector<Type *> AggParamTy;
+  ValueSet StructValues;
 
   // Add the types of the input values to the function's argument list
   for (Value *value : inputs) {
     LLVM_DEBUG(dbgs() << "value used in func: " << *value << "\n");
-    paramTy.push_back(value->getType());
+    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(value)) {
+      AggParamTy.push_back(value->getType());
+      StructValues.insert(value);
+    } else
+      ScalarParamTy.push_back(value->getType());
   }
 
   // Add the types of the output values to the function's argument list.
   for (Value *output : outputs) {
     LLVM_DEBUG(dbgs() << "instr used in func: " << *output << "\n");
-    if (AggregateArgs)
-      paramTy.push_back(output->getType());
-    else
-      paramTy.push_back(PointerType::getUnqual(output->getType()));
+    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(output)) {
+      AggParamTy.push_back(output->getType());
+      StructValues.insert(output);
+    } else
+      ScalarParamTy.push_back(PointerType::getUnqual(output->getType()));
+  }
+
+  assert(
+      (ScalarParamTy.size() + AggParamTy.size()) ==
+          (inputs.size() + outputs.size()) &&
+      "Number of scalar and aggregate params does not match inputs, outputs");
+
+  paramTy = ScalarParamTy;
+  StructType *StructTy = nullptr;
+  if (AggregateArgs && !AggParamTy.empty()) {
+    StructTy = StructType::get(M->getContext(), AggParamTy);
+    paramTy.push_back(PointerType::getUnqual(StructTy));
   }
 
   LLVM_DEBUG({
     dbgs() << "Function type: " << *RetTy << " f(";
     for (Type *i : paramTy)
       dbgs() << *i << ", ";
     dbgs() << ")\n";
   });
 
-  StructType *StructTy = nullptr;
-  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,
                                     AllowVarArgs && oldFunction->isVarArg());
 
   std::string SuffixToUse =
       Suffix.empty()
           ? (header->getName().empty() ? "extracted" : header->getName().str())
           : Suffix;
if (Attr.isStringAttribute()) {
       case Attribute::NoProfile:
         break;
       }
 
     newFunction->addFnAttr(Attr);
   }
   newFunction->getBasicBlockList().push_back(newRootNode);
 
-  // Create an iterator to name all of the arguments we inserted.
-  Function::arg_iterator AI = newFunction->arg_begin();
+  // Create scalar and aggregate iterators to name all of the arguments we
+  // inserted.
+  Function::arg_iterator ScalarAI = newFunction->arg_begin();
+  Function::arg_iterator AggAI = std::next(ScalarAI, ScalarParamTy.size());
 
   // Rewrite all users of the inputs in the extracted region to use the
   // arguments (or appropriate addressing into struct) instead.
-  for (unsigned i = 0, e = inputs.size(); i != e; ++i) {
+  for (unsigned i = 0, e = inputs.size(), aggIdx = 0; i != e; ++i) {
     Value *RewriteVal;
-    if (AggregateArgs) {
+    if (AggregateArgs && StructValues.contains(inputs[i])) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(header->getContext()));
-      Idx[1] = ConstantInt::get(Type::getInt32Ty(header->getContext()), i);
+      Idx[1] = ConstantInt::get(Type::getInt32Ty(header->getContext()), aggIdx);
       Instruction *TI = newFunction->begin()->getTerminator();
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
-          StructTy, &*AI, Idx, "gep_" + inputs[i]->getName(), TI);
-      RewriteVal = new LoadInst(StructTy->getElementType(i), GEP,
+          StructTy, &*AggAI, Idx, "gep_" + inputs[i]->getName(), TI);
+      RewriteVal = new LoadInst(StructTy->getElementType(aggIdx), GEP,
                                 "loadgep_" + inputs[i]->getName(), TI);
+      ++aggIdx;
     } else
-      RewriteVal = &*AI++;
+      RewriteVal = &*ScalarAI++;
 
     std::vector<User *> Users(inputs[i]->user_begin(), inputs[i]->user_end());
     for (User *use : Users)
       if (Instruction *inst = dyn_cast<Instruction>(use))
         if (Blocks.count(inst->getParent()))
           inst->replaceUsesOfWith(inputs[i], RewriteVal);
   }
 
   // Set names for input and output arguments.
-  if (!AggregateArgs) {
-    AI = newFunction->arg_begin();
-    for (unsigned i = 0, e = inputs.size(); i != e; ++i, ++AI)
-      AI->setName(inputs[i]->getName());
-    for (unsigned i = 0, e = outputs.size(); i != e; ++i, ++AI)
-      AI->setName(outputs[i]->getName()+".out");
+  if (!AggregateArgs || !ScalarParamTy.empty()) {
+    ScalarAI = newFunction->arg_begin();
+    for (unsigned i = 0, e = inputs.size(); i != e; ++i, ++ScalarAI)
+      if (!StructValues.contains(inputs[i]))
+        ScalarAI->setName(inputs[i]->getName());
+    for (unsigned i = 0, e = outputs.size(); i != e; ++i, ++ScalarAI)
+      if (!StructValues.contains(outputs[i]))
+        ScalarAI->setName(outputs[i]->getName() + ".out");
   }
 
   // Rewrite branches to basic blocks outside of the loop to new dummy blocks
   // within the new function. This must be done before we lose track of which
   // blocks were originally in the code region.
   std::vector<User *> Users(header->user_begin(), header->user_end());
   for (auto &U : Users)
     // The BasicBlock which contains the branch is not in the region

 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
 CallInst *CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
                                                     BasicBlock *codeReplacer,
                                                     ValueSet &inputs,
                                                     ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
-  std::vector<Value *> params, StructValues, ReloadOutputs, Reloads;
+  std::vector<Value *> params, ReloadOutputs, Reloads;
+  ValueSet StructValues;
 
   Module *M = newFunction->getParent();
   LLVMContext &Context = M->getContext();
   const DataLayout &DL = M->getDataLayout();
   CallInst *call = nullptr;
 
   // Add inputs as params, or to be filled into the struct
-  unsigned ArgNo = 0;
+  unsigned ScalarInputArgNo = 0;
   SmallVector<unsigned, 1> SwiftErrorArgs;
   for (Value *input : inputs) {
-    if (AggregateArgs)
-      StructValues.push_back(input);
+    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(input))
+      StructValues.insert(input);
     else {
       params.push_back(input);
       if (input->isSwiftError())
-        SwiftErrorArgs.push_back(ArgNo);
+        SwiftErrorArgs.push_back(ScalarInputArgNo);
     }
-    ++ArgNo;
+    ++ScalarInputArgNo;
   }
 
   // Create allocas for the outputs
+  unsigned ScalarOutputArgNo = 0;
   for (Value *output : outputs) {
-    if (AggregateArgs) {
-      StructValues.push_back(output);
+    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(output)) {
+      StructValues.insert(output);
     } else {
       AllocaInst *alloca =
         new AllocaInst(output->getType(), DL.getAllocaAddrSpace(),
                        nullptr, output->getName() + ".loc",
                        &codeReplacer->getParent()->front().front());
       ReloadOutputs.push_back(alloca);
       params.push_back(alloca);
+      ++ScalarOutputArgNo;
     }
   }
 
   StructType *StructArgTy = nullptr;
   AllocaInst *Struct = nullptr;
-  if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
+  unsigned NumAggregatedInputs = 0;
+  if (AggregateArgs && !StructValues.empty()) {
     std::vector<Type *> ArgTypes;
     for (Value *V : StructValues)
       ArgTypes.push_back(V->getType());
 
     // Allocate a struct at the beginning of this function
     StructArgTy = StructType::get(newFunction->getContext(), ArgTypes);
     Struct = new AllocaInst(StructArgTy, DL.getAllocaAddrSpace(), nullptr,
                             "structArg",
                             &codeReplacer->getParent()->front().front());
     params.push_back(Struct);
 
-    for (unsigned i = 0, e = inputs.size(); i != e; ++i) {
+    // Store aggregated inputs in the struct.
+    for (unsigned i = 0, e = StructValues.size(); i != e; ++i) {
+      if (inputs.contains(StructValues[i])) {
         Value *Idx[2];
         Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
         Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), i);
         GetElementPtrInst *GEP = GetElementPtrInst::Create(
             StructArgTy, Struct, Idx, "gep_" + StructValues[i]->getName());
         codeReplacer->getInstList().push_back(GEP);
         new StoreInst(StructValues[i], GEP, codeReplacer);
+        NumAggregatedInputs++;
+      }
     }
   }
 
   // Emit the call to the function
   call = CallInst::Create(newFunction, params,
                           NumExitBlocks > 1 ? "targetBlock" : "");
   // Add debug location to the new call, if the original function has debug
   // info. In that case, the terminator of the entry block of the extracted
   // function contains the first debug location of the extracted function,
   // set in extractCodeRegion.
   if (codeReplacer->getParent()->getSubprogram()) {
     if (auto DL = newFunction->getEntryBlock().getTerminator()->getDebugLoc())
       call->setDebugLoc(DL);
   }
   codeReplacer->getInstList().push_back(call);
 
   // Set swifterror parameter attributes.
   for (unsigned SwiftErrArgNo : SwiftErrorArgs) {
     call->addParamAttr(SwiftErrArgNo, Attribute::SwiftError);
     newFunction->addParamAttr(SwiftErrArgNo, Attribute::SwiftError);
   }
 
-  Function::arg_iterator OutputArgBegin = newFunction->arg_begin();
-  unsigned FirstOut = inputs.size();
-  if (!AggregateArgs)
-    std::advance(OutputArgBegin, inputs.size());
-
-  // Reload the outputs passed in by reference.
-  for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
+  // Reload the outputs passed in by reference, use the struct if output is in
+  // the aggregate or reload from the scalar argument.
+  for (unsigned i = 0, e = outputs.size(), scalarIdx = 0,
+                aggIdx = NumAggregatedInputs;
+       i != e; ++i) {
     Value *Output = nullptr;
-    if (AggregateArgs) {
+    if (AggregateArgs && StructValues.contains(outputs[i])) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
-      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
+      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), aggIdx);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructArgTy, Struct, Idx, "gep_reload_" + outputs[i]->getName());
       codeReplacer->getInstList().push_back(GEP);
       Output = GEP;
+      ++aggIdx;
     } else {
-      Output = ReloadOutputs[i];
+      Output = ReloadOutputs[scalarIdx];
+      ++scalarIdx;
     }
     LoadInst *load = new LoadInst(outputs[i]->getType(), Output,
                                   outputs[i]->getName() + ".reload",
                                   codeReplacer);
     Reloads.push_back(load);
     std::vector<User *> Users(outputs[i]->user_begin(), outputs[i]->user_end());
     for (unsigned u = 0, e = Users.size(); u != e; ++u) {
       Instruction *inst = cast<Instruction>(Users[u]);
for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) {
       // rewrite the original branch instruction with this new target
       TI->setSuccessor(i, NewTarget);
    }
   }
 
   // Store the arguments right after the definition of output value.
   // This should be proceeded after creating exit stubs to be ensure that invoke
   // result restore will be placed in the outlined function.
-  Function::arg_iterator OAI = OutputArgBegin;
-  for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
+  Function::arg_iterator ScalarOutputArgBegin = newFunction->arg_begin();
+  std::advance(ScalarOutputArgBegin, ScalarInputArgNo);
+  Function::arg_iterator AggOutputArgBegin = newFunction->arg_begin();
+  std::advance(AggOutputArgBegin, ScalarInputArgNo + ScalarOutputArgNo);
+
+  for (unsigned i = 0, e = outputs.size(), aggIdx = NumAggregatedInputs; i != e;
+       ++i) {
     auto *OutI = dyn_cast<Instruction>(outputs[i]);
     if (!OutI)
       continue;
 
     // Find proper insertion point.
     BasicBlock::iterator InsertPt;
     // In case OutI is an invoke, we insert the store at the beginning in the
     // 'normal destination' BB. Otherwise we insert the store right after OutI.
     if (auto *InvokeI = dyn_cast<InvokeInst>(OutI))
       InsertPt = InvokeI->getNormalDest()->getFirstInsertionPt();
     else if (auto *Phi = dyn_cast<PHINode>(OutI))
       InsertPt = Phi->getParent()->getFirstInsertionPt();
     else
       InsertPt = std::next(OutI->getIterator());
 
     Instruction *InsertBefore = &*InsertPt;
     assert((InsertBefore->getFunction() == newFunction ||
             Blocks.count(InsertBefore->getParent())) &&
            "InsertPt should be in new function");
-    assert(OAI != newFunction->arg_end() &&
-           "Number of output arguments should match "
-           "the amount of defined values");
-    if (AggregateArgs) {
+    if (AggregateArgs && StructValues.contains(outputs[i])) {
+      assert(AggOutputArgBegin != newFunction->arg_end() &&
+             "Number of aggregate output arguments should match "
+             "the number of defined values");
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
-      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
+      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), aggIdx);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
-          StructArgTy, &*OAI, Idx, "gep_" + outputs[i]->getName(),
+          StructArgTy, &*AggOutputArgBegin, Idx, "gep_" + outputs[i]->getName(),
           InsertBefore);
       new StoreInst(outputs[i], GEP, InsertBefore);
+      ++aggIdx;
       // Since there should be only one struct argument aggregating
-      // all the output values, we shouldn't increment OAI, which always
-      // points to the struct argument, in this case.
+      // all the output values, we shouldn't increment AggOutputArgBegin, which
+      // always points to the struct argument, in this case.
     } else {
-      new StoreInst(outputs[i], &*OAI, InsertBefore);
-      ++OAI;
+      assert(ScalarOutputArgBegin != newFunction->arg_end() &&
+             "Number of scalar output arguments should match "
+             "the number of defined values");
+      new StoreInst(outputs[i], &*ScalarOutputArgBegin, InsertBefore);
+      ++ScalarOutputArgBegin;
     }
   }
 
   // Now that we've done the deed, simplify the switch instruction.
   Type *OldFnRetTy = TheSwitch->getParent()->getParent()->getReturnType();
   switch (NumExitBlocks) {
   case 0:
     // There are no successors (the block containing the switch itself), which
for (auto AffectedValVH : AC->assumptionsFor(I->getOperand(0))) {
         return true;
       auto *AssumedInst = cast<Instruction>(AffectedCI->getOperand(0));
       if (AssumedInst->getFunction() != &OldFunc)
         return true;
     }
   }
   return false;
 }
+
+void CodeExtractor::excludeArgFromAggregate(Value *Arg) {
+  ExcludeArgsFromAggregate.insert(Arg);
+}

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

TEST(CodeExtractor, ExitBlockOrderingPhis) {
   ConstantInt *CIFirst = dyn_cast<ConstantInt>(FirstReturn->getReturnValue());
   EXPECT_TRUE(CIFirst->getLimitedValue() == 1u);
 
   Instruction *NextTerm = NextExitStub->getTerminator();
   ReturnInst *NextReturn = dyn_cast<ReturnInst>(NextTerm);
   EXPECT_TRUE(NextReturn);
   ConstantInt *CINext = dyn_cast<ConstantInt>(NextReturn->getReturnValue());
   EXPECT_TRUE(CINext->getLimitedValue() == 0u);
-  
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
 
 TEST(CodeExtractor, ExitBlockOrdering) {
   LLVMContext Ctx;
   SMDiagnostic Err;
   std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
TEST(CodeExtractor, ExitBlockOrdering) {
   ConstantInt *CIFirst = dyn_cast<ConstantInt>(FirstReturn->getReturnValue());
   EXPECT_TRUE(CIFirst->getLimitedValue() == 1u);
 
   Instruction *NextTerm = NextExitStub->getTerminator();
   ReturnInst *NextReturn = dyn_cast<ReturnInst>(NextTerm);
   EXPECT_TRUE(NextReturn);
   ConstantInt *CINext = dyn_cast<ConstantInt>(NextReturn->getReturnValue());
   EXPECT_TRUE(CINext->getLimitedValue() == 0u);
-  
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
 
 TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
   LLVMContext Ctx;
   SMDiagnostic Err;
   std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
TEST(CodeExtractor, RemoveBitcastUsesFromOuterLifetimeMarkers) {
   CE.findInputsOutputs(Inputs, Outputs, SinkingCands);
   EXPECT_EQ(Outputs.size(), 0U);
 
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
+
+TEST(CodeExtractor, PartialAggregateArgs) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"ir(
+    target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+    target triple = "x86_64-unknown-linux-gnu"
+
+    declare void @use(i32)
+
+    define void @foo(i32 %a, i32 %b, i32 %c) {
+    entry:
+      br label %extract
+
+    extract:
+      call void @use(i32 %a)
+      call void @use(i32 %b)
+      call void @use(i32 %c)
+      br label %exit
+
+    exit:
+      ret void
+    }
+  )ir",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 1> Blocks{getBlockByName(Func, "extract")};
+
+  // Create the CodeExtractor with arguments aggregation enabled.
+  CodeExtractor CE(Blocks, /* DominatorTree */ nullptr,
+                   /* AggregateArgs */ true);
+  EXPECT_TRUE(CE.isEligible());
+
+  CodeExtractorAnalysisCache CEAC(*Func);
+  SetVector<Value *> Inputs, Outputs, SinkingCands, HoistingCands;
+  BasicBlock *CommonExit = nullptr;
+  CE.findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
+  CE.findInputsOutputs(Inputs, Outputs, SinkingCands);
+  // Exclude the first input from the argument aggregate.
+  CE.excludeArgFromAggregate(Inputs[0]);
+
+  Function *Outlined = CE.extractCodeRegion(CEAC, Inputs, Outputs);
+  EXPECT_TRUE(Outlined);
+  // Expect 2 arguments in the outlined function: the excluded input and the
+  // struct aggregate for the remaining inputs.
+  EXPECT_EQ(Outlined->arg_size(), 2U);
+  EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
+}
 } // end anonymous namespace