Skip to content

Commit 857d1bd

Browse files
DougGregorDougGregor
authored
Merge pull request #62963 from ktoso/pick-discarding-task-group
2 parents 83388d3 + 3d837dd commit 857d1bd

30 files changed

+2339
-428
lines changed

‎include/swift/ABI/MetadataValues.h

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2383,6 +2383,24 @@ enum class TaskOptionRecordKind : uint8_t {
23832383
RunInline = UINT8_MAX,
23842384
};
23852385

2386+
/// Flags for TaskGroup.
2387+
class TaskGroupFlags : public FlagSet<uint32_t> {
2388+
public:
2389+
enum {
2390+
// 8 bits are reserved for future use
2391+
/// Request the TaskGroup to immediately release completed tasks,
2392+
/// and not store their results. This also effectively disables `next()`.
2393+
TaskGroup_DiscardResults = 8,
2394+
};
2395+
2396+
explicit TaskGroupFlags(uint32_t bits) : FlagSet(bits) {}
2397+
constexpr TaskGroupFlags() {}
2398+
2399+
FLAGSET_DEFINE_FLAG_ACCESSORS(TaskGroup_DiscardResults,
2400+
isDiscardResults,
2401+
setIsDiscardResults)
2402+
};
2403+
23862404
/// Flags for cancellation records.
23872405
class TaskStatusRecordFlags : public FlagSet<size_t> {
23882406
public:

‎include/swift/ABI/TaskGroup.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,14 @@ class alignas(Alignment_TaskGroup) TaskGroup {
5454

5555
// Provide accessor for task group's status record
5656
TaskGroupTaskStatusRecord *getTaskRecord();
57+
58+
/// The group is a `TaskGroup` that accumulates results.
59+
bool isAccumulatingResults() {
60+
return !isDiscardingResults();
61+
}
62+
63+
/// The group is a `DiscardingTaskGroup` that discards results.
64+
bool isDiscardingResults();
5765
};
5866

5967
} // end namespace swift

‎include/swift/ABI/TaskStatus.h

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,9 @@ class TaskStatusRecord {
5151
TaskStatusRecord(const TaskStatusRecord &) = delete;
5252
TaskStatusRecord &operator=(const TaskStatusRecord &) = delete;
5353

54-
TaskStatusRecordKind getKind() const { return Flags.getKind(); }
54+
TaskStatusRecordKind getKind() const {
55+
return Flags.getKind();
56+
}
5557

5658
TaskStatusRecord *getParent() const { return Parent; }
5759

@@ -179,8 +181,7 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
179181
TaskGroupTaskStatusRecord()
180182
: TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
181183
FirstChild(nullptr),
182-
LastChild(nullptr) {
183-
}
184+
LastChild(nullptr) {}
184185

185186
TaskGroupTaskStatusRecord(AsyncTask *child)
186187
: TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
@@ -189,7 +190,8 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
189190
assert(!LastChild || !LastChild->childFragment()->getNextChild());
190191
}
191192

192-
TaskGroup *getGroup() { return reinterpret_cast<TaskGroup *>(this); }
193+
/// Get the task group this record is associated with.
194+
TaskGroup *getGroup();
193195

194196
/// Return the first child linked by this record. This may be null;
195197
/// if not, it (and all of its successors) are guaranteed to satisfy

