Skip to content

Commit 93a0a6d

Browse files
authored
Merge pull request #62914 from ktoso/wip-task-group-twotypes
[Concurrency] DiscardingTaskGroup (rev 3)
2 parents d8c8ab3 + 1d1cda7 commit 93a0a6d

34 files changed

+2350
-439
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: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,13 +45,16 @@ class TaskStatusRecord {
4545
TaskStatusRecord(TaskStatusRecordKind kind,
4646
TaskStatusRecord *parent = nullptr)
4747
: Flags(kind) {
48+
getKind();
4849
resetParent(parent);
4950
}
5051

5152
TaskStatusRecord(const TaskStatusRecord &) = delete;
5253
TaskStatusRecord &operator=(const TaskStatusRecord &) = delete;
5354

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

5659
TaskStatusRecord *getParent() const { return Parent; }
5760

@@ -172,15 +175,14 @@ class ChildTaskStatusRecord : public TaskStatusRecord {
172175
/// Group child tasks DO NOT have their own `ChildTaskStatusRecord` entries,
173176
/// and are only tracked by their respective `TaskGroupTaskStatusRecord`.
174177
class TaskGroupTaskStatusRecord : public TaskStatusRecord {
178+
public:
175179
AsyncTask *FirstChild;
176180
AsyncTask *LastChild;
177181

178-
public:
179182
TaskGroupTaskStatusRecord()
180183
: TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
181184
FirstChild(nullptr),
182-
LastChild(nullptr) {
183-
}
185+
LastChild(nullptr) {}
184186

185187
TaskGroupTaskStatusRecord(AsyncTask *child)
186188
: TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
@@ -189,7 +191,8 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
189191
assert(!LastChild || !LastChild->childFragment()->getNextChild());
190192
}
191193

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

194197
/// Return the first child linked by this record. This may be null;
195198
/// 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
@@ -2658,6 +2658,9 @@ FunctionPointer::Kind irgen::classifyFunctionPointerKind(SILFunction *fn) {
26582658
if (name.equals("swift_taskGroup_wait_next_throwing"))
26592659
return SpecialKind::TaskGroupWaitNext;
26602660

2661+
if (name.equals("swift_taskGroup_waitAll"))
2662+
return SpecialKind::TaskGroupWaitAll;
2663+
26612664
if (name.equals("swift_distributed_execute_target"))
26622665
return SpecialKind::DistributedExecuteTarget;
26632666
}

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
@@ -2505,6 +2505,7 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
25052505
case BuiltinValueKind::EndAsyncLet:
25062506
case BuiltinValueKind::EndAsyncLetLifetime:
25072507
case BuiltinValueKind::CreateTaskGroup:
2508+
case BuiltinValueKind::CreateTaskGroupWithFlags:
25082509
case BuiltinValueKind::DestroyTaskGroup:
25092510
return;
25102511

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/BackDeployConcurrency/CompatibilityOverrideConcurrency.def

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,10 @@ OVERRIDE_TASK_GROUP(taskGroup_initialize, void,
241241
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
242242
swift::, (TaskGroup *group, const Metadata *T), (group, T))
243243

244+
OVERRIDE_TASK_GROUP(taskGroup_initializeWithFlags, void,
245+
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
246+
swift::, (size_t flags, TaskGroup *group, const Metadata *T), (flags, group, T))
247+
244248
OVERRIDE_TASK_STATUS(taskGroup_attachChild, void,
245249
SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
246250
swift::, (TaskGroup *group, AsyncTask *child),

0 commit comments

Comments
 (0)