Skip to content

[mypyc] feat: specialize any and all using for loop helpers if possible [1/2] #19948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 14 commits into
base: master
Choose a base branch
from
Open

[mypyc] feat: specialize any and all using for loop helpers if possible [1/2] #19948

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
41c7c9b
[mypyc] feat: specialize any and all using for loop helpers if possible
BobTheBuidler Oct 7, 2025
50bd3ce
Merge branch 'master' into any-all-for-helper
BobTheBuidler Oct 10, 2025
f636c95
Update specialize.py
BobTheBuidler Oct 13, 2025
1af0a05
Update specialize.py
BobTheBuidler Oct 13, 2025
df669fe
unique names for index
BobTheBuidler Oct 13, 2025
2b43f62
Update specialize.py
BobTheBuidler Oct 13, 2025
e0dc704
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 13, 2025
1028f8c
Update irbuild-basic.test
BobTheBuidler Oct 13, 2025
70fb118
Update irbuild-basic.test
BobTheBuidler Oct 13, 2025
1e22947
Update irbuild-statements.test
BobTheBuidler Oct 13, 2025
6a2c8a0
Merge branch 'master' into any-all-for-helper
BobTheBuidler Oct 13, 2025
a2ab6ba
Update irbuild-basic.test
BobTheBuidler Oct 14, 2025
f4c1c77
Merge branch 'master' into any-all-for-helper
BobTheBuidler Oct 18, 2025
ded8359
Merge branch 'master' into any-all-for-helper
BobTheBuidler Oct 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 26 additions & 1 deletion mypyc/irbuild/for_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

