Skip to content

GH-132554: Fix tier2 FOR_ITER implementation and optimizations #135137

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

Merged
merged 4 commits into from
Jun 5, 2025
+155 −173
Merged

GH-132554: Fix tier2 FOR_ITER implementation and optimizations #135137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
eedbe8b
Fix FOR_ITER for tier 2
markshannon Jun 4, 2025
a10eadb
Add test
markshannon Jun 4, 2025
5725821
Update debug stackref
markshannon Jun 4, 2025
cc09774
Address review comments
markshannon Jun 5, 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
3 changes: 2 additions & 1 deletion Include/internal/pycore_ceval.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,8 @@ PyAPI_FUNC(_PyStackRef) _PyFloat_FromDouble_ConsumeInputs(_PyStackRef left, _PyS
extern int _PyRunRemoteDebugger(PyThreadState *tstate);
#endif

_PyStackRef _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index);
PyAPI_FUNC(_PyStackRef)
_PyForIter_VirtualIteratorNext(PyThreadState* tstate, struct _PyInterpreterFrame* frame, _PyStackRef iter, _PyStackRef *index_ptr);

#ifdef __cplusplus
}
Expand Down
42 changes: 38 additions & 4 deletions Include/internal/pycore_stackref.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,21 +62,35 @@ PyAPI_FUNC(void) _Py_stackref_record_borrow(_PyStackRef ref, const char *filenam
extern void _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref);

static const _PyStackRef PyStackRef_NULL = { .index = 0 };
static const _PyStackRef PyStackRef_ERROR = { .index = 2 };

// Use the first 3 even numbers for None, True and False.
// Odd numbers are reserved for (tagged) integers
#define PyStackRef_None ((_PyStackRef){ .index = 2 } )
#define PyStackRef_False ((_PyStackRef){ .index = 4 })
#define PyStackRef_True ((_PyStackRef){ .index = 6 })
#define PyStackRef_None ((_PyStackRef){ .index = 4 } )
#define PyStackRef_False ((_PyStackRef){ .index = 6 })
#define PyStackRef_True ((_PyStackRef){ .index = 8 })

#define INITIAL_STACKREF_INDEX 8
#define INITIAL_STACKREF_INDEX 10

static inline int
PyStackRef_IsNull(_PyStackRef ref)
{
return ref.index == 0;
}

static inline bool
PyStackRef_IsError(_PyStackRef ref)
{
return ref.index == 2;
}

static inline bool
PyStackRef_IsValid(_PyStackRef ref)
{
/* Invalid values are ERROR and NULL */
return !PyStackRef_IsError(ref) && !PyStackRef_IsNull(ref);
}

