[mypyc] Call generator helper method directly in await expression #19376

JukkaL · 2025-07-04T13:52:48Z

Previously calls like await foo() were compiled to code that included code like this (in Python-like pseudocode):

a = foo()
...
b = get_coro(a)
...
c = next(b)

In the above code, get_coro(a) just returns a if foo is a native async function, so we now optimize this call away. Also next(b) calls b.__next__(), which calls the generated generator helper method __mypyc_generator_helper__. Now we call the helper method directly, which saves some unnecessary calls.

More importantly, in a follow-up PR I can easily change the way __mypyc_generator_helper__ is called, since we now call it directly. This makes it possible to avoid raising a StopIteration exception in many await expressions. The goal of this PR is to prepare for the latter optimization. This PR doesn't help performance significantly by itself.

In order to call the helper method directly, I had to generate the declaration of this method and the generated generator class before the main irbuild pass, since otherwise a call site could be processed before we have processed the called generator.

I also improved test coverage of related functionality. We don't have an IR test for async calls, since the IR is very verbose. I manually inspected the generated IR to verify that the new code path works both when calling a top-level function and when calling a method. I'll later add a mypyc benchmark to ensure that we will notice if the performance of async calls is degraded.

Also use generator class as the return type of generators.

…rn type

JukkaL · 2025-07-04T13:54:26Z

mypyc/irbuild/mapper.py

+                # Give a more precise type for generators, so that we can optimize
+                # code that uses them. They return a generator object, which has a
+                # specific class. Without this, the type would have to be 'object'.
+                ret: RType = RInstance(self.fdef_to_generator[fdef])


This is the first part of the optimization.

JukkaL · 2025-07-04T13:55:18Z

mypyc/irbuild/statement.py

+        and val.type.class_ir.has_method(helper_method)
+    ):
+        # This is a generated generator class, and we can use a fast path.
+        iter_val: Value = val


The changes in this function implemented the bulk of the optimization.

Maybe add a bit more detailed comment (essentially summarizing the PR description) so we will not forget the motivation.

ilevkivskyi

LG, just one suggestion.

ilevkivskyi · 2025-07-07T12:50:02Z

mypyc/irbuild/statement.py

+        and val.type.class_ir.has_method(helper_method)
+    ):
+        # This is a generated generator class, and we can use a fast path.
+        iter_val: Value = val


Maybe add a bit more detailed comment (essentially summarizing the PR description) so we will not forget the motivation.

When calling a native async function using `await`, e.g. `await foo()`, avoid raising `StopIteration` to pass the return value, since this is expensive. Instead, pass an extra `PyObject **` argument to the generator helper method and use that to return the return value. This is mostly helpful when there are many calls using await that don't block (e.g. there is a fast path that is usually taken that doesn't block). When awaiting from non-compiled code, the slow path is still taken. This builds on top of #19376. This PR makes this microbenchmark about 3x faster, which is about the ideal scenario for this optimization: ``` import asyncio from time import time async def inc(x: int) -> int: return x + 1 async def bench(n: int) -> int: x = 0 for i in range(n): x = await inc(x) return x asyncio.run(bench(1000)) t0 = time() asyncio.run(bench(1000 * 1000 * 200)) print(time() - t0) ```

Chainfire · 2025-09-16T09:42:34Z

from typing import Type

import asyncio
import traceback


class Parent:
    async def test(self) -> int:
        raise NotImplementedError()

    async def run(self) -> int:
        return await self.test()


class ChildTypeError1(Parent):
    async def test(self) -> int:
        return 1


class ChildTypeError2(Parent):
    async def run(self) -> int:
        return await self.test()


def test(cls: Type[Parent]) -> None:
    print(str(cls.__name__))
    try:
        print("- ", asyncio.run(cls().run()))
    except NotImplementedError:
        pass
    except Exception:
        traceback.print_exc()
        exit(1)


test(ChildTypeError1)
test(ChildTypeError2)
exit(0)

Results in:

ada0d2a3b
running build_ext
building 'test' extension
creating build/temp.linux-x86_64-cpython-312/build
x86_64-linux-gnu-gcc -fno-strict-overflow -Wsign-compare -DNDEBUG -g -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -g -fwrapv -O2 -fPIC -I/home/jorrit/python/omnes/mypy/mypyc/lib-rt -I/home/jorrit/python/omnes/omnes-router-services/.venv/include -I/usr/include/python3.12 -c build/__native.c -o build/temp.linux-x86_64-cpython-312/build/__native.o -O3 -g1 -Werror -Wno-unused-function -Wno-unused-label -Wno-unreachable-code -Wno-unused-variable -Wno-unused-command-line-argument -Wno-unknown-warning-option -Wno-unused-but-set-variable -Wno-ignored-optimization-argument -Wno-cpp
creating build/lib.linux-x86_64-cpython-312
x86_64-linux-gnu-gcc -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -g -fwrapv -O2 build/temp.linux-x86_64-cpython-312/build/__native.o -L/usr/lib/x86_64-linux-gnu -o build/lib.linux-x86_64-cpython-312/test.cpython-312-x86_64-linux-gnu.so
copying build/lib.linux-x86_64-cpython-312/test.cpython-312-x86_64-linux-gnu.so ->
ChildTypeError1
Traceback (most recent call last):
  File "test.py", line 28, in test
    print("- ", asyncio.run(cls().run()))
  File "/usr/lib/python3.12/asyncio/runners.py", line 195, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "test.py", line 12, in run
    return await self.test()
