[Concurrency] DiscardingTaskGroup

prepare flags

wip on options

implement discardResults as a flag passed to grout init

initial complete impl

prepare for cancellation handling

fix abi

move to DiscardingTaskGroup

more docs

complete impl, except body throws

implement simple rethrowing logic, however body throw always wins

DiscardingTaskGroup now shares some implementation with "Accumulating" TaskGroup

remove duplicated status logic

more cleanup

remove even more duplicated methods

undo any changes in backdeploy lib, we don't do any changes there

fix cast type in getting task record

remove changes in backdeploy lib, not needed

cleanup

remove import-darwin from test

Author: ktoso

include/swift/ABI/MetadataValues.h

@@ -2383,6 +2383,24 @@ enum class TaskOptionRecordKind : uint8_t {
   RunInline = UINT8_MAX,
 };
 
+/// Flags for TaskGroup.
+class TaskGroupFlags : public FlagSet<uint32_t> {
+public:
+  enum {
+    // 8 bits are reserved for future use
+    /// Request the TaskGroup to immediately release completed tasks,
+    /// and not store their results. This also effectively disables `next()`.
+    TaskGroup_DiscardResults = 8,
+  };
+
+  explicit TaskGroupFlags(uint32_t bits) : FlagSet(bits) {}
+  constexpr TaskGroupFlags() {}
+
+  FLAGSET_DEFINE_FLAG_ACCESSORS(TaskGroup_DiscardResults,
+                                isDiscardResults,
+                                setIsDiscardResults)
+};
+
 /// Flags for cancellation records.
 class TaskStatusRecordFlags : public FlagSet<size_t> {
 public:

include/swift/ABI/TaskGroup.h

@@ -54,6 +54,14 @@ class alignas(Alignment_TaskGroup) TaskGroup {
 
   // Provide accessor for task group's status record
   TaskGroupTaskStatusRecord *getTaskRecord();
+
+  /// The group is a `TaskGroup` that accumulates results.
+  bool isAccumulatingResults() {
+    return !isDiscardingResults();
+  }
+
+  /// The group is a `DiscardingTaskGroup` that discards results.
+  bool isDiscardingResults();
 };
 
 } // end namespace swift

include/swift/ABI/TaskStatus.h

@@ -51,7 +51,9 @@ class TaskStatusRecord {
   TaskStatusRecord(const TaskStatusRecord &) = delete;
   TaskStatusRecord &operator=(const TaskStatusRecord &) = delete;
 
-  TaskStatusRecordKind getKind() const { return Flags.getKind(); }
+  TaskStatusRecordKind getKind() const {
+    return Flags.getKind();
+  }
 
   TaskStatusRecord *getParent() const { return Parent; }
 
@@ -179,8 +181,7 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
   TaskGroupTaskStatusRecord()
       : TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
         FirstChild(nullptr),
-        LastChild(nullptr) {
-  }
+        LastChild(nullptr) {}
 
   TaskGroupTaskStatusRecord(AsyncTask *child)
       : TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
@@ -189,7 +190,8 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
     assert(!LastChild || !LastChild->childFragment()->getNextChild());
   }
 
-  TaskGroup *getGroup() { return reinterpret_cast<TaskGroup *>(this); }
+  /// Get the task group this record is associated with.
+  TaskGroup *getGroup();
 
   /// Return the first child linked by this record.  This may be null;
   /// if not, it (and all of its successors) are guaranteed to satisfy

include/swift/AST/Builtins.def

@@ -793,6 +793,10 @@ BUILTIN_MISC_OPERATION(ResumeThrowingContinuationThrowing,
 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)

include/swift/Runtime/Concurrency.h

@@ -185,13 +185,15 @@ void swift_task_future_wait_throwing(
 /// func swift_taskGroup_wait_next_throwing(
 ///     waitingTask: Builtin.NativeObject, // current task
 ///     group: Builtin.RawPointer
-/// ) async -> T
+/// ) async throws -> T
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swiftasync)
 void swift_taskGroup_wait_next_throwing(
-    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
-    TaskGroup *group, ThrowingTaskFutureWaitContinuationFunction *resumeFn,
+    OpaqueValue *resultPointer,
+    SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    TaskGroup *group,
+    ThrowingTaskFutureWaitContinuationFunction *resumeFn,
     AsyncContext *callContext);
 
 /// Initialize a `TaskGroup` in the passed `group` memory location.
