Merge pull request #37832 from aschwaighofer/change_async_task_wait_abi

Change _wait(_throwing) ABIs to reduce code size

Author: aschwaighofer

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>

include/swift/ABI/Task.h

@@ -481,7 +481,13 @@ class AsyncTask : public Job {
   /// \c Executing, then \c waitingTask has been added to the
   /// wait queue and will be scheduled when the future completes. Otherwise,
   /// the future has completed and can be queried.
-  FutureFragment::Status waitFuture(AsyncTask *waitingTask);
+  /// The waiting task's async context will be intialized with the parameters if
+  /// the current's task state is executing.
+  FutureFragment::Status waitFuture(AsyncTask *waitingTask,
+                                    AsyncContext *waitingTaskContext,
+                                    TaskContinuationFunction *resumeFn,
+                                    AsyncContext *callerContext,
+                                    OpaqueValue *result);
 
   /// Complete this future.
   ///

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
 ///

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);

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);
 }
 

lib/IRGen/CallEmission.h

@@ -119,6 +119,9 @@ class CallEmission {
 
     return result;
   }
+
+  virtual llvm::Value *getResumeFunctionPointer() = 0;
+  virtual llvm::Value *getAsyncContext() = 0;
 };
 
 std::unique_ptr<CallEmission>

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::AsyncLetWaitThrowing:
         case SpecialKind::TaskFutureWaitThrowing:
+        case SpecialKind::TaskFutureWait:
         case SpecialKind::AsyncLetWait:
-        case SpecialKind::AsyncLetWaitThrowing:
         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
 

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;
   }