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

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)
```

Author: JukkaL

mypyc/irbuild/context.py

@@ -167,6 +167,11 @@ def __init__(self, ir: ClassIR) -> None:
         # Holds the arg passed to send
         self.send_arg_reg: Value | None = None
 
+        # Holds the PyObject ** pointer through which return value can be passed
+        # instead of raising StopIteration(ret_value) (only if not NULL). This
+        # is used for faster native-to-native calls.
+        self.stop_iter_value_reg: Value | None = None
+
         # The switch block is used to decide which instruction to go using the value held in the
         # next-label register.
         self.switch_block = BasicBlock()

mypyc/irbuild/generator.py

@@ -32,7 +32,12 @@
     Unreachable,
     Value,
 )
-from mypyc.ir.rtypes import RInstance, int32_rprimitive, object_rprimitive
+from mypyc.ir.rtypes import (
+    RInstance,
+    int32_rprimitive,
+    object_pointer_rprimitive,
+    object_rprimitive,
+)
 from mypyc.irbuild.builder import IRBuilder, calculate_arg_defaults, gen_arg_defaults
 from mypyc.irbuild.context import FuncInfo, GeneratorClass
 from mypyc.irbuild.env_class import (
@@ -256,7 +261,14 @@ def add_next_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl:
         result = builder.add(
             Call(
                 fn_decl,
-                [builder.self(), none_reg, none_reg, none_reg, none_reg],
+                [
+                    builder.self(),
+                    none_reg,
+                    none_reg,
+                    none_reg,
+                    none_reg,
+                    Integer(0, object_pointer_rprimitive),
+                ],
                 fn_info.fitem.line,
             )
         )
@@ -272,7 +284,14 @@ def add_send_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl:
         result = builder.add(
             Call(
                 fn_decl,
-                [builder.self(), none_reg, none_reg, none_reg, builder.read(arg)],
+                [
+                    builder.self(),
+                    none_reg,
+                    none_reg,
+                    none_reg,
+                    builder.read(arg),
+                    Integer(0, object_pointer_rprimitive),
+                ],
                 fn_info.fitem.line,
             )
         )
@@ -297,7 +316,14 @@ def add_throw_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl:
         result = builder.add(
             Call(
                 fn_decl,
-                [builder.self(), builder.read(typ), builder.read(val), builder.read(tb), none_reg],
+                [
+                    builder.self(),
+                    builder.read(typ),
+                    builder.read(val),
+                    builder.read(tb),
+                    none_reg,
+                    Integer(0, object_pointer_rprimitive),
+                ],
                 fn_info.fitem.line,
             )
         )
@@ -377,8 +403,15 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None:
     # TODO: Use the right type here instead of object?
     exc_arg = builder.add_local(Var("arg"), object_rprimitive, is_arg=True)
 
+    # Parameter that can used to pass a pointer which can used instead of
+    # raising StopIteration(value). If the value is NULL, this won't be used.
+    stop_iter_value_arg = builder.add_local(
+        Var("stop_iter_ptr"), object_pointer_rprimitive, is_arg=True
+    )
+
     cls.exc_regs = (exc_type, exc_val, exc_tb)
     cls.send_arg_reg = exc_arg
+    cls.stop_iter_value_reg = stop_iter_value_arg
 
     cls.self_reg = builder.read(self_target, fitem.line)
     if builder.fn_info.can_merge_generator_and_env_classes():

mypyc/irbuild/nonlocalcontrol.py

@@ -16,9 +16,11 @@
     Integer,
     Register,
     Return,
+    SetMem,
     Unreachable,
     Value,
 )
+from mypyc.ir.rtypes import object_rprimitive
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.primitives.exc_ops import restore_exc_info_op, set_stop_iteration_value
 
@@ -108,10 +110,27 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
         # StopIteration instead of using RaiseStandardError because
         # the obvious thing doesn't work if the value is a tuple
         # (???).
+
+        true, false = BasicBlock(), BasicBlock()
+        stop_iter_reg = builder.fn_info.generator_class.stop_iter_value_reg
+        assert stop_iter_reg is not None
+
+        builder.add(Branch(stop_iter_reg, true, false, Branch.IS_ERROR))
+
+        builder.activate_block(true)
+        # The default/slow path is to raise a StopIteration exception with
+        # return value.
         builder.call_c(set_stop_iteration_value, [value], NO_TRACEBACK_LINE_NO)
         builder.add(Unreachable())
         builder.builder.pop_error_handler()
 
+        builder.activate_block(false)
+        # The fast path is to store return value via caller-provided pointer
+        # instead of raising an exception. This can only be used when the
+        # caller is a native function.
+        builder.add(SetMem(object_rprimitive, stop_iter_reg, value))
+        builder.add(Return(Integer(0, object_rprimitive)))
+
 
 class CleanupNonlocalControl(NonlocalControl):
     """Abstract nonlocal control that runs some cleanup code."""

mypyc/irbuild/prepare.py

@@ -56,6 +56,7 @@
     RType,
     dict_rprimitive,
     none_rprimitive,
+    object_pointer_rprimitive,
     object_rprimitive,
     tuple_rprimitive,
 )
@@ -220,6 +221,8 @@ def create_generator_class_if_needed(
                 RuntimeArg("value", object_rprimitive),
                 RuntimeArg("traceback", object_rprimitive),
                 RuntimeArg("arg", object_rprimitive),
+                # If non-NULL, used to store return value instead of raising StopIteration(retv)
+                RuntimeArg("stop_iter_ptr", object_pointer_rprimitive),
             ),
             object_rprimitive,
         )

mypyc/irbuild/statement.py

@@ -103,6 +103,7 @@
     get_exc_info_op,
     get_exc_value_op,
     keep_propagating_op,
+    propagate_if_error_op,
     raise_exception_op,
     reraise_exception_op,
     restore_exc_info_op,
@@ -958,21 +959,34 @@ def emit_yield_from_or_await(
 
     if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method):
         # Second fast path optimization: call helper directly (see also comment above).
+        #
+        # Calling a generated generator, so avoid raising StopIteration by passing
+        # an extra PyObject ** argument to helper where the stop iteration value is stored.
+        fast_path = True
         obj = builder.read(iter_reg)
         nn = builder.none_object()
-        m = MethodCall(obj, helper_method, [nn, nn, nn, nn], line)
+        stop_iter_val = Register(object_rprimitive)
+        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
+        builder.assign(stop_iter_val, err, line)
+        ptr = builder.add(LoadAddress(object_pointer_rprimitive, stop_iter_val))
+        m = MethodCall(obj, helper_method, [nn, nn, nn, nn, ptr], line)
         # Generators have custom error handling, so disable normal error handling.
         m.error_kind = ERR_NEVER
         _y_init = builder.add(m)
     else:
+        fast_path = False
         _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], line)
 
     builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
 
-    # Try extracting a return value from a StopIteration and return it.
-    # If it wasn't, this reraises the exception.
     builder.activate_block(stop_block)
-    builder.assign(result, builder.call_c(check_stop_op, [], line), line)
+    if fast_path:
+        builder.primitive_op(propagate_if_error_op, [stop_iter_val], line)
+        builder.assign(result, stop_iter_val, line)
+    else:
+        # Try extracting a return value from a StopIteration and return it.
+        # If it wasn't, this reraises the exception.
+        builder.assign(result, builder.call_c(check_stop_op, [], line), line)
     # Clear the spilled iterator/coroutine so that it will be freed.
     # Otherwise, the freeing of the spilled register would likely be delayed.
     err = builder.add(LoadErrorValue(iter_reg.type))

mypyc/lower/misc_ops.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from mypyc.ir.ops import GetElementPtr, LoadMem, Value
-from mypyc.ir.rtypes import PyVarObject, c_pyssize_t_rprimitive
+from mypyc.ir.ops import ComparisonOp, GetElementPtr, Integer, LoadMem, Value
+from mypyc.ir.rtypes import PyVarObject, c_pyssize_t_rprimitive, object_rprimitive
 from mypyc.irbuild.ll_builder import LowLevelIRBuilder
 from mypyc.lower.registry import lower_primitive_op
 
@@ -10,3 +10,9 @@
 def var_object_size(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
     elem_address = builder.add(GetElementPtr(args[0], PyVarObject, "ob_size"))
     return builder.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
+
+
+@lower_primitive_op("propagate_if_error")
+def propagate_if_error_op(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
+    # Return False on NULL. The primitive uses ERR_FALSE, so this is an error.
+    return builder.add(ComparisonOp(args[0], Integer(0, object_rprimitive), ComparisonOp.NEQ))

mypyc/primitives/exc_ops.py

@@ -4,7 +4,7 @@
 
 from mypyc.ir.ops import ERR_ALWAYS, ERR_FALSE, ERR_NEVER
 from mypyc.ir.rtypes import bit_rprimitive, exc_rtuple, object_rprimitive, void_rtype
-from mypyc.primitives.registry import custom_op
+from mypyc.primitives.registry import custom_op, custom_primitive_op
 
 # If the argument is a class, raise an instance of the class. Otherwise, assume
 # that the argument is an exception object, and raise it.
@@ -62,6 +62,16 @@
     error_kind=ERR_FALSE,
 )
 
+# If argument is NULL, propagate currently raised exception (in this case
+# an exception must have been raised). If this can be used, it's faster
+# than using PyErr_Occurred().
+propagate_if_error_op = custom_primitive_op(
+    "propagate_if_error",
+    arg_types=[object_rprimitive],
+    return_type=bit_rprimitive,
+    error_kind=ERR_FALSE,
+)
+
 # Catches a propagating exception and makes it the "currently
 # handled exception" (by sticking it into sys.exc_info()). Returns the
 # exception that was previously being handled, which must be restored