@@ -205,6 +207,17 @@ void swift_taskGroup_wait_next_throwing(
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_taskGroup_initialize(TaskGroup *group, const Metadata *T);
 
+/// Initialize a `TaskGroup` in the passed `group` memory location.
+/// The caller is responsible for retaining and managing the group's lifecycle.
+///
+/// Its Swift signature is
+///
+/// \code
+/// func swift_taskGroup_initialize(flags: Int, group: Builtin.RawPointer)
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_taskGroup_initializeWithFlags(size_t flags, TaskGroup *group, const Metadata *T);
+
 /// Attach a child task to the parent task's task group record.
 ///
 /// This function MUST be called from the AsyncTask running the task group.
@@ -276,6 +289,28 @@ void swift_taskGroup_cancelAll(TaskGroup *group);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_taskGroup_isCancelled(TaskGroup *group);
 
+/// Wait until all pending tasks from the task group have completed.
+/// If this task group is accumulating results, this also discards all those results.
+///
+/// This can be called from any thread. Its Swift signature is
+///
+/// \code
+/// func swift_taskGroup_waitAll(
+///     waitingTask: Builtin.NativeObject, // current task
+///     group: Builtin.RawPointer,
+///     bodyError: Swift.Error?
+/// ) async throws
+/// \endcode
+  SWIFT_EXPORT_FROM(swift_Concurrency)
+  SWIFT_CC(swiftasync)
+  void swift_taskGroup_waitAll(
+      OpaqueValue *resultPointer,
+      SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+      TaskGroup *group,
+      SwiftError *bodyError,
+      ThrowingTaskFutureWaitContinuationFunction *resumeFn,
+      AsyncContext *callContext);
+
 /// Check the readyQueue of a task group, return true if it has no pending tasks.
 ///
 /// This can be called from any thread. Its Swift signature is

include/swift/Runtime/RuntimeFunctions.def

@@ -2060,6 +2060,18 @@ FUNCTION(TaskGroupInitialize,
          ATTRS(NoUnwind),
          EFFECT(Concurrency))
 
+// void swift_taskGroup_initializeWithFlags(size_t flags, TaskGroup *group);
+FUNCTION(TaskGroupInitializeWithFlags,
+         swift_taskGroup_initializeWithFlags, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(SizeTy,           // flags
+              Int8PtrTy,        // group
+              TypeMetadataPtrTy // T.Type
+         ),
+         ATTRS(NoUnwind),
+         EFFECT(Concurrency))
+
 // void swift_taskGroup_destroy(TaskGroup *group);
 FUNCTION(TaskGroupDestroy,
          swift_taskGroup_destroy, SwiftCC,

lib/AST/Builtins.cpp

@@ -1539,6 +1539,24 @@ static ValueDecl *getCreateTaskGroup(ASTContext &ctx, Identifier id) {
                             _rawPointer);
 }
 
+static ValueDecl *getCreateTaskGroupWithFlags(ASTContext &ctx, Identifier id) {
+  ModuleDecl *M = ctx.TheBuiltinModule;
+  DeclContext *DC = &M->getMainFile(FileUnitKind::Builtin);
+  SynthesisContext SC(ctx, DC);
+
+  BuiltinFunctionBuilder builder(ctx);
+
+  // int
+  builder.addParameter(makeConcrete(ctx.getIntType())); // 0 flags
+
+  // T.self
+  builder.addParameter(makeMetatype(makeGenericParam(0))); // 1 ChildTaskResult.Type
+
+  // -> Builtin.RawPointer
+  builder.setResult(makeConcrete(synthesizeType(SC, _rawPointer)));
+  return builder.build(id);
+}
+
 static ValueDecl *getDestroyTaskGroup(ASTContext &ctx, Identifier id) {
   return getBuiltinFunction(ctx, id, _thin,
                             _parameters(_rawPointer),
@@ -2908,6 +2926,8 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::CreateTaskGroup:
     return getCreateTaskGroup(Context, Id);
+  case BuiltinValueKind::CreateTaskGroupWithFlags:
+    return getCreateTaskGroupWithFlags(Context, Id);
 
   case BuiltinValueKind::DestroyTaskGroup:
     return getDestroyTaskGroup(Context, Id);

lib/IRGen/Callee.h

@@ -175,6 +175,7 @@ namespace irgen {
       AsyncLetGetThrowing,
       AsyncLetFinish,
       TaskGroupWaitNext,
+      TaskGroupWaitAll,
       DistributedExecuteTarget,
     };
 
@@ -247,6 +248,7 @@ namespace irgen {
       case SpecialKind::AsyncLetGetThrowing:
       case SpecialKind::AsyncLetFinish:
       case SpecialKind::TaskGroupWaitNext:
+      case SpecialKind::TaskGroupWaitAll:
         return true;
       case SpecialKind::DistributedExecuteTarget:
         return false;
@@ -277,6 +279,7 @@ namespace irgen {
       case SpecialKind::AsyncLetGetThrowing:
       case SpecialKind::AsyncLetFinish:
       case SpecialKind::TaskGroupWaitNext:
+      case SpecialKind::TaskGroupWaitAll:
         return true;
       case SpecialKind::DistributedExecuteTarget:
         return false;

lib/IRGen/GenBuiltin.cpp

@@ -277,9 +277,22 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
   }
 
   if (Builtin.ID == BuiltinValueKind::CreateTaskGroup) {
+    llvm::Value *groupFlags = nullptr;
     // Claim metadata pointer.
     (void)args.claimAll();
-    out.add(emitCreateTaskGroup(IGF, substitutions));
+    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");
+    }
+    out.add(emitCreateTaskGroup(IGF, substitutions, groupFlags));
     return;
   }
 

lib/IRGen/GenCall.cpp

@@ -162,6 +162,7 @@ FunctionPointerKind::getStaticAsyncContextSize(IRGenModule &IGM) const {
   case SpecialKind::AsyncLetGetThrowing:
   case SpecialKind::AsyncLetFinish:
   case SpecialKind::TaskGroupWaitNext:
+  case SpecialKind::TaskGroupWaitAll:
   case SpecialKind::DistributedExecuteTarget:
     // The current guarantee for all of these functions is the same.
     // See TaskFutureWaitAsyncContext.