gh-143414: Implement unique reference tracking for JIT, optimize unpacking of such tuples by reidenong · Pull Request #144300 · python/cpython

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!

bedevere-app · 2026-02-05T13:30:15Z

Thanks for making the requested changes!

@markshannon: please review the changes made to this pull request.

markshannon

I think we should be very conservative in the analysis, and only track references on the evaluation stack. It simplifies the code a bit, and means we don't need to worry about escapes.

FTR, this optimizations is going to be rendered obsolete by #120619 which will eliminate many temporary tuples entirely.
BUILD_TUPLE 2; _RETURN_VALUE; UNPACK 2 -> _RETURN_PAIR

markshannon · 2026-02-06T11:02:09Z

Python/optimizer_bytecodes.c

    }

    op(_UNPACK_SEQUENCE_TWO_TUPLE, (seq -- val1, val0)) {
+        if (PyJitRef_IsUnique(seq)) {


Do the length check here, and not in the uop at runtime. If we have tracked the uniqueness of the reference, we most likely know the size of the tuple

markshannon · 2026-02-06T11:09:32Z

Python/bytecodes.c

+            assert(oparg == 2);
+            PyObject *seq_o = PyStackRef_AsPyObjectSteal(seq);
+            assert(PyTuple_CheckExact(seq_o));
+            DEOPT_IF(PyTuple_GET_SIZE(seq_o) != 2);


Do this check in the optimizer

markshannon · 2026-02-06T11:09:39Z

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));


If we only track references on the evaluation stack, then there is no need for this check as the reference cannot escape.

As soon as a value is stored in a local, it is unlikely to only have one reference to it when we actually use it.
Only tracking references on the stack means we don't need to worry about the various forms of LOAD_FAST, just STORE_FAST.

bedevere-app · 2026-02-06T11:15:48Z

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.

reidenong · 2026-02-06T14:28:53Z

Thanks for the explanation on only tracking references on the stack. Since we remove uniqueness on STORE_FAST, I'm wondering if it is still necessary to perform PyJitRef_RemoveUnique on each LOAD_FAST*?

I have made the requested changes; please review again.

bedevere-app · 2026-02-06T14:28:59Z

Thanks for making the requested changes!

@markshannon: please review the changes made to this pull request.

Fidget-Spinner · 2026-02-06T19:05:08Z

wondering if it is still necessary to perform PyJitRef_RemoveUnique on each LOAD_FAST*?

Probably not.

reidenong · 2026-02-20T08:59:37Z

@markshannon Hi hope you're well, I’ve implemented the suggested changes, could you please let me know if everything looks good, or if there’s anything further I should adjust? Specifically I suspect PyJitRef_RemoveUnique on LOAD_FAST* is redundant but would like some confirmation on the matter.

Thanks!

markshannon

Just two places where uniqueness removals should be asserts, as the ref should not be unique at that point. Otherwise, this looks good.

markshannon · 2026-02-20T10:59:42Z

Python/optimizer_bytecodes.c

        if (sym_is_null(value)) {
            ctx->done = true;
        }
+        value = PyJitRef_RemoveUnique(value);


Suggested change

value = PyJitRef_RemoveUnique(value);

assert(!PyJitRef_IsUnique(value));

Any ref stored in a local should not be treated as unique

markshannon · 2026-02-20T10:59:57Z

Python/optimizer_bytecodes.c


    op(_LOAD_FAST, (-- value)) {
        value = GETLOCAL(oparg);
+        value = PyJitRef_RemoveUnique(value);


bedevere-app · 2026-02-20T11:06: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.

reidenong · 2026-02-21T05:18:05Z

I have made the requested changes; please review again. Thanks!

bedevere-app · 2026-02-21T05:18:11Z

Thanks for making the requested changes!

@markshannon: please review the changes made to this pull request.

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

bedevere-app bot requested a review from markshannon February 5, 2026 13:30

markshannon requested changes Feb 6, 2026

View reviewed changes

bedevere-app bot removed the awaiting change review label Feb 6, 2026

bedevere-app bot added the awaiting changes label Feb 6, 2026

reidenong added 4 commits February 6, 2026 21:20

Move length checks to JIT

fb1d227

Refactor JIT reference remove unique

c7e53fb

Remove JIT reference uniqueness on store

6851962

Remove unnecessary DEOPT

8c8e0a9

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

bedevere-app bot requested a review from markshannon February 6, 2026 14:29

markshannon requested changes Feb 20, 2026

View reviewed changes

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

reidenong and others added 4 commits February 21, 2026 12:43

Remove redundant make unique ops

a6fffdc

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

78978aa

Regen cases

0b07ac4

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

14e5d3d

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

bedevere-app bot requested a review from markshannon February 21, 2026 05:18

	value = PyJitRef_RemoveUnique(value);
	assert(!PyJitRef_IsUnique(value));

Uh oh!

Comments

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!

bedevere-app bot commented Feb 5, 2026

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!

bedevere-app bot commented Feb 6, 2026

Uh oh!

reidenong commented Feb 6, 2026

Uh oh!

bedevere-app bot commented Feb 6, 2026

Uh oh!

Fidget-Spinner commented Feb 6, 2026

Uh oh!

reidenong commented Feb 20, 2026

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!

bedevere-app bot commented Feb 20, 2026

Uh oh!

reidenong commented Feb 21, 2026

Uh oh!

bedevere-app bot commented Feb 21, 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