Move coroutine upvars into locals for better memory economy

[dingxiangfei2009]

View all comments

Replace #127522
Related to #62958

The problem statement

#62958 demonstrates two problems. One is that upvars are always unconditionally promoted to prefix data fields of the state machine; the other is that the opportunity to achieve a more compact data layout is lost because captured upvars are not subjected to liveness analysis, in the sense that the memory space at one point occupied by upvars is never reclaimed and made available for other saved data across certain yield points, even when they are dead at those suspension locations.

The second problem is better demonstrated with this code snippet.

async fn work(another_fut: impl Future) {
    let _ = another_fut.await;
    // now `another_fut` is consumed
    let next_fut = async { .. };
    next_fut.await;
}

// `work`'s layout needs to reserve space for both `another_fut` and `next_fut`, while there is a clear missed opportunity
// to overlap the memory for `another_fut` and `next_fut` for better memory economy.

The difficulty lies with the fact that captured upvars do not receive their own locals inside a coroutine body. If we can assign locals to them somehow, we can run the layout scheme as usual and the optimisation on the data layout comes into effect out of the box in most cases.

Proposed changes

This is an initial work to improve memory economy of coroutine and async futures, by reducing the unnecessary of promotion of captured upvars into state prefix. In a nutshell, this patch works along the idea in this comment and this comment.

The patch contains the following changes.

1. Introduction of a RelocateUpvar MIR pass that inserts a MIR gadget, through which captured values by coroutine or async bodies or closures are moved into the inner MIR locals. This opens opportunities to subject the captured upvars to the same liveness analysis right before the StateTransform rewrites and determine which are the necessary ones to be stored in the coroutine state during suspension.
2. With this gadget, it means that we do not have to keep all upvars in the so-called prefix data regions of coroutine states. Instead, they are moved into the Unresumed state, or by convention the first variants of the state ADTs.
3. In addition, in case that some upvars are eventually used across more than one suspension point, which leads to their promotion into the prefix after all, we further arrange the coroutine state data layout, so that their offsets in the Unresumed state coincide with their memory slots after promotion. This means that during codegen, the additional moves introduced by the RelocateUpvar gadget are actually elided. The relevant change is implemented in rustc_abi.
4. We then have to pay the lip service to translate direct field access to the upvars into access behind the Unresumed variant.
5. We have to update diagnostics so that they are more informed about captured values and they make more sense in view of this change.
6. As requested by the review comments, the relocation only applies behind an unstable compiler flag -Z pack-coroutine-layout=captures-only. The default is pack-coroutine-layout=no, so that we keep the layout aligned with the stable.
7. Thanks to @ywxt, we spot that coroutine prefix accidentally picked up C-style enum prefix layout unnecessarily which mandates the variants to be aligned to the prefix enum. We relaxed this rule by resetting the alignment to 1 for variants.

Other than upvars, the coroutine state data layout scheme remains largely the same.

Test coverage

We have added the following tests.

• MIR transform diffs to show how the upvars are relocated.
• Run-pass tests to demonstrate that coroutine sizes can be compressed.
• Correctness assurance with additional revisions to existing run-pass tests, which switch on this new layout calculation.

Further optimisation to be implemented behind a feature gate

Point 4 mentions that any local to be saved across suspensions will be promoted whenever they are alive across two or more yield locations. We would like to run an experiment behind a feature gate on improvements of the layout scheme. For ease of reviewing, it is better to drop this part of work from this PR. Nevertheless, the idea runs along the implementation in #127522 and we intend to propose a second PR just for that.

Old PR description

Good day, this PR is related to #127522 and it is made easier to the public to test out a new coroutine/`async` state machine directly.

Prepare the compiler for tests

For starter, you may build the compiler as prescribed in the rustc-dev-guide instruction. If a test in the docker container is desirable, you may build this compiler with src/ci/docker/run.sh dist-x86_64-linux --dev for x86_64 and package the compiler with ../x dist to produce the artifacts in obj/dist-x86_64-linux/build/dist. This Dockerfile gets you a working Rust builder image which allows you to build your Rust applications in bookworm.

The state of performance

So far with this patch, I have been studying the performance impact on the cases of tokio's single- and multi-threaded runtime, as well as a simple axum HTTP service. As far as I can see, I can find a change in performance characteristics that are statistically significant, one-sided p = 0.05.