TypeError: test.test_Parent_gen object expected; got test.test_ChildTypeError1_gen

Ping #19764

CoolCat467 · 2025-09-16T13:57:37Z

Error message and highlevel reproduction looks awfully similar to #19558, I expect they might both be stemming from the same issue.

Chainfire · 2025-09-16T15:09:53Z

Error message and highlevel reproduction looks awfully similar to #19558, I expect they might both be stemming from the same issue.

They do not. Your issue exists prior to this merge - the code from your opening comment errors out if I go further back in history than this commit.

BobTheBuidler · 2025-09-16T17:37:15Z

At a quick glance it looks like if Parent is not @Final and Parent.test is not @Final, it will need to use object_rprimitive

edit: oh crap, sorry guy named @final I didn't mean to tag you

JukkaL · 2025-09-17T11:12:24Z

I created an issue about the regression: mypyc/mypyc#1141

Let's move further discussion there.

JukkaL added 21 commits July 4, 2025 13:10

WIP populate generator classes before main irbuild phase

d92c437

Also use generator class as the return type of generators.

Omit CPy_GetCoro call when calling native coro

4c44b4d

WIP better

ab44852

WIP attempt to fix naming

74db611

Fix segfault

9959e24

Simplify

baca816

Call __next__ directly if possible in await

e1282a7

Call helper method directly

5e6b575

Generator class is final (fix after rebase)

e3d3578

Fix helper method return type -- it should be object

7bdc1f5

Don't apply to decorated functions, since we don't know concrete retu…

a925378

…rn type

Fix nested generators

3328d18

Add test case

f293150

Support method calls

3048d03

Add comments and docstrings

2c184dc

Fix generator class detection

bcf9bbc

Lint

ce74b2f

Test calling async method

0ab4b19

Add comment

11a1e75

Add some type annotations

fff29df

Fix self check

c33a683

JukkaL commented Jul 4, 2025

View reviewed changes

ilevkivskyi approved these changes Jul 7, 2025

View reviewed changes

Add comments based on review feedback

1fa32b2

JukkaL merged commit ada0d2a into master Jul 7, 2025
13 checks passed

JukkaL deleted the mypyc-await-optimize-1 branch July 7, 2025 14:43

JukkaL mentioned this pull request Jul 8, 2025

[mypyc] Speed up native-to-native calls using await #19398

Merged

Chainfire mentioned this pull request Sep 16, 2025

1.18 Release Planning #19764

Closed

JukkaL mentioned this pull request Sep 17, 2025

TypeError with async methods and inheritance mypyc/mypyc#1141

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[mypyc] Call generator helper method directly in await expression #19376

[mypyc] Call generator helper method directly in await expression #19376

Uh oh!

JukkaL commented Jul 4, 2025 •

edited

Loading

Uh oh!

JukkaL Jul 4, 2025

Uh oh!

JukkaL Jul 4, 2025

Uh oh!

ilevkivskyi Jul 7, 2025

Uh oh!

ilevkivskyi left a comment

Uh oh!

ilevkivskyi Jul 7, 2025

Uh oh!

Uh oh!

Chainfire commented Sep 16, 2025

Uh oh!

CoolCat467 commented Sep 16, 2025

Uh oh!

Chainfire commented Sep 16, 2025 •

edited

Loading

Uh oh!

BobTheBuidler commented Sep 16, 2025 •

edited

Loading

Uh oh!

JukkaL commented Sep 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

[mypyc] Call generator helper method directly in await expression #19376

[mypyc] Call generator helper method directly in await expression #19376

Uh oh!

Conversation

JukkaL commented Jul 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JukkaL Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

JukkaL Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

ilevkivskyi Jul 7, 2025

Choose a reason for hiding this comment

Uh oh!

ilevkivskyi left a comment

Choose a reason for hiding this comment

Uh oh!

ilevkivskyi Jul 7, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Chainfire commented Sep 16, 2025

Uh oh!

CoolCat467 commented Sep 16, 2025

Uh oh!

Chainfire commented Sep 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

BobTheBuidler commented Sep 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JukkaL commented Sep 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

JukkaL commented Jul 4, 2025 •

edited

Loading

Chainfire commented Sep 16, 2025 •

edited

Loading

BobTheBuidler commented Sep 16, 2025 •

edited

Loading