Unsafe `withTaskGroup` optimization

[za-creature]

Description

I'm filing this under compiler because I get different results between -Onone and -O. It might be an issue with the soon-to-be-deprecated swift-atomics package but my guess is otherwise.

Reproduction

import Atomics

func inner(_ group: inout TaskGroup<Void>) async -> ManagedAtomic<Int> {
    let count = ManagedAtomic<Int>(314) // guard value
    for _ in 0..<10 {
        group.addTask { count.wrappingIncrement(ordering: .relaxed) }
    }
    return count
}

func outer1() async -> Int {
    async let count = withTaskGroup(of: Void.self, body: inner)
    return await count.load(ordering: .relaxed)
}

func outer2() async -> Int {
    let count = await withTaskGroup(of: Void.self, body: inner)
    return count.load(ordering: .relaxed)
}

await print(outer1(), outer2())

Expected behavior

While I'm using .relaxed loads and stores, there is a data dependency between outer1() / outer2() and inner(): neither outer func can know what address to load the atomic value from until withTaskGroup returns which is documented as:

A group waits for all of its child tasks to complete or be canceled before it returns. After this function returns, the task group is always empty.

When compiled with -Onone, I get the expected result: 324, 324

With -O or -Osize, the result is a worrying 324, 0 on both x86 and M1

It seems that in optimized builds, when not using async let, the atomic value is loaded not just before the return of withTaskGroup(), but before it was initialized to its 314 guard value.

Environment

swift-driver version: 1.87.3 Apple Swift version 5.9.2 (swiftlang-5.9.2.2.56 clang-1500.1.0.2.5)
Target: x86_64-apple-macosx13.0

Additional information

Initially reported at https://forums.swift.org/t/unsafe-withtaskgroup-optimization/69372

[AnthonyLatsis]

@ktoso Do you reckon this is a runtime bug, not an optimizer bug?

[ktoso]

No idea, it doesn't look like a runtime bug per se... I'd like to ask some optimizer folks to see what they think

[mikeash]

Thank you for the report! This is an interesting interaction with the unsafe pointer code inside swift-atomics getting inlined into the caller, which can cause the compiler to release the ManagedAtomic object before actually performing that final load. I've turned my investigation over to the swift-atomics folks, but for now you can work around the issue like this:

    return withExtendedLifetime(count) { count.load(ordering: .relaxed) }