This time, I would like to call for pooling in your valuable assessments and thoughts on this patch. I kindly request experiments from you and hopefully you can provide regression cases with perf record -e cycles:u,instructions:u,cache-misses:u reports.

Thank you all so much! 🙇

[rustbot]

r? @BoxyUwU

rustbot has assigned @BoxyUwU.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

[rustbot]

Some changes occurred in compiler/rustc_codegen_cranelift

cc @bjorn3

Some changes occurred to the CTFE / Miri interpreter

cc @rust-lang/miri

Some changes occurred to MIR optimizations

cc @rust-lang/wg-mir-opt

Some changes occurred to the CTFE machinery

cc @rust-lang/wg-const-eval

[bors]

☔ The latest upstream changes (presumably #135715) made this pull request unmergeable. Please resolve the merge conflicts.

[BoxyUwU]

I don't think this needs a reviewer?

[traviscross]

cc @Darksonn @tmandry @eholk @rust-lang/wg-async

Ding here is reworking the layout of coroutines to try to reduce their memory footprint (and that of Futures). He's curious to find whether this introduces any performance or other regressions. In this own testing, he's not been able to find any, but he's curious in more data and experience here to help inform whether this is a worthwhile change.

What do people think?

[tmandry]

For anyone searching for a description of what this PR changes, it's summarized at the top of compiler/rustc_mir_transform/src/coroutine/relocate_upvars.rs.

[tmandry]

Don't we decide the offsets of upvars in Unresumed in the same place as we decide the offset of saved locals? Couldn't we then "backpropagate" the field offsets for each upvar's local as the offset for the corresponding upvar?

[dingxiangfei2009]

Thank you for reviewing! I had a backlog of things due to sickness.

True indeed. This statement is completely voided by the work in the second commit. I will reword this section in the following way.

---

By enabling the feature gate coroutine_new_layout the field offsets of the upvars in Unresumed state are further exactly placed in the same place as their corresponding saved locals, which is guaranteed by the alternative coroutine layout calculator that enters in effect. <... quote the relevant comment/file/etc. ...>

[tmandry]

I don't personally have any means of performance testing this at the moment. It would be much easier if it landed behind a feature gate.

[bors]

☔ The latest upstream changes (presumably #135318) made this pull request unmergeable. Please resolve the merge conflicts.

[nikomatsakis]

Cc @arielb1 who was also investigated this

…

On Wed, Jan 29, 2025, at 7:56 PM, Tyler Mandry wrote: ***@***.**** commented on this pull request. In compiler/rustc_mir_transform/src/coroutine/relocate_upvars.rs <#135527 (comment)>: > +//! The reason is that it is possible that coroutine layout may change and the source memory location of +//! an upvar may not necessarily be mapped exactly to the same place as in the `Unresumed` state. Don't we decide the offsets of upvars in `Unresumed` in the same place as we decide the offset of saved locals? Couldn't we then "backpropagate" the field offsets for each upvar's local as the offset for the corresponding upvar? — Reply to this email directly, view it on GitHub <#135527 (review)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABF4ZTFDPQDUNGH5L6MGSL2NF2CHAVCNFSM6AAAAABVG4UUZ2VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDKOBSGY4TKMZUHA>. You are receiving this because you are on a team that was mentioned.Message ID: ***@***.***>

[dingxiangfei2009]

@tmandry

I don't personally have any means of performance testing this at the moment. It would be much easier if it landed behind a feature gate.

I think it is fair to land with a feature gate so that we can get to play with it. The PR has temporarily disabled the check on the feature gate. However, given that coroutine layout data is keyed individually by their DefId, I think it is still safe to allow code to link to each other even when the feature gate status varies among the crates.

[eholk]

I don't personally have any means of performance testing this at the moment. It would be much easier if it landed behind a feature gate.

Would this be better as a #[feature(...)] gate, or as -Z new_coroutine_layout? I think the compiler flag feels like a better fit for something like this.

[oli-obk]

Are there any issues if only one crate activates it but others do not? if there are no issues, a feature gate seems ok (and easier to use ^^)

[bors]

☔ The latest upstream changes (presumably #137030) made this pull request unmergeable. Please resolve the merge conflicts.

[Dirbaio]

A feature doesn't allow turning it on for the whole build, you'd have to fork every single crate that uses async. A -Z flag would be better IMO.

[tmandry]

Agreed on a -Z flag being better for testing for the reason @Dirbaio gave.

If my understanding is correct, we shouldn't expect any regression from this approach (only upside), but since we currently rely on later passes eliding copies there might be some regression. We could be more aggressive in eliding the copies ourselves, but maybe this is hard.

[dingxiangfei2009]

Thanks for looking into this!

I will have time this week to clean this up a bit and I will ask rustbot to set it to ready-for-review.

[rustbot]

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

[dingxiangfei2009]

Replying to @ywxt

Now layout.fields no longer matches upvar_tys.

I posted a patch to fix it up. It is still not perfect because it does not report prefix alignment value.

And good catch on the mysterious padding! @ywxt posted an example where prefix alignment can lead to unnecessary padding.

To quote:

async fn bar1(_: [u32; 3] ) {}
async fn bar2(_: [u8; 2])   {}

let a = [1, 2];
let b = [1, 2, 3,];
let c = async move {
    bar2(a).await;
    bar1(b).await;

};

With or without -Zpack-coroutine-layout, there is unnecessary padding of 4 bytes so that the non-prefix part of the coroutine state has to be padded to the alignment of 4, due to b whose alignment is 4. For pack-coroutine-layout=captures-only case, we have the following layout, UNRESUMED only.

field	kind	type	offset from absolute start	size	align
b	Upvar	[u32; 3]	0	12	4
discriminant	Tag	i8	12	1	1
-- prefix --			0	16, but could have been just 13	4
#variant0.a	Upvar	[u8; 2]	16, because we need to align to 4	2	1
#variant0.b	Upvar	[u32; 3]	12	12	4

Now the fix is apparently set the prefix alignment to Align::ONE when using StructKind::Prefixed, while the overall coroutine should still be aligned to at least 4 surely. We neither need our coroutine to work like a C-style tagged union, nor there would ever be the case.

For this matter, is there anyone who can approve or reject my suggestion?

[ywxt]

Replying to @ywxt

Now layout.fields no longer matches upvar_tys.

I posted a patch to fix it up. It is still not perfect because it does not report prefix alignment value.

And good catch on the mysterious padding! @ywxt posted an example where prefix alignment can lead to unnecessary padding.

To quote:

async fn bar1(_: [u32; 3] ) {}
async fn bar2(_: [u8; 2])   {}

let a = [1, 2];
let b = [1, 2, 3,];
let c = async move {
    bar2(a).await;
    bar1(b).await;

};

With or without -Zpack-coroutine-layout, there is unnecessary padding of 4 bytes so that the non-prefix part of the coroutine state has to be padded to the alignment of 4, due to b whose alignment is 4. For pack-coroutine-layout=captures-only case, we have the following layout, UNRESUMED only.

field kind type offset from absolute start size align
b Upvar [u32; 3] 0 12 4
discriminant Tag i8 12 1 1
-- prefix -- 0 16, but could have been just 13 4
#variant0.a Upvar [u8; 2] 16, because we need to align to 4 2 1
#variant0.b Upvar [u32; 3] 12 12 4
Now the fix is apparently set the prefix alignment to Align::ONE when using StructKind::Prefixed, while the overall coroutine should still be aligned to at least 4 surely. We neither need our coroutine to work like a C-style tagged union, nor there would ever be the case.

For this matter, is there anyone who can approve or reject my suggestion?

Is it possible to make the discriminant postposed to the last so that we can remove paddings?

[ywxt]

fn main() {
    struct Pointer{
        p: Box<i32>,
    }
    async fn inc(pointer: &mut Pointer) {
        *pointer.p += 1;
    }
    async fn task(mut pointer: Pointer) {
        inc(&mut pointer).await;
        inc(&mut pointer).await;
    }
    let pointer = Pointer { p: Box::new(0) };
    let task = task(pointer);
    println!("size of future: {}", std::mem::size_of_val(&task));
}

print-type-size type: `{async fn body of main::task()}`: 32 bytes, alignment: 8 bytes
print-type-size     discriminant: 1 bytes
print-type-size     variant `Unresumed`: 23 bytes
print-type-size         padding: 15 bytes
print-type-size         upvar `.pointer`: 8 bytes, alignment: 8 bytes
print-type-size     variant `Suspend0`: 31 bytes
print-type-size         local `.pointer`: 8 bytes, offset: 0 bytes, alignment: 8 bytes
print-type-size         padding: 8 bytes
print-type-size         local `.__awaitee`: 16 bytes, alignment: 8 bytes, type: {async fn body of main::inc()}
print-type-size     variant `Suspend1`: 31 bytes
print-type-size         local `.pointer`: 8 bytes, offset: 0 bytes, alignment: 8 bytes
print-type-size         padding: 8 bytes
print-type-size         local `.__awaitee`: 16 bytes, alignment: 8 bytes, type: {async fn body of main::inc()}
print-type-size     variant `Returned`: 0 bytes
print-type-size     variant `Panicked`: 0 bytes

I think we can eliminate paddings in Suspend0 and Suspend1 when putting the discriminant to the end.

[dingxiangfei2009]

Replying to @ywxt

I think we can eliminate paddings in Suspend0 and Suspend1 when putting the discriminant to the end.

Yes, I agree. To the extreme we could even slot it in some padding location in the middle, in case some types demand big alignment and leave a huge gap in between.

I would like to land the current patch as-is. This enhancement can go with a second packing scheme that actually depends on this.

[dingxiangfei2009]

Sorry @ytausky, wrong tagging.

I was replying to @ywxt

[dingxiangfei2009]

Oh can anyone explain why this is considered a trait system refactor? This PR does not affect trait solver at all. 😕

[rustbot]

This PR was rebased onto a different main commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

[rust-bors]

☔ The latest upstream changes (presumably #151842) made this pull request unmergeable. Please resolve the merge conflicts.

[ywxt]

Sorry, Anyone can explain what's the status of this pr right now?

[dingxiangfei2009]

I will be actively maintaining it, but I do not know anyone who can make the decision for go/no-go.

cc @traviscross what should we do about it? Do we still need more reviews? Who else can give us the verdict?

[dingxiangfei2009]

If there are concerns that would trigger immediate regressions in correctness, regressions in MIR semantics or regressions in performance with major async crates, we should fix it in this PR and I will be in charge.

If there are only minor concerns such as technicality, we could resolve it progressively in the later patch series. I will take the responsibility for further improvements.

This PR has unfortunately turned into a head-of-line block on a few other works at the moment. I hope the situation could be improved.

[traviscross]

Looking at the history, we had signaled lang approval for this to land under a nightly feature flag, and we had said that to later stabilize this would need lang FCP.

Do you want to land this under a feature flag, or would you prefer to propose stabilization at this point? If the latter, then please draft a summary of what's happened and changed over the course of this long thread, and please ensure that the PR description is the complete stabilization report you want the lang team to review, and then we'll nominate and discuss this for stabilization.

Have a look at the dev guide stabilization report template. The checklist items regarding test coverage and implementation maturity might be helpful in this case. Beyond everyone on lang feeling comfortable, it will of course need someone on compiler to r+ it, and these items would likely be helpful there.

[dingxiangfei2009]

@traviscross thanks for weighing in. I should be clearer. I would still prefer to land it behind the nightly compiler flag -Z just as planned, not stabilisation.

There are interested parties who would like to try out the flag on nightly first because the build system does not really allow building a custom rustc. This is the pre-requisite to collect the data to prove the efficacy of the flag, which will be one of the deliverables to the eventual stabilisation report and to cut down corner cases in general. Right now it is still too early to say anything about stabilising it but at least it will allow us to experiment with it in a wider ecosystem.

cc @ywxt @gmarcosb

[ywxt]

Thanks for your explanation! @traviscross @dingxiangfei2009

Landing on nightly first is absolutely an appropriate way to enables the current solution to be widely validated.

I was wondering if there’s a way we could help move this PR forward more smoothly, as it’s been ongoing for a while now (including the previous PRs).

[dingxiangfei2009]

@traviscross I have updated the PR front cover on how test coverage has been. One question is, should we also update the dev-guide?

Anyway, I will try...

r? compiler