from mypy.nodes import (
ARG_POS,
LDEF,
BytesExpr,
CallExpr,
DictionaryComprehension,
Expand All @@ -28,7 +29,7 @@
TypeAlias,
Var,
)
from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types
from mypy.types import LiteralType, TupleType, Type, get_proper_type, get_proper_types
from mypyc.ir.ops import (
ERR_NEVER,
BasicBlock,
Expand Down Expand Up @@ -1241,3 +1242,27 @@ def get_expr_length_value(
return builder.builder.builtin_len(expr_reg, line, use_pyssize_t=use_pyssize_t)
# The expression result is known at compile time, so we can use a constant.
return Integer(length, c_pyssize_t_rprimitive if use_pyssize_t else short_int_rprimitive)


def expr_has_specialized_for_helper(builder: IRBuilder, expr: Expression) -> bool:
if is_sequence_rprimitive(builder.node_type(expr)):
return True
if not isinstance(expr, CallExpr):
return False
if isinstance(expr.callee, RefExpr):
return expr.callee.fullname in {"builtins.range", "builtins.enumerate", "builtins.zip"}
elif isinstance(expr.callee, MemberExpr):
return expr.callee.fullname in {
"builtins.dict.keys",
"builtins.dict.values",
"builtins.dict.items",
}
return False


def create_synthetic_nameexpr(index_name: str, index_type: Type) -> NameExpr:
"""This helper spoofs a NameExpr to use as the lvalue in one of the for loop helpers."""
index = NameExpr(index_name)
Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I probably need something here to ensure no name collisions globally, but not sure what

Copy link
Collaborator

@JukkaL JukkaL Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In my other comment I suggested adding _<line>_column suffix. Another thing we use is to use the temp_counter defined in IRBuilder to get an increasing suffix.

Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should create_synthetic_nameexpr be an IRBuilder method instead of going here then?

I think it could be used to refactor a few other places within the codebase where NameExpr are artificially created

Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've implemented the temp_counter but left this function where it is for now

index.kind = LDEF
index.node = Var(index_name, index_type)
return index
89 changes: 65 additions & 24 deletions mypyc/irbuild/specialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,9 @@
from mypyc.irbuild.constant_fold import constant_fold_expr
from mypyc.irbuild.for_helpers import (
comprehension_helper,
create_synthetic_nameexpr,
expr_has_specialized_for_helper,
for_loop_helper,
sequence_from_generator_preallocate_helper,
translate_list_comprehension,
translate_set_comprehension,
Expand Down Expand Up @@ -412,29 +415,70 @@ def translate_safe_generator_call(

@specialize_function("builtins.any")
def translate_any_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
if (
len(expr.args) == 1
and expr.arg_kinds == [ARG_POS]
and isinstance(expr.args[0], GeneratorExpr)
):
return any_all_helper(builder, expr.args[0], builder.false, lambda x: x, builder.true)
if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
arg = expr.args[0]
if isinstance(arg, GeneratorExpr):
return any_all_helper(builder, arg, builder.false, lambda x: x, builder.true)
elif expr_has_specialized_for_helper(builder, arg):
Copy link
Collaborator

@JukkaL JukkaL Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It took me a while to understand why this check is useful. The idea is that if we'd have a generic for loop (PyObject_GetIter etc.) we can as well call the stdlib function?

What would happen if we'd always created the for loop here, since it would likely avoid the overhead of calling any using generic call op, and the namespace dict lookup? If that sounds feasible, can you run some microbenchmarks to check if it's useful? I think it might be worth it especially if the iterable is small (0 or 1 items), which is quite common in practice.

Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Oct 13, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"It took me a while to understand why this check is useful. The idea is that if we'd have a generic for loop (PyObject_GetIter etc.) we can as well call the stdlib function?"

Uh. No. While I appreciate your assessment, it wasn't exactly so thought-out. My only intent was to fall back to the existing logic, whatever that may be, if there is no supported ForHelper.

I need to think about how to implement this

Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wait I want to confirm, you're saying if we just remove the if check, we're good to go without any other changes? I suppose that makes sense since the for_loop_helper does have the generic ForGenerator class but I wonder if we're missing any edge cases where that might cause failure

retval = Register(bool_rprimitive)
builder.assign(retval, builder.false(), -1)
loop_exit = BasicBlock()
index_name = "__mypyc_any_item__"

def body_insts() -> None:
true_block = BasicBlock()
false_block = BasicBlock()
builder.add_bool_branch(builder.read(index_reg), true_block, false_block)
builder.activate_block(true_block)
builder.assign(retval, builder.true(), -1)
builder.goto(loop_exit)
builder.activate_block(false_block)

index_type = builder._analyze_iterable_item_type(arg)
index = create_synthetic_nameexpr(index_name, index_type)
index_reg = builder.add_local_reg(index.node, builder.type_to_rtype(index_type)) # type: ignore [arg-type]

for_loop_helper(builder, index, arg, body_insts, None, is_async=False, line=expr.line)
builder.goto_and_activate(loop_exit)
return retval
return None


@specialize_function("builtins.all")
def translate_all_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
if (
len(expr.args) == 1
and expr.arg_kinds == [ARG_POS]
and isinstance(expr.args[0], GeneratorExpr)
):
return any_all_helper(
builder,
expr.args[0],
builder.true,
lambda x: builder.unary_op(x, "not", expr.line),
builder.false,
)
if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
arg = expr.args[0]
if isinstance(arg, GeneratorExpr):
return any_all_helper(
builder,
arg,
builder.true,
lambda x: builder.unary_op(x, "not", expr.line),
builder.false,
)

elif expr_has_specialized_for_helper(builder, arg):
retval = Register(bool_rprimitive)
builder.assign(retval, builder.true(), -1)
loop_exit = BasicBlock()
index_name = "__mypyc_all_item__"

def body_insts() -> None:
true_block = BasicBlock()
false_block = BasicBlock()
builder.add_bool_branch(builder.read(index_reg), true_block, false_block)
builder.activate_block(false_block)
builder.assign(retval, builder.false(), -1)
builder.goto(loop_exit)
builder.activate_block(true_block)

index_type = builder._analyze_iterable_item_type(arg)
Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should I deduplicate this block or is this fine?

Copy link
Collaborator

@JukkaL JukkaL Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might be useful to deduplicate, since the only difference seems to be true/false switch, so a helper would just need a bool flag that indicates whether it's any or all. Not a big deal though.

Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Oct 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lets just keep it as-is, the intent of the code is explicit and clear

index = create_synthetic_nameexpr(index_name, index_type)
index_reg = builder.add_local_reg(index.node, builder.type_to_rtype(index_type)) # type: ignore [arg-type]

for_loop_helper(builder, index, arg, body_insts, None, is_async=False, line=expr.line)
builder.goto_and_activate(loop_exit)
return retval
return None


Expand Down Expand Up @@ -610,12 +654,9 @@ def translate_isinstance(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
return builder.builder.isinstance_helper(obj, irs, expr.line)

if isinstance(type_expr, RefExpr):
node = type_expr.node
if node:
desc = isinstance_primitives.get(node.fullname)
if desc:
obj = builder.accept(obj_expr)
return builder.primitive_op(desc, [obj], expr.line)
if node := type_expr.node:
Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Oct 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was definitely not necessary but it bothered me how ugly that snippet was

if desc := isinstance_primitives.get(node.fullname):
return builder.primitive_op(desc, [builder.accept(obj_expr)], expr.line)

elif isinstance(type_expr, TupleExpr):
node_names: list[str] = []
Expand Down
118 changes: 118 additions & 0 deletions mypyc/test-data/irbuild-basic.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2812,6 +2812,12 @@ def call_any(l: Iterable[int]) -> bool:
def call_all(l: Iterable[int]) -> bool:
return all(i == 0 for i in l)

def call_any_helper(l: list[Iterable[int]]) -> bool:
return any([str(i) for i in l])

def call_all_helper(l: list[Iterable[int]]) -> bool:
return all([str(i) for i in l])

[out]
def call_any(l):
l :: object
Expand Down Expand Up @@ -2870,6 +2876,118 @@ L6:
L7:
L8:
return r0
def call_any_helper(l):
l :: list
r0 :: bool
r1 :: native_int
r2 :: list
r3, r4 :: native_int
r5 :: bit
r6, i :: object
r7 :: str
r8, r9, r10 :: native_int
r11 :: bit
r12 :: object
r13, __mypyc_any_item__ :: str
r14 :: bit
r15 :: native_int
L0:
r0 = 0
r1 = var_object_size l
r2 = PyList_New(r1)
r3 = 0
L1:
r4 = var_object_size l
r5 = r3 < r4 :: signed
if r5 goto L2 else goto L4 :: bool
L2:
r6 = list_get_item_unsafe l, r3
i = r6
r7 = PyObject_Str(i)
CPyList_SetItemUnsafe(r2, r3, r7)
L3:
r8 = r3 + 1
r3 = r8
goto L1
L4:
r9 = 0
L5:
r10 = var_object_size r2
r11 = r9 < r10 :: signed
if r11 goto L6 else goto L10 :: bool
L6:
r12 = list_get_item_unsafe r2, r9
r13 = cast(str, r12)
__mypyc_any_item__ = r13
r14 = CPyStr_IsTrue(__mypyc_any_item__)
if r14 goto L7 else goto L8 :: bool
L7:
r0 = 1
goto L11
L8:
L9:
r15 = r9 + 1
r9 = r15
goto L5
L10:
L11:
return r0
def call_all_helper(l):
l :: list
r0 :: bool
r1 :: native_int
r2 :: list
r3, r4 :: native_int
r5 :: bit
r6, i :: object
r7 :: str
r8, r9, r10 :: native_int
r11 :: bit
r12 :: object
r13, __mypyc_all_item__ :: str
r14 :: bit
r15 :: native_int
L0:
r0 = 1
r1 = var_object_size l
r2 = PyList_New(r1)
r3 = 0
L1:
r4 = var_object_size l
r5 = r3 < r4 :: signed
if r5 goto L2 else goto L4 :: bool
L2:
r6 = list_get_item_unsafe l, r3
i = r6
r7 = PyObject_Str(i)
CPyList_SetItemUnsafe(r2, r3, r7)
L3:
r8 = r3 + 1
r3 = r8
goto L1
L4:
r9 = 0
L5:
r10 = var_object_size r2
r11 = r9 < r10 :: signed
if r11 goto L6 else goto L10 :: bool
L6:
r12 = list_get_item_unsafe r2, r9
r13 = cast(str, r12)
__mypyc_all_item__ = r13
r14 = CPyStr_IsTrue(__mypyc_all_item__)
Copy link
Contributor Author

@BobTheBuidler BobTheBuidler Sep 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The boolean check for the all call is now specialized per the dtype of the iterable, and we no longer have to use python's iterator protocol for the input

Copy link
Collaborator

@JukkaL JukkaL Oct 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On master this seemed to just call the stdlib all using a vector call? So the change is that we now generate a for loop for this, right?

if r14 goto L8 else goto L7 :: bool
L7:
r0 = 0
goto L11
L8:
L9:
r15 = r9 + 1
r9 = r15
goto L5
L10:
L11:
return r0

[case testSum]
from typing import Callable, Iterable
Expand Down
30 changes: 29 additions & 1 deletion mypyc/test-data/run-misc.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -794,8 +794,18 @@ def call_all(l: Iterable[int], val: int = 0) -> int:
res = all(i == val for i in l)
return 0 if res else 1

def call_any_with_for_helper(l: Iterable[int], val: int = 0) -> int:
# this listcomp isnt a reasonable real world use but proves the any for loop specializer is good
res = any([i == val for i in l])
return 0 if res else 1

def call_all_with_for_helper(l: Iterable[int], val: int = 0) -> int:
# this listcomp isnt a reasonable real world use but proves the all for loop specializer is good
res = all([i == val for i in l])
return 0 if res else 1

[file driver.py]
from native import call_any, call_all, call_any_nested
from native import call_any, call_all, call_any_nested, call_any_with_for_helper, call_all_with_for_helper

zeros = [0, 0, 0]
ones = [1, 1, 1]
Expand All @@ -807,24 +817,42 @@ mixed_101 = [1, 0, 1]
mixed_110 = [1, 1, 0]

assert call_any([]) == 1
assert call_any_with_for_helper([]) == 1
assert call_any(zeros) == 0
assert call_any_with_for_helper(zeros) == 0
assert call_any(ones) == 1
assert call_any_with_for_helper(ones) == 1
assert call_any(mixed_001) == 0
assert call_any_with_for_helper(mixed_001) == 0
assert call_any(mixed_010) == 0
assert call_any_with_for_helper(mixed_010) == 0
assert call_any(mixed_100) == 0
assert call_any_with_for_helper(mixed_100) == 0
assert call_any(mixed_011) == 0
assert call_any_with_for_helper(mixed_011) == 0
assert call_any(mixed_101) == 0
assert call_any_with_for_helper(mixed_101) == 0
assert call_any(mixed_110) == 0
assert call_any_with_for_helper(mixed_110) == 0

assert call_all([]) == 0
assert call_all_with_for_helper([]) == 0
assert call_all(zeros) == 0
assert call_all_with_for_helper(zeros) == 0
assert call_all(ones) == 1
assert call_all_with_for_helper(ones) == 1
assert call_all(mixed_001) == 1
assert call_all_with_for_helper(mixed_001) == 1
assert call_all(mixed_010) == 1
assert call_all_with_for_helper(mixed_010) == 1
assert call_all(mixed_100) == 1
assert call_all_with_for_helper(mixed_100) == 1
assert call_all(mixed_011) == 1
assert call_all_with_for_helper(mixed_011) == 1
assert call_all(mixed_101) == 1
assert call_all_with_for_helper(mixed_101) == 1
assert call_all(mixed_110) == 1
assert call_all_with_for_helper(mixed_110) == 1

assert call_any_nested([[1, 1, 1], [1, 1], []]) == 1
assert call_any_nested([[1, 1, 1], [0, 1], []]) == 0
Expand Down