Skip to content

Commit 0dfe587

Browse files
committed
Specialize FOR_ITER for ranges using tagged ints
1 parent 4dc28c9 commit 0dfe587

File tree

12 files changed

+258
-162
lines changed

12 files changed

+258
-162
lines changed

Include/internal/pycore_long.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,15 @@ _PyLong_IsNonNegativeCompact(const PyLongObject* op) {
207207
}
208208

209209

210+
/* Return the value of a non-negative compact as a machine int */
211+
static inline Py_ssize_t
212+
_PyLong_GetNonNegativeCompactValue(const PyLongObject* op) {
213+
assert(PyLong_Check(op));
214+
assert (_PyLong_IsNonNegativeCompact(op));
215+
return op->long_value.ob_digit[0];
216+
}
217+
218+
210219
static inline int
211220
_PyLong_BothAreCompact(const PyLongObject* a, const PyLongObject* b) {
212221
assert(PyLong_Check(a));

Include/internal/pycore_opcode_metadata.h

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Include/internal/pycore_range.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@ typedef struct {
1515
long len;
1616
} _PyRangeIterObject;
1717

18+
// Does this range have start == 0, step == 1 and step in compact int range?
19+
int _PyRange_IsSimpleCompact(PyObject *range);
20+
Py_ssize_t _PyRange_GetStopIfCompact(PyObject *range);
21+
1822
#ifdef __cplusplus
1923
}
2024
#endif

Include/internal/pycore_stackref.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,7 @@ PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
272272
return (_PyStackRef){ .bits = ref.bits + 4 };
273273
}
274274

275+
extern _PyStackRef PyStackRef_BoxInt(_PyStackRef i);
275276

276277
#ifdef Py_GIL_DISABLED
277278

Include/internal/pycore_uop_metadata.h

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Objects/longobject.c

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#include "pycore_long.h" // _Py_SmallInts
1111
#include "pycore_object.h" // _PyObject_Init()
1212
#include "pycore_runtime.h" // _PY_NSMALLPOSINTS
13+
#include "pycore_stackref.h"
1314
#include "pycore_structseq.h" // _PyStructSequence_FiniBuiltin()
1415
#include "pycore_unicodeobject.h" // _PyUnicode_Equal()
1516

@@ -6879,3 +6880,19 @@ PyLongWriter_Finish(PyLongWriter *writer)
68796880

68806881
return (PyObject*)obj;
68816882
}
6883+
6884+
// Tagged int support
6885+
6886+
_PyStackRef
6887+
PyStackRef_BoxInt(_PyStackRef i)
6888+
{
6889+
assert((i.bits & Py_INT_TAG) == Py_INT_TAG);
6890+
intptr_t val = (intptr_t)i.bits;
6891+
val = Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, val, 2);
6892+
PyObject *boxed = PyLong_FromSsize_t(val);
6893+
if (boxed == NULL) {
6894+
return PyStackRef_NULL;
6895+
}
6896+
return PyStackRef_FromPyObjectSteal(boxed);
6897+
}
6898+

Objects/rangeobject.c

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,26 @@ range_vectorcall(PyObject *rangetype, PyObject *const *args,
156156
return range_from_array((PyTypeObject *)rangetype, args, nargs);
157157
}
158158

