[GlobalOpt] Remove preallocated calls when possible

[aeubanks]

Summary

When possible (e.g. internal linkage), strip preallocated attribute off
parameters/arguments.
This requires removing the "preallocated" operand bundle from the call
site, replacing @llvm.call.preallocated.arg() with an alloca and a
bitcast to i8*, and removing the @llvm.call.preallocated.setup().

This sort of transformation may need to be moved to somewhere more
accessible to accomodate similar transformations like inlining in the
future.

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2289

I couldn't figure out how to remove the preallocated operand bundle any better than this, suggestions greatly appreciated

[aeubanks]

Remove obsolete comment

[efriedma]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2282

Maybe assert any non-callbase use is a BlockAddress

2289

Looked briefly; I don't think there's any other way.

2305

Is the extractProfTotalWeight really necessary? I would have thought copying a call using CallInst::Create() would copy that as well.

2312

I think it's possible to have multiple llvm.call.preallocated.arg() calls with the same index?

If you're creating dynamic allocas, you probably want to llvm.stacksave/stackrestore in appropriate places.

2450

Do you also need to check for musttail calls?

[aeubanks]

Don't remove preallocated when there are any musttail callers

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2282

Done in newly added hasMustTailCallers().

2305

Did a little test, yes it's required. Do you think it makes sense to fix this globally?

2312

I think it's possible to have multiple llvm.call.preallocated.arg() calls with the same index?

Hmm, that's true. I guess for each index that is used, create an alloca right after the llvm.call.preallocated.setup() and replace the corresponding llvm.call.preallocated.arg with the result of that. That way the alloca will dominate all its uses.

If you're creating dynamic allocas, you probably want to llvm.stacksave/stackrestore in appropriate places.

Oh I though llvm automatically handled dynamic allocas, I'll look more into this.

2450

Good point, done

[efriedma]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2305

Maybe worth doing a brief survey of other callers. I expect usually, you'd want to preserve almost all most metadata.

2312

If you don't do anything, the "automatic" handling for dynamic allocas is that they remain live until the function returns, similar to alloca() in C. That's probably not what you want here; you want the memory to be deallocated after the call.

[aeubanks]

Move more comprehensive tests to new file
Handle multiple @llvm.call.preallocated.arg() calls with same arg index
Add @llvm.stack{save,restore}

[aeubanks]

ping :)
I've gotten rid of the metadata cloning here. Originally I had seen the profile metadata copying somewhere else and I mostly copied that code, but I looked back and it was actually copying it from a different instruction.

[efriedma]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

"This doesn't handle invoke" is rough...

In the normal destination, you can stick a stackrestore as the first instruction (breaking the critical edge if necessary). In the unwind destination, I'm not sure. For Itanium-style unwinding, you could just stick it immediately after the landingpad, but I'm not sure what you can do on Windows; you can't stackrestore in a funclet.

I guess ultimately, we really want to use a static alloca (in the entry block) in most cases, but it's not safe in all cases, as we've discussed before.

At the very least, this probably should bail out somehow, not crash.

[aeubanks]

Bail out if any invoke callers

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

Bailed out for any invokes.

I've been trying to understand how inalloca handles this and still don't quite understand.

void bar2() {
 for (int i = 0; i < 1000000; ++i) {
  try {
   foo(A(), 5);
  } catch (int) {}
 }
}

becomes

define dso_local void @"?bar2@@YAXXZ"() local_unnamed_addr #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
entry:
  %agg.tmp.ensured = alloca %struct.A, align 8
  br label %for.body

for.cond.cleanup:                                 ; preds = %for.inc
  ret void

for.body:                                         ; preds = %for.inc, %entry
  %i.05 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
  %inalloca.save = call i8* @llvm.stacksave()
  %argmem = alloca inalloca <{ %struct.A*, %struct.A, i32 }>, align 4
  %0 = getelementptr inbounds <{ %struct.A*, %struct.A, i32 }>, <{ %struct.A*, %struct.A, i32 }>* %argmem, i32 0, i32 1
  %call = call x86_thiscallcc %struct.A* @"??0A@@QAE@XZ"(%struct.A* nonnull %0) #2
  %1 = getelementptr inbounds <{ %struct.A*, %struct.A, i32 }>, <{ %struct.A*, %struct.A, i32 }>* %argmem, i32 0, i32 0
  store %struct.A* %agg.tmp.ensured, %struct.A** %1, align 4
  %2 = getelementptr inbounds <{ %struct.A*, %struct.A, i32 }>, <{ %struct.A*, %struct.A, i32 }>* %argmem, i32 0, i32 2
  store i32 5, i32* %2, align 4, !tbaa !3
  %call1 = invoke %struct.A* @"?foo@@YA?AUA@@U1@H@Z"(<{ %struct.A*, %struct.A, i32 }>* inalloca nonnull %argmem)
          to label %invoke.cont unwind label %catch.dispatch

