Skip to content

Commit b90ecea

Browse files
markshannonmarkshannon
authored
GH-132554: Fix tier2 FOR_ITER implementation and optimizations (GH-135137)
1 parent d9cad07 commit b90ecea

File tree

12 files changed

+155
-173
lines changed

12 files changed

+155
-173
lines changed

‎Include/internal/pycore_ceval.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,8 @@ PyAPI_FUNC(_PyStackRef) _PyFloat_FromDouble_ConsumeInputs(_PyStackRef left, _PyS
353353
extern int _PyRunRemoteDebugger(PyThreadState *tstate);
354354
#endif
355355

356-
_PyStackRef _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index);
356+
PyAPI_FUNC(_PyStackRef)
357+
_PyForIter_VirtualIteratorNext(PyThreadState* tstate, struct _PyInterpreterFrame* frame, _PyStackRef iter, _PyStackRef *index_ptr);
357358

358359
#ifdef __cplusplus
359360
}

‎Include/internal/pycore_stackref.h

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,21 +62,35 @@ PyAPI_FUNC(void) _Py_stackref_record_borrow(_PyStackRef ref, const char *filenam
6262
extern void _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref);
6363

6464
static const _PyStackRef PyStackRef_NULL = { .index = 0 };
65+
static const _PyStackRef PyStackRef_ERROR = { .index = 2 };
6566

6667
// Use the first 3 even numbers for None, True and False.
6768
// Odd numbers are reserved for (tagged) integers
68-
#define PyStackRef_None ((_PyStackRef){ .index = 2 } )
69-
#define PyStackRef_False ((_PyStackRef){ .index = 4 })
70-
#define PyStackRef_True ((_PyStackRef){ .index = 6 })
69+
#define PyStackRef_None ((_PyStackRef){ .index = 4 } )
70+
#define PyStackRef_False ((_PyStackRef){ .index = 6 })
71+
#define PyStackRef_True ((_PyStackRef){ .index = 8 })
7172

72-
#define INITIAL_STACKREF_INDEX 8
73+
#define INITIAL_STACKREF_INDEX 10
7374

7475
static inline int
7576
PyStackRef_IsNull(_PyStackRef ref)
7677
{
7778
return ref.index == 0;
7879
}
7980