159+
int
160+
_PyRange_IsSimpleCompact(PyObject *range) {
161+
assert(PyRange_Check(range));
162+
rangeobject *r = (rangeobject*)range;
163+
if (r->start == _PyLong_GetZero() && r->step == _PyLong_GetOne() &&
164+
_PyLong_IsNonNegativeCompact((PyLongObject *)r->stop)
165+
) {
166+
return 1;
167+
}
168+
return 0;
169+
}
170+
171+
Py_ssize_t
172+
_PyRange_GetStopIfCompact(PyObject *range) {
173+
assert(PyRange_Check(range));
174+
rangeobject *r = (rangeobject*)range;
175+
assert(_PyLong_IsNonNegativeCompact((PyLongObject *)r->stop));
176+
return _PyLong_GetNonNegativeCompactValue((PyLongObject *)r->stop);
177+
}
178+
159179
PyDoc_STRVAR(range_doc,
160180
"range(stop) -> range object\n\
161181
range(start, stop[, step]) -> range object\n\

Python/bytecodes.c

Lines changed: 68 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -3122,11 +3122,20 @@ dummy_func(
31223122
index_or_null = PyStackRef_TagInt(0);
31233123
}
31243124
else {
3125-
PyObject *iter_o = PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable));
3126-
PyStackRef_CLOSE(iterable);
3127-
ERROR_IF(iter_o == NULL);
3128-
iter = PyStackRef_FromPyObjectSteal(iter_o);
3129-
index_or_null = PyStackRef_NULL;
3125+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iterable);
3126+
if (tp == &PyRange_Type && _PyRange_IsSimpleCompact(iter_o)) {
3127+
Py_ssize_t stop = _PyRange_GetStopIfCompact(iter_o);
3128+
PyStackRef_CLOSE(iterable);
3129+
iter = PyStackRef_TagInt(stop);
3130+
index_or_null = PyStackRef_TagInt(0);
3131+
}
3132+
else {
3133+
iter_o = PyObject_GetIter(iter_o);
3134+
PyStackRef_CLOSE(iterable);
3135+
ERROR_IF(iter_o == NULL);
3136+
iter = PyStackRef_FromPyObjectSteal(iter_o);
3137+
index_or_null = PyStackRef_NULL;
3138+
}
31303139
}
31313140
}
31323141

