Asynchronous Programming in Mojo

The Mojo team is resuming work on asynchronous programming in Mojo. The proposal can be found modular/Mojo/proposals/async-design.md at main · modular/modular · GitHub. Please comment on the proposal in this forum.

Over the next few weeks you can expect to see the experimental example published as well as compiler changes necessary to support the example. The goal is to (1) enhance the example as more compiler support is added and (2) migrate the example from an experimental directory to a stable test directory as we add components to the standard library.

I think that among the stated goals, the following are important:

  • Low overhead, so async is usable in performance-critical code, and ideally on constrained targets.
  • Async without malloc. Async should work on targets with no heap allocator.

As for open-ended questions:

  • Type erasure. Storing different futures in one collection needs some form of existential or boxing, like Box<dyn Future> in Rust. How much does this depend on existentials landing in Mojo?

The fastest type erased state machine is the one storing it’s state as function pointer, so driving it forward translates to a single level of indirection - https://github.com/newpavlov/fsm-bench This was also confirmed by my own benchmarks. It might be possible to make the implementation more efficient here than in Rust. It would be great if dyn Coroutine didn’t add another level of indirection.

Administrative Things

First, thanks Steffi for tackling async. It’s a big task with a lot sitting on top of it.

Since this is so important, I’ll ask everyone to try to keep this thread on track since Steffi is going to have to read the collective output of everyone with opinions on this.

Secondly, I’ll ask everyone to please respect that the compiler team is going to need substantial reasons to deviate from the points in the Decisions we have committed to section. The lowest levels of this interface are unlikely to be pretty, since they will have to handle every possible use-case for Mojo. There is interest in models of asynchronous computation which are strictly better than the proposed design, but I mean strictly in an academic sense. This means that if the model currently proposed can do express all of the same concepts, even if it requires 20 linear zero-sized types, a horrible abuse of origins, then you need to show that either your new model is better in at least one of the following categories and equivalent or extremely close to it in others; Runtime performance, compilation cost, ease of use for a particular use-case (performance may not suffer for this unless it’s otherwise impossible), memory usage, composability, and portability. I’ll ask that discussions of other designs which receive push-back under those criteria move to other threads where a resolution can be hashed out and, if the result is positive, then the model can be returned back to the main thread.

Open Questions

Async on the GPU. Should async def be usable inside GPU kernels, or only on the host to coordinate device work?

I think it absolutely should. async def and friends all form an abstraction over fairly primitive programming constructs (while and switch for one common mapping) which currently work on almost all of the MAX’s targets. Over time we may discover some targets are too limited to allow for many forms of async constructs, similar to those targets which do not allow for dynamic control flow the language feature will not be available there.

Type erasure. Storing different futures in one collection needs some form of existential or boxing, like Box<dyn Future> in Rust. How much does this depend on existentials landing in Mojo?

I think that, if we phrase existentials as the “make me a vtable please” feature, then I’d say we’d need some form of existentials for true dynamism. However, in many cases there is a compile-time known set of top-level coroutines such that a tagged union may be sufficient to de-virtualize everything. I think this means we can land “MVP” async without existentials, although its usage will be a bit un-ergonomic until existentials arrive, potentially depending on large hand-written TypeList values.

One other thing I’d like to poke at is the idea of something like a inline dyn Coroutine + size_bytes_leq[64]. Effectively, have the vtable be inline instead of behind another layer of indirection (yes, more binary size and memory consumption which isn’t great but it should be fine), and potentially the coroutine body as well, since this provides the guarantee that the actual data which matters is no larger than some amount, meaning that there is a safe way to have inline data without knowing the type. The “struct size” trait should be pretty cheap for the compiler to compute, given that it would be the amount of space needed for the vtable (potentially just a pointer) and space for the data, which is just size_of.

Function coloring. Given stackless coroutines, what can we do to reduce the cost of coloring? For example, a blocking wait() bridge, or APIs that are generic over sync and async.

My vote, although it might be a bit un-ergonomic in some cases, would be the following:

First, create def, sync def and async def. def means “I don’t care”, meaning that the function trait def() -> None is the union of async def() -> None and sync def -> None. Most functions in both user and library code are expected to be def functions, since most code doesn’t need to do fancy things with async and doesn’t really care about whether or not it makes a blocking syscall. sync def is for functions which are explicitly not async, and you may not call them in an async context. Similarly, async def is explicitly async and must be called in an async context.

From there, we have two options. In a def, either any function call which returns an awaitable has it automatically awaited (+ergonomics, -explicitness, potentially +confusion), or def must be written as “async correct” code where function calls to def or async def functions occur on await (potentially a lot of compiler transformation here). This second method is less ergonomic but is likely to cause less surprises.

Finally, at the lowest levels of code, you’ll run into why async def and sync def exist. These are for functions which are fundamentally one or the other, either due to making blocking syscalls themselves, or due to doing things like awaiting multiple things at once to increase productive CPU time (see Rust’s FuturesUnordered). Right before them, you have a bridge, of some form like the following:

def simple_print(s: String):
  comptime if async:
      await async_print(s)
  else:
      sync_print(s)

This async can be implemented as an “invisible parameter” passed down the call stack. It starts out sync, and then async executors can provide a primitive to “set” the “flag”. Additionally, awaiting a def function would make it definitely async since that means you’re in an async context. This means that, while we do technically end up with whole-program monomorphization to support this, it should be relatively cheap to support since it’s a boolean flag and I expect most programs mostly be written in one “dialect” (sync or async). This enables libraries to mostly be “synchness”-agnostic with a low-level IO layer handling the rest. This has the additional benefit of allowing the standard library to make std.atomic.Mutex or wherever we put it to behave “correctly” given the context, getting rid of a problem that Rust has where executors need to provide unique async synchronization primitives even for simple operations (although I expect executors will need to provide hooks in order to make this performant).

