Skip to content

Commit 6f23e47

Browse files
authored
[mypyc] Speed up for loop over native generator (#19415)
Call the generator helper method directly instead of calling `PyIter_Next` when calling a native generator from a native function. This way we can avoid raising StopIteration when the generator is exhausted. The approach is similar to what I used to speed up calls using await in #19398. Refer to that PR for a more detailed explanation. This helps mostly when a generator produces a small number of values, which is quite common. This PR improves the performance of this microbenchmark, which is a close to the ideal use case, by about 2.6x (now 5.7x faster than interpreted): ``` from typing import Iterator def foo(x: int) -> Iterator[int]: for a in range(x): yield a def bench(n: int) -> None: for i in range(n): for a in foo(1): pass from time import time bench(1000 * 1000) t0 = time() bench(50 * 1000 * 1000) print(time() - t0) ```
1 parent a79e85e commit 6f23e47

File tree

2 files changed

+76
-4
lines changed

2 files changed

+76
-4
lines changed

mypyc/irbuild/for_helpers.py

Lines changed: 74 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,19 +24,23 @@
2424
TypeAlias,
2525
)
2626
from mypyc.ir.ops import (
27+
ERR_NEVER,
2728
BasicBlock,
2829
Branch,
2930
Integer,
3031
IntOp,
3132
LoadAddress,
33+
LoadErrorValue,
3234
LoadMem,
35+
MethodCall,
3336
RaiseStandardError,
3437
Register,
3538
TupleGet,
3639
TupleSet,
3740
Value,
3841
)
3942
from mypyc.ir.rtypes import (
43+
RInstance,
4044
RTuple,
4145
RType,
4246
bool_rprimitive,
@@ -48,10 +52,13 @@
4852
is_short_int_rprimitive,
4953
is_str_rprimitive,
5054
is_tuple_rprimitive,
55+
object_pointer_rprimitive,
56+
object_rprimitive,
5157
pointer_rprimitive,
5258
short_int_rprimitive,
5359
)
5460
from mypyc.irbuild.builder import IRBuilder
61+
from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
5562
from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
5663
from mypyc.primitives.dict_ops import (
5764
dict_check_size_op,
@@ -62,7 +69,7 @@
6269
dict_next_value_op,
6370
dict_value_iter_op,
6471
)
65-
from mypyc.primitives.exc_ops import no_err_occurred_op
72+
from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op
6673
from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
6774
from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op, new_list_set_item_op
6875
from mypyc.primitives.misc_ops import stop_async_iteration_op
@@ -511,7 +518,15 @@ def make_for_loop_generator(
511518
# Default to a generic for loop.
512519
if iterable_expr_reg is None:
513520
iterable_expr_reg = builder.accept(expr)
514-
for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
521+
522+
it = iterable_expr_reg.type
523+
for_obj: ForNativeGenerator | ForIterable
524+
if isinstance(it, RInstance) and it.class_ir.has_method(GENERATOR_HELPER_NAME):
525+
# Directly call generator object methods if iterating over a native generator.
526+
for_obj = ForNativeGenerator(builder, index, body_block, loop_exit, line, nested)
527+
else:
528+
# Generic implementation that works of arbitrary iterables.
529+
for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
515530
item_type = builder._analyze_iterable_item_type(expr)
516531
item_rtype = builder.type_to_rtype(item_type)
517532
for_obj.init(iterable_expr_reg, item_rtype)
@@ -623,6 +638,63 @@ def gen_cleanup(self) -> None:
623638
self.builder.call_c(no_err_occurred_op, [], self.line)
624639

625640

641+
class ForNativeGenerator(ForGenerator):
642+
"""Generate IR for a for loop over a native generator."""
643+
644+
def need_cleanup(self) -> bool:
645+
# Create a new cleanup block for when the loop is finished.
646+
return True
647+
648+
def init(self, expr_reg: Value, target_type: RType) -> None:
649+
# Define target to contains the generator expression. It's also the iterator.
650+
# If we are inside a generator function, spill these into the environment class.
651+
builder = self.builder
652+
self.iter_target = builder.maybe_spill(expr_reg)
653+
self.target_type = target_type
654+
655+
def gen_condition(self) -> None:
656+
builder = self.builder
657+
line = self.line
658+
self.return_value = Register(object_rprimitive)
659+
err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
660+
builder.assign(self.return_value, err, line)
661+
662+
# Call generated generator helper method, passing a PyObject ** as the final
663+
# argument that will be used to store the return value in the return value
664+
# register. We ignore the return value but the presence of a return value
665+
# indicates that the generator has finished. This is faster than raising
666+
# and catching StopIteration, which is the non-native way of doing this.
667+
ptr = builder.add(LoadAddress(object_pointer_rprimitive, self.return_value))
668+
nn = builder.none_object()
669+
helper_call = MethodCall(
670+
builder.read(self.iter_target), GENERATOR_HELPER_NAME, [nn, nn, nn, nn, ptr], line
671+
)
672+
# We provide custom handling for error values.
673+
helper_call.error_kind = ERR_NEVER
674+
675+
self.next_reg = builder.add(helper_call)
676+
builder.add(Branch(self.next_reg, self.loop_exit, self.body_block, Branch.IS_ERROR))
677+
678+
def begin_body(self) -> None:
679+
# Assign the value obtained from the generator helper method to the
680+
# lvalue so that it can be referenced by code in the body of the loop.
681+
builder = self.builder
682+
line = self.line
683+
# We unbox here so that iterating with tuple unpacking generates a tuple based
684+
# unpack instead of an iterator based one.
685+
next_reg = builder.coerce(self.next_reg, self.target_type, line)
686+
builder.assign(builder.get_assignment_target(self.index), next_reg, line)
687+
688+
def gen_step(self) -> None:
689+
# Nothing to do here, since we get the next item as part of gen_condition().
690+
pass
691+
692+
def gen_cleanup(self) -> None:
693+
# If return value is NULL (it wasn't assigned to by the generator helper method),
694+
# an exception was raised that we need to propagate.
695+
self.builder.primitive_op(propagate_if_error_op, [self.return_value], self.line)
696+
697+
626698
class ForAsyncIterable(ForGenerator):
627699
"""Generate IR for an async for loop."""
628700

mypyc/irbuild/prepare.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from collections import defaultdict
1717
from collections.abc import Iterable
18-
from typing import NamedTuple
18+
from typing import Final, NamedTuple
1919

2020
from mypy.build import Graph
2121
from mypy.nodes import (
@@ -71,7 +71,7 @@
7171
from mypyc.options import CompilerOptions
7272
from mypyc.sametype import is_same_type
7373

74-
GENERATOR_HELPER_NAME = "__mypyc_generator_helper__"
74+
GENERATOR_HELPER_NAME: Final = "__mypyc_generator_helper__"
7575

7676

7777
def build_type_map(

0 commit comments

Comments
 (0)