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

reidenong · 2026-01-28T12:57:13Z

Implements unique reference tracking to references in the JIT Optimizer.
Uses the unique reference tracking to eliminate unnecessary reference counting for uniquely referenced tuples. ex:

x, y, z = f()

Some notes:

Only implements the unique reference tracking (and relevant optimization) for tuples in _BUILD_TUPLE, to be extended to other objects in the future
Objects are deemed as not uniquely referenced when they are made aliases or duplicated, in uops of _LOAD_FAST_* and _COPY.

Would appreciate any feedback.
Thanks

Issue: JIT: Implement unique reference tracking in Tier 2 for reference count optimizations #143414

Python/optimizer_bytecodes.c

reidenong · 2026-01-28T18:37:52Z

Discussed with @Fidget-Spinner, will reopen the PR after implementing PyStackRef_CLOSE_SPECIALIZED for _UNPACK_SEQUENCE_UNIQUE_TUPLE

Fidget-Spinner · 2026-01-29T17:03:08Z

Windows Ci failure looks possibly related, can you please look into it?

Fidget-Spinner

This is excellent. Just 2 comments, then I think it's safe to merge.

Include/internal/pycore_optimizer.h

Fidget-Spinner · 2026-01-31T17:44:59Z

Python/optimizer_bytecodes.c


    op(_COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
        assert(oparg > 0);
+        bottom = PyJitRef_IsUnique(bottom) ? PyJitRef_StripReferenceInfo(bottom) : bottom;


We only need to strip unique info here, not the borrow info too. This way we still get refcount elimination with the borrow.

I think this is consistent with the code?

If the ref is unique, we remove the unique tag, otherwise we leave it as is. Since a ref cannot be both unique and borrowed at the same time, borrowed/invalid refs will be maintained.

Sounds good, thanks!

Python/bytecodes.c

markshannon

We need to be very careful about tracking the uniqueness of a reference.
Uniqueness of a reference is a property of the whole VM, if we miss any operation that can duplicate a reference, we could introduce a serious bug.

With that in mind, I'd recommend only tracking references that remain on the stack, as it is too easy to make new references to objects stored in local variables:

Through away the uniqueness in STORE_FAST
Assert that references are not unique in LOAD_FAST[_BORROW]

Other than that, I think the analysis is sound.

This could be considerable more efficient, see comments below:

markshannon · 2026-02-02T12:27:36Z

Python/bytecodes.c

            DECREF_INPUTS();
        }

+        op(_UNPACK_SEQUENCE_UNIQUE_TUPLE, (seq -- values[oparg])) {


This is might end up being less efficient than _UNPACK_SEQUENCE_TWO_TUPLE due to the looping and having to spill the values to memory.

I think it is worth making two variants of this for 2 and 3 length tuples.

markshannon · 2026-02-02T12:29:05Z

Python/bytecodes.c

+            PyObject **items = _PyTuple_ITEMS(seq_o);
+            for (int i = oparg; --i >= 0; ) {
+                *values++ = PyStackRef_FromPyObjectSteal(items[i]);
+                items[i] = NULL;


If we free the tuple directly, we don't need to NULL out the fields.

markshannon · 2026-02-02T12:33:36Z

Python/bytecodes.c

+            PyObject *seq_o = PyStackRef_AsPyObjectBorrow(seq);
+            assert(PyTuple_CheckExact(seq_o));
+            assert(PyTuple_GET_SIZE(seq_o) == oparg);
+            DEOPT_IF(!_PyObject_IsUniquelyReferenced(seq_o));


Suggested change

DEOPT_IF(!_PyObject_IsUniquelyReferenced(seq_o));

assert(_PyObject_IsUniquelyReferenced(seq_o));

If our analysis is correct, we don't need a runtime check.

markshannon · 2026-02-02T12:34:39Z

Python/bytecodes.c

+                *values++ = PyStackRef_FromPyObjectSteal(items[i]);
+                items[i] = NULL;
+            }
+            PyStackRef_CLOSE_SPECIALIZED(seq, _PyTuple_EmptyExactDealloc);


Suggested change

PyStackRef_CLOSE_SPECIALIZED(seq, _PyTuple_EmptyExactDealloc);

PyObject_GC_UnTrack(seq_o);

_PyTuple_Free(seq_o);

markshannon · 2026-02-02T12:40:01Z

Objects/tupleobject.c

 /* Methods */

+void
+_PyTuple_EmptyExactDealloc(PyObject *obj)


Rather than a deallocing the tuple, we just need to free it.

Something like _PyTuple_Free.
Just assert that Py_SIZE(op) != 0 and free the tuple.

markshannon · 2026-02-02T12:42:16Z

Python/bytecodes.c

        }