catch.dispatch:                                   ; preds = %for.body
  %3 = catchswitch within none [label %catch] unwind to caller

catch:                                            ; preds = %catch.dispatch
  %4 = catchpad within %3 [%rtti.TypeDescriptor2* @"??_R0H@8", i32 0, i8* null]
  catchret from %4 to label %for.inc

for.inc:                                          ; preds = %invoke.cont, %catch
  %inc = add nuw nsw i32 %i.05, 1
  %exitcond = icmp eq i32 %inc, 1000000
  br i1 %exitcond, label %for.cond.cleanup, label %for.body

invoke.cont:                                      ; preds = %for.body
  call void @llvm.stackrestore(i8* %inalloca.save)
  call x86_thiscallcc void @"??1A@@QAE@XZ"(%struct.A* nonnull %agg.tmp.ensured) #2
  br label %for.inc
}

The @llvm.stackrestore() only gets called in the non-exceptional case, but I verified that this doesn't leak stack memory when foo() always throws. I'm not sure how that happens.

[efriedma]

LGTM

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

You're not testing what you want to test; the problem isn't when the function with the inalloca argument itself throws. The behavior is obvious in that case: the alloca gets consumed by the call. The issue is what happens if an exception gets thrown before that. Expanding out your example:

struct A { A() { throw 1; } ~A(); };
void foo(A a, int z); // Never gets called
void bar() {
 for (int i = 0; i < 1000000; ++i) {
  try {
   foo(A(), 5);
  } catch (int) {}
 }
}

I'm pretty sure with inalloca, the memory just leaks. I guess preallocated has the same problem,

Not sure how we solve it with preallocated. I guess in the backend, we need to apply some sort of adjustment after a catchret. Not sure how we compute the size of that adjustment, though. I guess we need to note on each catchret/cleanupret which allocations die at that point.

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

I realized that at some point and I've been trying out different variations of A() and foo() both randomly throwing, plus nested calls to foo(), but it seems like there's never a stack leak with inalloca with @llvm.stackrestores removed. The @llvm.stackrestore only appears in the non-exceptional path, so it can't possibly help prevent stack leaks due to exceptions? And in the non-exceptional case the stack will get cleaned up as expected. So I'm not sure that the stack restores are necessary for inalloca?

I am still trying to understand how stack cleanups work in an exception handler though.

void bar() {
        for (int i = 0; i < 1000000; ++i) {
                try {
                        foo(foo(A(), A(), 5), A(), 6);
                } catch (int) {
                }
        }
}

[efriedma]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

Try the following for some interesting results at -O0.

#include <cstdio>
#define INLINE __attribute((always_inline))
//#define INLINE __attribute((noinline))
struct A {
__attribute((optnone)) A(int) {z[1]=10;}
__attribute((optnone)) A() { z[1]=5; throw 1; }
__attribute((optnone)) ~A() { printf("%d\n", z[1]); }
int z[10000];
};
__attribute((optnone)) void foo(int z, A a) { a.z[0] = 100; }
INLINE static int inner() { try {foo(1, A());}catch(int){} return 3; }
void bar() {
 for (int i = 0; i < 10; ++i) {
   foo(inner(), A(1));
 }
}
int main() { bar(); }

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

Interesting...
noinline is fine (with and without removing the stack restores), but always_inline produces incorrect results. Keeping the stack restores, the printfs print a random value (consistent within a run), and removing the outer stack restore makes it crash after the first iteration.

So this is a bug with the current implementation of inalloca? I can't tell if it's specifically an inalloca/exception/inlining issue (the LLVM IR looks fine at a quick glance) or a deeper exception handling issue.

This reminds me of https://gcc.gnu.org/onlinedocs/gcc/Statement-Exprs.html which Reid has brought up to me before, where you can apparently jump out while in the middle of constructing arguments.

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

