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. 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.

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

yanked

# Design decisions

Why does this patch not perform relocation as part of the StateTransform pass?

This idea is explored in #120168 already back in 2023. The conclusion then was that it does not interact well with MIR dataflow analysis. It requires StateTransform pass to assign a virtual "MIR local" to each upvars at the beginning. Apparently this created difficulty in reviewing the piece as soon as we overload this huge StateTransform pass with this additional renumbering work. The idea has always been that it is better to perform the renumbering in its own pass, to keep StateTransform simple.

This patch has gone further to carry out the re-write as early as possible, so that the passes in between can perform rewrites as per current MIR local semantics and optimisation rules.

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]

@rustbot ready

• I have reverted most of the irrelevant hunks and I will probably post them in other PRs when I get to it.
• I realised that most of the special handling of UNRESUMED variant and upvars are superfluous, so I have cleaned them up.
• I fixed the missing case for TerminatorKind::UnwindResume. StorageDeads are issued before unwind-resuming.

[tmandry]

opt nit: Can be combined into the let chain

[rust-bors]

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

[tmandry]

I think the example would be much clearer if you show two locals and show moving both of them out of the struct first, then moving both to their "destination local" second. And then talk about how the destination locals will be rewritten to a struct field in a later pass.

[dingxiangfei2009]

Thank you for clarification! I have expanded the example and explained how the layout would look like for this particular sample coroutine.

[ywxt]

Thanks for your work.

Is there a plan to fix -Zprint-type-sizes options?

rust/compiler/rustc_ty_utils/src/layout.rs

Line 930 in b765963

let offset = layout.fields.offset(field_idx);

Now layout.fields no longer matches upvar_tys.

[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?

[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.

[gmarcosb]

Wonder if it makes sense to have this test be stricter, i.e. maybe not exact number but if size goes from 21 to 550, this test would still pass, so maybe constraining to < 50 or something?

Or maybe an exact match of 21 is actually a good thing, to catch regressions?