remove inalloca parameters in globalopt and simplify argpromotion

[inglorion]

Summary

Inalloca parameters require special handling in some optimizations.
This change causes globalopt to strip the inalloca attribute from
function parameters when it is safe to do so, removes the special
handling for inallocas from argpromotion, and replaces it with a
simple check that causes argpromotion to skip functions that receive
inallocas (for when the pass is invoked on code that didn't run
through globalopt first). This also avoids a case where argpromotion
would incorrectly try to pass an inalloca in a register.

Fixes PR41658.

[inglorion]

Eli, I saw after I wrote this that you suggested a slightly different approach on the bug (skip argpromote for functions with inallocas). If you'd rather have me do that instead of the change here, I'd be happy to. I just figured I'd put up the one I had already written for consideration.

[efriedma]

Special-casing x86_thiscall seems too narrow; there are other x86 calling conventions that place arguments in registers, so I'm not confident this covers all the cases we care about. Scanning for an inalloca argument should comprehensively cover all the relevant cases.

Alternatively, we could change the x86 calling convention code to make the result valid (if an argument is marked inalloca, force it into memory even if it's in an unexpected position in the argument list).

[inglorion]

Updated in response to @efriedma's comments. We now don't run
argpromotion on functions with inallocas. I've deleted some code
(mostly comments) that refers to them, and modified the test
which previously checked that we do perform the optimization in some
cases to instead check that we don't perform the optimization in
the case I found to be unsafe.

[rnk]

Now that I see we already had an optimization for inalloca, I'm hesitant to lose it. I think there's a simple fix, though. We can just remove inalloca in any situation where we could do argument promotion.

Hoisting the dynamic alloca to the entry block and stripping the inalloca marker off of it are also nice improvements, but they aren't necessary for correctness.

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

152

I think the simplest fix that preserves the existing optimization would be to strip off inalloca from the parameter and argument attributes. I think it's just a matter of pulling off inalloca here, and on the call site.

255–257

This is where you would remove inalloca from the call site attribute list.

Optionally, if inalloca was present and you are removing it, you could check if the value passed into the call site is a pointer to a dynamic alloca marked inalloca, and if so turn it into a static alloca by stripping the inalloca marker and hoisting it to the entry block. That's not necessary for correctness, though.

[inglorion]

Yeah, I kind of wanted to do the least intrusive thing that would fix the problem I observed and then think about the larger picture. But given that we're thinking about removing the inalloca attribute, the way I would like to do that is remove it in globalopt and then we can perform further optimizations (in globalopt and here) predicated on no inalloca being present. That takes care of half the FIXME for inalloca in globalopt, so after that it becomes really tempting to do follow-up with the other half and hoist the alloca to the entry block.

[rnk]

In D61286#1485024, @inglorion wrote:

Yeah, I kind of wanted to do the least intrusive thing that would fix the problem I observed and then think about the larger picture. But given that we're thinking about removing the inalloca attribute, the way I would like to do that is remove it in globalopt and then we can perform further optimizations (in globalopt and here) predicated on no inalloca being present. That takes care of half the FIXME for inalloca in globalopt, so after that it becomes really tempting to do follow-up with the other half and hoist the alloca to the entry block.

I approve of this plan to do the removal in globalopt and make the later passes more conservative around inalloca. :)

llvm/test/Transforms/ArgumentPromotion/inalloca.ll

1–2

I would say, keep this test case as is, but once you've improved globalopt, test that the combination of globalopt + argpromotion + sroa is able to do this transform successfully. It's a bit of an integration test, but we need a few of those than we have today anyway.

[inglorion]

Remove inallocas in globalopt and simplify argpromotion by just
bailing on inallocas.

Tests verify that the optimizations we previously performed are still
performed when both globalopt and argpromotion run.

[inglorion]

Correcting earlier upload to preserve original inalloca.ll tests and
add new tests in a separate file.

[rnk]

lgtm, thanks for the improvement and simplification.

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