+        op(_UNPACK_SEQUENCE_UNIQUE_TUPLE, (seq -- values[oparg])) {
+            PyObject *seq_o = PyStackRef_AsPyObjectBorrow(seq);


Suggested change

PyObject *seq_o = PyStackRef_AsPyObjectBorrow(seq);

PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);

We want to free the tuple directly, so we need to steal the reference.

bedevere-app · 2026-02-02T12:45:33Z

A Python core developer has requested some changes be made to your pull request before we can consider merging it. If you could please address their requests along with any other requests in other reviews from core developers that would be appreciated.

Once you have made the requested changes, please leave a comment on this pull request containing the phrase I have made the requested changes; please review again. I will then notify any core developers who have left a review that you're ready for them to take another look at this pull request.

And if you don't make the requested changes, you will be put in the comfy chair!

Add unique reference tracking to JIT, optimize tuple unpacking

e4b7f51

reidenong requested review from Fidget-Spinner, markshannon, savannahostrowski and tomasr8 as code owners January 28, 2026 12:57

bedevere-app bot mentioned this pull request Jan 28, 2026

JIT: Implement unique reference tracking in Tier 2 for reference count optimizations #143414

Open

bedevere-app bot added the awaiting review label Jan 28, 2026

reidenong and others added 2 commits January 28, 2026 20:58

Merge branch 'main' into branch-unique-reference-tracking

b32a8e2

Regenerate cases

7aeac0e

cocolato reviewed Jan 28, 2026

View reviewed changes

Python/optimizer_bytecodes.c Outdated Show resolved Hide resolved

Fix unique tuple unpacking checks

f8b7e4e

reidenong marked this pull request as draft January 28, 2026 18:38

bedevere-app bot removed the awaiting review label Jan 28, 2026

reidenong added 2 commits January 29, 2026 18:47

Implement _PyTuple_ExactDealloc

37d7c4b

Rename symbol

7df2176

reidenong marked this pull request as ready for review January 29, 2026 11:04

bedevere-app bot added the awaiting review label Jan 29, 2026

reidenong and others added 4 commits January 30, 2026 10:59

Merge branch 'main' into branch-unique-reference-tracking

6f66c56

Merge branch 'main' into branch-unique-reference-tracking

b251825

Regenerate cases

bcd5654

Add alignment checking

acfa737

reidenong marked this pull request as draft January 31, 2026 06:39

bedevere-app bot removed the awaiting review label Jan 31, 2026

reidenong and others added 2 commits January 31, 2026 23:11

Change JIT reference tags to 2 bit scheme

a0fd9e2

📜🤖 Added by blurb_it.

645776a

reidenong marked this pull request as ready for review January 31, 2026 16:36

bedevere-app bot added the awaiting review label Jan 31, 2026

Fidget-Spinner reviewed Jan 31, 2026

View reviewed changes

Fix style

cdb4a05

Fidget-Spinner reviewed Feb 2, 2026

View reviewed changes

Python/bytecodes.c Outdated Show resolved Hide resolved

reidenong added 2 commits February 2, 2026 17:49

Remove unnecesssary deopt

275732b

Regenerate cases

dbde059

markshannon requested changes Feb 2, 2026

View reviewed changes

bedevere-app bot added awaiting changes and removed awaiting review labels Feb 2, 2026

	DEOPT_IF(!_PyObject_IsUniquelyReferenced(seq_o));
	assert(_PyObject_IsUniquelyReferenced(seq_o));

	PyStackRef_CLOSE_SPECIALIZED(seq, _PyTuple_EmptyExactDealloc);
	PyObject_GC_UnTrack(seq_o);
	_PyTuple_Free(seq_o);

	PyObject *seq_o = PyStackRef_AsPyObjectBorrow(seq);
	PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);

Uh oh!

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

Are you sure you want to change the base?

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

Conversation

reidenong commented Jan 28, 2026 • edited by bedevere-app bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

reidenong commented Jan 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fidget-Spinner commented Jan 29, 2026

Uh oh!

Fidget-Spinner left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

markshannon left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

markshannon Feb 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bedevere-app bot commented Feb 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

reidenong commented Jan 28, 2026 •

edited by bedevere-app bot

Loading

reidenong commented Jan 28, 2026 •

edited

Loading

markshannon Feb 2, 2026 •

edited

Loading