ArgumentPromotion: Drop sret attribute on functions that are only called directly.

[pcc]

Summary

If the first argument to a function is a 'this' argument and the second
has the sret attribute, the ArgumentPromotion pass may promote the 'this'
argument to more than one argument, violating the IR constraint that 'sret'
may only be applied to the first or second argument.

Although this IR constraint is arguably unnecessary, it highlighted the fact
that ArgPromotion does not need to preserve this attribute. Dropping the
attribute reduces register pressure in the backend by avoiding the register
copy required by sret. Because sret implies noalias, we also replace the
former with the latter.

[rnk]

First, the limitation that sret be parameter 1 or 2 is artificial. At the time, I didn't see how it could appear on any other parameter and I tried to make the most minimal change. Now we have a reason to generalize to other parameters. That said, all the backends handle sret appearing on any parameter, because the first parameter can be split up into multiple MVTs. Consider passing i64 as the first parameter on a 32-bit platform. So, we could lift the verifier limitation.

Second, could argpromotion just drop sret in this case? sret does two things: it's equivalent to noalias, and it ensures that the sret pointer is returned in eax. If all uses of the function are internal, dropping sret is conservatively correct and reduces register pressure. We don't leverage that eax return because it would be an ABI break with old LLVMs, so we can replace sret with noalias without losing any power.

Finally, if you like having argpromotion canonicalize sret to be the first parameter, I guess that's also kind of nice, so we could continue on this way. :)

lib/IR/Function.cpp

157–158

Oops, that's just a bug.

[pcc]

Second, could argpromotion just drop sret in this case? sret does two things: it's equivalent to noalias, and it ensures that the sret pointer is returned in eax. If all uses of the function are internal, dropping sret is conservatively correct and reduces register pressure. We don't leverage that eax return because it would be an ABI break with old LLVMs, so we can replace sret with noalias without losing any power.

I agree, that seems like a much better solution.

[pcc]

New patch

[rnk]

lgtm

This might be an interesting change, since we'll basically remove sret from all directly called functions in LTO builds, even if we weren't going to do arg promotion. I'm pretty confident it's safe, though.

lib/Transforms/IPO/ArgumentPromotion.cpp

260–264

I think this is inefficient because we end up with lots of temporary attribute sets, but whatever. It's kind of an API deficiency. =/

The alternative would be to compute the set of parameters that are losing sret and feed that into DoPromotion, which is ugly.

lib/IR/Function.cpp

bool Argument::hasNoCaptureAttr() const {
   return getParent()->getAttributes().
     hasAttribute(getArgNo()+1, Attribute::NoCapture);
 }
 
 /// hasSRetAttr - Return true if this argument has the sret attribute on
 /// it in its containing function.
 bool Argument::hasStructRetAttr() const {
   if (!getType()->isPointerTy()) return false;
-  if (this != getParent()->arg_begin())
-    return false; // StructRet param must be first param
   return getParent()->getAttributes().
-    hasAttribute(1, Attribute::StructRet);
+    hasAttribute(getArgNo()+1, Attribute::StructRet);
 }
 
 /// hasReturnedAttr - Return true if this argument has the returned attribute on
 /// it in its containing function.
 bool Argument::hasReturnedAttr() const {
   return getParent()->getAttributes().
     hasAttribute(getArgNo()+1, Attribute::Returned);
 }

lib/Transforms/IPO/ArgumentPromotion.cpp

namespace {
 
   private:
     bool isDenselyPacked(Type *type, const DataLayout &DL);
     bool canPaddingBeAccessed(Argument *Arg);
     CallGraphNode *PromoteArguments(CallGraphNode *CGN);
     bool isSafeToPromoteArgument(Argument *Arg, bool isByVal) const;
     CallGraphNode *DoPromotion(Function *F,
                               SmallPtrSetImpl<Argument*> &ArgsToPromote,
-                              SmallPtrSetImpl<Argument*> &ByValArgsToTransform);
+                              SmallPtrSetImpl<Argument*> &ByValArgsToTransform,
+                              unsigned StructRetArgIndex);
     
     using llvm::Pass::doInitialization;
     bool doInitialization(CallGraph &CG) override;
     /// The maximum number of elements to expand, or 0 for unlimited.
     unsigned maxElements;
     DenseMap<const Function *, DISubprogram *> FunctionDIs;
   };
 }
CallGraphNode *ArgPromotion::PromoteArguments(CallGraphNode *CGN) {
   }
   
   const DataLayout &DL = F->getParent()->getDataLayout();
 
   // Check to see which arguments are promotable.  If an argument is promotable,
   // add it to ArgsToPromote.
   SmallPtrSet<Argument*, 8> ArgsToPromote;
   SmallPtrSet<Argument*, 8> ByValArgsToTransform;