@@ -3204,17 +3213,33 @@ dummy_func(
32043213

32053214
replaced op(_FOR_ITER, (iter, null_or_index -- iter, null_or_index, next)) {
32063215
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
3207-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32083216
if (PyStackRef_IsTaggedInt(null_or_index)) {
3209-
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3210-
if (PyStackRef_IsNull(next)) {
3211-
null_or_index = PyStackRef_TagInt(-1);
3212-
JUMPBY(oparg + 1);
3213-
DISPATCH();
3217+
if (PyStackRef_IsTaggedInt(iter)) {
3218+
if (PyStackRef_Is(iter, null_or_index)) {
3219+
null_or_index = PyStackRef_TagInt(-1);
3220+
JUMPBY(oparg + 1);
3221+
DISPATCH();
3222+
3223+
}
3224+
next = PyStackRef_BoxInt(null_or_index);
3225+
if (PyStackRef_IsNull(next)) {
3226+
ERROR_NO_POP();
3227+
}
3228+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
3229+
}
3230+
else {
3231+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3232+
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3233+
if (PyStackRef_IsNull(next)) {
3234+
null_or_index = PyStackRef_TagInt(-1);
3235+
JUMPBY(oparg + 1);
3236+
DISPATCH();
3237+
}
3238+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
32143239
}
3215-
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
32163240
}
32173241
else {
3242+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32183243
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
32193244
if (next_o == NULL) {
32203245
if (_PyErr_Occurred(tstate)) {
@@ -3264,10 +3289,25 @@ dummy_func(
32643289
inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
32653290
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32663291
if (PyStackRef_IsTaggedInt(null_or_index)) {
3267-
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3268-
if (PyStackRef_IsNull(next)) {
3269-
JUMPBY(oparg + 1);
3270-
DISPATCH();
3292+
if (PyStackRef_IsTaggedInt(iter)) {
3293+
if (PyStackRef_Is(iter, null_or_index)) {
3294+
null_or_index = PyStackRef_TagInt(-1);
3295+
JUMPBY(oparg + 1);
3296+
DISPATCH();
3297+
3298+
}
3299+
next = PyStackRef_BoxInt(null_or_index);
3300+
if (PyStackRef_IsNull(next)) {
3301+
ERROR_NO_POP();
3302+
}
3303+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
3304+
}
3305+
else {
3306+
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3307+
if (PyStackRef_IsNull(next)) {
3308+
JUMPBY(oparg + 1);
3309+
DISPATCH();
3310+
}
32713311
}
32723312
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
32733313
}
@@ -3298,10 +3338,10 @@ dummy_func(
32983338

32993339

33003340
op(_ITER_CHECK_LIST, (iter, null_or_index -- iter, null_or_index)) {
3301-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3302-
EXIT_IF(Py_TYPE(iter_o) != &PyList_Type);
3341+
EXIT_IF(PyStackRef_TYPE(iter) != &PyList_Type);
33033342
assert(PyStackRef_IsTaggedInt(null_or_index));
33043343
#ifdef Py_GIL_DISABLED
3344+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
33053345
EXIT_IF(!_Py_IsOwnedByCurrentThread(iter_o) && !_PyObject_GC_IS_SHARED(iter_o));
33063346
#endif
33073347
}
@@ -3383,8 +3423,7 @@ dummy_func(
33833423
_ITER_NEXT_LIST;
33843424

33853425
op(_ITER_CHECK_TUPLE, (iter, null_or_index -- iter, null_or_index)) {
3386-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3387-
EXIT_IF(Py_TYPE(iter_o) != &PyTuple_Type);
3426+
EXIT_IF(PyStackRef_TYPE(iter) != &PyTuple_Type);
33883427
assert(PyStackRef_IsTaggedInt(null_or_index));
33893428
}
33903429

@@ -3424,21 +3463,11 @@ dummy_func(
34243463
_ITER_NEXT_TUPLE;
34253464

34263465
op(_ITER_CHECK_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3427-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3428-
EXIT_IF(Py_TYPE(r) != &PyRangeIter_Type);
3429-
#ifdef Py_GIL_DISABLED
3430-
EXIT_IF(!_PyObject_IsUniquelyReferenced((PyObject *)r));
3431-
#endif
3466+
EXIT_IF(!PyStackRef_IsTaggedInt(iter));
34323467
}
34333468

34343469
replaced op(_ITER_JUMP_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3435-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3436-
assert(Py_TYPE(r) == &PyRangeIter_Type);
3437-
#ifdef Py_GIL_DISABLED
3438-
assert(_PyObject_IsUniquelyReferenced((PyObject *)r));
3439-
#endif
3440-
STAT_INC(FOR_ITER, hit);
3441-
if (r->len <= 0) {
3470+
if (PyStackRef_Is(iter, null_or_index)) {
34423471
// Jump over END_FOR instruction.
34433472
JUMPBY(oparg + 1);
34443473
DISPATCH();
@@ -3447,24 +3476,15 @@ dummy_func(
34473476

34483477
// Only used by Tier 2
34493478
op(_GUARD_NOT_EXHAUSTED_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3450-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3451-
assert(Py_TYPE(r) == &PyRangeIter_Type);
3452-
EXIT_IF(r->len <= 0);
3479+
EXIT_IF(PyStackRef_Is(iter, null_or_index));
34533480
}
34543481

34553482
op(_ITER_NEXT_RANGE, (iter, null_or_index -- iter, null_or_index, next)) {
3456-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3457-
assert(Py_TYPE(r) == &PyRangeIter_Type);
3458-
#ifdef Py_GIL_DISABLED
3459-
assert(_PyObject_IsUniquelyReferenced((PyObject *)r));
3460-
#endif
3461-
assert(r->len > 0);
3462-
long value = r->start;
3463-
r->start = value + r->step;
3464-
r->len--;
3465-
PyObject *res = PyLong_FromLong(value);
3466-
ERROR_IF(res == NULL);
3467-
next = PyStackRef_FromPyObjectSteal(res);
3483+
next = PyStackRef_BoxInt(null_or_index);
3484+
if (PyStackRef_IsNull(next)) {
3485+
ERROR_NO_POP();
3486+
}
3487+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
34683488
}
34693489

34703490
macro(FOR_ITER_RANGE) =
@@ -3474,8 +3494,8 @@ dummy_func(
34743494
_ITER_NEXT_RANGE;
34753495

34763496
op(_FOR_ITER_GEN_FRAME, (iter, null -- iter, null, gen_frame: _PyInterpreterFrame*)) {
3497+
DEOPT_IF(PyStackRef_TYPE(iter) != &PyGen_Type);
34773498
PyGenObject *gen = (PyGenObject *)PyStackRef_AsPyObjectBorrow(iter);
3478-
DEOPT_IF(Py_TYPE(gen) != &PyGen_Type);
34793499
#ifdef Py_GIL_DISABLED
34803500
// Since generators can't be used by multiple threads anyway we
34813501
// don't need to deopt here, but this lets us work on making

0 commit comments

Comments
 (0)