[Concurrency] Enable TaskGroup/DiscardingTaskGroup in Embedded Swift

Author: kubamracek

include/swift/ABI/MetadataValues.h

@@ -2792,6 +2792,10 @@ class TaskGroupFlags : public FlagSet<uint32_t> {
     /// Request the TaskGroup to immediately release completed tasks,
     /// and not store their results. This also effectively disables `next()`.
     TaskGroup_DiscardResults = 8,
+    /// The `Metadata *T` in swift_taskGroup_initialize and in
+    /// swift_taskGroup_initializeWithFlags pointer is actually a
+    /// TaskOptionRecord pointer. Used only in Embedded Swift.
+    TaskGroup_MetadataAsOptionRecord = 9,
   };
 
   explicit TaskGroupFlags(uint32_t bits) : FlagSet(bits) {}
@@ -2800,6 +2804,10 @@ class TaskGroupFlags : public FlagSet<uint32_t> {
   FLAGSET_DEFINE_FLAG_ACCESSORS(TaskGroup_DiscardResults,
                                 isDiscardResults,
                                 setIsDiscardResults)
+
+  FLAGSET_DEFINE_FLAG_ACCESSORS(TaskGroup_MetadataAsOptionRecord,
+                                isMetadataAsOptionRecord,
+                                setIsMetadataAsOptionRecord)
 };
 
 /// Flags for cancellation records.

include/swift/AST/Builtins.def

@@ -800,18 +800,6 @@ BUILTIN_MISC_OPERATION(ResumeThrowingContinuationReturning,
 BUILTIN_MISC_OPERATION(ResumeThrowingContinuationThrowing,
                        "resumeThrowingContinuationThrowing", "", Special)
 
-/// Create a task group.
-BUILTIN_MISC_OPERATION(CreateTaskGroup,
-                       "createTaskGroup", "", Special)
-
-/// Create a task group, with options.
-BUILTIN_MISC_OPERATION(CreateTaskGroupWithFlags,
-                       "createTaskGroupWithFlags", "", Special)
-
-/// Destroy a task group.
-BUILTIN_MISC_OPERATION(DestroyTaskGroup,
-                       "destroyTaskGroup", "", Special)
-
 /// Unchecked pointer alignment assertion. Allows the compiler to assume
 /// alignment of the pointer to emit more efficient code.
 ///
@@ -1054,6 +1042,15 @@ BUILTIN_SIL_OPERATION(CreateAsyncDiscardingTaskInGroupWithExecutor,
 BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildOrdinaryTaskExecutorRef,
     "buildOrdinaryTaskExecutorRef", "n", Special)
 
+/// Create a task group.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateTaskGroup, "createTaskGroup", "", Special)
+
+/// Create a task group, with options.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateTaskGroupWithFlags, "createTaskGroupWithFlags", "", Special)
+
+/// Destroy a task group.
+BUILTIN_MISC_OPERATION(DestroyTaskGroup, "destroyTaskGroup", "", Special)
+
 /// globalStringTablePointer has type String -> Builtin.RawPointer.
 /// It returns an immortal, global string table pointer for strings constructed
 /// from string literals. We consider it effects as readnone meaning that it

lib/IRGen/GenBuiltin.cpp

@@ -295,20 +295,14 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
 
   if (Builtin.ID == BuiltinValueKind::CreateTaskGroup) {
     llvm::Value *groupFlags = nullptr;
-    // Claim metadata pointer.
-    (void)args.claimAll();
+    assert(args.size() == 0);
     out.add(emitCreateTaskGroup(IGF, substitutions, groupFlags));
     return;
   }
 
   if (Builtin.ID == BuiltinValueKind::CreateTaskGroupWithFlags) {
     auto groupFlags = args.claimNext();
-    // Claim the remaining metadata pointer.
-    if (args.size() == 1) {
-      (void)args.claimNext();
-    } else if (args.size() > 1) {
-      llvm_unreachable("createTaskGroupWithFlags expects 1 or 2 arguments");
-    }
+    assert(args.size() == 0);
     out.add(emitCreateTaskGroup(IGF, substitutions, groupFlags));
     return;
   }

lib/IRGen/GenConcurrency.cpp

@@ -335,15 +335,39 @@ llvm::Value *irgen::emitCreateTaskGroup(IRGenFunction &IGF,
   assert(subs.getReplacementTypes().size() == 1 &&
          "createTaskGroup should have a type substitution");
   auto resultType = subs.getReplacementTypes()[0]->getCanonicalType();
-  auto resultTypeMetadata = IGF.emitAbstractTypeMetadataRef(resultType);
+
+  // In desktop Swift, we pass a Metadata pointer as the last argument. In
+  // Embedded Swift, we pass a TaskOptionRecord list, and mark a bit in
+  // groupFlags.
+  llvm::Value *metadataOrTypeOptionRecord;
+  if (IGF.IGM.Context.LangOpts.hasFeature(Feature::Embedded)) {
+    metadataOrTypeOptionRecord =
+        llvm::ConstantPointerNull::get(IGF.IGM.Int8PtrTy);
+    metadataOrTypeOptionRecord = maybeAddEmbeddedSwiftResultTypeInfo(
+        IGF, metadataOrTypeOptionRecord, resultType);
+    llvm::IntegerType *IntPtrTy = IGF.IGM.IntPtrTy;
+    TaskGroupFlags taskgroupFlags(0);
+    taskgroupFlags.setIsMetadataAsOptionRecord(true);
+    auto flagValue = llvm::ConstantInt::get(IGF.IGM.IntPtrTy,
+                                            taskgroupFlags.getOpaqueValue());
+    if (!groupFlags) {
+      groupFlags = flagValue;
+    } else {
+      groupFlags = IGF.Builder.CreateOr(groupFlags, flagValue);
+    }
+  } else {
+    metadataOrTypeOptionRecord = IGF.emitAbstractTypeMetadataRef(resultType);
+  }
 
   llvm::CallInst *call;
   if (groupFlags) {
-    call = IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeWithFlagsFunctionPointer(),
-                                  {groupFlags, group, resultTypeMetadata});
+    call = IGF.Builder.CreateCall(
+        IGF.IGM.getTaskGroupInitializeWithFlagsFunctionPointer(),
+        {groupFlags, group, metadataOrTypeOptionRecord});
   } else {
-    call = IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeFunctionPointer(),
-                                  {group, resultTypeMetadata});
+    call =
+        IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeFunctionPointer(),
+                               {group, metadataOrTypeOptionRecord});
   }
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);

