await f(...) where f is a coroutine function currently executes CALL (which runs RETURN_GENERATOR in the callee and returns to the caller), GET_AWAITABLE, PUSH_NULL, a load of None and SEND before the coroutine body starts.
This PR adds four specializations to the CALL family, ASYNC_CALL_PY_EXACT_ARGS, ASYNC_CALL_BOUND_METHOD_EXACT_ARGS, ASYNC_CALL_PY_GENERAL and ASYNC_CALL_BOUND_METHOD_GENERAL, and two to the CALL_KW family, ASYNC_CALL_KW_PY and ASYNC_CALL_KW_BOUND_METHOD. The specializer picks them when the callee is a coroutine function and the call is awaited at once (await f(...), async with, async comprehensions). They are superinstructions: they merge the call with the first execution of the four instructions that follow it, creating the coroutine and pushing its frame directly, with the caller left as if it were executing the SEND.
Design notes:
- The specializations reuse the uops of the matching
CALL_* or CALL_KW_* specialization; the only new uop is _ASYNC_CALL_GEN_FRAME, plus a tier-1 guard that the skipped instructions are not instrumented. The specializer reuses specialize_py_call() and specialize_py_call_kw(). The frame set-up shared with RETURN_GENERATOR, SEND_GEN and FOR_ITER_GEN is factored into gen_take_frame() and gen_link_frame(), and _Py_MakeCoroWithFrame() is _Py_MakeCoro() with the frame passed explicitly.
- Every other
CALL and CALL_KW path is untouched, so calls that cannot be merged (non-coroutine callables, types.coroutine functions, CALL_FUNCTION_EX) keep their current specializations.
RETURN_GENERATOR only pushes the None that the following POP_TOP discards when its frame already belongs to a generator, so coroutine functions whose prologue contains COPY_FREE_VARS or MAKE_CELL use the fast path too. The JIT trace recorder records such a RETURN_GENERATOR as a load of None, so the trace continues into the coroutine body.- The tier-2 interpreter traces through the new uop (
_ASYNC_CALL_GEN_FRAME + _PUSH_FRAME).
Micro-benchmarks (release build, no PGO, ns per awaited call including loop overhead): await func(x) 50.5 → 44.2, await obj.method(x) 52.2 → 47.0, await closure(x) 60.0 → 55.0, a 4-deep chain that suspends 278 → 251. Keyword calls: await func(x, b=1) 69.3 → 63.0, await obj.method(x, b=1) 71.5 → 65.6. Awaits of plain functions or classes are unchanged.
Tested with the full test suite on a debug build with --enable-experimental-jit=interpreter, and the async-related tests on a free-threaded debug build. Also tested with the async-related and JIT tests on a release build with the LLVM JIT (--enable-experimental-jit), where await f(i) in a loop goes from 54 to 47 ns per iteration, an awaited closure from 64 to 55 ns, and a 4-deep suspending chain from 341 to 290 ns. Not tested: Windows, the tail-calling interpreter.
Also in this PR: a trace function can no longer jump to the def line of a generator frame, which re-executed RETURN_GENERATOR in a frame that already belonged to the generator (undefined behaviour before this change).
Still to do: file an issue and add the NEWS entry.
This PR was developed with the help of an AI tool, which drafted the implementation and tests; the design, measurements and review are mine.
🤖 Generated with Claude Code