Use taskAlloc for dynamic allocas

aschwaighofer · aschwaighofer · commit 2271d4df55c0 · 2020-11-13T10:41:42.000-08:00
Async coroutines don't support dynamic allocas.
diff --git a/lib/IRGen/Address.h b/lib/IRGen/Address.h
@@ -112,6 +112,7 @@ class StackAddress {
 
   /// In a normal function, the result of llvm.stacksave or null.
   /// In a coroutine, the result of llvm.coro.alloca.alloc.
+  /// In an async function, the result of the taskAlloc call.
   llvm::Value *ExtraInfo;
 
 public:
diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp
@@ -330,10 +330,6 @@ AsyncContextLayout::AsyncContextLayout(
 #endif
 }
 
-static Size getAsyncContextSize(AsyncContextLayout layout) {
-  return layout.getSize();
-}
-
 static Alignment getAsyncContextAlignment(IRGenModule &IGM) {
   return IGM.getPointerAlignment();
 }
@@ -2240,7 +2236,6 @@ class AsyncCallEmission final : public CallEmission {
   using super = CallEmission;
 
   Address contextBuffer;
-  Size contextSize;
   Address context;
   llvm::Value *calleeFunction = nullptr;
   llvm::Value *currentResumeFn = nullptr;
@@ -2285,8 +2280,7 @@ class AsyncCallEmission final : public CallEmission {
         CurCallee.getFunctionPointer(), thickContext);
     auto *dynamicContextSize =
         IGF.Builder.CreateZExt(dynamicContextSize32, IGF.IGM.SizeTy);
-    std::tie(contextBuffer, contextSize) = emitAllocAsyncContext(
-        IGF, layout, dynamicContextSize, getAsyncContextSize(layout));
+    contextBuffer = emitAllocAsyncContext(IGF, dynamicContextSize);
     context = layout.emitCastTo(IGF, contextBuffer.getAddress());
     if (layout.canHaveError()) {
       auto fieldLayout = layout.getErrorLayout();
@@ -2299,7 +2293,7 @@ class AsyncCallEmission final : public CallEmission {
   void end() override {
     assert(contextBuffer.isValid());
     assert(context.isValid());
-    emitDeallocAsyncContext(IGF, contextBuffer, contextSize);
+    emitDeallocAsyncContext(IGF, contextBuffer);
     super::end();
   }
   void setFromCallee() override {
@@ -3523,8 +3517,7 @@ void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
   auto &IGM = IGF.IGM;
   auto size = getAsyncContextLayout(IGM, asyncFunction).getSize();
   auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
-      IGM.getAddrOfAsyncFunctionPointer(asyncFunction, NotForDefinition),
-      IGM.Int8PtrTy);
+      IGM.getAddrOfAsyncFunctionPointer(asyncFunction), IGM.Int8PtrTy);
   auto *id = IGF.Builder.CreateIntrinsicCall(
       llvm::Intrinsic::coro_id_async,
       {llvm::ConstantInt::get(IGM.Int32Ty, size.getValue()),
@@ -3590,28 +3583,6 @@ void irgen::emitDeallocYieldManyCoroutineBuffer(IRGenFunction &IGF,
   IGF.Builder.CreateLifetimeEnd(buffer, bufferSize);
 }
 
-Address irgen::emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
-                             Alignment alignment) {
-  auto *call = IGF.Builder.CreateCall(IGF.IGM.getTaskAllocFn(),
-                                      {IGF.getAsyncTask(), size});
-  call->setDoesNotThrow();
-  call->setCallingConv(IGF.IGM.SwiftCC);
-  call->addAttribute(llvm::AttributeList::FunctionIndex,
-                     llvm::Attribute::ReadNone);
-  auto address = Address(call, alignment);
-  return address;
-}
-
-void irgen::emitTaskDealloc(IRGenFunction &IGF, Address address,
-                            llvm::Value *size) {
-  auto *call = IGF.Builder.CreateCall(
-      IGF.IGM.getTaskDeallocFn(), {IGF.getAsyncTask(), address.getAddress()});
-  call->setDoesNotThrow();
-  call->setCallingConv(IGF.IGM.SwiftCC);
-  call->addAttribute(llvm::AttributeList::FunctionIndex,
-                     llvm::Attribute::ReadNone);
-}
-
 void irgen::emitTaskCancel(IRGenFunction &IGF, llvm::Value *task) {
   if (task->getType() != IGF.IGM.SwiftTaskPtrTy) {
     task = IGF.Builder.CreateBitCast(task, IGF.IGM.SwiftTaskPtrTy);
@@ -3663,28 +3634,17 @@ llvm::Value *irgen::emitTaskCreate(
   return result;
 }
 
-std::pair<Address, Size> irgen::emitAllocAsyncContext(IRGenFunction &IGF,
-                                                      AsyncContextLayout layout,
-                                                      llvm::Value *sizeValue,
-                                                      Size sizeLowerBound) {
+Address irgen::emitAllocAsyncContext(IRGenFunction &IGF,
+                                     llvm::Value *sizeValue) {
   auto alignment = getAsyncContextAlignment(IGF.IGM);
-  auto address = emitTaskAlloc(IGF, sizeValue, alignment);
-  IGF.Builder.CreateLifetimeStart(address, sizeLowerBound);
-  return {address, sizeLowerBound};
-}
-
-std::pair<Address, Size>
-irgen::emitAllocAsyncContext(IRGenFunction &IGF, AsyncContextLayout layout) {
-  auto size = getAsyncContextSize(layout);
-  auto *sizeValue = llvm::ConstantInt::get(IGF.IGM.SizeTy, size.getValue());
-  return emitAllocAsyncContext(IGF, layout, sizeValue, size);
+  auto address = IGF.emitTaskAlloc(sizeValue, alignment);
+  IGF.Builder.CreateLifetimeStart(address, Size(-1) /*dynamic size*/);
+  return address;
 }
 
-void irgen::emitDeallocAsyncContext(IRGenFunction &IGF, Address context,
-                                    Size size) {
-  auto *sizeValue = llvm::ConstantInt::get(IGF.IGM.SizeTy, size.getValue());
-  emitTaskDealloc(IGF, context, sizeValue);
-  IGF.Builder.CreateLifetimeEnd(context, size);
+void irgen::emitDeallocAsyncContext(IRGenFunction &IGF, Address context) {
+  IGF.emitTaskDealloc(context);
+  IGF.Builder.CreateLifetimeEnd(context, Size(-1) /*dynamic size*/);
 }
 
 llvm::Value *irgen::emitYield(IRGenFunction &IGF,
diff --git a/lib/IRGen/GenCall.h b/lib/IRGen/GenCall.h
@@ -397,9 +397,6 @@ namespace irgen {
                               CanSILFunctionType coroutineType,
                               NativeCCEntryPointArgumentEmission &emission);
 
-  Address emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
-                        Alignment alignment);
-  void emitTaskDealloc(IRGenFunction &IGF, Address address, llvm::Value *size);
   void emitTaskCancel(IRGenFunction &IGF, llvm::Value *task);
 
   /// Emit a class to swift_task_create[_f] with the given flags, parent task,
@@ -408,19 +405,9 @@ namespace irgen {
     IRGenFunction &IGF, llvm::Value *flags, llvm::Value *parentTask,
     llvm::Value *taskFunction, llvm::Value *localContextInfo);
 
-  /// Allocate task local storage for the specified layout but using the
-  /// provided dynamic size.  Allowing the size to be specified dynamically is
-  /// necessary for applies of thick functions the sizes of whose async contexts
-  /// are dependent on the underlying, already partially applied, called
-  /// function.  The provided sizeLowerBound will be used to track the lifetime
-  /// of the allocation that is known statically.
-  std::pair<Address, Size> emitAllocAsyncContext(IRGenFunction &IGF,
-                                                 AsyncContextLayout layout,
-                                                 llvm::Value *sizeValue,
-                                                 Size sizeLowerBound);
-  std::pair<Address, Size> emitAllocAsyncContext(IRGenFunction &IGF,
-                                                 AsyncContextLayout layout);
-  void emitDeallocAsyncContext(IRGenFunction &IGF, Address context, Size size);
+  /// Allocate task local storage for the provided dynamic size.
+  Address emitAllocAsyncContext(IRGenFunction &IGF, llvm::Value *sizeValue);
+  void emitDeallocAsyncContext(IRGenFunction &IGF, Address context);
 
   void emitAsyncFunctionEntry(IRGenFunction &IGF, SILFunction *asyncFunc);
 
diff --git a/lib/IRGen/GenOpaque.cpp b/lib/IRGen/GenOpaque.cpp
@@ -536,8 +536,19 @@ StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
                                               llvm::Value *arraySize,
                                               Alignment align,
                                               const llvm::Twine &name) {
+  // Async functions call task alloc.
+  if (isAsync()) {
+    llvm::Value *byteCount;
+    auto eltSize = IGM.DataLayout.getTypeAllocSize(eltTy);
+    if (eltSize == 1) {
+      byteCount = arraySize;
+    } else {
+      byteCount = Builder.CreateMul(arraySize, IGM.getSize(Size(eltSize)));
+    }
+    auto address = emitTaskAlloc(byteCount, align);
+    return {address, address.getAddress()};
   // In coroutines, call llvm.coro.alloca.alloc.
-  if (isCoroutine()) {
+  } else if (isCoroutine()) {
     // Compute the number of bytes to allocate.
     llvm::Value *byteCount;
     auto eltSize = IGM.DataLayout.getTypeAllocSize(eltTy);
@@ -589,11 +600,16 @@ StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
 /// Deallocate dynamic alloca's memory if requested by restoring the stack
 /// location before the dynamic alloca's call.
 void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address) {
+  // Async function use taskDealloc.
+  if (isAsync() && address.getAddress().isValid()) {
+    emitTaskDealloc(Address(address.getExtraInfo(), address.getAlignment()));
+    return;
+  }
   // In coroutines, unconditionally call llvm.coro.alloca.free.
   // Except if the address is invalid, this happens when this is a StackAddress
   // for a partial_apply [stack] that did not need a context object on the
   // stack.
-  if (isCoroutine() && address.getAddress().isValid()) {
+  else if (isCoroutine() && address.getAddress().isValid()) {
     auto allocToken = address.getExtraInfo();
     assert(allocToken && "dynamic alloca in coroutine without alloc token?");
     auto freeFn = llvm::Intrinsic::getDeclaration(
diff --git a/lib/IRGen/IRGenFunction.cpp b/lib/IRGen/IRGenFunction.cpp
@@ -490,3 +490,20 @@ void IRGenFunction::emitTrap(StringRef failureMessage, bool EmitUnreachable) {
   if (EmitUnreachable)
     Builder.CreateUnreachable();
 }
+
+Address IRGenFunction::emitTaskAlloc(llvm::Value *size, Alignment alignment) {
+  auto *call = Builder.CreateCall(IGM.getTaskAllocFn(), {getAsyncTask(), size});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGM.SwiftCC);
+  call->addAttribute(llvm::AttributeList::FunctionIndex,
+                     llvm::Attribute::ReadNone);
+  auto address = Address(call, alignment);
+  return address;
+}
+
+void IRGenFunction::emitTaskDealloc(Address address) {
+  auto *call = Builder.CreateCall(IGM.getTaskDeallocFn(),
+                                  {getAsyncTask(), address.getAddress()});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGM.SwiftCC);
+}
diff --git a/lib/IRGen/IRGenFunction.h b/lib/IRGen/IRGenFunction.h
@@ -456,6 +456,10 @@ class IRGenFunction {
   llvm::Value *emitIsEscapingClosureCall(llvm::Value *value, SourceLoc loc,
                                          unsigned verificationType);
 
+  Address emitTaskAlloc(llvm::Value *size,
+                        Alignment alignment);
+  void emitTaskDealloc(Address address);
+
   //--- Expression emission
   //------------------------------------------------------
 public:
