Change _wait(_throwing) ABIs to reduce code size

Changes the task, taskGroup, asyncLet wait funtion call ABIs. To reduce code size pass the context parameters and resumption function as arguments to the wait function. This means that the suspend point does not need to store parent context and resumption to the suspend point's context. ``` void swift_task_future_wait_throwing( OpaqueValue * result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext, AsyncTask *task, ThrowingTaskFutureWaitContinuationFunction *resume, AsyncContext *callContext); ``` The runtime passes the caller context to the resume entry point saving the load of the parent context in the resumption function. This patch adds a `Metadata *` field to `GroupImpl`. The await entry pointer no longer pass the metadata pointer and there is a path through the runtime where the task future is no longer available.
2025-12-14 20:36:38 +01:00 · 2021-06-07 10:55:19 -07:00
parent 6c51fc4707
commit 10e3d2e3af
26 changed files with 313 additions and 273 deletions
--- a/include/swift/ABI/Executor.h
+++ b/include/swift/ABI/Executor.h
@@ -128,6 +128,11 @@ using TaskContinuationFunction =
  SWIFT_CC(swiftasync)
  void (SWIFT_ASYNC_CONTEXT AsyncContext *);

+using ThrowingTaskFutureWaitContinuationFunction =
+  SWIFT_CC(swiftasync)
+  void (SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT void *);
+
+
 template <class AsyncSignature>
 class AsyncFunctionPointer;
 template <class AsyncSignature>
--- a/include/swift/Runtime/Concurrency.h
+++ b/include/swift/Runtime/Concurrency.h
@@ -125,12 +125,6 @@ swift_task_escalate(AsyncTask *task, JobPriority newPriority);
 // TODO: "async let wait" and "async let destroy" would be expressed
 //       similar to like TaskFutureWait;

-/// This matches the ABI of a closure `<T>(Builtin.NativeObject) async -> T`
-using TaskFutureWaitSignature =
-    SWIFT_CC(swiftasync)
-    void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
-
 /// Wait for a non-throwing future task to complete.
 ///
 /// This can be called from any thread. Its Swift signature is
@@ -141,12 +135,9 @@ using TaskFutureWaitSignature =
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_task_future_wait(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
-
-using TaskFutureWaitThrowingSignature =
-    SWIFT_CC(swiftasync)
-    void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
+         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *,
+         TaskContinuationFunction *,
+         AsyncContext *);

 /// Wait for a potentially-throwing future task to complete.
 ///
@@ -157,15 +148,12 @@ using TaskFutureWaitThrowingSignature =
 ///    async throws -> Success
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
-void swift_task_future_wait_throwing(OpaqueValue *,
-                                     SWIFT_ASYNC_CONTEXT AsyncContext *,
-                                     AsyncTask *, Metadata *);
-
-using TaskGroupFutureWaitThrowingSignature =
-SWIFT_CC(swiftasync)
-  void(OpaqueValue *,
-       SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, TaskGroup *,
-       const Metadata *successType);
+void swift_task_future_wait_throwing(
+  OpaqueValue *,
+  SWIFT_ASYNC_CONTEXT AsyncContext *,
+  AsyncTask *,
+  ThrowingTaskFutureWaitContinuationFunction *,
+  AsyncContext *);

 /// Wait for a readyQueue of a Channel to become non empty.
 ///
@@ -180,8 +168,9 @@ SWIFT_CC(swiftasync)
 SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swiftasync)
 void swift_taskGroup_wait_next_throwing(
-    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-    TaskGroup *group, const Metadata *successType);
+    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    TaskGroup *group, ThrowingTaskFutureWaitContinuationFunction *resumeFn,
+    AsyncContext *callContext);

 /// Initialize a `TaskGroup` in the passed `group` memory location.
 /// The caller is responsible for retaining and managing the group's lifecycle.
@@ -192,7 +181,7 @@ void swift_taskGroup_wait_next_throwing(
 /// func swift_taskGroup_initialize(group: Builtin.RawPointer)
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_taskGroup_initialize(TaskGroup *group);
+void swift_taskGroup_initialize(TaskGroup *group, const Metadata *T);

 /// Attach a child task to the parent task's task group record.
 ///
@@ -310,7 +299,8 @@ using AsyncLetWaitSignature =
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_asyncLet_wait(OpaqueValue *,
                         SWIFT_ASYNC_CONTEXT AsyncContext *,
-                         AsyncLet *, Metadata *);
+                         AsyncLet *, TaskContinuationFunction *,
+                         AsyncContext *);

 /// Wait for a potentially-throwing async-let to complete.
 ///
@@ -324,7 +314,9 @@ void swift_asyncLet_wait(OpaqueValue *,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_asyncLet_wait_throwing(OpaqueValue *,
                                  SWIFT_ASYNC_CONTEXT AsyncContext *,
-                                  AsyncLet *, Metadata *);
+                                  AsyncLet *,
+                                  ThrowingTaskFutureWaitContinuationFunction *,
+                                  AsyncContext *);

 /// Its Swift signature is
 ///
--- a/include/swift/Runtime/RuntimeFunctions.def
+++ b/include/swift/Runtime/RuntimeFunctions.def
@@ -1744,7 +1744,7 @@ FUNCTION(TaskGroupInitialize,
         swift_taskGroup_initialize, SwiftCC,
         ConcurrencyAvailability,
         RETURNS(VoidTy),
-         ARGS(Int8PtrTy),
+         ARGS(Int8PtrTy, TypeMetadataPtrTy),
         ATTRS(NoUnwind))

 // void swift_taskGroup_destroy(TaskGroup *group);
--- a/lib/AST/Builtins.cpp
+++ b/lib/AST/Builtins.cpp
@@ -1437,7 +1437,7 @@ static ValueDecl *getCreateAsyncTaskFuture(ASTContext &ctx, Identifier id) {

 static ValueDecl *getCreateAsyncTaskGroupFuture(ASTContext &ctx, Identifier id) {
  BuiltinFunctionBuilder builder(ctx);
-  auto genericParam = makeGenericParam().build(builder);
+  auto genericParam = makeGenericParam().build(builder); // <T>
  builder.addParameter(makeConcrete(ctx.getIntType())); // flags
  builder.addParameter(
      makeConcrete(OptionalType::get(ctx.TheRawPointerType))); // group
@@ -1520,7 +1520,8 @@ static ValueDecl *getEndAsyncLet(ASTContext &ctx, Identifier id) {

 static ValueDecl *getCreateTaskGroup(ASTContext &ctx, Identifier id) {
  return getBuiltinFunction(ctx, id, _thin,
-                            _parameters(),
+                            _generics(_unrestricted),
+                            _parameters(_metatype(_typeparam(0))),
                            _rawPointer);
 }

--- a/lib/IRGen/CallEmission.h
+++ b/lib/IRGen/CallEmission.h
@@ -119,6 +119,9 @@ public:

    return result;
  }
+
+  virtual llvm::Value *getResumeFunctionPointer() = 0;
+  virtual llvm::Value *getAsyncContext() = 0;
 };

 std::unique_ptr<CallEmission>
--- a/lib/IRGen/Callee.h
+++ b/lib/IRGen/Callee.h
@@ -209,24 +209,16 @@ namespace irgen {
      ///
      /// This is a micro-optimization we apply to certain special functions
      /// that we know don't need generics.
-      bool suppressGenerics() const {
+      bool useSpecialConvention() const {
        if (!isSpecial()) return false;

        switch (getSpecialKind()) {
-        case SpecialKind::TaskFutureWait:
-        case SpecialKind::TaskFutureWaitThrowing:
-        case SpecialKind::AsyncLetWait:
        case SpecialKind::AsyncLetWaitThrowing:
+        case SpecialKind::TaskFutureWaitThrowing:
+        case SpecialKind::TaskFutureWait:
+        case SpecialKind::AsyncLetWait:
        case SpecialKind::TaskGroupWaitNext:
-          // FIXME: I have disabled this optimization, if we bring it back we
-          // need to debug why it currently does not work (call emission
-          // computes an undef return pointer) and change the runtime entries to
-          // remove the extra type parameter.
-          //
-          // We suppress generics from these as a code-size optimization
-          // because the runtime can recover the success type from the
-          // future.
-          return false;
+          return true;
        }
        llvm_unreachable("covered switch");
      }
@@ -372,9 +364,7 @@ namespace irgen {
      return !kind.isAsyncFunctionPointer();
    }

-    bool suppressGenerics() const {
-      return kind.suppressGenerics();
-    }
+    bool useSpecialConvention() const { return kind.useSpecialConvention(); }
  };

  class Callee {
@@ -438,9 +428,7 @@ namespace irgen {
      return Fn.getSignature();
    }

-    bool suppressGenerics() const {
-      return Fn.suppressGenerics();
-    }
+    bool useSpecialConvention() const { return Fn.useSpecialConvention(); }

    /// If this callee has a value for the Swift context slot, return
    /// it; otherwise return non-null.
@@ -458,6 +446,7 @@ namespace irgen {
    llvm::Value *getObjCMethodSelector() const;
  };

+  FunctionPointer::Kind classifyFunctionPointerKind(SILFunction *fn);
 } // end namespace irgen
 } // end namespace swift

--- a/lib/IRGen/GenBuiltin.cpp
+++ b/lib/IRGen/GenBuiltin.cpp
@@ -250,7 +250,9 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
  }

  if (Builtin.ID == BuiltinValueKind::CreateTaskGroup) {
-    out.add(emitCreateTaskGroup(IGF));
+    // Claim metadata pointer.
+    (void)args.claimAll();
+    out.add(emitCreateTaskGroup(IGF, substitutions));
    return;
  }

