diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index c8f87f4fe17ac7..415abf6372f820 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -20,6 +20,10 @@ #For test of issue 136154 GLOBAL_136154 = 42 +# For frozendict JIT tests +FROZEN_DICT_CONST = frozendict(x=1, y=2) + + @contextlib.contextmanager def clear_executors(func): # Clear executors in func before and after running a block @@ -4155,6 +4159,35 @@ def testfunc(n): self.assertLessEqual(count_ops(ex, "_POP_TOP_INT"), 1) self.assertIn("_POP_TOP_NOP", uops) + def test_binary_subscr_frozendict_lowering(self): + def testfunc(n): + x = 0 + for _ in range(n): + x += FROZEN_DICT_CONST['x'] + return x + + res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + self.assertIn("_INSERT_2_LOAD_CONST_INLINE_BORROW", uops) + self.assertNotIn("_BINARY_OP_SUBSCR_DICT", uops) + + def test_binary_subscr_frozendict_const_fold(self): + def testfunc(n): + x = 0 + for _ in range(n): + if FROZEN_DICT_CONST['x'] == 1: + x += 1 + return x + + res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + # lookup result is folded to constant 1, so comparison is optimized away + self.assertNotIn("_COMPARE_OP_INT", uops) + def test_binary_subscr_list_slice(self): def testfunc(n): x = 0 diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index 2959d1d0bb4e6f..2fdf186d2a3a46 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -485,6 +485,9 @@ dummy_func(void) { res = sym_new_not_null(ctx); ds = dict_st; ss = sub_st; + if (sym_matches_type(dict_st, &PyFrozenDict_Type)) { + REPLACE_OPCODE_IF_EVALUATES_PURE(dict_st, sub_st, res); + } } op(_BINARY_OP_SUBSCR_LIST_SLICE, (list_st, sub_st -- res, ls, ss)) { diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index 745ee6b3ab9498..63a7e4b0ed5f71 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -1327,6 +1327,54 @@ res = sym_new_not_null(ctx); ds = dict_st; ss = sub_st; + if (sym_matches_type(dict_st, &PyFrozenDict_Type)) { + if ( + sym_is_safe_const(ctx, dict_st) && + sym_is_safe_const(ctx, sub_st) + ) { + JitOptRef dict_st_sym = dict_st; + JitOptRef sub_st_sym = sub_st; + _PyStackRef dict_st = sym_get_const_as_stackref(ctx, dict_st_sym); + _PyStackRef sub_st = sym_get_const_as_stackref(ctx, sub_st_sym); + _PyStackRef res_stackref; + _PyStackRef ds_stackref; + _PyStackRef ss_stackref; + /* Start of uop copied from bytecodes for constant evaluation */ + PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st); + PyObject *dict = PyStackRef_AsPyObjectBorrow(dict_st); + assert(PyAnyDict_CheckExact(dict)); + STAT_INC(BINARY_OP, hit); + PyObject *res_o; + int rc = PyDict_GetItemRef(dict, sub, &res_o); + if (rc == 0) { + _PyErr_SetKeyError(sub); + } + if (rc <= 0) { + JUMP_TO_LABEL(error); + } + res_stackref = PyStackRef_FromPyObjectSteal(res_o); + ds_stackref = dict_st; + ss_stackref = sub_st; + /* End of uop copied from bytecodes for constant evaluation */ + (void)ds_stackref; + (void)ss_stackref; + res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectSteal(res_stackref)); + if (sym_is_const(ctx, res)) { + PyObject *result = sym_get_const(ctx, res); + if (_Py_IsImmortal(result)) { + // Replace with _INSERT_2_LOAD_CONST_INLINE_BORROW since we have two inputs and an immortal result + ADD_OP(_INSERT_2_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)result); + } + } + CHECK_STACK_BOUNDS(1); + stack_pointer[-2] = res; + stack_pointer[-1] = ds; + stack_pointer[0] = ss; + stack_pointer += 1; + ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__); + break; + } + } CHECK_STACK_BOUNDS(1); stack_pointer[-2] = res; stack_pointer[-1] = ds; diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c index 2a8d8c45c588ba..7593d5b159620d 100644 --- a/Python/optimizer_symbols.c +++ b/Python/optimizer_symbols.c @@ -282,7 +282,8 @@ _Py_uop_sym_is_safe_const(JitOptContext *ctx, JitOptRef sym) return (typ == &PyUnicode_Type) || (typ == &PyFloat_Type) || (typ == &_PyNone_Type) || - (typ == &PyBool_Type); + (typ == &PyBool_Type) || + (typ == &PyFrozenDict_Type); } void