if (ByValArgsToTransform.count(&*I)) {
       StructType *STy = cast<StructType>(AgTy);
       Params.insert(Params.end(), STy->element_begin(), STy->element_end());
       ArgAttrVec.insert(ArgAttrVec.end(), STy->getNumElements(),
                         AttributeSet());
       ++NumByValArgsPromoted;
     } else if (!ArgsToPromote.count(&*I)) {
       // Unchanged argument
       Params.push_back(I->getType());
       ArgAttrVec.push_back(PAL.getParamAttributes(ArgNo));
     } else if (I->use_empty()) {
       // Dead argument (which are always marked as promotable)
       ++NumArgumentsDead;
 
       // There may be remaining metadata uses of the argument for things like
       // llvm.dbg.value. Replace them with undef.
       I->replaceAllUsesWith(UndefValue::get(I->getType()));
     } else {
while (!F->use_empty()) {
     IRBuilder<NoFolder> IRB(Call);
 
     // Loop over the operands, inserting GEP and loads in the caller as
     // appropriate.
     CallSite::arg_iterator AI = CS.arg_begin();
     ArgNo = 0;
     for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
          ++I, ++AI, ++ArgNo)
       if (!ArgsToPromote.count(&*I) && !ByValArgsToTransform.count(&*I)) {
         Args.push_back(*AI); // Unmodified argument
         ArgAttrVec.push_back(CallPAL.getParamAttributes(ArgNo));
       } else if (ByValArgsToTransform.count(&*I)) {
         // Emit a GEP and load for each element of the struct.
         Type *AgTy = cast<PointerType>(I->getType())->getElementType();
         StructType *STy = cast<StructType>(AgTy);
         Value *Idxs[2] = {
             ConstantInt::get(Type::getInt32Ty(F->getContext()), 0), nullptr};
         for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i) {
           Idxs[1] = ConstantInt::get(Type::getInt32Ty(F->getContext()), i);
static void markIndicesSafe(const IndicesVector &ToMark,
   }
 }
 
 /// isSafeToPromoteArgument - As you might guess from the name of this method,
 /// it checks to see if it is both safe and useful to promote the argument.
 /// This method limits promotion of aggregates to only promote up to three
 /// elements of the aggregate in order to avoid exploding the number of
 /// arguments passed in.
-static bool isSafeToPromoteArgument(Argument *Arg, bool isByValOrInAlloca,
-                                    AAResults &AAR, unsigned MaxElements) {
+static bool isSafeToPromoteArgument(Argument *Arg, bool isByVal, AAResults &AAR,
+                                    unsigned MaxElements) {
   using GEPIndicesSet = std::set<IndicesVector>;
 
   // Quick exit for unused arguments
   if (Arg->use_empty())
     return true;
 
   // We can only promote this argument if all of the uses are loads, or are GEP
   // instructions (with constant indices) that are subsequently loaded.
   //
   // Promoting the argument causes it to be loaded in the caller
   // unconditionally. This is only safe if we can prove that either the load
   // would have happened in the callee anyway (ie, there is a load in the entry
   // block) or the pointer passed in at every call site is guaranteed to be
   // valid.
   // In the former case, invalid loads can happen, but would have happened
   // anyway, in the latter case, invalid loads won't happen. This prevents us
   // from introducing an invalid load that wouldn't have happened in the
   // original code.
   //
   // This set will contain all sets of indices that are loaded in the entry
   // block, and thus are safe to unconditionally load in the caller.
-  //
-  // This optimization is also safe for InAlloca parameters, because it verifies
-  // that the address isn't captured.
   GEPIndicesSet SafeToUnconditionallyLoad;
 
   // This set contains all the sets of indices that we are planning to promote.
   // This makes it possible to limit the number of arguments added.
   GEPIndicesSet ToPromote;
 
   // If the pointer is always valid, any load with first index 0 is valid.
-  if (isByValOrInAlloca || allCallersPassInValidPointerForArgument(Arg))
+  if (isByVal || allCallersPassInValidPointerForArgument(Arg))
     SafeToUnconditionallyLoad.insert(IndicesVector(1, 0));
 
   // First, iterate the entry block and mark loads of (geps of) arguments as
   // safe.
   BasicBlock &EntryBlock = Arg->getParent()->front();
   // Declare this here so we can reuse it
   IndicesVector Indices;
   for (Instruction &I : EntryBlock)
for (Use &U : Arg->uses()) {
     } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(UR)) {
       if (GEP->use_empty()) {
         // Dead GEP's cause trouble later.  Just remove them if we run into
         // them.
         GEP->eraseFromParent();
         // TODO: This runs the above loop over and over again for dead GEPs
         // Couldn't we just do increment the UI iterator earlier and erase the
         // use?
-        return isSafeToPromoteArgument(Arg, isByValOrInAlloca, AAR,
-                                       MaxElements);
+        return isSafeToPromoteArgument(Arg, isByVal, AAR, MaxElements);
       }
 
       // Ensure that all of the indices are constants.
       for (User::op_iterator i = GEP->idx_begin(), e = GEP->idx_end(); i != e;
            ++i)
         if (ConstantInt *C = dyn_cast<ConstantInt>(*i))
           Operands.push_back(C->getSExtValue());
         else
promoteArguments(Function *F, function_ref<AAResults &(Function &F)> AARGetter,
   // Don't promote arguments for variadic functions. Adding, removing, or
   // changing non-pack parameters can change the classification of pack
   // parameters. Frontends encode that classification at the call site in the
   // IR, while in the callee the classification is determined dynamically based
   // on the number of registers consumed so far.
   if (F->isVarArg())
     return nullptr;
 
+  // Don't transform functions that receive inallocas, as the transformation may
+  // not be safe depending on calling convention.
+  if (F->getAttributes().hasAttrSomewhere(Attribute::InAlloca))
+    return nullptr;
+
   // First check: see if there are any pointer arguments!  If not, quick exit.
   SmallVector<Argument *, 16> PointerArgs;
   for (Argument &I : F->args())
     if (I.getType()->isPointerTy())
       PointerArgs.push_back(&I);
   if (PointerArgs.empty())
     return nullptr;
 
if (PtrArg->hasStructRetAttr()) {
         CallSite CS(U.getUser());
         CS.removeParamAttr(ArgNo, Attribute::StructRet);
         CS.addParamAttr(ArgNo, Attribute::NoAlias);
       }
     }
 
     // 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.
+    // packed or if we can prove the padding bytes are never accessed.
     bool isSafeToPromote =
         PtrArg->hasByValAttr() &&
         (isDenselyPacked(AgTy, DL) || !canPaddingBeAccessed(PtrArg));
     if (isSafeToPromote) {
       if (StructType *STy = dyn_cast<StructType>(AgTy)) {
         if (MaxElements > 0 && STy->getNumElements() > MaxElements) {
           LLVM_DEBUG(dbgs() << "argpromotion disable promoting argument '"
                             << PtrArg->getName()
if (isSelfRecursive) {
           }
         }
         if (RecursiveType)
           continue;
       }
     }
 
     // Otherwise, see if we can promote the pointer to its value.
-    if (isSafeToPromoteArgument(PtrArg, PtrArg->hasByValOrInAllocaAttr(), AAR,
+    if (isSafeToPromoteArgument(PtrArg, PtrArg->hasByValAttr(), AAR,
                                 MaxElements))
       ArgsToPromote.insert(PtrArg);
   }
 
   // No promotable pointer arguments.
   if (ArgsToPromote.empty() && ByValArgsToTransform.empty())
     return nullptr;
 

llvm/lib/Transforms/IPO/GlobalOpt.cpp

static void ChangeCalleesToFastCall(Function *F) {
   for (User *U : F->users()) {
     if (isa<BlockAddress>(U))
       continue;
     CallSite CS(cast<Instruction>(U));
     CS.setCallingConv(CallingConv::Fast);
   }
 }
 
-static AttributeList StripNest(LLVMContext &C, AttributeList Attrs) {
-  // There can be at most one attribute set with a nest attribute.
-  unsigned NestIndex;
-  if (Attrs.hasAttrSomewhere(Attribute::Nest, &NestIndex))
-    return Attrs.removeAttribute(C, NestIndex, Attribute::Nest);
+static AttributeList StripAttr(LLVMContext &C, AttributeList Attrs,
+                               Attribute::AttrKind A) {
+  unsigned AttrIndex;
+  if (Attrs.hasAttrSomewhere(A, &AttrIndex))
+    return Attrs.removeAttribute(C, AttrIndex, A);
   return Attrs;
 }
 
-static void RemoveNestAttribute(Function *F) {
-  F->setAttributes(StripNest(F->getContext(), F->getAttributes()));
+static void RemoveAttribute(Function *F, Attribute::AttrKind A) {
+  F->setAttributes(StripAttr(F->getContext(), F->getAttributes(), A));
   for (User *U : F->users()) {
     if (isa<BlockAddress>(U))
       continue;
     CallSite CS(cast<Instruction>(U));
-    CS.setAttributes(StripNest(F->getContext(), CS.getAttributes()));
+    CS.setAttributes(StripAttr(F->getContext(), CS.getAttributes(), A));
   }
 }
 
 /// Return true if this is a calling convention that we'd like to change.  The
 /// idea here is that we don't want to mess with the convention if the user
 /// explicitly requested something with performance implications like coldcc,
 /// GHC, or anyregcc.
 static bool hasChangeableCC(Function *F) {
   CallingConv::ID CC = F->getCallingConv();
 
   // FIXME: Is it worth transforming x86_stdcallcc and x86_fastcallcc?
   if (CC != CallingConv::C && CC != CallingConv::X86_ThisCall)
     return false;
 
-  // Don't break the invariant that the inalloca parameter is the only parameter
-  // passed in memory.
-  // FIXME: GlobalOpt should remove inalloca when possible and hoist the dynamic
-  // alloca it uses to the entry block if possible.
-  if (F->getAttributes().hasAttrSomewhere(Attribute::InAlloca))
-    return false;
-
   // FIXME: Change CC for the whole chain of musttail calls when possible.
   //
   // Can't change CC of the function that either has musttail calls, or is a
   // musttail callee itself
   for (User *U : F->users()) {
     if (isa<BlockAddress>(U))
       continue;
     CallInst* CI = dyn_cast<CallInst>(U);
if (!F->isDeclaration()) {
       }
     }
 
     Changed |= processGlobal(*F, TLI, LookupDomTree);
 
     if (!F->hasLocalLinkage())
       continue;
 
+    // If we have an inalloca parameter that we can safely remove the
+    // inalloca attribute from, do so. This unlocks optimizations that
+    // wouldn't be safe in the presence of inalloca.
+    // FIXME: We should also hoist alloca affected by this to the entry
+    // block if possible.
+    if (F->getAttributes().hasAttrSomewhere(Attribute::InAlloca) &&
+        !F->hasAddressTaken()) {
+      RemoveAttribute(F, Attribute::InAlloca);
+      Changed = true;
+    }
+
     if (hasChangeableCC(F) && !F->isVarArg() && !F->hasAddressTaken()) {
       NumInternalFunc++;
       TargetTransformInfo &TTI = GetTTI(*F);
       // Change the calling convention to coldcc if either stress testing is
       // enabled or the target would like to use coldcc on functions which are
       // cold at all call sites and the callers contain no other non coldcc
       // calls.
       if (EnableColdCCStressTest ||
if (hasChangeableCC(F) && !F->isVarArg() &&
       ++NumFastCallFns;
       Changed = true;
     }
 
     if (F->getAttributes().hasAttrSomewhere(Attribute::Nest) &&
         !F->hasAddressTaken()) {
       // The function is not used by a trampoline intrinsic, so it is safe
       // to remove the 'nest' attribute.
-      RemoveNestAttribute(F);
+      RemoveAttribute(F, Attribute::Nest);
       ++NumNestRemoved;
       Changed = true;
     }
   }
   return Changed;
 }
 
 static bool

llvm/test/Transforms/ArgumentPromotion/X86/thiscall.ll

+; In PR41658, argpromotion put an inalloca in a position that per the
+; calling convention is passed in a register. This test verifies that
+; we don't do that anymore. It also verifies that the combination of
+; globalopt and argpromotion is able to optimize the call safely.
+;
+; RUN: opt -S -argpromotion %s | FileCheck --check-prefix=THIS %s
+; RUN: opt -S -globalopt -argpromotion %s | FileCheck --check-prefix=OPT %s
+; THIS: define internal x86_thiscallcc void @internalfun(%struct.a* %this, <{ %struct.a
+; OPT: define internal fastcc void @internalfun(<{ %struct.a }>*)
+
+target datalayout = "e-m:x-p:32:32-i64:64-f80:32-n8:16:32-a:0:32-S32"
+target triple = "i386-pc-windows-msvc19.11.0"
+
+%struct.a = type { i8 }
+
+define internal x86_thiscallcc void @internalfun(%struct.a* %this, <{ %struct.a }>* inalloca) {
+entry:
+  %a = getelementptr inbounds <{ %struct.a }>, <{ %struct.a }>* %0, i32 0, i32 0
+  %argmem = alloca inalloca <{ %struct.a }>, align 4
+  %1 = getelementptr inbounds <{ %struct.a }>, <{ %struct.a }>* %argmem, i32 0, i32 0
+  %call = call x86_thiscallcc %struct.a* @copy_ctor(%struct.a* %1, %struct.a* dereferenceable(1) %a)
+  call void @ext(<{ %struct.a }>* inalloca %argmem)
+  ret void
+}
+
+; This is here to ensure @internalfun is live.
+define void @exportedfun(%struct.a* %a) {
+  %inalloca.save = tail call i8* @llvm.stacksave()
+  %argmem = alloca inalloca <{ %struct.a }>, align 4
+  call x86_thiscallcc void @internalfun(%struct.a* %a, <{ %struct.a }>* inalloca %argmem)
+  call void @llvm.stackrestore(i8* %inalloca.save)
+  ret void
+}
+
+declare x86_thiscallcc %struct.a* @copy_ctor(%struct.a* returned, %struct.a* dereferenceable(1))
+declare void @ext(<{ %struct.a }>* inalloca)
+declare i8* @llvm.stacksave()
+declare void @llvm.stackrestore(i8*)

llvm/test/Transforms/ArgumentPromotion/inalloca.ll

-; RUN: opt %s -argpromotion -sroa -S | FileCheck %s
-; RUN: opt %s -passes='argpromotion,function(sroa)' -S | FileCheck %s
+; RUN: opt %s -globalopt -argpromotion -sroa -S | FileCheck %s
+; RUN: opt %s -passes='module(globalopt),cgscc(argpromotion),function(sroa)' -S | FileCheck %s
 
 target datalayout = "E-p:64:64:64-a0:0:8-f32:32:32-f64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-v64:64:64-v128:128:128"
 
 %struct.ss = type { i32, i32 }
 
 ; Argpromote + sroa should change this to passing the two integers by value.
 define internal i32 @f(%struct.ss* inalloca  %s) {
 entry:
   %f0 = getelementptr %struct.ss, %struct.ss* %s, i32 0, i32 0
   %f1 = getelementptr %struct.ss, %struct.ss* %s, i32 0, i32 1
   %a = load i32, i32* %f0, align 4
   %b = load i32, i32* %f1, align 4
   %r = add i32 %a, %b
   ret i32 %r
 }
-; CHECK-LABEL: define internal i32 @f
+; CHECK-LABEL: define internal fastcc i32 @f
 ; CHECK-NOT: load
 ; CHECK: ret
 
 define i32 @main() {
 entry:
   %S = alloca inalloca %struct.ss
   %f0 = getelementptr %struct.ss, %struct.ss* %S, i32 0, i32 0
   %f1 = getelementptr %struct.ss, %struct.ss* %S, i32 0, i32 1
   store i32 1, i32* %f0, align 4
   store i32 2, i32* %f1, align 4
   %r = call i32 @f(%struct.ss* inalloca %S)
   ret i32 %r
 }
 ; CHECK-LABEL: define i32 @main
 ; CHECK-NOT: load
 ; CHECK: ret
 
 ; Argpromote can't promote %a because of the icmp use.
 define internal i1 @g(%struct.ss* %a, %struct.ss* inalloca %b) nounwind  {
-; CHECK: define internal i1 @g(%struct.ss* %a, %struct.ss* inalloca %b)
+; CHECK: define internal fastcc i1 @g(%struct.ss* %a, %struct.ss* %b)
 entry:
   %c = icmp eq %struct.ss* %a, %b
   ret i1 %c
 }
 
 define i32 @test() {
 entry:
   %S = alloca inalloca %struct.ss
   %c = call i1 @g(%struct.ss* %S, %struct.ss* inalloca %S)
-; CHECK: call i1 @g(%struct.ss* %S, %struct.ss* inalloca %S)
+; CHECK: call fastcc i1 @g(%struct.ss* %S, %struct.ss* %S)
   ret i32 0
 }

llvm/test/Transforms/GlobalOpt/fastcc.ll

 
 define internal i32 @j(i32* %m) {
 ; CHECK-LABEL: define internal i32 @j
   %v = load i32, i32* %m
   ret i32 %v
 }
 
 define internal i32 @inalloca(i32* inalloca %p) {
-; CHECK-LABEL: define internal i32 @inalloca(i32* inalloca %p)
+; CHECK-LABEL: define internal fastcc i32 @inalloca(i32* %p)
   %rv = load i32, i32* %p
   ret i32 %rv
 }
 
 define void @call_things() {
   %m = alloca i32
   call i32 @f(i32* %m)
   call x86_thiscallcc i32 @g(i32* %m)
   call coldcc i32 @h(i32* %m)
   call i32 @j(i32* %m)
   %args = alloca inalloca i32
   call i32 @inalloca(i32* inalloca %args)
   ret void
 }
 
 @llvm.used = appending global [1 x i8*] [
    i8* bitcast (i32(i32*)* @j to i8*)
 ], section "llvm.metadata"
 
 ; CHECK-LABEL: define void @call_things()
 ; CHECK: call fastcc i32 @f
 ; CHECK: call fastcc i32 @g
 ; CHECK: call coldcc i32 @h
 ; CHECK: call i32 @j
-; CHECK: call i32 @inalloca(i32* inalloca %args)
+; CHECK: call fastcc i32 @inalloca(i32* %args)