gh-100239: Specialize binary operations using BINARY_OP_EXTEND

[eendebakpt]

• We add list and tuple concatenation to BINARY_OP_EXTEND
• We pass type information in BINARY_OP_EXTEND in tier 2. This allows the jit to perform better optimizations.
• In the jit we can now eliminate the _GUARD_BINARY_OP_EXTEND if type information is known.
• The fraction of code specialized for BINARY_OP increases from 70% to 90%.

Benchmark are performance neutral (in the +- 1% range) is seems.

Benchmark script

"""Benchmark for BINARY_OP_EXTEND type propagation.
Tests whether the tier 2 optimizer can eliminate guards when types are
known from previous BINARY_OP_EXTEND results.
Usage:
./python bench_binary_op_extend.py
./python bench_binary_op_extend.py --save result.json
./python bench_binary_op_extend.py --compare a.json b.json
"""
import sys
import pyperf
INNER = 2000
def bench_list_concat_subscr(n):
"""list + list followed by subscript — tests list type propagation."""
a = [1, 2, 3]
b = [4, 5, 6]
total = 0
for _ in range(n):
c = a + b
total += c[0] + c[3]
return total
def bench_tuple_concat_unpack(n):
"""tuple + tuple followed by unpack — tests tuple type propagation."""
t1 = (1, 2)
t2 = (3, 4)
total = 0
for _ in range(n):
a, b, c, d = t1 + t2
total += a + d
return total
def bench_str_repeat(n):
"""str * int in a loop — tests str type propagation."""
s = "ab"
total = 0
for i in range(n):
r = s * (i % 5)
total += len(r)
return total
def bench_bytes_concat(n):
"""bytes + bytes in a loop — tests bytes type propagation."""
a = b"hello"
b_ = b" world"
total = 0
for _ in range(n):
c = a + b_
total += len(c)
return total
def bench_bytes_repeat(n):
"""bytes * int in a loop — tests bytes type propagation."""
b = b"ab"
total = 0
for i in range(n):
r = b * (i % 3)
total += len(r)
return total
def bench_tuple_repeat(n):
"""tuple * int in a loop — tests tuple type propagation."""
t = (1, 2, 3)
total = 0
for i in range(n):
r = t * (i % 3)
total += len(r)
return total
def bench_dict_merge(n):
"""dict | dict in a loop — tests dict type propagation."""
d1 = {"a": 1, "b": 2}
d2 = {"c": 3, "d": 4}
total = 0
for _ in range(n):
d = d1 | d2
total += len(d)
return total
def bench_chained_list_ops(n):
"""Multiple list ops chained — tests guard elimination across ops."""
a = [1, 2]
b = [3, 4]
total = 0
for _ in range(n):
c = a + b
d = c + a
total += d[0] + d[4]
return total
def bench_mixed_float_int(n):
"""float + int and int + float — existing EXTEND specializations."""
x = 1.5
total = 0.0
for i in range(n):
a = x + i
total += a return total
def float_mix_mul(n):
"""float + int then float * float — tests unique flag for inplace mul."""
x = 1.5
total = 0.0
for i in range(n):
a = (x + i) * 2.0 # result of x+i should be unique -> inplace multiply
total += a
return total
BENCHMARKS = [
("list_concat_subscr", bench_list_concat_subscr),
("tuple_concat_unpack", bench_tuple_concat_unpack),
("str_repeat", bench_str_repeat),
("bytes_concat", bench_bytes_concat),
("bytes_repeat", bench_bytes_repeat),
("tuple_repeat", bench_tuple_repeat),
("dict_merge", bench_dict_merge),
("chained_list_ops", bench_chained_list_ops),
("mixed_float_int", bench_mixed_float_int),
("float_mix_mul", float_mix_mul),
]
def main():
args = sys.argv[1:]
if "--compare" in args:
idx = args.index("--compare")
file_a = args[idx + 1]
file_b = args[idx + 2]
import subprocess
subprocess.run([sys.executable, "-m", "pyperf", "compare_to",
file_a, file_b, "--table"])
return
save_file = None
if "--save" in args:
idx = args.index("--save")
save_file = args[idx + 1]
runner = pyperf.Runner()
for name, func in BENCHMARKS:
# Warm up
func(INNER)
runner.bench_func(name, func, INNER)
if save_file and runner.args.output:
import shutil
shutil.copy(runner.args.output, save_file)
if __name__ == "__main__":
main()

