Skip to content

Commit 658a422

Browse files
markshannonmarkshannon
committed
Specialize FOR_ITER for ranges using tagged ints
1 parent 37757c6 commit 658a422

File tree

12 files changed

+255
-159
lines changed

12 files changed

+255
-159
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: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,8 @@ PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
276276

277277
#define PyStackRef_IsDeferredOrTaggedInt(ref) (((ref).bits & Py_TAG_REFCNT) != 0)
278278

279+
extern _PyStackRef PyStackRef_BoxInt(_PyStackRef i);
280+
279281
#ifdef Py_GIL_DISABLED
280282

281283
#define Py_TAG_DEFERRED Py_TAG_REFCNT

‎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

@@ -6872,3 +6873,19 @@ PyLongWriter_Finish(PyLongWriter *writer)
68726873

68736874
return (PyObject*)obj;
68746875
}
6876+
6877+
// Tagged int support
6878+
6879+
_PyStackRef
6880+
PyStackRef_BoxInt(_PyStackRef i)
6881+
{
6882+
assert((i.bits & Py_INT_TAG) == Py_INT_TAG);
6883+
intptr_t val = (intptr_t)i.bits;
6884+
val = Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, val, 2);
6885+
PyObject *boxed = PyLong_FromSsize_t(val);
6886+
if (boxed == NULL) {
6887+
return PyStackRef_NULL;
6888+
}
6889+
return PyStackRef_FromPyObjectSteal(boxed);
6890+
}
6891+

‎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: 67 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -3078,11 +3078,20 @@ dummy_func(
30783078
index_or_null = PyStackRef_TagInt(0);
30793079
}
30803080
else {
3081-
PyObject *iter_o = PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable));
3082-
PyStackRef_CLOSE(iterable);
3083-
ERROR_IF(iter_o == NULL);
3084-
iter = PyStackRef_FromPyObjectSteal(iter_o);
3085-
index_or_null = PyStackRef_NULL;
3081+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iterable);
3082+
if (tp == &PyRange_Type && _PyRange_IsSimpleCompact(iter_o)) {
3083+
Py_ssize_t stop = _PyRange_GetStopIfCompact(iter_o);
3084+
PyStackRef_CLOSE(iterable);
3085+
iter = PyStackRef_TagInt(stop);
3086+
index_or_null = PyStackRef_TagInt(0);
3087+
}
3088+
else {
3089+
iter_o = PyObject_GetIter(iter_o);
3090+
PyStackRef_CLOSE(iterable);
3091+
ERROR_IF(iter_o == NULL);
3092+
iter = PyStackRef_FromPyObjectSteal(iter_o);
3093+
index_or_null = PyStackRef_NULL;
3094+
}
30863095
}
30873096
}
30883097

