[Preallocated] Add @llvm.call.preallocated.teardown

[aeubanks]

Summary

This cleans up the stack allocated by a @llvm.call.preallocated.setup.
Should either call the teardown or the preallocated call to clean up the
stack. Calling both is UB.

Add LangRef.

Add verifier check that the token argument is a @llvm.call.preallocated.setup.

[efriedma]

LGTM

llvm/docs/LangRef.rst

 to DCE), where otherwise we cannot know how large the arguments are.
 
 It is undefined behavior if this is called with a token from an
 '``llvm.call.preallocated.setup``' if another
 '``llvm.call.preallocated.setup``' has already been called or if the
 preallocated call corresponding to the '``llvm.call.preallocated.setup``'
 has already been called.
 
+.. _int_call_preallocated_teardown:
+
+'``llvm.call.preallocated.teardown``' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+::
+
+      declare i8* @llvm.call.preallocated.teardown(token %setup_token)
+
+Overview:
+"""""""""
+
+The '``llvm.call.preallocated.teardown``' intrinsic cleans up the stack
+created by a '``llvm.call.preallocated.setup``'.
+
+Semantics:
+""""""""""
+
+The token argument must be a '``llvm.call.preallocated.setup``'.
+
+The '``llvm.call.preallocated.teardown``' intrinsic cleans up the stack
+allocated by the corresponding '``llvm.call.preallocated.setup``'. Exactly
+one of this or the preallocated call must be called to prevent stack leaks.
+It is undefined behavior to call both a '``llvm.call.preallocated.teardown``'
+and the preallocated call for a given '``llvm.call.preallocated.setup``'.
+
+For example, if the stack is allocated for a preallocated call by a
+'``llvm.call.preallocated.setup``', then an initializer function called on an
+allocated argument throws an exception, there should be a
+'``llvm.call.preallocated.teardown``' in the exception handler to prevent
+stack leaks.
+
+Following the nesting rules in '``llvm.call.preallocated.setup``', nested
+calls to '``llvm.call.preallocated.setup``' and
+'``llvm.call.preallocated.teardown``' are allowed but must be properly
+nested.
+
+Example:
+""""""""
+
+.. code-block:: llvm
+
+        %cs = call token @llvm.call.preallocated.setup(i32 1)
+        %x = call i8* @llvm.call.preallocated.arg(token %cs, i32 0) preallocated(i32)
+        %y = bitcast i8* %x to i32*
+        invoke void @constructor(i32* %y) to label %conta unwind label %contb
+    conta:
+        call void @foo1(i32* preallocated(i32) %y) ["preallocated"(token %cs)]
+        ret void
+    contb:
+        %s = catchswitch within none [label %catch] unwind to caller
+    catch:
+        %p = catchpad within %s []
+        call void @llvm.call.preallocated.teardown(token %cs)
+        ret void
+
 Standard C Library Intrinsics
 -----------------------------
 
 LLVM provides intrinsics for a few important standard C library
 functions. These intrinsics allow source-language front-ends to pass
 information about the alignment of the pointer arguments to the code
 generator, providing opportunity for more efficient code generation.
 

llvm/include/llvm/IR/Intrinsics.td

 def int_instrprof_value_profile : Intrinsic<[],
                                             [llvm_ptr_ty, llvm_i64_ty,
                                              llvm_i64_ty, llvm_i32_ty,
                                              llvm_i32_ty],
                                             []>;
 
 def int_call_preallocated_setup : Intrinsic<[llvm_token_ty], [llvm_i32_ty]>;
 def int_call_preallocated_arg : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_i32_ty]>;