• Issue: Specialize long tail of binary operations using a table. #100239

[markshannon]

This looks good overall, but I've not done a detailed review.

I have a couple of general concerns about the BINARY_OP_EXTEND optimization in general, not in this PR, but something to keep in mind:

• How do we ensure the robustness of the VM and optimizations when we expose this to 3rd party code as we intend to do at some point in the future?
• As binaryop_extend_descrs gets larger, specialization of binary ops will get slower. Can we sort the array, or use a mapping to reduce the overhead? It shouldn't be a problem yet, but could be in the future if there were 100s of entries.

[markshannon]

I note that you had to add a lot of new guard functions, most of which just check that the operands types are the same as those specified in the new lhs_type and rhs_type fields.
Rather than have the duplication, how about doing the guard by either calling the guard or checking the given types.

In other words, the check becomes:
descr->guard == NULL ? (ltype == descr->lhstype && rtype == descr->rhstype) : descr->guard(lobj, robj)

This looks slower and more complex but it is more robust, since we don't need to worry about the guard function and the types being out of sync. The additional specializations should more than compensate in the interpreter, and we can easily eliminate the additional check in the JIT.

If the guard function is not NULL, then both lhs_type and rhs_type should be NULL.
This will reduce the amount of information available to the JIT, but is much less error prone.
We can special case the compact long and float guards.

[markshannon]

A few minor suggestions

[eendebakpt]

In other words, the check becomes: descr->guard == NULL ? (ltype == descr->lhstype && rtype == descr->rhstype) : descr->guard(lobj, robj)

That was a nice suggestion. Less guards, and faster for most cases (no function call needed). Updated benchmarks:

Benchmark	main	branch	Change
list_concat_subscr	95.9 us	94.5 us	1.01x faster
tuple_concat_unpack	90.6 us	89.0 us	1.02x faster
str_repeat	108 us	97.6 us	1.10x faster
bytes_repeat	97.2 us	86.5 us	1.12x faster
tuple_repeat	93.3 us	82.0 us	1.14x faster
dict_merge	193 us	189 us	1.02x faster
chained_list_ops	141 us	137 us	1.03x faster
mixed_float_int	45.8 us	44.9 us	1.02x faster
float_mix_mul	46.4 us	46.2 us	1.00x faster
Geometric mean	(ref)		1.05x faster

[kumaraditya303]

Would you check if #148384 helps in this?

[kumaraditya303]

This would still lookup the function pointer each time this gets called, have you tried exposing the function and using that directly?

[eendebakpt]

I considered it, but did not do it yet. With the current PR we get a performance increase and better type information in tier2 and I did not want to make too many changes. For some other ops exposing the function meant not having to create an extra method in specialize.c. Even when exposing PyUnicode_Type.tp_as_sequence->sq_repeat (which is unicode_repeat from unicodeobject.c) we would still need the str_int_multiply as unicode_repeat takes an int.

Exposing it and using it here would be another minor performance improvement though. So this let me know if you want me to make the change.

[markshannon]

It would make sense to expose the functions, like you have for _PyBytes_Concat for example, but it can wait for a future PR.

[eendebakpt]

Would you check if #148384 helps in this?

I'll run benchmarks later (will have to be on a stable machine, I suspect the difference will be small)

[markshannon]

There is potential to optimize out the float/compact int and compact int/float guards.
Can be done in another PR.

[markshannon]

This looks good now.
There more optimizations we can do, but let's get this in first.

Can you fix the merge conflicts, and then I can merge it.

[markshannon]

Further optimizations that can be done (in future PRs):

• Remove _GUARD_BINARY_OP_EXTEND for guards on float/compact_int pairs
• Apply Low effort, precise specialization in tier 2 using lookup tables, replace and replicate. faster-cpython/ideas#660 to generate optimum uops for each variant.
• Avoid indirections like PyUnicode_Type.tp_as_sequence->sq_repeat\ by exposing _PyUnicode_Repeat`

[markshannon]

The aarch64-apple-darwin machine seems to be generally flaky lately.