I tried putting the outer stackrestore in the exception path instead of the non-exceptional path and that makes it produce the right output without stack leaks.

Filed https://bugs.llvm.org/show_bug.cgi?id=46386.

[efriedma]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

One, it looks like llvm::CloneBasicBlock doesn't recognize inalloca allocations as dynamic in some cases. That's pretty clearly a bug, and I think fixing that might be enough to deal with my earlier testcase. But really, it shouldn't matter: the frontend should wrap inalloca allocations with appropriate stacksave/stackrestore operations anyway.

Two, it looks like clang has some general issues with the interaction between dynamic allocation and exceptions. Take the following testcase:

int z = 10000;
__attribute((optnone)) int bar(int*z) { int a = *z; throw a;}
int main() {
    for (int i = 0; i < 1000; ++i) {
    try {
        int a[z];
        bar(a);
        __builtin_trap();
    } catch (...) {
    }
    }
}

This currently overflows the stack on Linux because there isn't an appropriate stackrestore anywhere.

Three, I suspect that unwinding on 32-bit Windows is freeing dynamic allocations in cases where it shouldn't, and this is hiding some inalloca issues. Don't have a good testcase for this; this is just what I suspect from not seeing stack overflows in cases where the IR indicates there should be stack overflows.

[aeubanks]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2334

Fixing CloneBasicBlock doesn't seem to fix the example you gave.

Your example makes more sense as to why the stackrestore is necessary, thanks.

(also I really do appreciate you taking the time to look into things and explain things to me)

[hans]

llvm/lib/Transforms/IPO/GlobalOpt.cpp

2450

Actually, I think the old code is also missing that check :-) Sent D83300

llvm/lib/Transforms/IPO/GlobalOpt.cpp

 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
for (Instruction &I : BB) {
         if (!isColdCallSite(*CI, CallerBFI))
           return false;
       }
     }
   }
   return true;
 }
 