@@ -3160,16 +3169,32 @@ dummy_func(
31603169

31613170
replaced op(_FOR_ITER, (iter, null_or_index -- iter, null_or_index, next)) {
31623171
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
3163-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
31643172
if (PyStackRef_IsTaggedInt(null_or_index)) {
3165-
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3166-
if (PyStackRef_IsNull(next)) {
3167-
JUMPBY(oparg + 1);
3168-
DISPATCH();
3173+
if (PyStackRef_IsTaggedInt(iter)) {
3174+
if (PyStackRef_Is(iter, null_or_index)) {
3175+
null_or_index = PyStackRef_TagInt(-1);
3176+
JUMPBY(oparg + 1);
3177+
DISPATCH();
3178+
3179+
}
3180+
next = PyStackRef_BoxInt(null_or_index);
3181+
if (PyStackRef_IsNull(next)) {
3182+
ERROR_NO_POP();
3183+
}
3184+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
3185+
}
3186+
else {
3187+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
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);
31693194
}
3170-
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
31713195
}
31723196
else {
3197+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
31733198
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
31743199
if (next_o == NULL) {
31753200
if (_PyErr_Occurred(tstate)) {
@@ -3219,10 +3244,25 @@ dummy_func(
32193244
inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
32203245
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32213246
if (PyStackRef_IsTaggedInt(null_or_index)) {
3222-
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3223-
if (PyStackRef_IsNull(next)) {
3224-
JUMPBY(oparg + 1);
3225-
DISPATCH();
3247+
if (PyStackRef_IsTaggedInt(iter)) {
3248+
if (PyStackRef_Is(iter, null_or_index)) {
3249+
null_or_index = PyStackRef_TagInt(-1);
3250+
JUMPBY(oparg + 1);
3251+
DISPATCH();
3252+
3253+
}
3254+
next = PyStackRef_BoxInt(null_or_index);
3255+
if (PyStackRef_IsNull(next)) {
3256+
ERROR_NO_POP();
3257+
}
3258+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
3259+
}
3260+
else {
3261+
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
3262+
if (PyStackRef_IsNull(next)) {
3263+
JUMPBY(oparg + 1);
3264+
DISPATCH();
3265+
}
32263266
}
32273267
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
32283268
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
@@ -3254,10 +3294,10 @@ dummy_func(
32543294

32553295

32563296
op(_ITER_CHECK_LIST, (iter, null_or_index -- iter, null_or_index)) {
3257-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3258-
EXIT_IF(Py_TYPE(iter_o) != &PyList_Type);
3297+
EXIT_IF(PyStackRef_TYPE(iter) != &PyList_Type);
32593298
assert(PyStackRef_IsTaggedInt(null_or_index));
32603299
#ifdef Py_GIL_DISABLED
3300+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32613301
EXIT_IF(!_Py_IsOwnedByCurrentThread(iter_o) && !_PyObject_GC_IS_SHARED(iter_o));
32623302
#endif
32633303
}
@@ -3339,8 +3379,7 @@ dummy_func(
33393379
_ITER_NEXT_LIST;
33403380

33413381
op(_ITER_CHECK_TUPLE, (iter, null_or_index -- iter, null_or_index)) {
3342-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
3343-
EXIT_IF(Py_TYPE(iter_o) != &PyTuple_Type);
3382+
EXIT_IF(PyStackRef_TYPE(iter) != &PyTuple_Type);
33443383
assert(PyStackRef_IsTaggedInt(null_or_index));
33453384
}
33463385

@@ -3380,21 +3419,11 @@ dummy_func(
33803419
_ITER_NEXT_TUPLE;
33813420

33823421
op(_ITER_CHECK_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3383-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3384-
EXIT_IF(Py_TYPE(r) != &PyRangeIter_Type);
3385-
#ifdef Py_GIL_DISABLED
3386-
EXIT_IF(!_PyObject_IsUniquelyReferenced((PyObject *)r));
3387-
#endif
3422+
EXIT_IF(!PyStackRef_IsTaggedInt(iter));
33883423
}
33893424

33903425
replaced op(_ITER_JUMP_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3391-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3392-
assert(Py_TYPE(r) == &PyRangeIter_Type);
3393-
#ifdef Py_GIL_DISABLED
3394-
assert(_PyObject_IsUniquelyReferenced((PyObject *)r));
3395-
#endif
3396-
STAT_INC(FOR_ITER, hit);
3397-
if (r->len <= 0) {
3426+
if (PyStackRef_Is(iter, null_or_index)) {
33983427
// Jump over END_FOR instruction.
33993428
JUMPBY(oparg + 1);
34003429
DISPATCH();
@@ -3403,24 +3432,15 @@ dummy_func(
34033432

34043433
// Only used by Tier 2
34053434
op(_GUARD_NOT_EXHAUSTED_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3406-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3407-
assert(Py_TYPE(r) == &PyRangeIter_Type);
3408-
EXIT_IF(r->len <= 0);
3435+
EXIT_IF(PyStackRef_Is(iter, null_or_index));
34093436
}
34103437

34113438
op(_ITER_NEXT_RANGE, (iter, null_or_index -- iter, null_or_index, next)) {
3412-
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
3413-
assert(Py_TYPE(r) == &PyRangeIter_Type);
3414-
#ifdef Py_GIL_DISABLED
3415-
assert(_PyObject_IsUniquelyReferenced((PyObject *)r));
3416-
#endif
3417-
assert(r->len > 0);
3418-
long value = r->start;
3419-
r->start = value + r->step;
3420-
r->len--;
3421-
PyObject *res = PyLong_FromLong(value);
3422-
ERROR_IF(res == NULL);
3423-
next = PyStackRef_FromPyObjectSteal(res);
3439+
next = PyStackRef_BoxInt(null_or_index);
3440+
if (PyStackRef_IsNull(next)) {
3441+
ERROR_NO_POP();
3442+
}
3443+
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
34243444
}
34253445

34263446
macro(FOR_ITER_RANGE) =
@@ -3430,8 +3450,8 @@ dummy_func(
34303450
_ITER_NEXT_RANGE;
34313451

34323452
op(_FOR_ITER_GEN_FRAME, (iter, null -- iter, null, gen_frame: _PyInterpreterFrame*)) {
3453+
DEOPT_IF(PyStackRef_TYPE(iter) != &PyGen_Type);
34333454
PyGenObject *gen = (PyGenObject *)PyStackRef_AsPyObjectBorrow(iter);
3434-
DEOPT_IF(Py_TYPE(gen) != &PyGen_Type);
34353455
#ifdef Py_GIL_DISABLED
34363456
// Since generators can't be used by multiple threads anyway we
34373457
// don't need to deopt here, but this lets us work on making

0 commit comments

Comments
 (0)