Async Rust never left the MVP state

Async Rust never left the MVP state

I've previously explained async bloat and some work-arounds for it, but would much prefer to solve the issue at the root, in the compiler. I've submitted a Project Goal, and am looking for help to fund the effort.

I love me some async Rust! It's amazing how we can write executor agnostic code that can run concurrently on huge servers and tiny microcontrollers.

But especially on those tiny microcontrollers we notice that async Rust is far from the zero cost abstractions we were promised. That's because every byte of binary size counts and async introduces a lot of bloat. This bloat exists on desktops and servers as well, but it's much less noticable when you have substantially more memory and compute available.

I've previously explained some work-arounds for this issue, but would much prefer to get to the root of the problem, and work on improving async bloat in the compiler. As such I have submitted a Project Goal.

This is part 2 of my blog series on this topic. In this second part we'll dive into the internals and translate the methods of blog 1 into optimizations for the compiler.

What I won't be talking about is the often discussed problem of futures becoming bigger than necessary and them doing a lot of copying. People are aware of that already. In fact, there is an open PR that tackles part of it.

Anatomy of a generated future

We're going to be looking at this code:

fn foo() -> impl Future<Output = i32> {
    async { 5 }
}
fn bar() -> impl Future<Output = i32> {
    async {
        foo().await + foo().await
    }
}

We're using the desugared syntax for futures because it's easier to see what's happening. So what does the bar future look like?

There are two await points, so the state machine must have at least two states, right? Well, yes. But there's more.

Luckily we can ask the compiler to dump MIR for us at the coroutine_resume pass. This is the last async-specific MIR pass. Why is this important? Well, async is a language feature that still exists in MIR, but not in LLVM IR. So the transformation of async to state machine happens as a MIR pass.

The bar function generates 360 lines of MIR. Pretty crazy, right? Although this gets optimized somewhat later on, the non-async version uses only 23 lines for this.

The compiler also outputs the CoroutineLayout. It's basically an enum with these states:

variant_fields: {
    Unresumed(0): [], // Starting state
    Returned (1): [],
    Panicked (2): [],
    Suspend0 (3): [_s1], // At await point 1, _s1 = the foo future
    Suspend1 (4): [_s0, _s2], // At await point 2, _s0 = result of _s1, s2 = the second foo future
},

Why panic?

But is it reasonable? Futures in the Returned state will panic. But they don't have to. The only thing we can't do is cause UB to happen.

Panics are relatively expensive. They introduce a path with a side-effect that's not easily optimized out. What if instead, we just return Pending again? Nothing unsafe going on, so we fulfill the contract of the Future type.

I've hacked this in the compiler to try it out and saw a 2%-5% reduction in binary size for async embedded firmware. So I propose this should be a switch, just like overflow-checks = false is for integer overflow. In debug builds it would still panic, but in release builds we get smaller futures.

Always a state machine

We've looked at bar, but not yet at foo.

fn foo() -> impl Future<Output = i32> {
    async { 5 }
}

If we implement it manually, we don't need any state. We just return the number. However, the generated MIR for the version the compiler gives us still has the 3 default states and switches on them. There's a big optimization opportunity here to have no states and always return Poll::Ready(5) on every poll.

LLVM to the rescue?

LLVM will pick up the pieces, right? Well, sometimes. But only when the futures are simple enough and you're running opt-level=3. If the future grows too complex or you're optimizing for size, LLVM doesn't optimize this all away.

LLVM is not our savior here sadly. We really do need to give it good inputs. It eventually can't keep up when the code gets less trivial because we're relying on LLVM to spot that it should optimize out the things we're asking it to create.