+def int_call_preallocated_teardown : Intrinsic<[], [llvm_token_ty]>;
 
 //===------------------- Standard C Library Intrinsics --------------------===//
 //
 
 def int_memcpy  : Intrinsic<[],
                             [llvm_anyptr_ty, llvm_anyptr_ty, llvm_anyint_ty,
                              llvm_i1_ty],
                             [IntrArgMemOnly, IntrWillReturn,

llvm/lib/IR/Verifier.cpp

for (User *U : Call.users()) {
         Assert(AllocArgIndex != nullptr,
                "llvm.call.preallocated.alloc arg index must be a constant");
         auto AllocArgIndexInt = AllocArgIndex->getValue();
         Assert(AllocArgIndexInt.sge(0) &&
                    AllocArgIndexInt.slt(NumArgs->getValue()),
                "llvm.call.preallocated.alloc arg index must be between 0 and "
                "corresponding "
                "llvm.call.preallocated.setup's argument count");
+      } else if (Fn && Fn->getIntrinsicID() ==
+                           Intrinsic::call_preallocated_teardown) {
+        // nothing to do
       } else {
         Assert(!FoundCall, "Can have at most one call corresponding to a "
                            "llvm.call.preallocated.setup");
         FoundCall = true;
         size_t NumPreallocatedArgs = 0;
         for (unsigned i = 0; i < UseCall->getNumArgOperands(); i++) {
           if (UseCall->paramHasAttr(i, Attribute::Preallocated)) {
             ++NumPreallocatedArgs;
Assert(Token && Token->getCalledFunction()->getIntrinsicID() ==
                         Intrinsic::call_preallocated_setup,
            "llvm.call.preallocated.arg token argument must be a "
            "llvm.call.preallocated.setup");
     Assert(Call.hasFnAttr(Attribute::Preallocated),
            "llvm.call.preallocated.arg must be called with a \"preallocated\" "
            "call site attribute");
     break;
   }
+  case Intrinsic::call_preallocated_teardown: {
+    auto *Token = dyn_cast<CallBase>(Call.getArgOperand(0));
+    Assert(Token && Token->getCalledFunction()->getIntrinsicID() ==
+                        Intrinsic::call_preallocated_setup,
+           "llvm.call.preallocated.teardown token argument must be a "
+           "llvm.call.preallocated.setup");
+    break;
+  }
   case Intrinsic::gcroot:
   case Intrinsic::gcwrite:
   case Intrinsic::gcread:
     if (ID == Intrinsic::gcroot) {
       AllocaInst *AI =
           dyn_cast<AllocaInst>(Call.getArgOperand(0)->stripPointerCasts());
       Assert(AI, "llvm.gcroot parameter #1 must be an alloca.", Call);
       Assert(isa<Constant>(Call.getArgOperand(1)),

llvm/test/Verifier/preallocated-invalid.ll

 ; RUN: not opt -S %s -verify 2>&1 | FileCheck %s
 
 declare token @llvm.call.preallocated.setup(i32)
 declare i8* @llvm.call.preallocated.arg(token, i32)
+declare void @llvm.call.preallocated.teardown(token)
 
 ; Fake LLVM intrinsic to return a token
 declare token @llvm.what()
 
 declare void @foo0()
 declare void @foo1(i32* preallocated(i32))
 declare void @foo2(i32* preallocated(i32), i32*, i32* preallocated(i32))
 declare i32 @blackbox()
define void @musttail_and_bundle(i32* preallocated(i32) %a) {
     ret void
 }
 
 ; CHECK: cannot guarantee tail call due to mismatched ABI impacting function attributes
 define void @musttail_attr_no_match(i32* preallocated(i32) %a) {
     musttail call void @musttail_and_bundle(i32* %a)
     ret void
 }
+
+; CHECK: token argument must be a llvm.call.preallocated.setup
+define void @teardown_token_not_from_setup() {
+    %cs = call token @llvm.what()
+    call void @llvm.call.preallocated.teardown(token %cs)
+    ret void
+}

llvm/test/Verifier/preallocated-valid.ll

 ; RUN: opt -S %s -verify
 
 declare token @llvm.call.preallocated.setup(i32)
 declare i8* @llvm.call.preallocated.arg(token, i32)
+declare void @llvm.call.preallocated.teardown(token)
+
+declare i32 @__CxxFrameHandler3(...)
 
 declare void @foo1(i32* preallocated(i32))
 declare i64 @foo1_i64(i32* preallocated(i32))
 declare void @foo2(i32* preallocated(i32), i32*, i32* preallocated(i32))
 
+declare void @constructor(i32*)
+
 define void @preallocated() {
     %cs = call token @llvm.call.preallocated.setup(i32 1)
     %x = call i8* @llvm.call.preallocated.arg(token %cs, i32 0) preallocated(i32)
     %y = bitcast i8* %x to i32*
     call void @foo1(i32* preallocated(i32) %y) ["preallocated"(token %cs)]
     ret void
 }
 
define void @preallocated_num_args() {
     %x1 = bitcast i8* %x to i32*
     %y = call i8* @llvm.call.preallocated.arg(token %cs, i32 1) preallocated(i32)
     %y1 = bitcast i8* %y to i32*
     %a = inttoptr i32 0 to i32*
     call void @foo2(i32* preallocated(i32) %x1, i32* %a, i32* preallocated(i32) %y1) ["preallocated"(token %cs)]
     ret void
 }
 
-define void @preallocate_musttail(i32* preallocated(i32) %a) {
+define void @preallocated_musttail(i32* preallocated(i32) %a) {
     musttail call void @foo1(i32* preallocated(i32) %a)
     ret void
 }
 
-define i64 @preallocate_musttail_i64(i32* preallocated(i32) %a) {
+define i64 @preallocated_musttail_i64(i32* preallocated(i32) %a) {
     %r = musttail call i64 @foo1_i64(i32* preallocated(i32) %a)
     ret i64 %r
 }
+
+define void @preallocated_teardown() {
+    %cs = call token @llvm.call.preallocated.setup(i32 1)
+    call void @llvm.call.preallocated.teardown(token %cs)
+    ret void
+}
+
+define void @preallocated_teardown_invoke() personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+    %cs = call token @llvm.call.preallocated.setup(i32 1)
+    %x = call i8* @llvm.call.preallocated.arg(token %cs, i32 0) preallocated(i32)
+    %y = bitcast i8* %x to i32*
+    invoke void @constructor(i32* %y) to label %conta unwind label %contb
+conta:
+    call void @foo1(i32* preallocated(i32) %y) ["preallocated"(token %cs)]
+    ret void
+contb:
+    %s = catchswitch within none [label %catch] unwind to caller
+catch:
+    %p = catchpad within %s []
+    call void @llvm.call.preallocated.teardown(token %cs)
+    ret void
+}