gh-143414: Implement unique reference tracking for JIT, optimize unpacking of such tuples (GH-144300)

Co-authored-by: Ken Jin <kenjin4096@gmail.com>
This commit is contained in:
reiden 2026-03-23 00:57:23 +08:00 committed by GitHub
parent 1ceb1fb284
commit e36f8db7e5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1414 additions and 911 deletions

View file

@ -33,7 +33,7 @@
#include "pycore_sliceobject.h" // _PyBuildSlice_ConsumeRefs
#include "pycore_stackref.h"
#include "pycore_template.h" // _PyTemplate_Build()
#include "pycore_tuple.h" // _PyTuple_ITEMS()
#include "pycore_tuple.h" // _PyStolenTuple_Free(), _PyTuple_ITEMS()
#include "pycore_typeobject.h" // _PySuper_Lookup()
#include "pycore_dict.h"
@ -1739,6 +1739,25 @@ dummy_func(
PyStackRef_CLOSE(seq);
}
op(_UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE, (seq -- val1, val0)) {
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
}
op(_UNPACK_SEQUENCE_UNIQUE_THREE_TUPLE, (seq -- val2, val1, val0)) {
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
val2 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 2));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
}
macro(UNPACK_SEQUENCE_TUPLE) =
_GUARD_TOS_TUPLE + unused/1 + _UNPACK_SEQUENCE_TUPLE;
@ -1754,6 +1773,20 @@ dummy_func(
DECREF_INPUTS();
}
op(_UNPACK_SEQUENCE_UNIQUE_TUPLE, (seq -- values[oparg])) {
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
assert(PyTuple_CheckExact(seq_o));
assert(PyTuple_GET_SIZE(seq_o) == oparg);
assert(_PyObject_IsUniquelyReferenced(seq_o));
STAT_INC(UNPACK_SEQUENCE, hit);
PyObject **items = _PyTuple_ITEMS(seq_o);
for (int i = oparg; --i >= 0; ) {
*values++ = PyStackRef_FromPyObjectSteal(items[i]);
}
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
}
macro(UNPACK_SEQUENCE_LIST) =
_GUARD_TOS_LIST + unused/1 + _UNPACK_SEQUENCE_LIST;

View file

