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

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