lib/SILGen/SILGenBuiltin.cpp

@@ -1750,6 +1750,38 @@ static ManagedValue emitBuiltinCreateDiscardingTask(
         CreateTaskOptions::Discarding });
 }
 
+// Emit SIL for the named builtin: createTaskGroup.
+// These formally take a metatype argument that's never actually used, so
+// we ignore it.
+static ManagedValue emitBuiltinCreateTaskGroup(SILGenFunction &SGF,
+                                               SILLocation loc,
+                                               SubstitutionMap subs,
+                                               ArrayRef<ManagedValue> args,
+                                               SGFContext C) {
+  auto &ctx = SGF.getASTContext();
+  auto resultType = SILType::getRawPointerType(ctx);
+  auto value = SGF.B.createBuiltin(
+      loc, ctx.getIdentifier(getBuiltinName(BuiltinValueKind::CreateTaskGroup)),
+      resultType, subs, {});
+  return ManagedValue::forObjectRValueWithoutOwnership(value);
+}
+
+// Emit SIL for the named builtin: createTaskGroupWithFlags.
+// These formally take a metatype argument that's never actually used, so
+// we ignore it.
+static ManagedValue emitBuiltinCreateTaskGroupWithFlags(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, SGFContext C) {
+  auto &ctx = SGF.getASTContext();
+  auto resultType = SILType::getRawPointerType(ctx);
+  auto value = SGF.B.createBuiltin(
+      loc,
+      ctx.getIdentifier(
+          getBuiltinName(BuiltinValueKind::CreateTaskGroupWithFlags)),
+      resultType, subs, {args[0].getValue()});
+  return ManagedValue::forObjectRValueWithoutOwnership(value);
+}
+
 ManagedValue
 SILGenFunction::emitCreateAsyncMainTask(SILLocation loc, SubstitutionMap subs,
                                         ManagedValue flags,

stdlib/public/Concurrency/DiscardingTaskGroup.swift

@@ -67,7 +67,9 @@ import Swift
 ///
 /// - SeeAlso: ``withThrowingDiscardingTaskGroup(returning:body:)``
 @available(SwiftStdlib 5.9, *)
+#if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
+#endif
 @inlinable
 public func withDiscardingTaskGroup<GroupResult>(
   returning returnType: GroupResult.Type = GroupResult.self,
@@ -609,7 +611,9 @@ extension DiscardingTaskGroup: Sendable { }
 /// }
 /// ```
 @available(SwiftStdlib 5.9, *)
+#if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
+#endif
 @inlinable
 public func withThrowingDiscardingTaskGroup<GroupResult>(
     returning returnType: GroupResult.Type = GroupResult.self,

stdlib/public/Concurrency/TaskGroup.cpp

@@ -26,6 +26,7 @@
 #include "swift/ABI/Metadata.h"
 #include "swift/ABI/Task.h"
 #include "swift/ABI/TaskGroup.h"
+#include "swift/ABI/TaskOptions.h"
 #include "swift/Basic/RelativePointer.h"
 #include "swift/Basic/STLExtras.h"
 #include "swift/Runtime/Concurrency.h"
@@ -964,16 +965,37 @@ static void swift_taskGroup_initializeImpl(TaskGroup *group, const Metadata *T)
 SWIFT_CC(swift)
 static void swift_taskGroup_initializeWithFlagsImpl(size_t rawGroupFlags,
                                                     TaskGroup *group, const Metadata *T) {
-
-#if !SWIFT_CONCURRENCY_EMBEDDED
-  ResultTypeInfo resultType;
-  resultType.metadata = T;
-
   TaskGroupFlags groupFlags(rawGroupFlags);
   SWIFT_TASK_GROUP_DEBUG_LOG_0(group, "create group, from task:%p; flags: isDiscardingResults=%d",
                                swift_task_getCurrent(),
                                groupFlags.isDiscardResults());
 
+  ResultTypeInfo resultType;
+#if !SWIFT_CONCURRENCY_EMBEDDED
+  resultType.metadata = T;
+  assert(!groupFlags.isMetadataAsOptionRecord());
+#else
+  assert(groupFlags.isMetadataAsOptionRecord());
+  TaskOptionRecord *options = (TaskOptionRecord *)T;
+  for (auto option = options; option; option = option->getParent()) {
+    switch (option->getKind()) {
+    case TaskOptionRecordKind::ResultTypeInfo: {
+      auto *typeInfo = cast<ResultTypeInfoTaskOptionRecord>(option);
+      resultType = {
+          .size = typeInfo->size,
+          .alignMask = typeInfo->alignMask,
+          .initializeWithCopy = typeInfo->initializeWithCopy,
+          .storeEnumTagSinglePayload = typeInfo->storeEnumTagSinglePayload,
+          .destroy = typeInfo->destroy,
+      };
+      break;
+    }
+    default:
+      swift_unreachable("only ResultTypeInfo expected");
+    }
+  }
+#endif
+
   TaskGroupBase *impl;
   if (groupFlags.isDiscardResults()) {
     impl = ::new(group) DiscardingTaskGroup(resultType);
@@ -993,9 +1015,6 @@ static void swift_taskGroup_initializeWithFlagsImpl(size_t rawGroupFlags,
     }
     return true;
   });
-#else
-  swift_unreachable("task groups not supported yet in embedded Swift");
-#endif
 }
 
 // =============================================================================

stdlib/public/Concurrency/TaskGroup.swift

@@ -63,7 +63,9 @@ import Swift
 /// For tasks that need to handle cancellation by throwing an error,
 /// use the `withThrowingTaskGroup(of:returning:body:)` method instead.
 @available(SwiftStdlib 5.1, *)
+#if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
+#endif
 @inlinable
 public func withTaskGroup<ChildTaskResult, GroupResult>(
   of childTaskResultType: ChildTaskResult.Type = ChildTaskResult.self,
@@ -198,7 +200,9 @@ public func _unsafeInheritExecutor_withTaskGroup<ChildTaskResult, GroupResult>(
 /// which gives you a chance to handle the individual error
 /// or to let the group rethrow the error.
 @available(SwiftStdlib 5.1, *)
+#if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
+#endif
 @inlinable
 public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
   of childTaskResultType: ChildTaskResult.Type = ChildTaskResult.self,
@@ -592,7 +596,9 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
   ///
   /// - Returns: The value returned by the next child task that completes.
   @available(SwiftStdlib 5.1, *)
+  #if !$Embedded
   @backDeployed(before: SwiftStdlib 6.0)
+  #endif
   public mutating func next(isolation: isolated (any Actor)? = #isolation) async -> ChildTaskResult? {
     // try!-safe because this function only exists for Failure == Never,
     // and as such, it is impossible to spawn a throwing child task.
@@ -610,7 +616,9 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
   /// Await all of the pending tasks added this group.
   @usableFromInline
   @available(SwiftStdlib 5.1, *)
+  #if !$Embedded
   @backDeployed(before: SwiftStdlib 6.0)
+  #endif
   internal mutating func awaitAllRemainingTasks(isolation: isolated (any Actor)? = #isolation) async {
     while let _ = await next(isolation: isolation) {}
   }
@@ -750,7 +758,9 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   /// Await all the remaining tasks on this group.
   @usableFromInline
   @available(SwiftStdlib 5.1, *)
+  #if !$Embedded
   @backDeployed(before: SwiftStdlib 6.0)
+  #endif
   internal mutating func awaitAllRemainingTasks(isolation: isolated (any Actor)? = #isolation) async {
     while true {
       do {
@@ -1028,7 +1038,9 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   ///
   /// - SeeAlso: `nextResult()`
   @available(SwiftStdlib 5.1, *)
+  #if !$Embedded
   @backDeployed(before: SwiftStdlib 6.0)
+  #endif
   public mutating func next(isolation: isolated (any Actor)? = #isolation) async throws -> ChildTaskResult? {
     return try await _taskGroupWaitNext(group: _group)
   }

stdlib/public/Concurrency/TaskPrivate.h

@@ -322,8 +322,12 @@ class TaskFutureWaitAsyncContext : public AsyncContext {
     fillWithError(future->getError());
   }
   void fillWithError(SwiftError *error) {
+    #if SWIFT_CONCURRENCY_EMBEDDED
+    swift_unreachable("untyped error used in embedded Swift");
+    #else
     errorResult = error;
     swift_errorRetain(error);
+    #endif
   }
 };
 

test/embedded/concurrency-taskgroup.swift

@@ -0,0 +1,61 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -enable-experimental-feature Embedded -target %target-cpu-apple-macos14 -parse-as-library %s -c -o %t/a.o
+// RUN: %target-clang -x c -c %S/Inputs/print.c -o %t/print.o
+// RUN: %target-clang %t/a.o %t/print.o -o %t/a.out %swift_obj_root/lib/swift/embedded/%target-cpu-apple-macos/libswift_Concurrency.a -dead_strip
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: swift_in_compiler
+// REQUIRES: optimized_stdlib
+// REQUIRES: OS=macosx
+
+import _Concurrency
+
+protocol Go: Actor {
+  var name: String { get }
+  func go(times: Int) async -> Int
+}
+
+extension Go {
+  func go(times: Int) async -> Int {
+    for i in 0..<times {
+      print("\(name) @ \(i)")
+      await Task.yield()
+    }
+    return times
+  }
+}
+
+actor One: Go { var name = "One" }
+actor Two: Go { var name = "Two" }
+
+func yielding() async {
+  let one = One()
+  let two = Two()
+  await withTaskGroup(of: Int.self) { group in
+    group.addTask {
+      await one.go(times: 5)
+    }
+    group.addTask {
+      await two.go(times: 5)
+    }
+  }
+}
+
+@main struct Main {
+    static func main() async {
+        await yielding()
+        print("All done!")
+        // CHECK: One @ 0
+        // CHECK: Two @ 0
+        // CHECK: One @ 1
+        // CHECK: Two @ 1
+        // CHECK: One @ 2
+        // CHECK: Two @ 2
+        // CHECK: One @ 3
+        // CHECK: Two @ 3
+        // CHECK: One @ 4
+        // CHECK: Two @ 4
+        // CHECK: All done!
+    }
+}