@ -7485,6 +7485,121 @@
break;
}
case _UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE_r02: {
CHECK_CURRENT_CACHED_VALUES(0);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
_PyStackRef seq;
_PyStackRef val1;
_PyStackRef val0;
seq = stack_pointer[-1];
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
_tos_cache1 = val0;
_tos_cache0 = val1;
SET_CURRENT_CACHED_VALUES(2);
stack_pointer += -1;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
break;
}
case _UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE_r12: {
CHECK_CURRENT_CACHED_VALUES(1);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
_PyStackRef seq;
_PyStackRef val1;
_PyStackRef val0;
_PyStackRef _stack_item_0 = _tos_cache0;
seq = _stack_item_0;
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
_tos_cache1 = val0;
_tos_cache0 = val1;
SET_CURRENT_CACHED_VALUES(2);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
break;
}
case _UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE_r23: {
CHECK_CURRENT_CACHED_VALUES(2);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
_PyStackRef seq;
_PyStackRef val1;
_PyStackRef val0;
_PyStackRef _stack_item_0 = _tos_cache0;
_PyStackRef _stack_item_1 = _tos_cache1;
seq = _stack_item_1;
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
_tos_cache2 = val0;
_tos_cache1 = val1;
_tos_cache0 = _stack_item_0;
SET_CURRENT_CACHED_VALUES(3);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
break;
}
case _UNPACK_SEQUENCE_UNIQUE_THREE_TUPLE_r03: {
CHECK_CURRENT_CACHED_VALUES(0);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
_PyStackRef seq;
_PyStackRef val2;
_PyStackRef val1;
_PyStackRef val0;
seq = stack_pointer[-1];
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
val2 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 2));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
_tos_cache2 = val0;
_tos_cache1 = val1;
_tos_cache0 = val2;
SET_CURRENT_CACHED_VALUES(3);
stack_pointer += -1;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
break;
}
case _UNPACK_SEQUENCE_UNIQUE_THREE_TUPLE_r13: {
CHECK_CURRENT_CACHED_VALUES(1);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
_PyStackRef seq;
_PyStackRef val2;
_PyStackRef val1;
_PyStackRef val0;
_PyStackRef _stack_item_0 = _tos_cache0;
seq = _stack_item_0;
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
STAT_INC(UNPACK_SEQUENCE, hit);
val0 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 0));
val1 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 1));
val2 = PyStackRef_FromPyObjectSteal(PyTuple_GET_ITEM(seq_o, 2));
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
_tos_cache2 = val0;
_tos_cache1 = val1;
_tos_cache0 = val2;
SET_CURRENT_CACHED_VALUES(3);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
break;
}
case _UNPACK_SEQUENCE_TUPLE_r10: {
CHECK_CURRENT_CACHED_VALUES(1);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
@ -7520,6 +7635,33 @@
break;
}
case _UNPACK_SEQUENCE_UNIQUE_TUPLE_r10: {
CHECK_CURRENT_CACHED_VALUES(1);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
_PyStackRef seq;
_PyStackRef *values;
_PyStackRef _stack_item_0 = _tos_cache0;
oparg = CURRENT_OPARG();
seq = _stack_item_0;
values = &stack_pointer[0];
PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
assert(PyTuple_CheckExact(seq_o));
assert(PyTuple_GET_SIZE(seq_o) == oparg);
assert(_PyObject_IsUniquelyReferenced(seq_o));
STAT_INC(UNPACK_SEQUENCE, hit);
PyObject **items = _PyTuple_ITEMS(seq_o);
for (int i = oparg; --i >= 0; ) {
*values++ = PyStackRef_FromPyObjectSteal(items[i]);
}
PyObject_GC_UnTrack(seq_o);
_PyStolenTuple_Free(seq_o);
SET_CURRENT_CACHED_VALUES(0);
stack_pointer += oparg;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());
break;
}
case _UNPACK_SEQUENCE_LIST_r10: {
CHECK_CURRENT_CACHED_VALUES(1);
assert(WITHIN_STACK_BOUNDS_IGNORING_CACHE());

View file

@ -93,7 +93,19 @@ dummy_func(void) {
if (sym_is_immortal(PyJitRef_Unwrap(value))) {
ADD_OP(_NOP, 0, 0);
}
value = PyJitRef_StripReferenceInfo(value);
value = PyJitRef_StripBorrowInfo(value);
}
op(_COPY_FREE_VARS, (--)) {
PyCodeObject *co = get_current_code_object(ctx);
if (co == NULL) {
ctx->done = true;
break;
}
int offset = co->co_nlocalsplus - oparg;
for (int i = 0; i < oparg; ++i) {
ctx->frame->locals[offset + i] = sym_new_not_null(ctx);
}
}
op(_LOAD_FAST_CHECK, (-- value)) {
@ -102,20 +114,24 @@ dummy_func(void) {
if (sym_is_null(value)) {
ctx->done = true;
}
assert(!PyJitRef_IsUnique(value));
}
op(_LOAD_FAST, (-- value)) {
value = GETLOCAL(oparg);
assert(!PyJitRef_IsUnique(value));
}
op(_LOAD_FAST_BORROW, (-- value)) {
value = PyJitRef_Borrow(GETLOCAL(oparg));
assert(!PyJitRef_IsUnique(value));
}
op(_LOAD_FAST_AND_CLEAR, (-- value)) {
value = GETLOCAL(oparg);
JitOptRef temp = sym_new_null(ctx);
GETLOCAL(oparg) = temp;
assert(!PyJitRef_IsUnique(value));
}
op(_STORE_ATTR_INSTANCE_VALUE, (offset/1, value, owner -- o)) {
@ -132,7 +148,7 @@ dummy_func(void) {
op(_SWAP_FAST, (value -- trash)) {
JitOptRef tmp = GETLOCAL(oparg);
GETLOCAL(oparg) = value;
GETLOCAL(oparg) = PyJitRef_RemoveUnique(value);
trash = tmp;
}
@ -713,6 +729,7 @@ dummy_func(void) {
op(_COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
assert(oparg > 0);
bottom = PyJitRef_RemoveUnique(bottom);
top = bottom;
}
@ -1311,6 +1328,7 @@ dummy_func(void) {
op(_BUILD_TUPLE, (values[oparg] -- tup)) {
tup = sym_new_tuple(ctx, oparg, values);
tup = PyJitRef_MakeUnique(tup);
}
op(_BUILD_LIST, (values[oparg] -- list)) {
@ -1339,11 +1357,20 @@ dummy_func(void) {
}
op(_UNPACK_SEQUENCE_TWO_TUPLE, (seq -- val1, val0)) {
if (PyJitRef_IsUnique(seq) && sym_tuple_length(seq) == 2) {
ADD_OP(_UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE, oparg, 0);
}
val0 = sym_tuple_getitem(ctx, seq, 0);
val1 = sym_tuple_getitem(ctx, seq, 1);
}
op(_UNPACK_SEQUENCE_TUPLE, (seq -- values[oparg])) {
if (PyJitRef_IsUnique(seq) && sym_tuple_length(seq) == 3) {
ADD_OP(_UNPACK_SEQUENCE_UNIQUE_THREE_TUPLE, oparg, 0);
}
else if (PyJitRef_IsUnique(seq) && sym_tuple_length(seq) == oparg) {
ADD_OP(_UNPACK_SEQUENCE_UNIQUE_TUPLE, oparg, 0);
}
for (int i = 0; i < oparg; i++) {
values[i] = sym_tuple_getitem(ctx, seq, oparg - i - 1);
}

View file

@ -31,6 +31,7 @@
if (sym_is_null(value)) {
ctx->done = true;
}
assert(!PyJitRef_IsUnique(value));
CHECK_STACK_BOUNDS(1);
stack_pointer[0] = value;
stack_pointer += 1;
@ -41,6 +42,7 @@
case _LOAD_FAST: {
JitOptRef value;
value = GETLOCAL(oparg);
assert(!PyJitRef_IsUnique(value));
CHECK_STACK_BOUNDS(1);
stack_pointer[0] = value;
stack_pointer += 1;
@ -51,6 +53,7 @@
case _LOAD_FAST_BORROW: {
JitOptRef value;
value = PyJitRef_Borrow(GETLOCAL(oparg));
assert(!PyJitRef_IsUnique(value));
CHECK_STACK_BOUNDS(1);
stack_pointer[0] = value;
stack_pointer += 1;
@ -63,6 +66,7 @@
value = GETLOCAL(oparg);
JitOptRef temp = sym_new_null(ctx);
GETLOCAL(oparg) = temp;
assert(!PyJitRef_IsUnique(value));
CHECK_STACK_BOUNDS(1);
stack_pointer[0] = value;
stack_pointer += 1;
@ -102,7 +106,7 @@
JitOptRef trash;
value = stack_pointer[-1];
JitOptRef tmp = GETLOCAL(oparg);
GETLOCAL(oparg) = value;
GETLOCAL(oparg) = PyJitRef_RemoveUnique(value);
trash = tmp;
stack_pointer[-1] = trash;
break;
@ -1321,7 +1325,7 @@
if (sym_is_immortal(PyJitRef_Unwrap(value))) {
ADD_OP(_NOP, 0, 0);
}
value = PyJitRef_StripReferenceInfo(value);
value = PyJitRef_StripBorrowInfo(value);
stack_pointer[-1] = value;
break;
}
@ -1487,6 +1491,9 @@
JitOptRef val1;
JitOptRef val0;
seq = stack_pointer[-1];
if (PyJitRef_IsUnique(seq) && sym_tuple_length(seq) == 2) {
ADD_OP(_UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE, oparg, 0);
}
val0 = sym_tuple_getitem(ctx, seq, 0);
val1 = sym_tuple_getitem(ctx, seq, 1);
CHECK_STACK_BOUNDS(1);
@ -1497,11 +1504,46 @@
break;
}
case _UNPACK_SEQUENCE_UNIQUE_TWO_TUPLE: {
JitOptRef val1;
JitOptRef val0;
val1 = sym_new_not_null(ctx);
val0 = sym_new_not_null(ctx);
CHECK_STACK_BOUNDS(1);
stack_pointer[-1] = val1;
stack_pointer[0] = val0;
stack_pointer += 1;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
break;
}
case _UNPACK_SEQUENCE_UNIQUE_THREE_TUPLE: {
JitOptRef val2;
JitOptRef val1;
JitOptRef val0;
val2 = sym_new_not_null(ctx);
val1 = sym_new_not_null(ctx);
val0 = sym_new_not_null(ctx);
CHECK_STACK_BOUNDS(2);
stack_pointer[-1] = val2;
stack_pointer[0] = val1;
stack_pointer[1] = val0;
stack_pointer += 2;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
break;
}
case _UNPACK_SEQUENCE_TUPLE: {
JitOptRef seq;
JitOptRef *values;
seq = stack_pointer[-1];
values = &stack_pointer[-1];
if (PyJitRef_IsUnique(seq) && sym_tuple_length(seq) == 3) {
ADD_OP(_UNPACK_SEQUENCE_UNIQUE_THREE_TUPLE, oparg, 0);
}
else if (PyJitRef_IsUnique(seq) && sym_tuple_length(seq) == oparg) {
ADD_OP(_UNPACK_SEQUENCE_UNIQUE_TUPLE, oparg, 0);
}
for (int i = 0; i < oparg; i++) {
values[i] = sym_tuple_getitem(ctx, seq, oparg - i - 1);
}
@ -1511,6 +1553,18 @@
break;
}
case _UNPACK_SEQUENCE_UNIQUE_TUPLE: {
JitOptRef *values;
values = &stack_pointer[-1];
for (int _i = oparg; --_i >= 0;) {
values[_i] = sym_new_not_null(ctx);
}
CHECK_STACK_BOUNDS(-1 + oparg);
stack_pointer += -1 + oparg;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
break;
}
case _UNPACK_SEQUENCE_LIST: {
JitOptRef *values;
values = &stack_pointer[-1];
@ -1763,6 +1817,15 @@
}
case _COPY_FREE_VARS: {
PyCodeObject *co = get_current_code_object(ctx);
if (co == NULL) {
ctx->done = true;
break;
}
int offset = co->co_nlocalsplus - oparg;
for (int i = 0; i < oparg; ++i) {
ctx->frame->locals[offset + i] = sym_new_not_null(ctx);
}
break;
}
@ -1801,6 +1864,7 @@
JitOptRef tup;
values = &stack_pointer[-oparg];
tup = sym_new_tuple(ctx, oparg, values);
tup = PyJitRef_MakeUnique(tup);
CHECK_STACK_BOUNDS(1 - oparg);
stack_pointer[-oparg] = tup;
stack_pointer += 1 - oparg;
@ -3793,8 +3857,10 @@
JitOptRef top;
bottom = stack_pointer[-1 - (oparg-1)];
assert(oparg > 0);
bottom = PyJitRef_RemoveUnique(bottom);
top = bottom;
CHECK_STACK_BOUNDS(1);
stack_pointer[-1 - (oparg-1)] = bottom;
stack_pointer[0] = top;
stack_pointer += 1;
ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);

View file

@ -1521,7 +1521,7 @@ _Py_uop_frame_new(
// Initialize with the initial state of all local variables
for (int i = 0; i < arg_len; i++) {
frame->locals[i] = args[i];
frame->locals[i] = PyJitRef_RemoveUnique(args[i]);
}
// If the args are known, then it's safe to just initialize