+  unsigned StructRetArgIndex = 0;
   for (unsigned i = 0, e = PointerArgs.size(); i != e; ++i) {
     Argument *PtrArg = PointerArgs[i];
     Type *AgTy = cast<PointerType>(PtrArg->getType())->getElementType();
 
+    // Keep track of the location of the sret argument.
+    if (PtrArg->hasStructRetAttr()) {
+      StructRetArgIndex = PtrArg->getArgNo() + 1;
+      continue;
+    }
+
     // If this is a byval argument, and if the aggregate type is small, just
     // pass the elements, which is always safe, if the passed value is densely
     // packed or if we can prove the padding bytes are never accessed. This does
     // not apply to inalloca.
     bool isSafeToPromote =
         PtrArg->hasByValAttr() &&
         (isDenselyPacked(AgTy, DL) || !canPaddingBeAccessed(PtrArg));
     if (isSafeToPromote) {
       if (StructType *STy = dyn_cast<StructType>(AgTy)) {
         if (maxElements > 0 && STy->getNumElements() > maxElements) {
           DEBUG(dbgs() << "argpromotion disable promoting argument '"
                 << PtrArg->getName() << "' because it would require adding more"
                 << " than " << maxElements << " arguments to the function.\n");
           continue;
         }
         
         // If all the elements are single-value types, we can promote it.
for (unsigned i = 0, e = PointerArgs.size(); i != e; ++i) {
     if (isSafeToPromoteArgument(PtrArg, PtrArg->hasByValOrInAllocaAttr()))
       ArgsToPromote.insert(PtrArg);
   }
 
   // No promotable pointer arguments.
   if (ArgsToPromote.empty() && ByValArgsToTransform.empty()) 
     return nullptr;
 
-  return DoPromotion(F, ArgsToPromote, ByValArgsToTransform);
+  return DoPromotion(F, ArgsToPromote, ByValArgsToTransform, StructRetArgIndex);
 }
 
 /// AllCallersPassInValidPointerForArgument - Return true if we can prove that
 /// all callees pass in a valid pointer for the specified function argument.
 static bool AllCallersPassInValidPointerForArgument(Argument *Arg) {
   Function *Callee = Arg->getParent();
   const DataLayout &DL = Callee->getParent()->getDataLayout();
 
bool ArgPromotion::isSafeToPromoteArgument(Argument *Arg,
   return true;
 }
 
 /// DoPromotion - This method actually performs the promotion of the specified
 /// arguments, and returns the new function.  At this point, we know that it's
 /// safe to do so.
 CallGraphNode *ArgPromotion::DoPromotion(Function *F,
                              SmallPtrSetImpl<Argument*> &ArgsToPromote,
-                             SmallPtrSetImpl<Argument*> &ByValArgsToTransform) {
+                             SmallPtrSetImpl<Argument*> &ByValArgsToTransform,
+                             unsigned StructRetArgIndex) {
 
   // Start by computing a new prototype for the function, which is the same as
   // the old function, but has modified arguments.
   FunctionType *FTy = F->getFunctionType();
   std::vector<Type*> Params;
 
   typedef std::set<std::pair<Type *, IndicesVector>> ScalarizeTable;
 
CallGraphNode *ArgPromotion::DoPromotion(Function *F,
   const AttributeSet &PAL = F->getAttributes();
 
   // Add any return attributes.
   if (PAL.hasAttributes(AttributeSet::ReturnIndex))
     AttributesVec.push_back(AttributeSet::get(F->getContext(),
                                               PAL.getRetAttributes()));
 
   // First, determine the new argument list
+
+  if (StructRetArgIndex) {
+    // Add the sret argument first now in case the first argument is being
+    // promoted to more than one argument, which would otherwise violate the IR
+    // constraint that 'sret' may only be applied to the first or second
+    // argument.
+    Params.push_back(F->getFunctionType()->getParamType(StructRetArgIndex - 1));
+    AttributeSet attrs = PAL.getParamAttributes(StructRetArgIndex);
+    AttrBuilder B(attrs, StructRetArgIndex);
+    AttributesVec.push_back(
+        AttributeSet::get(F->getContext(), Params.size(), B));
+  }
+
   unsigned ArgIndex = 1;
   for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
        ++I, ++ArgIndex) {
+    if (ArgIndex == StructRetArgIndex)
+      continue;
     if (ByValArgsToTransform.count(I)) {
       // Simple byval argument? Just add all the struct element types.
       Type *AgTy = cast<PointerType>(I->getType())->getElementType();
       StructType *STy = cast<StructType>(AgTy);
       Params.insert(Params.end(), STy->element_begin(), STy->element_end());
       ++NumByValArgsPromoted;
     } else if (!ArgsToPromote.count(I)) {
       // Unchanged argument
while (!F->use_empty()) {
     Instruction *Call = CS.getInstruction();
     const AttributeSet &CallPAL = CS.getAttributes();
 
     // Add any return attributes.
     if (CallPAL.hasAttributes(AttributeSet::ReturnIndex))
       AttributesVec.push_back(AttributeSet::get(F->getContext(),
                                                 CallPAL.getRetAttributes()));
 
+    // Add the sret argument.
+    if (StructRetArgIndex) {
+      Args.push_back(CS.getArgument(StructRetArgIndex - 1));
+      AttrBuilder B(CallPAL, StructRetArgIndex);
+      AttributesVec.
+          push_back(AttributeSet::get(F->getContext(), Args.size(), B));
+    }
+
     // Loop over the operands, inserting GEP and loads in the caller as
     // appropriate.
     CallSite::arg_iterator AI = CS.arg_begin();
     ArgIndex = 1;
     for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end();
-         I != E; ++I, ++AI, ++ArgIndex)
+         I != E; ++I, ++AI, ++ArgIndex) {
+      if (ArgIndex == StructRetArgIndex)
+        continue;
       if (!ArgsToPromote.count(I) && !ByValArgsToTransform.count(I)) {
         Args.push_back(*AI);          // Unmodified argument
 
         if (CallPAL.hasAttributes(ArgIndex)) {
           AttrBuilder B(CallPAL, ArgIndex);
           AttributesVec.
             push_back(AttributeSet::get(F->getContext(), Args.size(), B));
         }
for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end();
           AAMDNodes AAInfo;
           OrigLoad->getAAMetadata(AAInfo);
           newLoad->setAAMetadata(AAInfo);
 
           Args.push_back(newLoad);
           AA.copyValue(OrigLoad, Args.back());
         }
       }
+    }
 
     // Push any varargs arguments on the list.
     for (; AI != CS.arg_end(); ++AI, ++ArgIndex) {
       Args.push_back(*AI);
       if (CallPAL.hasAttributes(ArgIndex)) {
         AttrBuilder B(CallPAL, ArgIndex);
         AttributesVec.
           push_back(AttributeSet::get(F->getContext(), Args.size(), B));
CallGraphNode *ArgPromotion::DoPromotion(Function *F,
   }
 
   // Since we have now created the new function, splice the body of the old
   // function right into the new function, leaving the old rotting hulk of the
   // function empty.
   NF->getBasicBlockList().splice(NF->begin(), F->getBasicBlockList());
 
   // Loop over the argument list, transferring uses of the old arguments over to
-  // the new arguments, also transferring over the names as well.
-  //
-  for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(),
-       I2 = NF->arg_begin(); I != E; ++I) {
+  // the new arguments, also transferring over the names as well, but deal with
+  // the sret argument first.
+  Function::arg_iterator I2 = NF->arg_begin();
+  if (StructRetArgIndex) {
+    Function::arg_iterator I = F->arg_begin();
+    std::advance(I, StructRetArgIndex - 1);
+    I->replaceAllUsesWith(I2);
+    I2->takeName(I);
+    AA.replaceWithNewValue(I, I2);
+    ++I2;
+  }
+
+  for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E; ++I) {
+    if (I->hasStructRetAttr())
+      continue;
     if (!ArgsToPromote.count(I) && !ByValArgsToTransform.count(I)) {
       // If this is an unmodified argument, move the name and users over to the
       // new version.
       I->replaceAllUsesWith(I2);
       I2->takeName(I);
       AA.replaceWithNewValue(I, I2);
       ++I2;
       continue;

test/Transforms/ArgumentPromotion/sret.ll

+; RUN: opt < %s -argpromotion -S | FileCheck %s
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc"
+
+; CHECK: define internal void @add(i32* sret %[[SR:.*]], i32 %[[THIS1:.*]], i32 %[[THIS2:.*]])
+define internal void @add({i32, i32}* %this, i32* sret %r) {
+  %ap = getelementptr {i32, i32}, {i32, i32}* %this, i32 0, i32 0
+  %bp = getelementptr {i32, i32}, {i32, i32}* %this, i32 0, i32 1
+  %a = load i32, i32* %ap
+  %b = load i32, i32* %bp
+  ; CHECK: %[[AB:.*]] = add i32 %[[THIS1]], %[[THIS2]]
+  %ab = add i32 %a, %b
+  ; CHECK: store i32 %[[AB]], i32* %[[SR]]
+  store i32 %ab, i32* %r
+  ret void
+}
+
+; CHECK: define void @f()
+define void @f() {
+  ; CHECK: %[[R:.*]] = alloca i32
+  ; CHECK: %[[PAIR:.*]] = alloca { i32, i32 }
+  %r = alloca i32
+  %pair = alloca {i32, i32}
+
+  ; CHECK: %[[PAIR1P:.*]] = getelementptr { i32, i32 }, { i32, i32 }* %[[PAIR]], i64 0, i32 0
+  ; CHECK: %[[PAIR1:.*]] = load i32, i32* %[[PAIR1P]]
+  ; CHECK: %[[PAIR2P:.*]] = getelementptr { i32, i32 }, { i32, i32 }* %[[PAIR]], i64 0, i32 1
+  ; CHECK: %[[PAIR2:.*]] = load i32, i32* %[[PAIR2P]]
+  ; CHECK: call void @add(i32* sret %[[R]], i32 %[[PAIR1]], i32 %[[PAIR2]])
+  call void @add({i32, i32}* %pair, i32* sret %r)
+  ret void
+}