--- a/lib/IRGen/GenCall.cpp
+++ b/lib/IRGen/GenCall.cpp
@@ -89,17 +89,15 @@ AsyncContextLayout irgen::getAsyncContextLayout(IRGenModule &IGM,
      IGM.getMaximalTypeExpansionContext());
  auto layout = getAsyncContextLayout(
      IGM, originalType, substitutedType, forwardingSubstitutionMap,
-      /*suppressGenerics*/ false,
+      /*useSpecialConvention*/ false,
      FunctionPointer::Kind(FunctionPointer::BasicKind::AsyncFunctionPointer));
  return layout;
 }

-AsyncContextLayout
-irgen::getAsyncContextLayout(IRGenModule &IGM, CanSILFunctionType originalType,
-                             CanSILFunctionType substitutedType,
-                             SubstitutionMap substitutionMap,
-                             bool suppressGenerics,
-                             FunctionPointer::Kind kind) {
+AsyncContextLayout irgen::getAsyncContextLayout(
+    IRGenModule &IGM, CanSILFunctionType originalType,
+    CanSILFunctionType substitutedType, SubstitutionMap substitutionMap,
+    bool useSpecialConvention, FunctionPointer::Kind kind) {
  SmallVector<const TypeInfo *, 4> typeInfos;
  SmallVector<SILType, 4> valTypes;
  SmallVector<AsyncContextLayout::ArgumentInfo, 4> paramInfos;
@@ -137,11 +135,8 @@ irgen::getAsyncContextLayout(IRGenModule &IGM, CanSILFunctionType originalType,
  }

  // Add storage for data used by runtime entry points.
-  // See TaskFutureWaitAsyncContext and TaskGroupNextAsyncContext.
+  // See TaskFutureWaitAsyncContext.
  if (kind.isSpecial()) {
-    switch (kind.getSpecialKind()) {
-    case FunctionPointer::SpecialKind::TaskFutureWait:
-    case FunctionPointer::SpecialKind::TaskFutureWaitThrowing: {
      // This needs to match the layout of TaskFutureWaitAsyncContext.
      // Add storage for the waiting future's result pointer (OpaqueValue *).
      auto ty = SILType();
@@ -152,51 +147,6 @@ irgen::getAsyncContextLayout(IRGenModule &IGM, CanSILFunctionType originalType,
      // OpaqueValue *successResultPointer
      valTypes.push_back(ty);
      typeInfos.push_back(&ti);
-      // void (*, *) async  *asyncResumeEntryPoint;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      break;
-    }
-    case FunctionPointer::SpecialKind::AsyncLetWait:
-    case FunctionPointer::SpecialKind::AsyncLetWaitThrowing: {
-      // This needs to match the layout of TaskFutureWaitAsyncContext.
-      // Add storage for the waiting future's result pointer (OpaqueValue *).
-      auto ty = SILType();
-      auto &ti = IGM.getSwiftContextPtrTypeInfo();
-      // SwiftError *
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      // OpaqueValue *successResultPointer
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      // void (*, *) async  *asyncResumeEntryPoint;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      break;
-    }
-    case FunctionPointer::SpecialKind::TaskGroupWaitNext: {
-      // This needs to match the layout of TaskGroupNextAsyncContext.
-      // Add storage for the waiting future's result pointer (OpaqueValue *).
-      auto ty = SILType();
-      auto &ti = IGM.getSwiftContextPtrTypeInfo();
-      // SwiftError * errorResult;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      // OpaqueValue *successResultPointer;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      // void (*, *) async  *asyncResumeEntryPoint;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      // TaskGroup *group;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      // Metadata *successType;
-      valTypes.push_back(ty);
-      typeInfos.push_back(&ti);
-      break;
-    }
-    }
  }
  return AsyncContextLayout(IGM, LayoutStrategy::Optimal, valTypes, typeInfos,
                            originalType, substitutedType, substitutionMap);
@@ -442,13 +392,14 @@ namespace {
    bool CanUseSRet = true;
    bool CanUseError = true;
    bool CanUseSelf = true;
-    bool SuppressGenerics;
+    bool useSpecialConvention;
    unsigned AsyncContextIdx;
    unsigned AsyncResumeFunctionSwiftSelfIdx = 0;

    SignatureExpansion(IRGenModule &IGM, CanSILFunctionType fnType,
-                       bool suppressGenerics)
-      : IGM(IGM), FnType(fnType), SuppressGenerics(suppressGenerics) {}
+                       bool useSpecialConvention)
+        : IGM(IGM), FnType(fnType), useSpecialConvention(useSpecialConvention) {
+    }

    /// Expand the components of the primary entrypoint of the function type.
    void expandFunctionType();
@@ -1653,8 +1604,15 @@ void SignatureExpansion::expandParameters() {
  }

  // Next, the generic signature.
-  if (hasPolymorphicParameters(FnType) && !SuppressGenerics)
+  if (hasPolymorphicParameters(FnType) && !useSpecialConvention)
    expandPolymorphicSignature(IGM, FnType, ParamIRTypes);
+  if (useSpecialConvention) {
+    // Async waiting functions add the resume function pointer, and the context
+    // for the call.
+    // (But skip passing the metadata.)
+    ParamIRTypes.push_back(IGM.Int8PtrTy);
+    ParamIRTypes.push_back(IGM.SwiftContextPtrTy);
+  }

  // Context is next.
  if (hasSelfContext) {
@@ -1812,8 +1770,14 @@ void SignatureExpansion::expandAsyncEntryType() {
  }

  // Next, the generic signature.
-  if (hasPolymorphicParameters(FnType) && !SuppressGenerics)
+  if (hasPolymorphicParameters(FnType) && !useSpecialConvention)
    expandPolymorphicSignature(IGM, FnType, ParamIRTypes);
+  if (useSpecialConvention) {
+    // Async waiting functions add the resume function pointer.
+    // (But skip passing the metadata.)
+    ParamIRTypes.push_back(IGM.Int8PtrTy);
+    ParamIRTypes.push_back(IGM.SwiftContextPtrTy);
+  }

  // Context is next.
  if (hasSelfContext) {
@@ -1941,9 +1905,9 @@ Signature SignatureExpansion::getSignature() {

 Signature Signature::getUncached(IRGenModule &IGM,
                                 CanSILFunctionType formalType,
-                                 bool suppressGenerics) {
+                                 bool useSpecialConvention) {
  GenericContextScope scope(IGM, formalType->getInvocationGenericSignature());
-  SignatureExpansion expansion(IGM, formalType, suppressGenerics);
+  SignatureExpansion expansion(IGM, formalType, useSpecialConvention);
  expansion.expandFunctionType();
  return expansion.getSignature();
 }
@@ -1960,25 +1924,28 @@ Signature Signature::forAsyncReturn(IRGenModule &IGM,
                                    CanSILFunctionType fnType) {
  assert(fnType->isAsync());
  GenericContextScope scope(IGM, fnType->getInvocationGenericSignature());
-  SignatureExpansion expansion(IGM, fnType, /*suppress generics*/ false);
+  SignatureExpansion expansion(IGM, fnType,
+                               /*suppress generics*/ false);
  expansion.expandAsyncReturnType();
  return expansion.getSignature();
 }

-Signature Signature::forAsyncAwait(IRGenModule &IGM,
-                                   CanSILFunctionType fnType) {
+Signature Signature::forAsyncAwait(IRGenModule &IGM, CanSILFunctionType fnType,
+                                   bool useSpecialConvention) {
  assert(fnType->isAsync());
  GenericContextScope scope(IGM, fnType->getInvocationGenericSignature());
-  SignatureExpansion expansion(IGM, fnType, /*suppress generics*/ false);
+  SignatureExpansion expansion(IGM, fnType,
+                               /*suppress generics*/ useSpecialConvention);
  expansion.expandAsyncAwaitType();
  return expansion.getSignature();
 }

-Signature Signature::forAsyncEntry(IRGenModule &IGM,
-                                   CanSILFunctionType fnType) {
+Signature Signature::forAsyncEntry(IRGenModule &IGM, CanSILFunctionType fnType,
+                                   bool useSpecialConvention) {
  assert(fnType->isAsync());
  GenericContextScope scope(IGM, fnType->getInvocationGenericSignature());
-  SignatureExpansion expansion(IGM, fnType, /*suppress generics*/ false);
+  SignatureExpansion expansion(IGM, fnType,
+                               /*suppress generics*/ useSpecialConvention);
  expansion.expandAsyncEntryType();
  return expansion.getSignature();
 }
@@ -2284,6 +2251,14 @@ public:
  Address getCalleeErrorSlot(SILType errorType, bool isCalleeAsync) override {
    return IGF.getCalleeErrorResultSlot(errorType);
  };
+
+  llvm::Value *getResumeFunctionPointer() override {
+    llvm_unreachable("Should not call getResumeFunctionPointer on a sync call");
+  }
+
+  llvm::Value *getAsyncContext() override {
+    llvm_unreachable("Should not call getAsyncContext on a sync call");
+  }
 };

 class AsyncCallEmission final : public CallEmission {
@@ -2294,15 +2269,15 @@ class AsyncCallEmission final : public CallEmission {
  llvm::Value *calleeFunction = nullptr;
  llvm::Value *currentResumeFn = nullptr;
  llvm::Value *thickContext = nullptr;
+  Size initialContextSize = Size(0);
  Optional<AsyncContextLayout> asyncContextLayout;

  AsyncContextLayout getAsyncContextLayout() {
    if (!asyncContextLayout) {
      asyncContextLayout.emplace(::getAsyncContextLayout(
          IGF.IGM, getCallee().getOrigFunctionType(),
-          getCallee().getSubstFunctionType(),
-          getCallee().getSubstitutions(),
-          getCallee().suppressGenerics(), 
+          getCallee().getSubstFunctionType(), getCallee().getSubstitutions(),
+          getCallee().useSpecialConvention(),
          getCallee().getFunctionPointer().getKind()));
    }
    return *asyncContextLayout;
@@ -2332,7 +2307,7 @@ public:
    auto layout = getAsyncContextLayout();
    // Allocate space for the async context.

-    auto initialContextSize = Size(0);
+    initialContextSize = Size(0);
    // Only c++ runtime functions should use the initial context size.
    if (CurCallee.getFunctionPointer().getKind().isSpecial()) {
      initialContextSize = layout.getSize();
@@ -2344,13 +2319,18 @@ public:
        std::make_pair(true, true), initialContextSize);
    auto *dynamicContextSize =
        IGF.Builder.CreateZExt(dynamicContextSize32, IGF.IGM.SizeTy);
-    contextBuffer = emitAllocAsyncContext(IGF, dynamicContextSize);
+    contextBuffer = getCallee().useSpecialConvention()
+                        ? emitStaticAllocAsyncContext(IGF, initialContextSize)
+                        : emitAllocAsyncContext(IGF, dynamicContextSize);
    context = layout.emitCastTo(IGF, contextBuffer.getAddress());
  }
  void end() override {
    assert(contextBuffer.isValid());
    assert(context.isValid());
-    emitDeallocAsyncContext(IGF, contextBuffer);
+    if (getCallee().useSpecialConvention())
+      emitStaticDeallocAsyncContext(IGF, contextBuffer, initialContextSize);
+    else
+      emitDeallocAsyncContext(IGF, contextBuffer);
    super::end();
  }
  void setFromCallee() override {
@@ -2385,7 +2365,8 @@ public:
                                       .getCorrespondingCodeAuthInfo();
    return FunctionPointer(
        FunctionPointer::Kind::Function, calleeFunction, codeAuthInfo,
-        Signature::forAsyncAwait(IGF.IGM, getCallee().getOrigFunctionType()));
+        Signature::forAsyncAwait(IGF.IGM, getCallee().getOrigFunctionType(),
+                                 getCallee().useSpecialConvention()));
  }

  SILType getParameterType(unsigned index) override {
@@ -2407,7 +2388,11 @@ public:
    original.transferInto(asyncExplosion, fnConv.getNumIndirectSILResults());

    // Pass the async context.
-    asyncExplosion.add(contextBuffer.getAddress());
+    if (getCallee().useSpecialConvention()) {
+      // Pass the caller context.
+      asyncExplosion.add(IGF.getAsyncContext());
+    } else
+      asyncExplosion.add(contextBuffer.getAddress());

    // Pass along the coroutine buffer.
    switch (origCalleeType->getCoroutineKind()) {
@@ -2471,7 +2456,7 @@ public:

    auto layout = getAsyncContextLayout();
    // Set caller info into the context.
-    { // caller context
+    if (!getCallee().useSpecialConvention()) { // caller context
      Explosion explosion;
      auto fieldLayout = layout.getParentLayout();
      auto *context = IGF.getAsyncContext();
@@ -2485,7 +2470,8 @@ public:
      explosion.add(context);
      saveValue(fieldLayout, explosion, isOutlined);
    }
-    { // Return to caller function.
+    if (!getCallee().useSpecialConvention()) { // Return to caller function.
+      assert(currentResumeFn == nullptr);
      auto fieldLayout = layout.getResumeParentLayout();
      currentResumeFn = IGF.Builder.CreateIntrinsicCall(
          llvm::Intrinsic::coro_async_resume, {});
@@ -2613,7 +2599,9 @@ public:
    arguments.push_back(
        IGM.getInt32(paramAttributeFlags));
    arguments.push_back(currentResumeFn);
-    auto resumeProjFn = IGF.getOrCreateResumePrjFn();
+    auto resumeProjFn = getCallee().useSpecialConvention()
+                            ? IGF.getOrCreateResumeFromSuspensionFn()
+                            : IGF.getOrCreateResumePrjFn();
    arguments.push_back(
        Builder.CreateBitOrPointerCast(resumeProjFn, IGM.Int8PtrTy));
    auto dispatchFn = IGF.createAsyncDispatchFn(
@@ -2631,6 +2619,18 @@ public:
        cast<llvm::StructType>(signature.getType()->getReturnType());
    return IGF.emitSuspendAsyncCall(asyncContextIndex, resultTy, arguments);
  }
+  llvm::Value *getResumeFunctionPointer() override {
+    assert(getCallee().useSpecialConvention());
+    assert(currentResumeFn == nullptr);
+    currentResumeFn =
+        IGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_async_resume, {});
+    return currentResumeFn;
+  }
+
+
+  llvm::Value *getAsyncContext() override {
+    return contextBuffer.getAddress();
+  }
 };

 } // end anonymous namespace
@@ -3888,6 +3888,20 @@ void irgen::emitDeallocAsyncContext(IRGenFunction &IGF, Address context) {
  IGF.Builder.CreateLifetimeEnd(context, Size(-1) /*dynamic size*/);
 }

+Address irgen::emitStaticAllocAsyncContext(IRGenFunction &IGF,
+                                           Size size) {
+  auto alignment = IGF.IGM.getAsyncContextAlignment();
+  auto &IGM = IGF.IGM;
+  auto address = IGF.createAlloca(IGM.Int8Ty, IGM.getSize(size), alignment);
+  IGF.Builder.CreateLifetimeStart(address, size);
+  return address;
+}
+
+void irgen::emitStaticDeallocAsyncContext(IRGenFunction &IGF, Address context,
+                                          Size size) {
+  IGF.Builder.CreateLifetimeEnd(context, size);
+}
+
 llvm::Value *irgen::emitYield(IRGenFunction &IGF,
                              CanSILFunctionType coroutineType,
                              Explosion &substValues) {
--- a/lib/IRGen/GenCall.h
+++ b/lib/IRGen/GenCall.h
@@ -120,7 +120,7 @@ namespace irgen {
                                           CanSILFunctionType originalType,
                                           CanSILFunctionType substitutedType,
                                           SubstitutionMap substitutionMap,
-                                           bool suppressGenerics,
+                                           bool useSpecialConvention,
                                           FunctionPointer::Kind kind);

  /// Given an async function, get the pointer to the function to be called and
@@ -239,6 +239,9 @@ namespace irgen {
  /// Allocate task local storage for the provided dynamic size.
  Address emitAllocAsyncContext(IRGenFunction &IGF, llvm::Value *sizeValue);
  void emitDeallocAsyncContext(IRGenFunction &IGF, Address context);
+  Address emitStaticAllocAsyncContext(IRGenFunction &IGF, Size size);
+  void emitStaticDeallocAsyncContext(IRGenFunction &IGF, Address context,
+                                     Size size);

  void emitAsyncFunctionEntry(IRGenFunction &IGF,
                              const AsyncContextLayout &layout,
--- a/lib/IRGen/GenConcurrency.cpp
+++ b/lib/IRGen/GenConcurrency.cpp
@@ -206,15 +206,20 @@ void irgen::emitEndAsyncLet(IRGenFunction &IGF, llvm::Value *alet) {
  IGF.Builder.CreateLifetimeEnd(alet);
 }

-llvm::Value *irgen::emitCreateTaskGroup(IRGenFunction &IGF) {
+llvm::Value *irgen::emitCreateTaskGroup(IRGenFunction &IGF,
+                                        SubstitutionMap subs) {
  auto ty = llvm::ArrayType::get(IGF.IGM.Int8PtrTy, NumWords_TaskGroup);
  auto address = IGF.createAlloca(ty, Alignment(Alignment_TaskGroup));
  auto group = IGF.Builder.CreateBitCast(address.getAddress(),
                                         IGF.IGM.Int8PtrTy);
  IGF.Builder.CreateLifetimeStart(group);
+  assert(subs.getReplacementTypes().size() == 1 &&
+         "createTaskGroup should have a type substitution");
+  auto resultType = subs.getReplacementTypes()[0]->getCanonicalType();
+  auto resultTypeMetadata = IGF.emitAbstractTypeMetadataRef(resultType);

  auto *call = IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeFn(),
-                                      {group});
+                                      {group, resultTypeMetadata});
  call->setDoesNotThrow();
  call->setCallingConv(IGF.IGM.SwiftCC);

--- a/lib/IRGen/GenConcurrency.h
+++ b/lib/IRGen/GenConcurrency.h
@@ -67,7 +67,7 @@ llvm::Value *emitBuiltinStartAsyncLet(IRGenFunction &IGF,
 void emitEndAsyncLet(IRGenFunction &IGF, llvm::Value *alet);

 /// Emit the createTaskGroup builtin.
-llvm::Value *emitCreateTaskGroup(IRGenFunction &IGF);
+llvm::Value *emitCreateTaskGroup(IRGenFunction &IGF, SubstitutionMap subs);

 /// Emit the destroyTaskGroup builtin.
 void emitDestroyTaskGroup(IRGenFunction &IGF, llvm::Value *group);
--- a/lib/IRGen/GenDecl.cpp
+++ b/lib/IRGen/GenDecl.cpp
@@ -2648,7 +2648,9 @@ void IRGenModule::createReplaceableProlog(IRGenFunction &IGF, SILFunction *f) {
        asyncFnPtr.getAuthInfo().getCorrespondingCodeAuthInfo();
    auto newFnPtr = FunctionPointer(
        FunctionPointer::Kind::Function, asyncFnPtr.getPointer(IGF),
-        codeAuthInfo, Signature::forAsyncAwait(IGM, silFunctionType));
+        codeAuthInfo,
+        Signature::forAsyncAwait(IGM, silFunctionType,
+                                 /*useSpecialConvention*/ false));
    SmallVector<llvm::Value *, 16> forwardedArgs;
    for (auto &arg : IGF.CurFn->args())
      forwardedArgs.push_back(&arg);
@@ -3107,8 +3109,9 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
             isLazilyEmittedFunction(*f, getSILModule())) {
    IRGen.addLazyFunction(f);
  }
-
-  Signature signature = getSignature(f->getLoweredFunctionType());
+  auto fpKind = irgen::classifyFunctionPointerKind(f);
+  Signature signature =
+      getSignature(f->getLoweredFunctionType(), fpKind.useSpecialConvention());
  addLLVMFunctionAttributes(f, signature);

  LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
--- a/lib/IRGen/GenFunc.cpp
+++ b/lib/IRGen/GenFunc.cpp
@@ -593,12 +593,11 @@ getFuncSignatureInfoForLowered(IRGenModule &IGM, CanSILFunctionType type) {
  llvm_unreachable("bad function type representation");
 }

-Signature
-IRGenModule::getSignature(CanSILFunctionType type,
-                          bool suppressGenerics) {
+Signature IRGenModule::getSignature(CanSILFunctionType type,
+                                    bool useSpecialConvention) {
  // Don't bother caching if we've been asked to suppress generics.
-  if (suppressGenerics)
-    return Signature::getUncached(*this, type, suppressGenerics);
+  if (useSpecialConvention)
+    return Signature::getUncached(*this, type, useSpecialConvention);

  auto &sigInfo = getFuncSignatureInfoForLowered(*this, type);
  return sigInfo.getSignature(*this);
@@ -1071,7 +1070,7 @@ public:
            substType, outType, subs, layout, conventions),
        layout(getAsyncContextLayout(
            subIGF.IGM, origType, substType, subs,
-            staticFnPtr ? staticFnPtr->suppressGenerics() : false,
+            staticFnPtr ? staticFnPtr->useSpecialConvention() : false,
            FunctionPointer::Kind(
                FunctionPointer::BasicKind::AsyncFunctionPointer))),
        currentArgumentIndex(outType->getNumParameters()) {}
@@ -1160,7 +1159,8 @@ public:
        fnPtr.getAuthInfo().getCorrespondingCodeAuthInfo();
    auto newFnPtr = FunctionPointer(
        FunctionPointer::Kind::Function, fnPtr.getPointer(subIGF), newAuthInfo,
-        Signature::forAsyncAwait(subIGF.IGM, origType));
+        Signature::forAsyncAwait(subIGF.IGM, origType,
+                                 /*useSpecialConvention*/ false));
    auto &Builder = subIGF.Builder;

    auto argValues = args.claimAll();
@@ -1268,7 +1268,8 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
  IRGenFunction subIGF(IGM, fwd);
  if (origType->isAsync()) {
    auto asyncContextIdx =
-        Signature::forAsyncEntry(IGM, outType).getAsyncContextIndex();
+        Signature::forAsyncEntry(IGM, outType, /*useSpecialConvention*/ false)
+            .getAsyncContextIndex();
    asyncLayout.emplace(irgen::getAsyncContextLayout(
        IGM, origType, substType, subs, /*suppress generics*/ false,
        FunctionPointer::Kind(
@@ -1372,8 +1373,9 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
  // captured arguments which we will do later. Otherwise, we have to
  // potentially bind polymorphic arguments from the context if it was a
  // partially applied argument.
-  bool hasPolymorphicParams = hasPolymorphicParameters(origType) &&
-    (!staticFnPtr || !staticFnPtr->suppressGenerics());
+  bool hasPolymorphicParams =
+      hasPolymorphicParameters(origType) &&
+      (!staticFnPtr || !staticFnPtr->useSpecialConvention());
  if (!layout && hasPolymorphicParams) {
    assert(conventions.size() == 1);
    // We could have either partially applied an argument from the function
--- a/lib/IRGen/GenThunk.cpp
+++ b/lib/IRGen/GenThunk.cpp
@@ -252,8 +252,9 @@ void IRGenThunk::emit() {
  GenericContextScope scope(IGF.IGM, origTy->getInvocationGenericSignature());

  if (isAsync) {
-    auto asyncContextIdx =
-        Signature::forAsyncEntry(IGF.IGM, origTy).getAsyncContextIndex();
+    auto asyncContextIdx = Signature::forAsyncEntry(
+                               IGF.IGM, origTy, /*useSpecialConvention*/ false)
+                               .getAsyncContextIndex();
    IGF.setupAsync(asyncContextIdx);

    auto entity = LinkEntity::forDispatchThunk(declRef);
--- a/lib/IRGen/IRGenModule.h
+++ b/lib/IRGen/IRGenModule.h
@@ -1393,7 +1393,7 @@ public:
  void finishEmitAfterTopLevel();

  Signature getSignature(CanSILFunctionType fnType,
-                         bool suppressGenerics = false);
+                         bool useSpecialConvention = false);
  llvm::FunctionType *getFunctionType(CanSILFunctionType type,
                                      llvm::AttributeList &attrs,
                                      ForeignFunctionInfo *foreignInfo=nullptr);
--- a/lib/IRGen/IRGenSIL.cpp
+++ b/lib/IRGen/IRGenSIL.cpp
@@ -1732,7 +1732,8 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
  }

  if (f->getLoweredFunctionType()->isAsync()) {
-    setupAsync(Signature::forAsyncEntry(IGM, f->getLoweredFunctionType())
+    setupAsync(Signature::forAsyncEntry(IGM, f->getLoweredFunctionType(),
+                                        /*useSpecialConvention*/ false)
                   .getAsyncContextIndex());
  }
 }
@@ -1951,10 +1952,11 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
  }

  if (funcTy->isAsync()) {
-    emitAsyncFunctionEntry(
-        IGF, getAsyncContextLayout(IGF.IGM, IGF.CurSILFn),
-        LinkEntity::forSILFunction(IGF.CurSILFn),
-        Signature::forAsyncEntry(IGF.IGM, funcTy).getAsyncContextIndex());
+    emitAsyncFunctionEntry(IGF, getAsyncContextLayout(IGF.IGM, IGF.CurSILFn),
+                           LinkEntity::forSILFunction(IGF.CurSILFn),
+                           Signature::forAsyncEntry(
+                               IGF.IGM, funcTy, /*useSpecialConvention*/ false)
+                               .getAsyncContextIndex());
    if (IGF.CurSILFn->isDynamicallyReplaceable()) {
      IGF.IGM.createReplaceableProlog(IGF, IGF.CurSILFn);
      // Remap the entry block.
@@ -2503,7 +2505,7 @@ void IRGenSILFunction::visitDifferentiabilityWitnessFunctionInst(
  setLoweredFunctionPointer(i, FunctionPointer(fnType, diffWitness, signature));
 }

-static FunctionPointer::Kind classifyFunctionPointerKind(SILFunction *fn) {
+FunctionPointer::Kind irgen::classifyFunctionPointerKind(SILFunction *fn) {
  using SpecialKind = FunctionPointer::SpecialKind;

  // Check for some special cases, which are currently all async:
@@ -2530,9 +2532,9 @@ void IRGenSILFunction::visitFunctionRefBaseInst(FunctionRefBaseInst *i) {
  auto fn = i->getInitiallyReferencedFunction();
  auto fnType = fn->getLoweredFunctionType();

-  auto fpKind = classifyFunctionPointerKind(fn);
+  auto fpKind = irgen::classifyFunctionPointerKind(fn);

-  auto sig = IGM.getSignature(fnType, fpKind.suppressGenerics());
+  auto sig = IGM.getSignature(fnType, fpKind.useSpecialConvention());

  // Note that the pointer value returned by getAddrOfSILFunction doesn't
  // necessarily have element type sig.getType(), e.g. if it's imported.
@@ -3072,11 +3074,15 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
  }

  // Pass the generic arguments.
-  if (hasPolymorphicParameters(origCalleeType) &&
-      !emission->getCallee().getFunctionPointer().suppressGenerics()) {
+  auto useSpecialConvention =
+      emission->getCallee().getFunctionPointer().useSpecialConvention();
+  if (hasPolymorphicParameters(origCalleeType) && !useSpecialConvention) {
    SubstitutionMap subMap = site.getSubstitutionMap();
    emitPolymorphicArguments(*this, origCalleeType,
                             subMap, &witnessMetadata, llArgs);
+  } else if (useSpecialConvention) {
+    llArgs.add(emission->getResumeFunctionPointer());
+    llArgs.add(emission->getAsyncContext());
  }

  // Add all those arguments.
--- a/lib/IRGen/Signature.h
+++ b/lib/IRGen/Signature.h
@@ -124,9 +124,8 @@ public:
  /// This is a private detail of the implementation of
  /// IRGenModule::getSignature(CanSILFunctionType), which is what
  /// clients should generally be using.
-  static Signature getUncached(IRGenModule &IGM,
-                               CanSILFunctionType formalType,
-                               bool suppressGenerics);
+  static Signature getUncached(IRGenModule &IGM, CanSILFunctionType formalType,
+                               bool useSpecialConvention);

  /// Compute the signature of a coroutine's continuation function.
  static Signature forCoroutineContinuation(IRGenModule &IGM,
@@ -134,10 +133,10 @@ public:

  static Signature forAsyncReturn(IRGenModule &IGM,
                                  CanSILFunctionType asyncType);
-  static Signature forAsyncAwait(IRGenModule &IGM,
-                                 CanSILFunctionType asyncType);
-  static Signature forAsyncEntry(IRGenModule &IGM,
-                                 CanSILFunctionType asyncType);
+  static Signature forAsyncAwait(IRGenModule &IGM, CanSILFunctionType asyncType,
+                                 bool useSpecialConvention);
+  static Signature forAsyncEntry(IRGenModule &IGM, CanSILFunctionType asyncType,
+                                 bool useSpecialConvention);

  llvm::FunctionType *getType() const {
    assert(isValid());
--- a/stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def
+++ b/stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def
@@ -108,17 +108,19 @@ OVERRIDE_TASK(task_create_group_future_common, AsyncTaskAndContext, , , ,
 OVERRIDE_TASK(task_future_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
              SWIFT_CC(swiftasync), swift::,
              (OpaqueValue *result,
-               SWIFT_ASYNC_CONTEXT AsyncContext *rawContext, AsyncTask *task,
-               Metadata *T),
-              (result, rawContext, task, T))
+               SWIFT_ASYNC_CONTEXT AsyncContext *callerContext, AsyncTask *task,
+               TaskContinuationFunction *resumeFunction,
+               AsyncContext *callContext),
+              (result, callerContext, task, resumeFunction, callContext))

 OVERRIDE_TASK(task_future_wait_throwing, void,
              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
              swift::,
              (OpaqueValue *result,
-               SWIFT_ASYNC_CONTEXT AsyncContext *rawContext, AsyncTask *task,
-               Metadata *T),
-              (result, rawContext, task, T))
+               SWIFT_ASYNC_CONTEXT AsyncContext *callerContext, AsyncTask *task,
+               ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
+               AsyncContext *callContext),
+              (result, callerContext, task, resumeFunction, callContext))

 OVERRIDE_TASK(continuation_resume, void, SWIFT_EXPORT_FROM(swift_Concurrency),
              SWIFT_CC(swift), swift::,
@@ -169,17 +171,20 @@ OVERRIDE_ASYNC_LET(asyncLet_start, void,
 OVERRIDE_ASYNC_LET(asyncLet_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
                   SWIFT_CC(swiftasync), swift::,
                   (OpaqueValue *result,
-                       SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-                       AsyncLet *alet, Metadata *T),
-                   (result, rawContext, alet, T))
+                       SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+                       AsyncLet *alet, TaskContinuationFunction *resumeFn,
+                       AsyncContext *callContext),
+                   (result, callerContext, alet, resumeFn, callContext))

 OVERRIDE_ASYNC_LET(asyncLet_wait_throwing, void,
                   SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
                   swift::,
                   (OpaqueValue *result,
-                       SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-                       AsyncLet *alet, Metadata *T),
-                   (result, rawContext, alet, T))
+                       SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+                       AsyncLet *alet,
+                       ThrowingTaskFutureWaitContinuationFunction *resume,
+                       AsyncContext *callContext),
+                   (result, callerContext, alet, resume, callContext))

 OVERRIDE_ASYNC_LET(asyncLet_end, void,
                   SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
@@ -187,7 +192,7 @@ OVERRIDE_ASYNC_LET(asyncLet_end, void,

 OVERRIDE_TASK_GROUP(taskGroup_initialize, void,
                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-                    swift::, (TaskGroup *group), (group))
+                    swift::, (TaskGroup *group, const Metadata *T), (group, T))

 OVERRIDE_TASK_GROUP(taskGroup_attachChild, void,
                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
@@ -202,9 +207,12 @@ OVERRIDE_TASK_GROUP(taskGroup_wait_next_throwing, void,
                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
                    swift::,
                    (OpaqueValue *resultPointer,
-                     SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-                     TaskGroup *_group, const Metadata *successType),
-                    (resultPointer, rawContext, _group, successType))
+                     SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+                     TaskGroup *_group,
+                     ThrowingTaskFutureWaitContinuationFunction *resumeFn,
+                     AsyncContext *callContext),
+                    (resultPointer, callerContext, _group, resumeFn,
+                    callContext))

 OVERRIDE_TASK_GROUP(taskGroup_isEmpty, bool,
                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
--- a/stdlib/public/Concurrency/AsyncLet.cpp
+++ b/stdlib/public/Concurrency/AsyncLet.cpp
@@ -128,18 +128,23 @@ static void swift_asyncLet_startImpl(AsyncLet *alet,

 SWIFT_CC(swiftasync)
 static void swift_asyncLet_waitImpl(
-    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-    AsyncLet *alet, Metadata *T) {
+    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    AsyncLet *alet, TaskContinuationFunction *resumeFunction,
+    AsyncContext *callContext) {
  auto task = alet->getTask();
-  swift_task_future_wait(result, rawContext, task, T);
+  swift_task_future_wait(result, callerContext, task, resumeFunction,
+                         callContext);
 }

 SWIFT_CC(swiftasync)
 static void swift_asyncLet_wait_throwingImpl(
-    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-    AsyncLet *alet, Metadata *T) {
+    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    AsyncLet *alet,
+    ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
+    AsyncContext * callContext) {
  auto task = alet->getTask();
-  swift_task_future_wait_throwing(result, rawContext, task, T);
+  swift_task_future_wait_throwing(result, callerContext, task, resumeFunction,
+                                  callContext);
 }

 // =============================================================================
--- a/stdlib/public/Concurrency/Task.cpp
+++ b/stdlib/public/Concurrency/Task.cpp
@@ -371,7 +371,16 @@ SWIFT_CC(swiftasync)
 static void task_wait_throwing_resume_adapter(SWIFT_ASYNC_CONTEXT AsyncContext *_context) {

  auto context = static_cast<TaskFutureWaitAsyncContext *>(_context);
-  return context->asyncResumeEntryPoint(_context, context->errorResult);
+  auto resumeWithError =
+      reinterpret_cast<AsyncVoidClosureEntryPoint *>(context->ResumeParent);
+  return resumeWithError(context->Parent, context->errorResult);
+}
+
+SWIFT_CC(swiftasync)
+static void
+task_future_wait_resume_adapter(SWIFT_ASYNC_CONTEXT AsyncContext *_context) {
+
+  return _context->ResumeParent(_context->Parent);
 }

 /// All `swift_task_create*` variants funnel into this common implementation.
@@ -673,17 +682,21 @@ swift::swift_task_create_group_future(
 }

 SWIFT_CC(swiftasync)
-static void swift_task_future_waitImpl(OpaqueValue *result,
-                                       SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-                                       AsyncTask *task, Metadata *T) {
+static void swift_task_future_waitImpl(
+  OpaqueValue *result,
+  SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+  AsyncTask *task,
+  TaskContinuationFunction *resumeFn,
+  AsyncContext *callContext) {
  // Suspend the waiting task.
  auto waitingTask = swift_task_getCurrent();
-  waitingTask->ResumeTask = rawContext->ResumeParent;
-  waitingTask->ResumeContext = rawContext;
+  waitingTask->ResumeTask = task_future_wait_resume_adapter;
+  waitingTask->ResumeContext = callContext;

  // Stash the result pointer for when we resume later.
-  auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
-  context->asyncResumeEntryPoint = nullptr;
+  auto context = static_cast<TaskFutureWaitAsyncContext *>(callContext);
+  context->ResumeParent = resumeFn;
+  context->Parent = callerContext;
  context->successResultPointer = result;
  context->errorResult = nullptr;

@@ -708,20 +721,21 @@ static void swift_task_future_waitImpl(OpaqueValue *result,

 SWIFT_CC(swiftasync)
 void swift_task_future_wait_throwingImpl(
-    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-    AsyncTask *task, Metadata *T) {
+    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    AsyncTask *task,
+    ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
+    AsyncContext *callContext) {
  auto waitingTask = swift_task_getCurrent();
  // Suspend the waiting task.
-  auto originalResumeParent =
-      reinterpret_cast<AsyncVoidClosureResumeEntryPoint *>(
-          rawContext->ResumeParent);
  waitingTask->ResumeTask = task_wait_throwing_resume_adapter;
-  waitingTask->ResumeContext = rawContext;
+  waitingTask->ResumeContext = callContext;

  // Stash the result pointer for when we resume later.
-  auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
+  auto context = static_cast<TaskFutureWaitAsyncContext *>(callContext);
+  context->ResumeParent =
+      reinterpret_cast<TaskContinuationFunction *>(resumeFunction);
+  context->Parent = callerContext;
  context->successResultPointer = result;
-  context->asyncResumeEntryPoint = originalResumeParent;
  context->errorResult = nullptr;

  // Wait on the future.
--- a/stdlib/public/Concurrency/TaskGroup.cpp
+++ b/stdlib/public/Concurrency/TaskGroup.cpp
@@ -100,6 +100,8 @@ public:
    /// object itself.
    OpaqueValue *storage;

+    const Metadata *successType;
+
    /// The completed task, if necessary to keep alive until consumed by next().
    ///
    /// # Important: swift_release
@@ -124,6 +126,7 @@ public:
        /*storage*/ hadErrorResult ?
                    reinterpret_cast<OpaqueValue *>(fragment->getError()) :
                    fragment->getStoragePtr(),
+        /*successType*/fragment->getResultType(),
        /*task*/ asyncTask
      };
    }
@@ -292,15 +295,17 @@ private:
  /// or `nullptr` if no task is currently waiting.
  std::atomic<AsyncTask *> waitQueue;

+  const Metadata *successType;
+
  friend class ::swift::AsyncTask;

 public:
-  explicit TaskGroupImpl()
+  explicit TaskGroupImpl(const Metadata *T)
    : TaskGroupTaskStatusRecord(),
      status(GroupStatus::initial().status),
      readyQueue(),
 //          readyQueue(ReadyQueueItem::get(ReadyStatus::Empty, nullptr)),
-      waitQueue(nullptr) {}
+      waitQueue(nullptr), successType(T) {}


  TaskGroupTaskStatusRecord *getTaskRecord() {
@@ -451,8 +456,8 @@ static TaskGroup *asAbstract(TaskGroupImpl *group) {

 // Initializes into the preallocated _group an actual TaskGroupImpl.
 SWIFT_CC(swift)
-static void swift_taskGroup_initializeImpl(TaskGroup *group) {
-  TaskGroupImpl *impl = new (group) TaskGroupImpl();
+static void swift_taskGroup_initializeImpl(TaskGroup *group, const Metadata *T) {
+  TaskGroupImpl *impl = new (group) TaskGroupImpl(T);
  auto record = impl->getTaskRecord();
  assert(impl == record && "the group IS the task record");

@@ -503,7 +508,7 @@ void TaskGroup::offer(AsyncTask *completedTask, AsyncContext *context) {
  asImpl(this)->offer(completedTask, context);
 }

-static void fillGroupNextResult(TaskGroupNextAsyncContext *context,
+static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
                                PollResult result) {
  /// Fill in the result value
  switch (result.status) {
@@ -517,7 +522,7 @@ static void fillGroupNextResult(TaskGroupNextAsyncContext *context,

  case PollStatus::Success: {
    // Initialize the result as an Optional<Success>.
-    const Metadata *successType = context->successType;
+    const Metadata *successType = result.successType;
    OpaqueValue *destPtr = context->successResultPointer;
    // TODO: figure out a way to try to optimistically take the
    // value out of the finished task's future, if there are no
@@ -529,7 +534,7 @@ static void fillGroupNextResult(TaskGroupNextAsyncContext *context,

  case PollStatus::Empty: {
    // Initialize the result as a nil Optional<Success>.
-    const Metadata *successType = context->successType;
+    const Metadata *successType = result.successType;
    OpaqueValue *destPtr = context->successResultPointer;
    successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);
    return;
@@ -593,7 +598,7 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
        mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization

        auto waitingContext =
-            static_cast<TaskGroupNextAsyncContext *>(
+            static_cast<TaskFutureWaitAsyncContext *>(
                waitingTask->ResumeContext);

        fillGroupNextResult(waitingContext, result);
@@ -633,31 +638,32 @@ SWIFT_CC(swiftasync)
 static void
 task_group_wait_resume_adapter(SWIFT_ASYNC_CONTEXT AsyncContext *_context) {

-  auto context = static_cast<TaskGroupNextAsyncContext *>(_context);
-  return context->asyncResumeEntryPoint(_context, context->errorResult);
+  auto context = static_cast<TaskFutureWaitAsyncContext *>(_context);
+  auto resumeWithError =
+      reinterpret_cast<AsyncVoidClosureResumeEntryPoint *>(context->ResumeParent);
+  return resumeWithError(context->Parent, context->errorResult);
 }

 // =============================================================================
 // ==== group.next() implementation (wait_next and groupPoll) ------------------
 SWIFT_CC(swiftasync)
 static void swift_taskGroup_wait_next_throwingImpl(
-    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-    TaskGroup *_group, const Metadata *successType) {
+    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    TaskGroup *_group,
+    ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
+    AsyncContext *rawContext) {
  auto waitingTask = swift_task_getCurrent();
-  auto originalResumeParent =
-      reinterpret_cast<AsyncVoidClosureResumeEntryPoint *>(
-          rawContext->ResumeParent);
  waitingTask->ResumeTask = task_group_wait_resume_adapter;
  waitingTask->ResumeContext = rawContext;

-  auto context = static_cast<TaskGroupNextAsyncContext *>(rawContext);
+  auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
+  context->ResumeParent =
+      reinterpret_cast<TaskContinuationFunction *>(resumeFunction);
+  context->Parent = callerContext;
  context->errorResult = nullptr;
-  context->asyncResumeEntryPoint = originalResumeParent;
  context->successResultPointer = resultPointer;
-  context->group = _group;
-  context->successType = successType;

-  auto group = asImpl(context->group);
+  auto group = asImpl(_group);
  assert(group && "swift_taskGroup_wait_next_throwing was passed context without group!");

  PollResult polled = group->poll(waitingTask);
@@ -681,6 +687,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {

  PollResult result;
  result.storage = nullptr;
+  result.successType = nullptr;
  result.retainedTask = nullptr;

  // ==== 1) bail out early if no tasks are pending ----------------------------
@@ -690,6 +697,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
    // Bail out and return `nil` from `group.next()`.
    statusRemoveWaiting();
    result.status = PollStatus::Empty;
+    result.successType = this->successType;
    mutex.unlock(); // TODO: remove group lock, and use status for synchronization
    return result;
  }
@@ -711,6 +719,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
      if (!taskDequeued) {
        result.status = PollStatus::MustWait;
        result.storage = nullptr;
+        result.successType = nullptr;
        result.retainedTask = nullptr;
        mutex.unlock(); // TODO: remove group lock, and use status for synchronization
        return result;
@@ -728,6 +737,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
          // Immediately return the polled value
          result.status = PollStatus::Success;
          result.storage = futureFragment->getStoragePtr();
+          result.successType = futureFragment->getResultType();
          assert(result.retainedTask && "polled a task, it must be not null");
          _swift_tsan_acquire(static_cast<Job *>(result.retainedTask));
          mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
@@ -738,6 +748,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
          result.status = PollStatus::Error;
          result.storage =
              reinterpret_cast<OpaqueValue *>(futureFragment->getError());
+          result.successType = nullptr;
          assert(result.retainedTask && "polled a task, it must be not null");
          _swift_tsan_acquire(static_cast<Job *>(result.retainedTask));
          mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
@@ -747,6 +758,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
          result.status = PollStatus::Empty;
          result.storage = nullptr;
          result.retainedTask = nullptr;
+          result.successType = this->successType;
          mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
          return result;
      }
--- a/stdlib/public/Concurrency/TaskGroup.swift
+++ b/stdlib/public/Concurrency/TaskGroup.swift
@@ -70,7 +70,7 @@ public func withTaskGroup<ChildTaskResult, GroupResult>(
 ) async -> GroupResult {
  #if compiler(>=5.5) && $BuiltinTaskGroup

-  let _group = Builtin.createTaskGroup()
+  let _group = Builtin.createTaskGroup(ChildTaskResult.self)
  var group = TaskGroup<ChildTaskResult>(group: _group)

  // Run the withTaskGroup body.
@@ -149,7 +149,7 @@ public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
 ) async rethrows -> GroupResult {
  #if compiler(>=5.5) && $BuiltinTaskGroup

-  let _group = Builtin.createTaskGroup()
+  let _group = Builtin.createTaskGroup(ChildTaskResult.self)
  var group = ThrowingTaskGroup<ChildTaskResult, Error>(group: _group)

  do {
--- a/stdlib/public/Concurrency/TaskPrivate.h
+++ b/stdlib/public/Concurrency/TaskPrivate.h
@@ -132,15 +132,15 @@ namespace {
 ///   @_silgen_name("swift_asyncLet_waitThrowing")
 ///   func _asyncLetGetThrowing<T>(_ task: Builtin.RawPointer) async throws -> T
 ///
+///   @_silgen_name("swift_taskGroup_wait_next_throwing")
+///   func _taskGroupWaitNext<T>(group: Builtin.RawPointer) async throws -> T?
+///
 class TaskFutureWaitAsyncContext : public AsyncContext {
 public:
  SwiftError *errorResult;

  OpaqueValue *successResultPointer;

-  AsyncVoidClosureResumeEntryPoint *__ptrauth_swift_task_resume_function
-      asyncResumeEntryPoint;
-
  void fillWithSuccess(AsyncTask::FutureFragment *future) {
    fillWithSuccess(future->getStoragePtr(), future->getResultType(),
                    successResultPointer);
@@ -159,35 +159,6 @@ public:
  }
 };

-/// The layout of a frame to call the following function:
-///
-///   @_silgen_name("swift_taskGroup_wait_next_throwing")
-///   func _taskGroupWaitNext<T>(group: Builtin.RawPointer) async throws -> T?
-///
-class TaskGroupNextAsyncContext : public AsyncContext {
-public:
-  SwiftError *errorResult;
-
-  OpaqueValue *successResultPointer;
-
-  AsyncVoidClosureResumeEntryPoint *__ptrauth_swift_task_resume_function
-      asyncResumeEntryPoint;
-
-  // Arguments.
-  TaskGroup *group;
-
-  const Metadata *successType;
-
-  void fillWithSuccess(OpaqueValue *src, const Metadata *successType) {
-    successType->vw_initializeWithCopy(successResultPointer, src);
-  }
-
-  void fillWithError(SwiftError *error) {
-    errorResult = error;
-    swift_errorRetain(error);
-  }
-};
-
 } // end anonymous namespace

 /// The current state of a task's status records.
--- a/test/IRGen/async.swift
+++ b/test/IRGen/async.swift
@@ -13,15 +13,18 @@ public func h(_: @Sendable (Int) -> Int) { }

 public class SomeClass {}

-@_silgen_name("swift_task_future_wait")
-public func task_future_wait(_ task: __owned SomeClass) async throws -> Int
+//@_silgen_name("swift_task_future_wait")
+//public func task_future_wait(_ task: __owned SomeClass) async throws -> Int
+
+@_silgen_name("swift_task_future_wait_throwing")
+public func _taskFutureGetThrowing<T>(_ task: SomeClass) async throws -> T

 // CHECK: define{{.*}} swift{{(tail)?}}cc void @"$s5async8testThisyyAA9SomeClassCnYaF"(%swift.context* swiftasync %0{{.*}}
-// CHECK-64: call swiftcc i8* @swift_task_alloc(i64 48)
-// CHECK: {{(must)?}}tail call swift{{(tail)?}}cc void @swift_task_future_wait(
+// CHECK-NOT: @swift_task_alloc
+// CHECK: {{(must)?}}tail call swift{{(tail)?}}cc void @swift_task_future_wait_throwing(%swift.opaque* {{.*}}, %swift.context* {{.*}}, %T5async9SomeClassC* {{.*}}, i8* {{.*}}, %swift.context* {{.*}})
 public func testThis(_ task: __owned SomeClass) async {
  do {
-    let _ = try await task_future_wait(task)
+    let _ : Int = try await _taskFutureGetThrowing(task)
  } catch _ {
    print("error")
  }
--- a/test/IRGen/async/builtins.sil
+++ b/test/IRGen/async/builtins.sil
@@ -114,12 +114,13 @@ bb0:
  // CHECK:      [[TASKGROUP:%.*]] = alloca [32 x i8*], align 16
  // CHECK:      [[T0:%.*]] = bitcast [32 x i8*]* [[TASKGROUP]] to i8*
  // CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[T0]])
-  // CHECK-NEXT: call swiftcc void @swift_taskGroup_initialize(i8* [[T0]])
-  %0 = builtin "createTaskGroup"() : $Builtin.RawPointer
+  // CHECK-NEXT: call swiftcc void @swift_taskGroup_initialize(i8* [[T0]], %swift.type* {{.*}})
+  %0 = metatype $@thin Builtin.Int32.Type
+  %1 = builtin "createTaskGroup"<Builtin.Int32>(%0: $@thin Builtin.Int32.Type) : $Builtin.RawPointer

  // CHECK-NEXT: call swiftcc void @swift_taskGroup_destroy(i8* [[T0]])
  // CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[T0]])
-  builtin "destroyTaskGroup"(%0 : $Builtin.RawPointer) : $()
+  builtin "destroyTaskGroup"(%1 : $Builtin.RawPointer) : $()

  %21 = tuple ()
  return %21 : $()
--- a/unittests/runtime/CompatibilityOverrideConcurrency.cpp
+++ b/unittests/runtime/CompatibilityOverrideConcurrency.cpp
@@ -153,12 +153,12 @@ TEST_F(CompatibilityOverrideConcurrencyTest,
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_future_wait) {
-  swift_task_future_wait(nullptr, nullptr, nullptr, nullptr);
+  swift_task_future_wait(nullptr, nullptr, nullptr, nullptr, nullptr);
 }

 TEST_F(CompatibilityOverrideConcurrencyTest,
       test_swift_task_future_wait_throwing) {
-  swift_task_future_wait_throwing(nullptr, nullptr, nullptr, nullptr);
+  swift_task_future_wait_throwing(nullptr, nullptr, nullptr, nullptr, nullptr);
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_continuation_resume) {
@@ -180,11 +180,11 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_asyncLet_start) {
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_asyncLet_wait) {
-  swift_asyncLet_wait(nullptr, nullptr, nullptr, nullptr);
+  swift_asyncLet_wait(nullptr, nullptr, nullptr, nullptr, nullptr);
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_asyncLet_wait_throwing) {
-  swift_asyncLet_wait(nullptr, nullptr, nullptr, nullptr);
+  swift_asyncLet_wait(nullptr, nullptr, nullptr, nullptr, nullptr);
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_asyncLet_end) {
@@ -192,7 +192,7 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_asyncLet_end) {
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_taskGroup_initialize) {
-  swift_taskGroup_initialize(nullptr);
+  swift_taskGroup_initialize(nullptr, nullptr);
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_taskGroup_attachChild) {
@@ -205,7 +205,8 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_taskGroup_destroy) {

 TEST_F(CompatibilityOverrideConcurrencyTest,
       test_swift_taskGroup_wait_next_throwing) {
-  swift_taskGroup_wait_next_throwing(nullptr, nullptr, nullptr, nullptr);
+  swift_taskGroup_wait_next_throwing(nullptr, nullptr, nullptr, nullptr,
+                                     nullptr);
 }

 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_taskGroup_isEmpty) {