Skip to content

gh-152215: Add uop peepholes for folded stack shuffles and adjacent guards #153035

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
cocolato wants to merge 3 commits into
base: main
Choose a base branch
from
Open

gh-152215: Add uop peepholes for folded stack shuffles and adjacent guards #153035

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0be7147
add peephole helper
cocolato Jul 4, 2026
2dc5a6c
Merge branch 'main' into jit/uop_peepholes
cocolato Jul 21, 2026
aeab76c
Combine adjacent uop guards
cocolato Jul 21, 2026
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
24 changes: 22 additions & 2 deletions Include/internal/pycore_uop_ids.h
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

84 changes: 84 additions & 0 deletions Include/internal/pycore_uop_metadata.h
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

134 changes: 119 additions & 15 deletions Lib/test/test_capi/test_opt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -640,7 +640,6 @@ def _run_with_optimizer(self, testfunc, arg):
ex = get_first_executor(testfunc)
return res, ex


def test_int_type_propagation(self):
def testfunc(loops):
num = 0
Expand All @@ -656,9 +655,10 @@ def testfunc(loops):
binop_count = [opname for opname in iter_opnames(ex) if opname == "_BINARY_OP_ADD_INT"]
guard_tos_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_INT"]
guard_nos_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_NOS_INT"]
combined_guard_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_AND_NOS_INT"]
self.assertGreaterEqual(len(binop_count), 3)
self.assertLessEqual(len(guard_tos_int_count), 1)
self.assertLessEqual(len(guard_nos_int_count), 1)
self.assertLessEqual(len(guard_tos_int_count) + len(combined_guard_count), 1)
self.assertLessEqual(len(guard_nos_int_count) + len(combined_guard_count), 1)

def test_int_type_propagation_through_frame(self):
def double(x):
Expand All @@ -679,9 +679,10 @@ def testfunc(loops):
binop_count = [opname for opname in iter_opnames(ex) if opname == "_BINARY_OP_ADD_INT"]
guard_tos_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_INT"]
guard_nos_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_NOS_INT"]
combined_guard_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_AND_NOS_INT"]
self.assertGreaterEqual(len(binop_count), 3)
self.assertLessEqual(len(guard_tos_int_count), 1)
self.assertLessEqual(len(guard_nos_int_count), 1)
self.assertLessEqual(len(guard_tos_int_count) + len(combined_guard_count), 1)
self.assertLessEqual(len(guard_nos_int_count) + len(combined_guard_count), 1)

def test_int_type_propagation_from_frame(self):
def double(x):
Expand All @@ -702,9 +703,10 @@ def testfunc(loops):
binop_count = [opname for opname in iter_opnames(ex) if opname == "_BINARY_OP_ADD_INT"]
guard_tos_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_INT"]
guard_nos_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_NOS_INT"]
combined_guard_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_AND_NOS_INT"]
self.assertGreaterEqual(len(binop_count), 3)
self.assertLessEqual(len(guard_tos_int_count), 1)
self.assertLessEqual(len(guard_nos_int_count), 1)
self.assertLessEqual(len(guard_tos_int_count) + len(combined_guard_count), 1)
self.assertLessEqual(len(guard_nos_int_count) + len(combined_guard_count), 1)

def test_int_impure_region(self):
def testfunc(loops):
Expand Down Expand Up @@ -751,6 +753,7 @@ def testfunc(n):
uops = get_opnames(ex)
self.assertNotIn("_GUARD_TOS_INT", uops)
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertNotIn("_GUARD_TOS_AND_NOS_INT", uops)

def test_int_value_numbering(self):
def testfunc(n):
Expand All @@ -770,6 +773,7 @@ def testfunc(n):
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_INT", uops)
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertNotIn("_GUARD_TOS_AND_NOS_INT", uops)
guard_tos_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_TOS_INT"]
self.assertEqual(len(guard_tos_count), 1)

Expand Down Expand Up @@ -4041,7 +4045,7 @@ def test_float_truediv_speculative_guards_from_tracing(self):
# a, b are locals with no statically known type. _RECORD_TOS_TYPE /
# _RECORD_NOS_TYPE (added to the BINARY_OP macro) capture the observed
# operand types during tracing, and the optimizer then speculatively
# emits _GUARD_{TOS,NOS}_FLOAT and specializes the division.
# emits a combined float guard and specializes the division.
def testfunc(args):
a, b, n = args
total = 0.0
Expand All @@ -4053,14 +4057,13 @@ def testfunc(args):
self.assertAlmostEqual(res, TIER2_THRESHOLD * (10.0 / 3.0))
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_FLOAT", uops)
self.assertIn("_GUARD_NOS_FLOAT", uops)
self.assertIn("_GUARD_TOS_AND_NOS_FLOAT", uops)
self.assertIn("_BINARY_OP_TRUEDIV_FLOAT", uops)

