Skip to content

Commit b867b60

Browse files
[mypyc] feat: reorder gens for speed in ForZip
Not all generators have an equally fast gen_condition, some are quite fast while others quite the opposite This PR orders them in what I see to be a reasonably fastest -> slowest order
1 parent fb16e93 commit b867b60

File tree

1 file changed

+37
-1
lines changed

1 file changed

+37
-1
lines changed

mypyc/irbuild/for_helpers.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1185,7 +1185,43 @@ def init(self, indexes: list[Lvalue], exprs: list[Expression]) -> None:
11851185
self.gens.append(gen)
11861186

11871187
def gen_condition(self) -> None:
1188-
for i, gen in enumerate(self.gens):
1188+
# We don't necessarily need to check the gens in order,
1189+
# we just need to know which gen ends first. Some gens
1190+
# are quicker to check than others, so we will check the
1191+
# specialized ForHelpers before we check any generic
1192+
# ForIterable
1193+
gens = self.gens
1194+
1195+
def check_type(obj: Any, typ: Type[Any]) -> bool:
1196+
# ForEnumerate gen_condition is as fast as it's underlying generator's
1197+
return isinstance(obj, typ) or isinstance(obj, ForEnumerate) and isinstance(obj.gen, typ)
1198+
1199+
# these are slowest, they invoke Python's iteration protocol
1200+
for_iterable = [g for g in gens if check_type(g, ForSequence)]
1201+
1202+
# These aren't the slowest but they're slow, we need to pack an RTuple and then get and item and do a comparison
1203+
for_dict = [g for g in gens if check_type(g, ForDictionaryCommon)]
1204+
1205+
# These are faster than ForIterable but not as fast as others (faster than ForDict?)
1206+
for_native = [g for g in gens if check_type(g, ForNativeGenerator)]
1207+
1208+
# forward involves in the best case one pyssize_t comparison, else one length check + the comparison
1209+
# reverse is slightly slower than forward, with one extra check
1210+
for_sequence_reverse = [g for g in gens if check_type(g, ForSequence) and (g.gen if isinstance(g, ForEnumerate) else g).reverse]
1211+
for_sequence_forward = [g for g in gens if check_type(g, ForSequence) and not (g.gen if isinstance(g, ForEnumerate) else g).reverse]
1212+
1213+
# these are really fast, just a C int equality check
1214+
for_range = [g for g in gens if isinstance(g, ForRange)]
1215+
1216+
ordered = for_range + for_sequence_forward + for_sequence_reverse + for_native + for_dict
1217+
1218+
# this is a failsafe for ForHelper classes which might have been added after this commit but not added to this function's code
1219+
others = [g for g in gens if g not in ordered + for_iterable]
1220+
1221+
ordered += others
1222+
ordered += for_iterable
1223+
1224+
for i, gen in enumerate(ordered):
11891225
gen.gen_condition()
11901226
if i < len(self.gens) - 1:
11911227
self.builder.activate_block(self.cond_blocks[i])

0 commit comments

Comments
 (0)