Move coroutine upvars into locals for better memory economy

[dingxiangfei2009]

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.

The changes are broken into commits for reviews in isolation. Among them, the changes are as follows.

1. Changes to the UI and diagnostic messages, upfront for results
2. 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 and determine which are the necessary ones to be stored in the coroutine state during suspension.
3. 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.
4. 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.
5. We then have to pay the lip service to translate direct field access to the upvars into access behind the Unresumed variant.
6. We have to update diagnostics so that they are more informed about captured values and they make more sense in view of this change.

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

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

[rustbot]

Some changes occurred in src/tools/clippy

cc @rust-lang/clippy

[dingxiangfei2009]

@rustbot ready

• The new coroutine layout experimental calculator is temporarily dropped from this PR for better reviewing experience. It will resume in a second PR behind a feature gate.
• The changeset is broken into commits with logical sequence, furnished with more explanation.
• The PR description is updated with further explanation.

[bors]

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