‎include/swift/AST/Builtins.def

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -793,6 +793,10 @@ BUILTIN_MISC_OPERATION(ResumeThrowingContinuationThrowing,
793793
BUILTIN_MISC_OPERATION(CreateTaskGroup,
794794
"createTaskGroup", "", Special)
795795

796+
/// Create a task group, with options.
797+
BUILTIN_MISC_OPERATION(CreateTaskGroupWithFlags,
798+
"createTaskGroupWithFlags", "", Special)
799+
796800
/// Destroy a task group.
797801
BUILTIN_MISC_OPERATION(DestroyTaskGroup,
798802
"destroyTaskGroup", "", Special)

‎include/swift/Runtime/Concurrency.h

Lines changed: 38 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -185,13 +185,15 @@ void swift_task_future_wait_throwing(
185185
/// func swift_taskGroup_wait_next_throwing(
186186
/// waitingTask: Builtin.NativeObject, // current task
187187
/// group: Builtin.RawPointer
188-
/// ) async -> T
188+
/// ) async throws -> T
189189
/// \endcode
190190
SWIFT_EXPORT_FROM(swift_Concurrency)
191191
SWIFT_CC(swiftasync)
192192
void swift_taskGroup_wait_next_throwing(
193-
OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
194-
TaskGroup *group, ThrowingTaskFutureWaitContinuationFunction *resumeFn,
193+
OpaqueValue *resultPointer,
194+
SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
195+
TaskGroup *group,
196+
ThrowingTaskFutureWaitContinuationFunction *resumeFn,
195197
AsyncContext *callContext);
196198

197199
/// Initialize a `TaskGroup` in the passed `group` memory location.
@@ -205,6 +207,17 @@ void swift_taskGroup_wait_next_throwing(
205207
SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
206208
void swift_taskGroup_initialize(TaskGroup *group, const Metadata *T);
207209

210+
/// Initialize a `TaskGroup` in the passed `group` memory location.
211+
/// The caller is responsible for retaining and managing the group's lifecycle.
212+
///
213+
/// Its Swift signature is
214+
///
215+
/// \code
216+
/// func swift_taskGroup_initialize(flags: Int, group: Builtin.RawPointer)
217+
/// \endcode
218+
SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
219+
void swift_taskGroup_initializeWithFlags(size_t flags, TaskGroup *group, const Metadata *T);
220+
208221
/// Attach a child task to the parent task's task group record.
209222
///
210223
/// This function MUST be called from the AsyncTask running the task group.
@@ -276,6 +289,28 @@ void swift_taskGroup_cancelAll(TaskGroup *group);
276289
SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
277290
bool swift_taskGroup_isCancelled(TaskGroup *group);
278291

292+
/// Wait until all pending tasks from the task group have completed.
293+
/// If this task group is accumulating results, this also discards all those results.
294+
///
295+
/// This can be called from any thread. Its Swift signature is
296+
///
297+
/// \code
298+
/// func swift_taskGroup_waitAll(
299+
/// waitingTask: Builtin.NativeObject, // current task
300+
/// group: Builtin.RawPointer,
301+
/// bodyError: Swift.Error?
302+
/// ) async throws
303+
/// \endcode
304+
SWIFT_EXPORT_FROM(swift_Concurrency)
305+
SWIFT_CC(swiftasync)
306+
void swift_taskGroup_waitAll(
307+
OpaqueValue *resultPointer,
308+
SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
309+
TaskGroup *group,
310+
SwiftError *bodyError,
311+
ThrowingTaskFutureWaitContinuationFunction *resumeFn,
312+
AsyncContext *callContext);
313+
279314
/// Check the readyQueue of a task group, return true if it has no pending tasks.
280315
///
281316
/// This can be called from any thread. Its Swift signature is

‎include/swift/Runtime/RuntimeFunctions.def

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2060,6 +2060,18 @@ FUNCTION(TaskGroupInitialize,
20602060
ATTRS(NoUnwind),
20612061
EFFECT(Concurrency))
20622062

2063+
// void swift_taskGroup_initializeWithFlags(size_t flags, TaskGroup *group);
2064+
FUNCTION(TaskGroupInitializeWithFlags,
2065+
swift_taskGroup_initializeWithFlags, SwiftCC,
2066+
ConcurrencyAvailability,
2067+
RETURNS(VoidTy),
2068+
ARGS(SizeTy, // flags
2069+
Int8PtrTy, // group
2070+
TypeMetadataPtrTy // T.Type
2071+
),
2072+
ATTRS(NoUnwind),
2073+
EFFECT(Concurrency))
2074+
20632075
// void swift_taskGroup_destroy(TaskGroup *group);
20642076
FUNCTION(TaskGroupDestroy,
20652077
swift_taskGroup_destroy, SwiftCC,

‎lib/AST/Builtins.cpp

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1539,6 +1539,24 @@ static ValueDecl *getCreateTaskGroup(ASTContext &ctx, Identifier id) {
15391539
_rawPointer);
15401540
}
15411541

1542+
static ValueDecl *getCreateTaskGroupWithFlags(ASTContext &ctx, Identifier id) {
1543+
ModuleDecl *M = ctx.TheBuiltinModule;
1544+
DeclContext *DC = &M->getMainFile(FileUnitKind::Builtin);
1545+
SynthesisContext SC(ctx, DC);
1546+
1547+
BuiltinFunctionBuilder builder(ctx);
1548+
1549+
// int
1550+
builder.addParameter(makeConcrete(ctx.getIntType())); // 0 flags
1551+
1552+
// T.self
1553+
builder.addParameter(makeMetatype(makeGenericParam(0))); // 1 ChildTaskResult.Type
1554+
1555+
// -> Builtin.RawPointer
1556+
builder.setResult(makeConcrete(synthesizeType(SC, _rawPointer)));
1557+
return builder.build(id);
1558+
}
1559+
15421560
static ValueDecl *getDestroyTaskGroup(ASTContext &ctx, Identifier id) {
15431561
return getBuiltinFunction(ctx, id, _thin,
15441562
_parameters(_rawPointer),
@@ -2908,6 +2926,8 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
29082926

29092927
case BuiltinValueKind::CreateTaskGroup:
29102928
return getCreateTaskGroup(Context, Id);
2929+
case BuiltinValueKind::CreateTaskGroupWithFlags:
2930+
return getCreateTaskGroupWithFlags(Context, Id);
29112931

29122932
case BuiltinValueKind::DestroyTaskGroup:
29132933
return getDestroyTaskGroup(Context, Id);

‎lib/IRGen/Callee.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ namespace irgen {
175175
AsyncLetGetThrowing,
176176
AsyncLetFinish,
177177
TaskGroupWaitNext,
178+
TaskGroupWaitAll,
178179
DistributedExecuteTarget,
179180
};
180181

@@ -247,6 +248,7 @@ namespace irgen {
247248
case SpecialKind::AsyncLetGetThrowing:
248249
case SpecialKind::AsyncLetFinish:
249250
case SpecialKind::TaskGroupWaitNext:
251+
case SpecialKind::TaskGroupWaitAll:
250252
return true;
251253
case SpecialKind::DistributedExecuteTarget:
252254
return false;
@@ -277,6 +279,7 @@ namespace irgen {
277279
case SpecialKind::AsyncLetGetThrowing:
278280
case SpecialKind::AsyncLetFinish:
279281
case SpecialKind::TaskGroupWaitNext:
282+
case SpecialKind::TaskGroupWaitAll:
280283
return true;
281284
case SpecialKind::DistributedExecuteTarget:
282285
return false;

‎lib/IRGen/GenBuiltin.cpp

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -277,9 +277,22 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
277277
}
278278

279279
if (Builtin.ID == BuiltinValueKind::CreateTaskGroup) {
280+
llvm::Value *groupFlags = nullptr;
280281
// Claim metadata pointer.
281282
(void)args.claimAll();
282-
out.add(emitCreateTaskGroup(IGF, substitutions));
283+
out.add(emitCreateTaskGroup(IGF, substitutions, groupFlags));
284+
return;
285+
}
286+
287+
if (Builtin.ID == BuiltinValueKind::CreateTaskGroupWithFlags) {
288+
auto groupFlags = args.claimNext();
289+
// Claim the remaining metadata pointer.
290+
if (args.size() == 1) {
291+
(void)args.claimNext();
292+
} else if (args.size() > 1) {
293+
llvm_unreachable("createTaskGroupWithFlags expects 1 or 2 arguments");
294+
}
295+
out.add(emitCreateTaskGroup(IGF, substitutions, groupFlags));
283296
return;
284297
}
285298

‎lib/IRGen/GenCall.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ FunctionPointerKind::getStaticAsyncContextSize(IRGenModule &IGM) const {
162162
case SpecialKind::AsyncLetGetThrowing:
163163
case SpecialKind::AsyncLetFinish:
164164
case SpecialKind::TaskGroupWaitNext:
165+
case SpecialKind::TaskGroupWaitAll:
165166
case SpecialKind::DistributedExecuteTarget:
166167
// The current guarantee for all of these functions is the same.
167168
// See TaskFutureWaitAsyncContext.

‎lib/IRGen/GenConcurrency.cpp

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,8 @@ void irgen::emitEndAsyncLet(IRGenFunction &IGF, llvm::Value *alet) {
271271
}
272272

273273
llvm::Value *irgen::emitCreateTaskGroup(IRGenFunction &IGF,
274-
SubstitutionMap subs) {
274+
SubstitutionMap subs,
275+
llvm::Value *groupFlags) {
275276
auto ty = llvm::ArrayType::get(IGF.IGM.Int8PtrTy, NumWords_TaskGroup);
276277
auto address = IGF.createAlloca(ty, Alignment(Alignment_TaskGroup));
277278
auto group = IGF.Builder.CreateBitCast(address.getAddress(),
@@ -282,9 +283,14 @@ llvm::Value *irgen::emitCreateTaskGroup(IRGenFunction &IGF,
282283
auto resultType = subs.getReplacementTypes()[0]->getCanonicalType();
283284
auto resultTypeMetadata = IGF.emitAbstractTypeMetadataRef(resultType);
284285

285-
auto *call =
286-
IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeFunctionPointer(),
287-
{group, resultTypeMetadata});
286+
llvm::CallInst *call;
287+
if (groupFlags) {
288+
call = IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeWithFlagsFunctionPointer(),
289+
{groupFlags, group, resultTypeMetadata});
290+
} else {
291+
call = IGF.Builder.CreateCall(IGF.IGM.getTaskGroupInitializeFunctionPointer(),
292+
{group, resultTypeMetadata});
293+
}
288294
call->setDoesNotThrow();
289295
call->setCallingConv(IGF.IGM.SwiftCC);
290296

‎lib/IRGen/GenConcurrency.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ llvm::Value *emitBuiltinStartAsyncLet(IRGenFunction &IGF,
6969
void emitEndAsyncLet(IRGenFunction &IGF, llvm::Value *alet);
7070

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

7475
/// Emit the destroyTaskGroup builtin.
7576
void emitDestroyTaskGroup(IRGenFunction &IGF, llvm::Value *group);

‎lib/IRGen/IRGenSIL.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2655,6 +2655,9 @@ FunctionPointer::Kind irgen::classifyFunctionPointerKind(SILFunction *fn) {
26552655
if (name.equals("swift_taskGroup_wait_next_throwing"))
26562656
return SpecialKind::TaskGroupWaitNext;
26572657

2658+
if (name.equals("swift_taskGroup_waitAll"))
2659+
return SpecialKind::TaskGroupWaitAll;
2660+
26582661
if (name.equals("swift_distributed_execute_target"))
26592662
return SpecialKind::DistributedExecuteTarget;
26602663
}

‎lib/SIL/IR/OperandOwnership.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -821,6 +821,7 @@ BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, EndAsyncLet)
821821
BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, StartAsyncLetWithLocalBuffer)
822822
BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, EndAsyncLetLifetime)
823823
BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroup)
824+
BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroupWithFlags)
824825
BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, DestroyTaskGroup)
825826

826827
BUILTIN_OPERAND_OWNERSHIP(ForwardingConsume, COWBufferForReading)

‎lib/SIL/IR/ValueOwnership.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -559,6 +559,7 @@ CONSTANT_OWNERSHIP_BUILTIN(None, EndAsyncLet)
559559
CONSTANT_OWNERSHIP_BUILTIN(None, StartAsyncLetWithLocalBuffer)
560560
CONSTANT_OWNERSHIP_BUILTIN(None, EndAsyncLetLifetime)
561561
CONSTANT_OWNERSHIP_BUILTIN(None, CreateTaskGroup)
562+
CONSTANT_OWNERSHIP_BUILTIN(None, CreateTaskGroupWithFlags)
562563
CONSTANT_OWNERSHIP_BUILTIN(None, DestroyTaskGroup)
563564
CONSTANT_OWNERSHIP_BUILTIN(None, TaskRunInline)
564565
CONSTANT_OWNERSHIP_BUILTIN(None, Copy)

‎lib/SIL/Utils/MemAccessUtils.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2539,6 +2539,7 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
25392539
case BuiltinValueKind::EndAsyncLet:
25402540
case BuiltinValueKind::EndAsyncLetLifetime:
25412541
case BuiltinValueKind::CreateTaskGroup:
2542+
case BuiltinValueKind::CreateTaskGroupWithFlags:
25422543
case BuiltinValueKind::DestroyTaskGroup:
25432544
return;
25442545

‎lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ static bool isBarrier(SILInstruction *inst) {
150150
case BuiltinValueKind::EndAsyncLet:
151151
case BuiltinValueKind::EndAsyncLetLifetime:
152152
case BuiltinValueKind::CreateTaskGroup:
153+
case BuiltinValueKind::CreateTaskGroupWithFlags:
153154
case BuiltinValueKind::DestroyTaskGroup:
154155
case BuiltinValueKind::StackAlloc:
155156
case BuiltinValueKind::StackDealloc:

‎stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,10 @@ OVERRIDE_TASK_GROUP(taskGroup_initialize, void,
275275
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
276276
swift::, (TaskGroup *group, const Metadata *T), (group, T))
277277

278+
OVERRIDE_TASK_GROUP(taskGroup_initializeWithFlags, void,
279+
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
280+
swift::, (size_t flags, TaskGroup *group, const Metadata *T), (flags, group, T))
281+
278282
OVERRIDE_TASK_STATUS(taskGroup_attachChild, void,
279283
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
280284
swift::, (TaskGroup *group, AsyncTask *child),
@@ -312,6 +316,17 @@ OVERRIDE_TASK_GROUP(taskGroup_addPending, bool,
312316
swift::, (TaskGroup *group, bool unconditionally),
313317
(group, unconditionally))
314318

319+
OVERRIDE_TASK_GROUP(taskGroup_waitAll, void,
320+
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
321+
swift::,
322+
(OpaqueValue *resultPointer,
323+
SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
324+
TaskGroup *_group,
325+
SwiftError *bodyError,
326+
ThrowingTaskFutureWaitContinuationFunction *resumeFn,
327+
AsyncContext *callContext),
328+
(resultPointer, callerContext, _group, bodyError,
329+
resumeFn, callContext))
315330

316331
OVERRIDE_TASK_LOCAL(task_reportIllegalTaskLocalBindingWithinWithTaskGroup, void,
317332
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,

‎stdlib/public/Concurrency/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ add_swift_target_library(swift_Concurrency ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} I
113113
TaskStatus.cpp
114114
TaskGroup.cpp
115115
TaskGroup.swift
116+
DiscardingTaskGroup.swift
116117
TaskLocal.cpp
117118
TaskLocal.swift
118119
TaskSleep.swift

0 commit comments

Comments
 (0)