Skip to content

gh-144145: Cleanups for object property tracking in JIT optimizer #144366

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

Closed
Fidget-Spinner wants to merge 4 commits into from
Closed

gh-144145: Cleanups for object property tracking in JIT optimizer #144366

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
37f0530
Optimize NULL checks away completely in property tracking
Fidget-Spinner Jan 31, 2026
9e64c07
Add a test
Fidget-Spinner Jan 31, 2026
3d79c64
remove printf
Fidget-Spinner Jan 31, 2026
8361602
Update ASCII art, address review
Fidget-Spinner Feb 1, 2026
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_optimizer_types.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,8 @@ typedef struct {
typedef struct _jit_opt_descr {
uint8_t tag;
uint8_t num_descrs;
uint16_t last_modified_index; // Index in out_buffer when this object was last modified
// Index in out_buffer when this object was last escaped
uint16_t last_escape_index;
uint32_t type_version;
JitOptDescrMapping *descrs;
} JitOptDescrObject;
Expand Down
4 changes: 2 additions & 2 deletions Include/internal/pycore_uop_metadata.h
View file
Open in desktop

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

2 changes: 0 additions & 2 deletions Modules/_testinternalcapi/test_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.

6 changes: 2 additions & 4 deletions Python/bytecodes.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2677,8 +2677,7 @@ dummy_func(
STAT_INC(STORE_ATTR, hit);
assert(_PyObject_GetManagedDict(owner_o) == NULL);
PyObject **value_ptr = (PyObject**)(((char *)owner_o) + offset);
PyObject *old_value = *value_ptr;
DEOPT_IF(old_value != NULL);
assert(*value_ptr == NULL);
FT_ATOMIC_STORE_PTR_RELEASE(*value_ptr, PyStackRef_AsPyObjectSteal(value));
PyDictValues *values = _PyObject_InlineValues(owner_o);
Py_ssize_t index = value_ptr - values->values;
Expand Down Expand Up @@ -2756,8 +2755,7 @@ dummy_func(
DEOPT_IF(!LOCK_OBJECT(owner_o));
char *addr = (char *)owner_o + index;
STAT_INC(STORE_ATTR, hit);
PyObject *old_value = *(PyObject **)addr;
DEOPT_IF(old_value != NULL);
assert(*(PyObject **)addr == NULL);
FT_ATOMIC_STORE_PTR_RELEASE(*(PyObject **)addr, PyStackRef_AsPyObjectSteal(value));
UNLOCK_OBJECT(owner_o);
INPUTS_DEAD();
Expand Down
70 changes: 8 additions & 62 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.

2 changes: 0 additions & 2 deletions Python/generated_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.

4 changes: 2 additions & 2 deletions Python/optimizer_analysis.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -526,9 +526,9 @@ optimize_uops(
if (ctx->out_buffer.next == out_ptr) {
*(ctx->out_buffer.next++) = *this_instr;
}
// Track escapes - but skip when from init shim frame, since self hasn't escaped yet
bool is_init_shim = CURRENT_FRAME_IS_INIT_SHIM();
Copy link
Copy Markdown
Member

@cocolato cocolato Feb 1, 2026

Choose a reason for hiding this comment

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

Suggested change
bool is_init_shim = CURRENT_FRAME_IS_INIT_SHIM();

Perhaps we can remove the is_init_shim now.

Copy link
Copy Markdown
Member Author

@Fidget-Spinner Fidget-Spinner Feb 1, 2026

Choose a reason for hiding this comment

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

I think it's fine to just use the variable for now.

if ((_PyUop_Flags[out_ptr->opcode] & HAS_ESCAPES_FLAG) && !is_init_shim)
// Track escapes
if (_PyUop_Flags[out_ptr->opcode] & HAS_ESCAPES_FLAG)
{
ctx->last_escape_index = uop_buffer_length(&ctx->out_buffer) - 1;
}
Expand Down
96 changes: 50 additions & 46 deletions Python/optimizer_symbols.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -980,7 +980,7 @@ _Py_uop_sym_new_descr_object(JitOptContext *ctx, unsigned int type_version)
res->descr.num_descrs = 0;
res->descr.descrs = NULL;
res->descr.type_version = type_version;
res->descr.last_modified_index = uop_buffer_length(&ctx->out_buffer);
res->descr.last_escape_index = uop_buffer_length(&ctx->out_buffer);
return PyJitRef_Wrap(res);
}

Expand All @@ -995,7 +995,7 @@ _Py_uop_sym_get_attr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index)
// 1. Symbol has mappings allocated
// 2. No escape has occurred since last modification (state is fresh)
if (sym->descr.descrs != NULL &&
sym->descr.last_modified_index >= ctx->last_escape_index)
sym->descr.last_escape_index >= ctx->last_escape_index)
{
for (int i = 0; i < sym->descr.num_descrs; i++) {
if (sym->descr.descrs[i].slot_index == slot_index) {
Expand Down Expand Up @@ -1026,57 +1026,55 @@ JitOptRef
_Py_uop_sym_set_attr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index, JitOptRef value)
{
JitOptSymbol *sym = PyJitRef_Unwrap(ref);
int curr_index = uop_buffer_length(&ctx->out_buffer);

switch (sym->tag) {
case JIT_SYM_DESCR_TAG:
break;
default:
return _Py_uop_sym_new_not_null(ctx);
}

// Check escape
if (sym->descr.last_modified_index < ctx->last_escape_index) {
sym->descr.num_descrs = 0;
}
case JIT_SYM_DESCR_TAG: {
// Check escape
if (sym->descr.last_escape_index < ctx->last_escape_index) {
sym->descr.num_descrs = 0;
return _Py_uop_sym_new_unknown(ctx);
Copy link
Copy Markdown
Member

@cocolato cocolato Feb 1, 2026

Choose a reason for hiding this comment

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

Suggested change
return _Py_uop_sym_new_unknown(ctx);
sym->descr.num_descrs = 0;
sym->descr.last_escape_index = uop_buffer_length(&ctx->out_buffer);

Can we update the last_escape_index here to prevent subsequent set_attr calls from repeatedly clearing num_descrs?

Copy link
Copy Markdown
Member Author

@Fidget-Spinner Fidget-Spinner Feb 1, 2026

Choose a reason for hiding this comment

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

Done.

}

// Get old value before updating
JitOptRef old_value = PyJitRef_NULL;
if (sym->descr.descrs != NULL) {
for (int i = 0; i < sym->descr.num_descrs; i++) {
if (sym->descr.descrs[i].slot_index == slot_index) {
old_value = PyJitRef_Wrap(allocation_base(ctx) + sym->descr.descrs[i].symbol);
break;
// Get old value before updating
JitOptRef old_value = _Py_uop_sym_new_unknown(ctx);
if (sym->descr.descrs != NULL) {
for (int i = 0; i < sym->descr.num_descrs; i++) {
if (sym->descr.descrs[i].slot_index == slot_index) {
old_value = PyJitRef_Wrap(allocation_base(ctx) + sym->descr.descrs[i].symbol);
break;
}
}
}
}
}

// Update the last modified timestamp
sym->descr.last_modified_index = curr_index;
// Check if have arena space allocated
if (sym->descr.descrs == NULL) {
sym->descr.descrs = descr_arena_alloc(ctx);
if (sym->descr.descrs == NULL) {
ctx->done = true;
ctx->out_of_space = true;
return _Py_uop_sym_new_null(ctx);
}
}
// Check if the slot already exists
for (int i = 0; i < sym->descr.num_descrs; i++) {
if (sym->descr.descrs[i].slot_index == slot_index) {
sym->descr.descrs[i].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
assert(!_Py_uop_sym_is_null(old_value));
return old_value;
}
}
// Add new mapping if there's space
if (sym->descr.num_descrs < MAX_SYMBOLIC_DESCR_SIZE) {
int idx = sym->descr.num_descrs++;
sym->descr.descrs[idx].slot_index = slot_index;
sym->descr.descrs[idx].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
}

// Check if have arena space allocated
if (sym->descr.descrs == NULL) {
sym->descr.descrs = descr_arena_alloc(ctx);
if (sym->descr.descrs == NULL) {
return PyJitRef_IsNull(old_value) ? _Py_uop_sym_new_not_null(ctx) : old_value;
// No value there. Objects are initialized to NULL, so it's safe.
return _Py_uop_sym_new_null(ctx);
}
default:
return _Py_uop_sym_new_unknown(ctx);
}
// Check if the slot already exists
for (int i = 0; i < sym->descr.num_descrs; i++) {
if (sym->descr.descrs[i].slot_index == slot_index) {
sym->descr.descrs[i].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
assert(!PyJitRef_IsNull(old_value));
return old_value;
}
}
// Add new mapping if there's space
if (sym->descr.num_descrs < MAX_SYMBOLIC_DESCR_SIZE) {
int idx = sym->descr.num_descrs++;
sym->descr.descrs[idx].slot_index = slot_index;
sym->descr.descrs[idx].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
}

return PyJitRef_IsNull(old_value) ? PyJitRef_Borrow(_Py_uop_sym_new_null(ctx)) : old_value;
}

// 0 on success, -1 on error.
Expand Down Expand Up @@ -1698,6 +1696,12 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored))
retrieved = _Py_uop_sym_get_attr(ctx, descr_obj, 0);
TEST_PREDICATE(_Py_uop_sym_get_const(ctx, retrieved) == val_43,
"descr getattr(0) changed unexpectedly");
// Test setattr with escape
ctx->last_escape_index = INT_MAX;
retrieved = _Py_uop_sym_set_attr(ctx, descr_obj, 1, slot_val3);
TEST_PREDICATE(PyJitRef_Unwrap(retrieved)->tag == JIT_SYM_UNKNOWN_TAG,
"descr setattr should be unknown after escaping");
ctx->last_escape_index = 0;

// Test escape invalidation
JitOptRef descr_obj3 = _Py_uop_sym_new_descr_object(ctx, 100);
Expand Down
8 changes: 7 additions & 1 deletion Tools/cases_generator/analyzer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -710,8 +710,14 @@ def has_error_without_pop(op: parser.CodeDef) -> bool:
"_Py_set_eval_breaker_bit",
"trigger_backoff_counter",
"_PyThreadState_PopCStackRefSteal",
"_PyFrame_PushTrampolineUnchecked",
)

KNOWN_NON_ESCAPING_UOPS = {
# This only calls escaping functions in the error path.
"_CREATE_INIT_FRAME",
}


def check_escaping_calls(instr: parser.CodeDef, escapes: dict[SimpleStmt, EscapingCall]) -> None:
error: lexer.Token | None = None
Expand Down Expand Up @@ -965,7 +971,7 @@ def compute_properties(op: parser.CodeDef) -> Properties:
)
error_with_pop = has_error_with_pop(op)
error_without_pop = has_error_without_pop(op)
escapes = stmt_escapes(op.block)
escapes = not (isinstance(op, parser.InstDef) and op.name in KNOWN_NON_ESCAPING_UOPS) and stmt_escapes(op.block)
pure = False if isinstance(op, parser.LabelDef) else "pure" in op.annotations
no_save_ip = False if isinstance(op, parser.LabelDef) else "no_save_ip" in op.annotations
unpredictable, branches_seen = stmt_has_jump_on_unpredictable_path(op.block, 0)
Expand Down
Loading