Use-cases

DPUs

DPUs, although they are nominally a fancy NIC, typically come with normal CPU cores running Linux. They also exist pretty much entirely for doing IO. As such, in my opinion they represent an easy target for “MAX should probably let me shove async code here so I can do data loading.”

Accelerator Driven IO

NVMe is a fundamentally asynchronous protocol, and given that most disks are moving towards using NVMe (we even have NVMe hard drives), it’s unlikely to go anywhere for a bit. Even if it does, I don’t expect that we as an industry going to suddenly decide that sync IO is a good idea again.

Given that production-grade KV caching now involves tiering to local and network storage, giving the Accelerator the ability to express the async protocol that is NVMe directly is very useful, because while we do have a lot of warps, there’s no reason to stall an entire warp per block requested from disk. If we’re already going to do async IO, then I think we should probably plumb async/await up to it.

Networking is also fundamentally asynchronous on pretty much every network in common use today. Having interactions with your scale-out network be able to be represented the same way on both the host and accelerator is quite nice.

DMA Devices

Another place this is useful is in expressing async transfers inside of accelerators. Once off of CPUs, it becomes fairly common to have some of “data mover” or DMA device you can ask to move things around for you, and on some platforms it’s mandatory if you want to access the accelerator’s main memory, such as AMD XDNA and Tenstorrent Tensix. In the latter case, where you may only have a few KiB of local memory, I don’t think anything other than precise control over the memory locations of coroutines is going to work.

General use in embedded devices

Async is a natural abstraction for cooperative multi-tasking, which is great for embedded devices which may not have the processing power to run a full threading abstraction or which may not have a large number of spare cores to support many threads.

No-allocator on larger servers

In networking, as traffic volume goes up, tighter and tighter coordination with the NIC is desirable for performance reasons. As a result, once one starts to push the server hard enough, manual control over memory becomes desirable. No-allocator generally means you need to provide the memory for everything, which, while typically a feature for embedded, is also very nice when working on a large server where you want to have a few hundred million coroutines floating around without getting smacked by time complexity.

Hopefully this isn’t too premature, as I’m not qualified to speak on some of these lower-level design questions. But I’m wondering if structured concurrency is planned for the stdlib implementation? I’d guess that might present challenges for Python interop. (The tricky thing is that structured code doesn’t mix well with non-structured code, as you lose structured concurrency guarantees when it’s possible to spawn standalone coroutines outside of a task group. So it’s best if it can be enforced. anyio has been really successful at bringing structured concurrency to Python, but many popular libraries still use asyncio.create_task all over the place.) I know Swift takes a structured approach, is the plan to replicate that here?

Well I would expect there to be a Thread.spawn relatively soon, since async would also use that with a multi threaded executor, although I suppose a single threaded executor would be the default.

I have a few question:

  1. Will the mojo LLDB fork support async and generator, so you can track variables like in sync code?
  2. How about multithreaded Vs single threaded executors? And how is the default executor going to look like?
  3. async def is going to sugar for block_on right?
block_on(_run)
  1. And multithreaded would look something like this?
block_on(_run, multi_threaded=True)
  1. How will the executor handle blocking code? E.g.
async def run():
    thread.sleep(secs=1)

Or

for x, y in very_large_list():
    compute(x, y)
  1. Will it be possible to detect such cases, e.g. with debugging tool like Tokio has?

Thanks for the feedback Owen! I will respond to each section independently, starting with Async on GPU. Firstly, the case where a host coroutine that awaits device works (though is unstable) today. We plan to add an example in the experimental directory to stress that path next week.

The more interesting case is executing the coroutine on the device and suspending inside a kernel, which is further out. It needs a pluggable frame allocator (today the frame is a host malloc), an address space on the frame pointer so the frame can live in shared or local memory, a representation for the resume handle other than a host function pointer, and an answer for what suspension means when lanes in a warp diverge.

We do not have a date for this support but it is shaping our design decisions. Ongoing frame construction refactoring is making frame analysis and layout pluggable so a device specific version can be slotted in rather than forking the pass.

One other very interesting option is that AMD documents their command processor, so you could potentially run an async executor there since it has a natural “spawn subtask” primitive. Hardware acceleration of these is potentially something to look at. Intel DLB is also something to consider, although it looks like it might be telco skus only for next gen Intel.

Everything else sounds good!

Being able to execute coroutines on device would be extremely interesting for latency-limited HPC applications, where this could enable overlapping computation with NVSHMEM-like device-initiated network communication without the code turning into an incomprehensible mess.

I ran across this paper from about a week ago and thought I’d share in case the Modular team wasn’t already aware of it. Looks like a few researchers did a kind of taxonomy of async/await systems in different languages. Could be a useful reference as you’re trying to design your own system.

Hi Steffi, thanks for opening this proposal. I just read the proposal again, and while I agree with the goals and the technical design of coroutines, it occurs to me that it never really goes into any detail regarding how the async/coroutine model composes with the rest of Mojo.

For example, it doesn’t discuss data flow between coroutine chains, nor does it describe how control is handed back to the Task abstraction. In Python, this is done by having yield somewhere in the coroutine chain, but since Mojo doesn’t have generators, I’m interested in hearing whether control models have been discussed.

Generally, I’m a Python enjoyer and think Mojo should use the same async control flow model since I believe it is really well thought out and dovetails cleanly with other Python features (which Mojo has or is going to have) - iterators, generators, context managers. However, it is not clear to me if this works in a compiled model, so I would love to hear your thoughts.