static inline int
PyStackRef_IsTrue(_PyStackRef ref)
{
Expand Down Expand Up @@ -104,6 +118,7 @@ PyStackRef_IsTaggedInt(_PyStackRef ref)
static inline PyObject *
_PyStackRef_AsPyObjectBorrow(_PyStackRef ref, const char *filename, int linenumber)
{
assert(!PyStackRef_IsError(ref));
assert(!PyStackRef_IsTaggedInt(ref));
_Py_stackref_record_borrow(ref, filename, linenumber);
return _Py_stackref_get_object(ref);
Expand Down Expand Up @@ -155,6 +170,7 @@ _PyStackRef_CLOSE(_PyStackRef ref, const char *filename, int linenumber)
static inline void
_PyStackRef_XCLOSE(_PyStackRef ref, const char *filename, int linenumber)
{
assert(!PyStackRef_IsError(ref));
if (PyStackRef_IsNull(ref)) {
return;
}
Expand All @@ -165,6 +181,7 @@ _PyStackRef_XCLOSE(_PyStackRef ref, const char *filename, int linenumber)
static inline _PyStackRef
_PyStackRef_DUP(_PyStackRef ref, const char *filename, int linenumber)
{
assert(!PyStackRef_IsError(ref));
if (PyStackRef_IsTaggedInt(ref)) {
return ref;
}
Expand Down Expand Up @@ -241,9 +258,25 @@ PyStackRef_IsNullOrInt(_PyStackRef ref);
#else

#define Py_INT_TAG 3
#define Py_TAG_INVALID 2
#define Py_TAG_REFCNT 1
#define Py_TAG_BITS 3

static const _PyStackRef PyStackRef_ERROR = { .bits = Py_TAG_INVALID };

static inline bool
PyStackRef_IsError(_PyStackRef ref)
{
return ref.bits == Py_TAG_INVALID;
}

static inline bool
PyStackRef_IsValid(_PyStackRef ref)
{
/* Invalid values are ERROR and NULL */
return ref.bits >= Py_INT_TAG;
}

static inline bool
PyStackRef_IsTaggedInt(_PyStackRef i)
{
Expand Down Expand Up @@ -284,6 +317,7 @@ PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)


static const _PyStackRef PyStackRef_NULL = { .bits = Py_TAG_DEFERRED};

#define PyStackRef_IsNull(stackref) ((stackref).bits == PyStackRef_NULL.bits)
#define PyStackRef_True ((_PyStackRef){.bits = ((uintptr_t)&_Py_TrueStruct) | Py_TAG_DEFERRED })
#define PyStackRef_False ((_PyStackRef){.bits = ((uintptr_t)&_Py_FalseStruct) | Py_TAG_DEFERRED })
Expand Down
11 changes: 11 additions & 0 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 @@ -1183,6 +1183,17 @@ def testfunc(n):
self.assertIsNotNone(ex)
self.assertIn("_RETURN_GENERATOR", get_opnames(ex))

def test_for_iter(self):
def testfunc(n):
t = 0
for i in set(range(n)):
t += i
return t
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(res, TIER2_THRESHOLD * (TIER2_THRESHOLD - 1) // 2)
self.assertIsNotNone(ex)
self.assertIn("_FOR_ITER_TIER_TWO", get_opnames(ex))

@unittest.skip("Tracing into generators currently isn't supported.")
def test_for_iter_gen(self):
def gen(n):
Expand Down
97 changes: 23 additions & 74 deletions Python/bytecodes.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3125,100 +3125,49 @@ dummy_func(
}

replaced op(_FOR_ITER, (iter, null_or_index -- iter, null_or_index, next)) {
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
if (PyStackRef_IsTaggedInt(null_or_index)) {
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
if (PyStackRef_IsNull(next)) {
JUMPBY(oparg + 1);
DISPATCH();
}
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
}
else {
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
if (next_o == NULL) {
if (_PyErr_Occurred(tstate)) {
int matches = _PyErr_ExceptionMatches(tstate, PyExc_StopIteration);
if (!matches) {
ERROR_NO_POP();
}
_PyEval_MonitorRaise(tstate, frame, this_instr);
_PyErr_Clear(tstate);
}
/* iterator ended normally */
assert(next_instr[oparg].op.code == END_FOR ||
next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
/* Jump forward oparg, then skip following END_FOR */
JUMPBY(oparg + 1);
DISPATCH();
_PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
if (!PyStackRef_IsValid(item)) {
if (PyStackRef_IsError(item)) {
ERROR_NO_POP();
}
next = PyStackRef_FromPyObjectSteal(next_o);
// Jump forward by oparg and skip the following END_FOR
JUMPBY(oparg + 1);
Copy link
Member

Choose a reason for hiding this comment

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

Can you add back the comment that this is skipping the END_FOR?

Suggested change
JUMPBY(oparg + 1);
// Jump forward by oparg, then skip the following END_FOR:
JUMPBY(oparg + 1);

DISPATCH();
}
next = item;
}

op(_FOR_ITER_TIER_TWO, (iter, null_or_index -- iter, null_or_index, next)) {
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
EXIT_IF(!PyStackRef_IsNull(null_or_index));
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
if (next_o == NULL) {
if (_PyErr_Occurred(tstate)) {
int matches = _PyErr_ExceptionMatches(tstate, PyExc_StopIteration);
if (!matches) {
ERROR_NO_POP();
}
_PyEval_MonitorRaise(tstate, frame, frame->instr_ptr);
_PyErr_Clear(tstate);
_PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
if (!PyStackRef_IsValid(item)) {
if (PyStackRef_IsError(item)) {
ERROR_NO_POP();
}
/* iterator ended normally */
/* The translator sets the deopt target just past the matching END_FOR */
EXIT_IF(true);
}
next = PyStackRef_FromPyObjectSteal(next_o);
// Common case: no jump, leave it to the code generator
next = item;
}


macro(FOR_ITER) = _SPECIALIZE_FOR_ITER + _FOR_ITER;


inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
if (PyStackRef_IsTaggedInt(null_or_index)) {
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
if (PyStackRef_IsNull(next)) {
JUMPBY(oparg + 1);
DISPATCH();
}
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
}
else {
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
if (next_o != NULL) {
next = PyStackRef_FromPyObjectSteal(next_o);
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
}
else {
if (_PyErr_Occurred(tstate)) {
int matches = _PyErr_ExceptionMatches(tstate, PyExc_StopIteration);
if (!matches) {
ERROR_NO_POP();
}
_PyEval_MonitorRaise(tstate, frame, this_instr);
_PyErr_Clear(tstate);
}
/* iterator ended normally */
assert(next_instr[oparg].op.code == END_FOR ||
next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
/* Skip END_FOR */
JUMPBY(oparg + 1);
DISPATCH();
_PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
if (!PyStackRef_IsValid(item)) {
if (PyStackRef_IsError(item)) {
ERROR_NO_POP();
}
// Jump forward by oparg and skip the following END_FOR
JUMPBY(oparg + 1);
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
JUMPBY(oparg + 1);
// Jump forward by oparg, then skip the following END_FOR:
JUMPBY(oparg + 1);

DISPATCH();
}
next = item;
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
}


op(_ITER_CHECK_LIST, (iter, null_or_index -- iter, null_or_index)) {
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
EXIT_IF(Py_TYPE(iter_o) != &PyList_Type);
Expand Down
28 changes: 26 additions & 2 deletions Python/ceval.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3439,8 +3439,8 @@ _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject *na
return value;
}

_PyStackRef
_PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index)
static _PyStackRef
foriter_next(PyObject *seq, _PyStackRef index)
{
assert(PyStackRef_IsTaggedInt(index));
assert(PyTuple_CheckExact(seq) || PyList_CheckExact(seq));
Expand All @@ -3459,6 +3459,30 @@ _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index)
return PyStackRef_FromPyObjectSteal(item);
}

_PyStackRef _PyForIter_VirtualIteratorNext(PyThreadState* tstate, _PyInterpreterFrame* frame, _PyStackRef iter, _PyStackRef* index_ptr)
{
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
_PyStackRef index = *index_ptr;
if (PyStackRef_IsTaggedInt(index)) {
*index_ptr = PyStackRef_IncrementTaggedIntNoOverflow(index);
return foriter_next(iter_o, index);
}
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
if (next_o == NULL) {
if (_PyErr_Occurred(tstate)) {
if (_PyErr_ExceptionMatches(tstate, PyExc_StopIteration)) {
_PyEval_MonitorRaise(tstate, frame, frame->instr_ptr);
_PyErr_Clear(tstate);
}
else {
return PyStackRef_ERROR;
}
}
return PyStackRef_NULL;
}
return PyStackRef_FromPyObjectSteal(next_o);
}

/* Check if a 'cls' provides the given special method. */
static inline int
type_has_special_method(PyTypeObject *cls, PyObject *name)
Expand Down
25 changes: 6 additions & 19 deletions Python/executor_cases.c.h
View file
Open in desktop

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

Loading
Loading