81+
static inline bool
82+
PyStackRef_IsError(_PyStackRef ref)
83+
{
84+
return ref.index == 2;
85+
}
86+
87+
static inline bool
88+
PyStackRef_IsValid(_PyStackRef ref)
89+
{
90+
/* Invalid values are ERROR and NULL */
91+
return !PyStackRef_IsError(ref) && !PyStackRef_IsNull(ref);
92+
}
93+
8094
static inline int
8195
PyStackRef_IsTrue(_PyStackRef ref)
8296
{
@@ -104,6 +118,7 @@ PyStackRef_IsTaggedInt(_PyStackRef ref)
104118
static inline PyObject *
105119
_PyStackRef_AsPyObjectBorrow(_PyStackRef ref, const char *filename, int linenumber)
106120
{
121+
assert(!PyStackRef_IsError(ref));
107122
assert(!PyStackRef_IsTaggedInt(ref));
108123
_Py_stackref_record_borrow(ref, filename, linenumber);
109124
return _Py_stackref_get_object(ref);
@@ -155,6 +170,7 @@ _PyStackRef_CLOSE(_PyStackRef ref, const char *filename, int linenumber)
155170
static inline void
156171
_PyStackRef_XCLOSE(_PyStackRef ref, const char *filename, int linenumber)
157172
{
173+
assert(!PyStackRef_IsError(ref));
158174
if (PyStackRef_IsNull(ref)) {
159175
return;
160176
}
@@ -165,6 +181,7 @@ _PyStackRef_XCLOSE(_PyStackRef ref, const char *filename, int linenumber)
165181
static inline _PyStackRef
166182
_PyStackRef_DUP(_PyStackRef ref, const char *filename, int linenumber)
167183
{
184+
assert(!PyStackRef_IsError(ref));
168185
if (PyStackRef_IsTaggedInt(ref)) {
169186
return ref;
170187
}
@@ -241,9 +258,25 @@ PyStackRef_IsNullOrInt(_PyStackRef ref);
241258
#else
242259

243260
#define Py_INT_TAG 3
261+
#define Py_TAG_INVALID 2
244262
#define Py_TAG_REFCNT 1
245263
#define Py_TAG_BITS 3
246264

265+
static const _PyStackRef PyStackRef_ERROR = { .bits = Py_TAG_INVALID };
266+
267+
static inline bool
268+
PyStackRef_IsError(_PyStackRef ref)
269+
{
270+
return ref.bits == Py_TAG_INVALID;
271+
}
272+
273+
static inline bool
274+
PyStackRef_IsValid(_PyStackRef ref)
275+
{
276+
/* Invalid values are ERROR and NULL */
277+
return ref.bits >= Py_INT_TAG;
278+
}
279+
247280
static inline bool
248281
PyStackRef_IsTaggedInt(_PyStackRef i)
249282
{
@@ -284,6 +317,7 @@ PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
284317

285318

286319
static const _PyStackRef PyStackRef_NULL = { .bits = Py_TAG_DEFERRED};
320+
287321
#define PyStackRef_IsNull(stackref) ((stackref).bits == PyStackRef_NULL.bits)
288322
#define PyStackRef_True ((_PyStackRef){.bits = ((uintptr_t)&_Py_TrueStruct) | Py_TAG_DEFERRED })
289323
#define PyStackRef_False ((_PyStackRef){.bits = ((uintptr_t)&_Py_FalseStruct) | Py_TAG_DEFERRED })

‎Lib/test/test_capi/test_opt.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1183,6 +1183,17 @@ def testfunc(n):
11831183
self.assertIsNotNone(ex)
11841184
self.assertIn("_RETURN_GENERATOR", get_opnames(ex))
11851185

1186+
def test_for_iter(self):
1187+
def testfunc(n):
1188+
t = 0
1189+
for i in set(range(n)):
1190+
t += i
1191+
return t
1192+
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
1193+
self.assertEqual(res, TIER2_THRESHOLD * (TIER2_THRESHOLD - 1) // 2)
1194+
self.assertIsNotNone(ex)
1195+
self.assertIn("_FOR_ITER_TIER_TWO", get_opnames(ex))
1196+
11861197
@unittest.skip("Tracing into generators currently isn't supported.")
11871198
def test_for_iter_gen(self):
11881199
def gen(n):

‎Python/bytecodes.c

Lines changed: 23 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -3125,100 +3125,49 @@ dummy_func(
31253125
}
31263126

31273127
replaced op(_FOR_ITER, (iter, null_or_index -- iter, null_or_index, next)) {
3128-
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
3129-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3130-
if (PyStackRef_IsTaggedInt(null_or_index)) {
3131-
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3132-
if (PyStackRef_IsNull(next)) {
3133-
JUMPBY(oparg + 1);
3134-
DISPATCH();
3135-
}
3136-
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
3137-
}
3138-
else {
3139-
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
3140-
if (next_o == NULL) {
3141-
if (_PyErr_Occurred(tstate)) {
3142-
int matches = _PyErr_ExceptionMatches(tstate, PyExc_StopIteration);
3143-
if (!matches) {
3144-
ERROR_NO_POP();
3145-
}
3146-
_PyEval_MonitorRaise(tstate, frame, this_instr);
3147-
_PyErr_Clear(tstate);
3148-
}
3149-
/* iterator ended normally */
3150-
assert(next_instr[oparg].op.code == END_FOR ||
3151-
next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
3152-
/* Jump forward oparg, then skip following END_FOR */
3153-
JUMPBY(oparg + 1);
3154-
DISPATCH();
3128+
_PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
3129+
if (!PyStackRef_IsValid(item)) {
3130+
if (PyStackRef_IsError(item)) {
3131+
ERROR_NO_POP();
31553132
}
3156-
next = PyStackRef_FromPyObjectSteal(next_o);
3133+
// Jump forward by oparg and skip the following END_FOR
3134+
JUMPBY(oparg + 1);
3135+
DISPATCH();
31573136
}
3137+
next = item;
31583138
}
31593139

31603140
op(_FOR_ITER_TIER_TWO, (iter, null_or_index -- iter, null_or_index, next)) {
3161-
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
3162-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3163-
EXIT_IF(!PyStackRef_IsNull(null_or_index));
3164-
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
3165-
if (next_o == NULL) {
3166-
if (_PyErr_Occurred(tstate)) {
3167-
int matches = _PyErr_ExceptionMatches(tstate, PyExc_StopIteration);
3168-
if (!matches) {
3169-
ERROR_NO_POP();
3170-
}
3171-
_PyEval_MonitorRaise(tstate, frame, frame->instr_ptr);
3172-
_PyErr_Clear(tstate);
3141+
_PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
3142+
if (!PyStackRef_IsValid(item)) {
3143+
if (PyStackRef_IsError(item)) {
3144+
ERROR_NO_POP();
31733145
}
31743146
/* iterator ended normally */
31753147
/* The translator sets the deopt target just past the matching END_FOR */
31763148
EXIT_IF(true);
31773149
}
3178-
next = PyStackRef_FromPyObjectSteal(next_o);
3179-
// Common case: no jump, leave it to the code generator
3150+
next = item;
31803151
}
31813152

3153+
31823154
macro(FOR_ITER) = _SPECIALIZE_FOR_ITER + _FOR_ITER;
31833155

31843156

31853157
inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
3186-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3187-
if (PyStackRef_IsTaggedInt(null_or_index)) {
3188-
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3189-
if (PyStackRef_IsNull(next)) {
3190-
JUMPBY(oparg + 1);
3191-
DISPATCH();
3192-
}
3193-
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
3194-
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
3195-
}
3196-
else {
3197-
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
3198-
if (next_o != NULL) {
3199-
next = PyStackRef_FromPyObjectSteal(next_o);
3200-
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
3201-
}
3202-
else {
3203-
if (_PyErr_Occurred(tstate)) {
3204-
int matches = _PyErr_ExceptionMatches(tstate, PyExc_StopIteration);
3205-
if (!matches) {
3206-
ERROR_NO_POP();
3207-
}
3208-
_PyEval_MonitorRaise(tstate, frame, this_instr);
3209-
_PyErr_Clear(tstate);
3210-
}
3211-
/* iterator ended normally */
3212-
assert(next_instr[oparg].op.code == END_FOR ||
3213-
next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
3214-
/* Skip END_FOR */
3215-
JUMPBY(oparg + 1);
3216-
DISPATCH();
3158+
_PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
3159+
if (!PyStackRef_IsValid(item)) {
3160+
if (PyStackRef_IsError(item)) {
3161+
ERROR_NO_POP();
32173162
}
3163+
// Jump forward by oparg and skip the following END_FOR
3164+
JUMPBY(oparg + 1);
3165+
DISPATCH();
32183166
}
3167+
next = item;
3168+
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
32193169
}
32203170

3221-
32223171
op(_ITER_CHECK_LIST, (iter, null_or_index -- iter, null_or_index)) {
32233172
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32243173
EXIT_IF(Py_TYPE(iter_o) != &PyList_Type);

‎Python/ceval.c

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3439,8 +3439,8 @@ _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject *na
34393439
return value;
34403440
}
34413441

3442-
_PyStackRef
3443-
_PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index)
3442+
static _PyStackRef
3443+
foriter_next(PyObject *seq, _PyStackRef index)
34443444
{
34453445
assert(PyStackRef_IsTaggedInt(index));
34463446
assert(PyTuple_CheckExact(seq) || PyList_CheckExact(seq));
@@ -3459,6 +3459,30 @@ _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index)
34593459
return PyStackRef_FromPyObjectSteal(item);
34603460
}
34613461

3462+
_PyStackRef _PyForIter_VirtualIteratorNext(PyThreadState* tstate, _PyInterpreterFrame* frame, _PyStackRef iter, _PyStackRef* index_ptr)
3463+
{
3464+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3465+
_PyStackRef index = *index_ptr;
3466+
if (PyStackRef_IsTaggedInt(index)) {
3467+
*index_ptr = PyStackRef_IncrementTaggedIntNoOverflow(index);
3468+
return foriter_next(iter_o, index);
3469+
}
3470+
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
3471+
if (next_o == NULL) {
3472+
if (_PyErr_Occurred(tstate)) {
3473+
if (_PyErr_ExceptionMatches(tstate, PyExc_StopIteration)) {
3474+
_PyEval_MonitorRaise(tstate, frame, frame->instr_ptr);
3475+
_PyErr_Clear(tstate);
3476+
}
3477+
else {
3478+
return PyStackRef_ERROR;
3479+
}
3480+
}
3481+
return PyStackRef_NULL;
3482+
}
3483+
return PyStackRef_FromPyObjectSteal(next_o);
3484+
}
3485+
34623486
/* Check if a 'cls' provides the given special method. */
34633487
static inline int
34643488
type_has_special_method(PyTypeObject *cls, PyObject *name)

‎Python/executor_cases.c.h

Lines changed: 6 additions & 19 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)