+static bool hasMustTailCallers(Function *F) {
+  for (User *U : F->users()) {
+    CallBase *CB = dyn_cast<CallBase>(U);
+    if (!CB) {
+      assert(isa<BlockAddress>(U) &&
+             "Expected either CallBase or BlockAddress");
+      continue;
+    }
+    if (CB->isMustTailCall())
+      return true;
+  }
+  return false;
+}
+
+static void RemovePreallocated(Function *F) {
+  RemoveAttribute(F, Attribute::Preallocated);
+
+  auto *M = F->getParent();
+
+  IRBuilder<> Builder(M->getContext());
+
+  // Cannot modify users() while iterating over it, so make a copy.
+  SmallVector<User *, 4> PreallocatedCalls(F->users());
+  for (User *U : PreallocatedCalls) {
+    CallBase *CB = dyn_cast<CallBase>(U);
+    if (!CB)
+      continue;
+
+    assert(
+        !CB->isMustTailCall() &&
+        "Shouldn't call RemotePreallocated() on a musttail preallocated call");
+    // Create copy of call without "preallocated" operand bundle.
+    SmallVector<OperandBundleDef, 1> OpBundles;
+    CB->getOperandBundlesAsDefs(OpBundles);
+    CallBase *PreallocatedSetup = nullptr;
+    for (auto *It = OpBundles.begin(); It != OpBundles.end(); ++It) {
+      if (It->getTag() == "preallocated") {
+        PreallocatedSetup = cast<CallBase>(*It->input_begin());
+        OpBundles.erase(It);
+        break;
+      }
+    }
+    assert(PreallocatedSetup && "Did not find preallocated bundle");
+    uint64_t ArgCount =
+        cast<ConstantInt>(PreallocatedSetup->getArgOperand(0))->getZExtValue();
+    CallBase *NewCB = nullptr;
+    if (InvokeInst *II = dyn_cast<InvokeInst>(CB)) {
+      NewCB = InvokeInst::Create(II, OpBundles, CB);
+    } else {
+      CallInst *CI = cast<CallInst>(CB);
+      NewCB = CallInst::Create(CI, OpBundles, CB);
+    }
+    CB->replaceAllUsesWith(NewCB);
+    NewCB->takeName(CB);
+    CB->eraseFromParent();
+
+    Builder.SetInsertPoint(PreallocatedSetup);
+    auto *StackSave =
+        Builder.CreateCall(Intrinsic::getDeclaration(M, Intrinsic::stacksave));
+
+    // FIXME: This doesn't handle invoke
+    Builder.SetInsertPoint(NewCB->getNextNonDebugInstruction());
+    Builder.CreateCall(Intrinsic::getDeclaration(M, Intrinsic::stackrestore),
+                       StackSave);
+
+    // Replace @llvm.call.preallocated.arg() with alloca.
+    // Cannot modify users() while iterating over it, so make a copy.
+    // @llvm.call.preallocated.arg() can be called with the same index multiple
+    // times. So for each @llvm.call.preallocated.arg(), we see if we have
+    // already created a Value* for the index, and if not, create an alloca and
+    // bitcast right after the @llvm.call.preallocated.setup() so that it
+    // dominates all uses.
+    SmallVector<Value *, 2> ArgAllocas(ArgCount);
+    SmallVector<User *, 2> PreallocatedArgs(PreallocatedSetup->users());
+    for (auto *User : PreallocatedArgs) {
+      auto *UseCall = cast<CallBase>(User);
+      assert(UseCall->getCalledFunction()->getIntrinsicID() ==
+                 Intrinsic::call_preallocated_arg &&
+             "preallocated token use was not a llvm.call.preallocated.arg");
+      uint64_t AllocArgIndex =
+          cast<ConstantInt>(UseCall->getArgOperand(1))->getZExtValue();
+      Value *AllocaReplacement = ArgAllocas[AllocArgIndex];
+      if (!AllocaReplacement) {
+        auto AddressSpace = UseCall->getType()->getPointerAddressSpace();
+        auto *ArgType = UseCall
+                            ->getAttribute(AttributeList::FunctionIndex,
+                                           Attribute::Preallocated)
+                            .getValueAsType();
+        auto *InsertBefore = PreallocatedSetup->getNextNonDebugInstruction();
+        Builder.SetInsertPoint(InsertBefore);
+        auto *Alloca =
+            Builder.CreateAlloca(ArgType, AddressSpace, nullptr, "paarg");
+        auto *BitCast = Builder.CreateBitCast(
+            Alloca, Type::getInt8PtrTy(CB->getContext()), UseCall->getName());
+        ArgAllocas[AllocArgIndex] = BitCast;
+        AllocaReplacement = BitCast;
+      }
+
+      UseCall->replaceAllUsesWith(AllocaReplacement);
+      UseCall->eraseFromParent();
+    }
+    // Remove @llvm.call.preallocated.setup().
+    cast<Instruction>(PreallocatedSetup)->eraseFromParent();
+  }
+}
+
 static bool
 OptimizeFunctions(Module &M,
                   function_ref<TargetLibraryInfo &(Function &)> GetTLI,
                   function_ref<TargetTransformInfo &(Function &)> GetTTI,
                   function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
                   function_ref<DominatorTree &(Function &)> LookupDomTree,
                   SmallPtrSetImpl<const Comdat *> &NotDiscardableComdats) {
 
for (Module::iterator FI = M.begin(), E = M.end(); FI != E; ) {
     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.
-    // FIXME: handle preallocated
     if (F->getAttributes().hasAttrSomewhere(Attribute::InAlloca) &&
         !F->hasAddressTaken()) {
       RemoveAttribute(F, Attribute::InAlloca);
       Changed = true;
     }
 
+    if (F->getAttributes().hasAttrSomewhere(Attribute::Preallocated) &&
+        !F->hasAddressTaken() && !hasMustTailCallers(F)) {
+      RemovePreallocated(F);
+      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 ||

llvm/test/Transforms/GlobalOpt/fastcc.ll

 
 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 internal i32 @preallocated(i32* preallocated(i32) %p) {
-; TODO: handle preallocated:
-; CHECK-NOT-LABEL: define internal fastcc i32 @preallocated(i32* %p)
+; CHECK-LABEL: define internal fastcc i32 @preallocated(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)
-  ; TODO: handle preallocated
-  ;%c = call token @llvm.call.preallocated.setup(i32 1)
-  ;%N = call i8* @llvm.call.preallocated.arg(token %c, i32 0) preallocated(i32)
-  ;%n = bitcast i8* %N to i32*
-   ;call i32 @preallocated(i32* preallocated(i32) %n) ["preallocated"(token %c)]
+  %c = call token @llvm.call.preallocated.setup(i32 1)
+  %N = call i8* @llvm.call.preallocated.arg(token %c, i32 0) preallocated(i32)
+  %n = bitcast i8* %N to i32*
+  call i32 @preallocated(i32* preallocated(i32) %n) ["preallocated"(token %c)]
   ret void
 }
-
[e-mail] = 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 fastcc i32 @inalloca(i32* %args)
+; CHECK-NOT: llvm.call.preallocated
+; CHECK: call fastcc i32 @preallocated(i32* %n)
+
[e-mail] = appending global [1 x i8*] [
+   i8* bitcast (i32(i32*)* @j to i8*)
+], section "llvm.metadata"

llvm/test/Transforms/GlobalOpt/preallocated.ll

+; RUN: opt < %s -globalopt -S | FileCheck %s
+
+declare token @llvm.call.preallocated.setup(i32)
+declare i8* @llvm.call.preallocated.arg(token, i32)
+
+define internal i32 @preallocated_musttail(i32* preallocated(i32) %p) {
+; CHECK-LABEL: define internal i32 @preallocated_musttail(i32* preallocated(i32) %p)
+  %rv = load i32, i32* %p
+  ret i32 %rv
+}
+
+define i32 @call_preallocated_musttail(i32* preallocated(i32) %a) {
+  %r = musttail call i32 @preallocated_musttail(i32* preallocated(i32) %a)
+  ret i32 %r
+}
+; CHECK-LABEL: define i32 @call_preallocated_musttail(i32* preallocated(i32) %a)
+; CHECK: musttail call i32 @preallocated_musttail(i32* preallocated(i32) %a)
+
+define i32 @call_preallocated_musttail_without_musttail() {
+  %c = call token @llvm.call.preallocated.setup(i32 1)
+  %N = call i8* @llvm.call.preallocated.arg(token %c, i32 0) preallocated(i32)
+  %n = bitcast i8* %N to i32*
+  %r = call i32 @preallocated_musttail(i32* preallocated(i32) %n) ["preallocated"(token %c)]
+  ret i32 %r
+}
+; CHECK-LABEL: define i32 @call_preallocated_musttail_without_musttail()
+; CHECK: call i32 @preallocated_musttail(i32* preallocated(i32) %n)
+
+define internal i32 @preallocated(i32* preallocated(i32) %a) {
+; CHECK-LABEL: define internal fastcc i32 @preallocated(i32* %a)
+  %rv = load i32, i32* %a
+  ret i32 %rv
+}
+
+declare void @foo(i8*)
+
+define i32 @call_preallocated_multiple_args() {
+; CHECK-LABEL: define i32 @call_preallocated_multiple_args()
+; CHECK-NEXT: [[SS:%[0-9a-zA-Z_]+]] = call i8* @llvm.stacksave()
+; CHECK-NEXT: [[ARG0:%[0-9a-zA-Z_]+]] = alloca i32
+; CHECK-NEXT: [[ARG1:%[0-9a-zA-Z_]+]] = bitcast i32* [[ARG0]] to i8*
+; CHECK-NEXT: call void @foo(i8* [[ARG1]])
+; CHECK-NEXT: call void @foo(i8* [[ARG1]])
+; CHECK-NEXT: call void @foo(i8* [[ARG1]])
+; CHECK-NEXT: [[ARG2:%[0-9a-zA-Z_]+]] = bitcast i8* [[ARG1]] to i32*
+; CHECK-NEXT: call fastcc i32 @preallocated(i32* [[ARG2]])
+; CHECK-NEXT: call void @llvm.stackrestore(i8* [[SS]])
+; CHECK-NEXT: ret
+  %c = call token @llvm.call.preallocated.setup(i32 1)
+  %a1 = call i8* @llvm.call.preallocated.arg(token %c, i32 0) preallocated(i32)
+  call void @foo(i8* %a1)
+  %a2 = call i8* @llvm.call.preallocated.arg(token %c, i32 0) preallocated(i32)
+  call void @foo(i8* %a2)
+  %a3 = call i8* @llvm.call.preallocated.arg(token %c, i32 0) preallocated(i32)
+  call void @foo(i8* %a3)
+  %b = bitcast i8* %a3 to i32*
+  %r = call i32 @preallocated(i32* preallocated(i32) %b) ["preallocated"(token %c)]
+  ret i32 %r
+}