def test_float_remainder_speculative_guards_from_tracing(self):
# a, b are locals with no statically known type. Tracing records
# them as floats; the optimizer then speculatively emits
# _GUARD_{TOS,NOS}_FLOAT for NB_REMAINDER. That narrows both
# _GUARD_TOS_AND_NOS_FLOAT for NB_REMAINDER. That narrows both
# operands to float, and the _BINARY_OP handler marks the result
# as a unique float. Downstream, `* 2.0` therefore specializes
# to _BINARY_OP_MULTIPLY_FLOAT_INPLACE.
Expand All @@ -4075,8 +4078,7 @@ def testfunc(args):
self.assertAlmostEqual(res, TIER2_THRESHOLD * (10.0 % 3.0) * 2.0)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_FLOAT", uops)
self.assertIn("_GUARD_NOS_FLOAT", uops)
self.assertIn("_GUARD_TOS_AND_NOS_FLOAT", uops)
self.assertIn("_BINARY_OP_MULTIPLY_FLOAT_INPLACE", uops)

def test_float_truediv_type_propagation(self):
Expand Down Expand Up @@ -5582,8 +5584,7 @@ def testfunc(*args):
uops = get_opnames(ex)

self.assertIn("_BINARY_OP_SUBSCR_INIT_CALL", uops)
# _POP_TOP_NOP is a sign the optimizer ran and didn't hit contradiction.
self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
self.assertNotIn("_POP_TOP_NOP", uops)

def test_load_attr_property_frame(self):
class B:
Expand Down Expand Up @@ -6215,6 +6216,109 @@ def __exit__(self, e, v, t): ...
f1()
"""), PYTHON_JIT="1")

def test_peephole_cancels_adjacent_push_pop(self):
def negated_constant(n):
x = 0
for i in range(n):
a = 1
result = -a
if result < 0:
x += 1
return x

res, ex = self._run_with_optimizer(negated_constant, TIER2_THRESHOLD)
self.assertEqual(res, TIER2_THRESHOLD)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertNotIn("_SWAP", uops)
self.assertNotIn("_SWAP_2", uops)

def test_peephole_removes_redundant_rrot(self):
def add(a, b):
return a + b

def testfunc(n):
x = 0
for _ in range(n):
x += add(1, 2)
return x

res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(res, 3 * TIER2_THRESHOLD)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertNotIn("_RROT_3", uops)

def test_combine_adjacent_int_guards(self):
def testfunc(left, right, loops):
for _ in range(loops):
result = left + right
return result

result = testfunc(40, 2, TIER2_THRESHOLD)

self.assertEqual(result, 42)
ex = get_first_executor(testfunc)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_AND_NOS_INT", uops)
self.assertNotIn("_GUARD_TOS_INT", uops)
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertEqual(testfunc(40.0, 2.0, 1), 42.0)
self.assertEqual(testfunc(1 << 100, 2, 1), (1 << 100) + 2)

def test_combine_adjacent_overflowed_guards(self):
def testfunc(start, loops):
for value in range(start, start + loops):
inverted = ~value
result = value + inverted
return result

result = testfunc(0, TIER2_THRESHOLD)

self.assertEqual(result, -1)
ex = get_first_executor(testfunc)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_AND_NOS_OVERFLOWED", uops)
self.assertNotIn("_GUARD_TOS_OVERFLOWED", uops)
self.assertNotIn("_GUARD_NOS_OVERFLOWED", uops)
self.assertEqual(testfunc(1 << 100, 1), -1)

def test_combine_adjacent_float_guards(self):
def testfunc(left, right, loops):
for _ in range(loops):
result = left + right
return result

result = testfunc(1.5, 2.5, TIER2_THRESHOLD)

self.assertEqual(result, 4.0)
ex = get_first_executor(testfunc)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_AND_NOS_FLOAT", uops)
self.assertNotIn("_GUARD_TOS_FLOAT", uops)
self.assertNotIn("_GUARD_NOS_FLOAT", uops)
self.assertEqual(testfunc(1, 2, 1), 3)

def test_combine_adjacent_unicode_guards(self):
def testfunc(left, right, loops):
for _ in range(loops):
result = left + right
return result

result = testfunc("a", "b", TIER2_THRESHOLD)

self.assertEqual(result, "ab")
ex = get_first_executor(testfunc)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_GUARD_TOS_AND_NOS_UNICODE", uops)
self.assertNotIn("_GUARD_TOS_UNICODE", uops)
self.assertNotIn("_GUARD_NOS_UNICODE", uops)
self.assertEqual(testfunc(b"a", b"b", 1), b"ab")

def global_identity(x):
return x

Expand Down
Loading
Loading