gh-130704: Strength reduce `LOAD_FAST{_LOAD_FAST}` #130708

mpage · 2025-02-28T22:00:43Z

This PR eliminates most reference counting overhead for references pushed onto the operand stack using LOAD_FAST{_LOAD_FAST} when we can be sure that the reference in the frame outlives the reference that is pushed onto the operand stack. Instructions that meet this criteria are replaced with new variants (LOAD_FAST_BORROW{_LOAD_FAST_BORROW}) that push appropriately tagged borrowed references.

Performance on the benchmark suite looks good:

A ~3% improvement on the free-threaded build. I think this might actually be higher and will continue investigating, but I wanted to put this up for review.
A ~2.7% improvement on the default build.

This approach looks like its quite effective at optimization too, at least on the benchmark suite. Roughly 97% of LOAD_FAST{_LOAD_FAST} instructions are optimized according to pystats. Note that these stats were collected using fastbench, so may not match those collected using pyperformance exactly.

The main pieces of the PR are:

New bytecodes

This adds two new bytecode instructions: LOAD_FAST_BORROW and its superinstruction form, LOAD_FAST_BORROW_LOAD_FAST_BORROW.

A new optimization pass

This adds a new optimization pass, optimize_load_fast, to the bytecode compiler that identifies and optimizes eligible instructions. Please read the detailed comment in flowgraph.c for a description of how it works.

Runtime support changes

A new function, PyStackRef_Borrow, was added to the stackref API. It creates a new stackref from an existing stackref without incrementing the reference count on the underlying object.

There are a few places in the runtime where we need to convert borrowed references into owned references:

When a frame escapes into the heap (i.e. when it is copied into a generator or when its materialized and unwound).
When a reference flows up the call stack (i.e. in RETURN_VALUE or YIELD_VALUE).
When someone destroys a reference to the frame "out of band" by poking at f_locals. We place the old reference into a tuple owned by the frame object.

The default build also required:

Adding support for stackrefs in the GC.
Removing reuse of stackrefs in PyFloat_FromDoubleConsumeInputs.

Issue: Optimize reference counting overhead of LOAD_FAST variants #130704

📚 Documentation preview 📚: https://cpython-previews--130708.org.readthedocs.build/

derp

Ref will be 2 if borrowed

…r frame

Otherwise, it ends up being loaded using `LOAD_FAST_CHECK`, which increfs and causes the refcount check to fail when it uses `LOAD_FAST_BORROW`.

These need to be tagged appropriately, not just increfed, so that they are decrefed when the frame is destroyed.

This may be 1 if the `LOAD_FAST` is optimized to a `LOAD_FAST_BORROW`. It's not clear that this is testing anything useful, so remove it.

The initial value will differ depending on whether a owned or borrowed reference is loaded onto the operand stack.

These don't push enough values on the stack.

…unconditional_jump_threading` Make sure we have a statically known stack depth

…mized to borrowed variants

PyStackRef_AsPyObjectSteal creates a new reference if the stackref is deferred. This reference is leaked if we deopt before the corresponding decref.

These may provide support for borrowed references contained in frames closer to the top of the call stack. Add them to a list attached to the frame when they are overwritten, to be destroyed when the frame is destroyed.

`STORE_FAST_LOAD_FAST` and `LOAD_FAST_AND_CLEAR` both need to kill the local.

This ensures we hit all the blocks

mpage · 2025-03-18T23:19:17Z

I think the tracerefs buildbot failure is unrelated to this PR. It's failing in the same way for other, unrelated PRs.

markshannon · 2025-03-20T09:29:14Z

I think this introduces a bug, or rather turns an existing piece of code into a bug.

There is code that assumes that owning a reference to an object and that the refcount of that object is one means that that code is free to mutate and reuse the object.
Which is fine if that reference is a heap reference, as this PR relies on there being at least one strong reference on the stack, but I'm fairly sure that there is some code (annoyingly I can't remember what code) which makes this assumption for a stack reference.

We should fix that code, not this, but I'd like to have a test for it first.

colesbury · 2025-03-20T14:46:35Z

Python/bytecodes.c

+            _PyStackRef tmp = PyStackRef_MakeHeapSafe(v);
+            DEAD(v);


Can we do something like https://gist.github.com/colesbury/ce98264761c831d42d185f758395933e instead?

Having a _PyStackRef that only exists on the C stack worries me especially with the PyObject_CallMethodOneArg() call below.

Yeah, good call

This is a bug in the code-generator, that you can't write v = PyStackRef_MakeHeapSafe(v);
It's trying to prevent losing a reference in an input variable, but it is too strict.

#131513

Fixed, you should be able to write

_PyStackRef tmp = PyStackRef_MakeHeapSafe(v); DEAD(v); v = tmp;

now.

mpage · 2025-03-20T15:42:18Z

Which is fine if that reference is a heap reference, as this PR relies on there being at least one strong reference on the stack, but I'm fairly sure that there is some code (annoyingly I can't remember what code) which makes this assumption for a stack reference.

We should fix that code, not this, but I'd like to have a test for it first.

@markshannon - Are you maybe thinking of this code?

markshannon

This looks really good overall. A few percent speedup is impressive.

I have a few comments and questions, but nothing blocking.

markshannon · 2025-03-20T09:31:41Z

Include/internal/pycore_stackref.h

+static inline _PyStackRef
+PyStackRef_Borrow(_PyStackRef ref)
+{
+    return PyStackRef_DUP(ref)


Is there an issue for adding lifetime checking for this?
If not, could you make one and assign it to either yourself or me?

markshannon · 2025-03-20T14:59:24Z

Lib/test/test_importlib/test_abc.py

@@ -801,6 +801,9 @@ def verify_code(self, code_object, *, bytecode_written=False):
            data.extend(self.init._pack_uint32(0))
            data.extend(self.init._pack_uint32(self.loader.source_mtime))
            data.extend(self.init._pack_uint32(self.loader.source_size))
+            # Make sure theres > 1 reference to code_object so that the


Suggested change

# Make sure theres > 1 reference to code_object so that the

# Make sure there's > 1 reference to code_object so that the

or there is

markshannon · 2025-03-20T15:01:51Z

Lib/test/test_peepholer.py

@@ -1330,15 +1335,6 @@ def test_fold_tuple_of_constants(self):
        ]
        self.cfg_optimization_test(before, after, consts=[], expected_consts=[(1, 2, 3)])

-        # not enough consts


Why have you removed these tests?

The bytecode is invalid.

markshannon · 2025-03-20T15:07:47Z

Lib/test/test_peepholer.py

+            ("LOAD_FAST", 0, 1),
+            ("LOAD_CONST", 0, 2),
+            ("LOAD_CONST", 0, 3),
+            ("STORE_FAST_STORE_FAST", 0 << 4 | 1, 4),


Suggested change

("STORE_FAST_STORE_FAST", 0 << 4 | 1, 4),

("STORE_FAST_STORE_FAST", ((0 << 4) | 1), 4),

Mentally parsing shifts and bitwise operations is extra work for the reader. Be nice and add parentheses.

markshannon · 2025-03-20T15:08:12Z

Lib/test/test_peepholer.py

+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_CONST", 0, 3),
+            ("STORE_FAST_STORE_FAST", 0 << 4 | 1, 4),


Suggested change

("STORE_FAST_STORE_FAST", 0 << 4 | 1, 4),

("STORE_FAST_STORE_FAST", ((0 << 4) | 1), 4),

Misc/NEWS.d/next/Core_and_Builtins/2025-02-28-11-29-35.gh-issue-130704.7RDVLE.rst

markshannon · 2025-03-20T15:26:51Z

Python/flowgraph.c

+ * non-violating LOAD_FAST{_LOAD_FAST} can be optimized.
+ */
+static int
+optimize_load_fast(cfg_builder *g, bool compute_stackdepth)


Maybe drop the compute_stackdepth parameter and replace optimize_load_fast(g, true) with calculate_stackdepth(g); optimize_load_fast(g)

markshannon · 2025-03-20T15:29:59Z

Python/flowgraph.c

+}
+
+/*
+ * Strength reduce LOAD_FAST{_LOAD_FAST} instructions into weaker variants that


Suggested change

* Strength reduce LOAD_FAST{_LOAD_FAST} instructions into weaker variants that

* Strength reduce LOAD_FAST{_LOAD_FAST} instructions into faster variants that

It's true that they are weaker, but the reason we're doing this is that they're faster.

markshannon · 2025-03-20T15:56:56Z

Python/bytecodes.c

+            _PyStackRef tmp = PyStackRef_MakeHeapSafe(v);
+            DEAD(v);


This is a bug in the code-generator, that you can't write v = PyStackRef_MakeHeapSafe(v);
It's trying to prevent losing a reference in an input variable, but it is too strict.

#131513

iritkatriel · 2025-03-20T16:41:15Z

Python/flowgraph.c

+{
+    if (stack->size == stack->capacity) {
+        Py_ssize_t new_cap = Py_MAX(32, stack->capacity * 2);
+        ref *refs = PyMem_Realloc(stack->refs, sizeof(*stack->refs) * new_cap);


We have a utility for growing arrays in Include/internal/pycore_c_array.h (_Py_c_array_t).

I don't think the code ends up being clearer if we use that here.

iritkatriel · 2025-03-20T17:09:29Z

Include/internal/pycore_flowgraph.h

@@ -41,6 +41,9 @@ PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
        PyObject *consts,
        int nlocals);

+// Export for '_testinternalcapi' shared extension
+PyAPI_FUNC(PyObject*) _PyCompile_OptimizeLoadFast(PyObject *instructions);


We don't tend to expose individual peephole optimisations, we test them through OptimizeCFG (see DirectCfgOptimizerTests). Is there a reason why you can't do that for this optimization?

Sure, I can add a call to optimize_load_fast into _PyCompile_OptimizeCFG. I wasn't sure if _PyCompile_OptimizeCFG was intended to only test _PyCfg_OptimizeCodeUnit. We want this pass to be called as late as possible in compilation, which is why it is called in _PyCfg_OptimizedCfgToInstructionSequence instead of _PyCfg_OptimizeCodeUnit.

markshannon · 2025-03-21T17:22:41Z

Which is fine if that reference is a heap reference, as this PR relies on there being at least one strong reference on the stack, but I'm fairly sure that there is some code (annoyingly I can't remember what code) which makes this assumption for a stack reference.
We should fix that code, not this, but I'd like to have a test for it first.

@markshannon - Are you maybe thinking of this code?

No, but that needed fixing as well. Thanks.

I had a look and there doesn't seem to a way to break anything without using the C API.
There are a couple of places that check whether a parameter has refcnt == 1 and reuse the object.
It was partial_new that I was thinking of.
It would be unsafe, if the kwargs were passed directly from the interpreter, but it appears that is impossible as functions.partial(foo, **d) creates a copy of d before it makes the call.

Not enough items on stack

mpage · 2025-03-24T18:19:43Z

A note for reviewers: I realized that we need to special case opcodes that leave some of their operands on the stack. This doesn't appear to reduce the effectiveness of the optimization. Stats look roughly the same as the previous version. Performance is a little worse (2.1% for the free-threaded build, 2.5% for the default), but it's hard for me to tell if that's just from noise in the runner or from changes that landed in main since the last time I ran the benchmarks.

colesbury

LGTM

markshannon

Missing one hint for the cases generator, otherwise looks good.

We could make the analysis more robust by using the cases generator, but that's for a later PR.

markshannon · 2025-03-26T15:30:59Z

Python/bytecodes.c

@@ -1209,7 +1219,7 @@ dummy_func(
                PyGenObject *gen = (PyGenObject *)receiver_o;
                _PyInterpreterFrame *gen_frame = &gen->gi_iframe;
                STACK_SHRINK(1);
-                _PyFrame_StackPush(gen_frame, v);
+                _PyFrame_StackPush(gen_frame, PyStackRef_MakeHeapSafe(v));


I think you need a DEAD(v); here

Python/bytecodes.c

markshannon · 2025-03-26T15:51:05Z

Python/flowgraph.c

+                    break;
+                }
+
+                // We treat opcodes that do not consume all of their inputs on


This approach seems fine for now, but the code generator knows exactly how many values are popped and consumed, as opposed to peek at.

We should add a _PyOpcode_num_peeked function, the we'd have consumed = _PyOpcode_num_popped() - _PyOpcode_num_peeked() which would be more robust.

mpage added 30 commits February 28, 2025 11:41

Experiment with borrowing load_fast

d716faa

Checkpoint poc

3736923

Fix pyframe copy

7a14254

Strengthen refs when frame is copied

b1607aa

Cleanup

e765735

Consider all instructions when computing mutations

291ace9

derp

Add a super instruction

17d6dd6

Don't optimize during quickening

0a74052

Use abstract interpretation

afbfd88

Fix test_generators

696c630

Ref will be 2 if borrowed

Optimize returns

483ac7a

Remove unused arg

259d5db

Make sure we convert borrowed refs on frame

aeafa98

Don't test with malformed bytecode

85f9a64

Make sure we convert borrowed refs to func/code when copying generato…

b6ab2f7

…r frame

Add support for disassembling LOAD_FAST_BORROW_LOAD_FAST_BORROW

fd1ad3d

Make sure exc_obj is always defined

eee2195

Otherwise, it ends up being loaded using `LOAD_FAST_CHECK`, which increfs and causes the refcount check to fail when it uses `LOAD_FAST_BORROW`.

Make sure we store new stackrefs for frame executable/funcobj

d75ec9a

These need to be tagged appropriately, not just increfed, so that they are decrefed when the frame is destroyed.

Remove refcount check

66f5351

This may be 1 if the `LOAD_FAST` is optimized to a `LOAD_FAST_BORROW`. It's not clear that this is testing anything useful, so remove it.

Don't hardcode initial refcount in refcount tests

7ef6a0b

The initial value will differ depending on whether a owned or borrowed reference is loaded onto the operand stack.

Remove invalid bytecode from test_peepholer

2af2bbc

These don't push enough values on the stack.

Fix invalid bytecode in `test_peepholer.DirectCfgOptimizerTests.test_…

bf19b7d

…unconditional_jump_threading` Make sure we have a statically known stack depth

Fix tests that checked for LOAD_FAST instructions that are now opti…

a9bca03

…mized to borrowed variants

Update disassembly in test_dis to match new bytecode

293c317

Fix refleak in _BINARY_OP_INPLACE_ADD_UNICODE

a12ccd9

PyStackRef_AsPyObjectSteal creates a new reference if the stackref is deferred. This reference is leaked if we deopt before the corresponding decref.

Create new references to fast locals overwritten via f_locals

1ef26c5

These may provide support for borrowed references contained in frames closer to the top of the call stack. Add them to a list attached to the frame when they are overwritten, to be destroyed when the frame is destroyed.

Implement two missing opcodes in the static analysis

1eb9226

`STORE_FAST_LOAD_FAST` and `LOAD_FAST_AND_CLEAR` both need to kill the local.

Use g_block_list when resetting stack depth

7291c49

This ensures we hit all the blocks

Avoid reallocating state for each basic block

90bf8df

Generators

9bfa922

colesbury reviewed Mar 20, 2025

View reviewed changes

markshannon reviewed Mar 20, 2025

View reviewed changes

iritkatriel reviewed Mar 20, 2025

View reviewed changes

mpage mentioned this pull request Mar 20, 2025

Dynamic borrow checking for stackrefs #131527

Open

mpage added 7 commits March 20, 2025 15:37

Merge branch 'main' into load-fast-borrow-absinterp

85b0b00

Fix post-merge issues

5ff2dea

English is hard

0c1e67f

Improve readability of test cases

ae2ec65

Elaborate in the blurb

03c474e

Remove parameter to calculate stackdepth

60665c9

Update comment

ac8940b

mpage added 7 commits March 21, 2025 11:39

Test optimize_load_fast as part of OptimizeCfg

f12573f

Remove test with invalid bytecode

44f7ffc

Not enough items on stack

Add helper macro for pushing refs

818e94e

Handle opcodes that leave at least one input on the stack

c30e1e9

Merge branch 'main' into load-fast-borrow-absinterp

112cee6

Avoid having stackref only visible from the c stack

80fc5aa

Merge branch 'main' into load-fast-borrow-absinterp

492cce1

mpage requested review from colesbury, iritkatriel and markshannon March 24, 2025 18:19

colesbury approved these changes Mar 24, 2025

View reviewed changes

bedevere-app bot added awaiting merge and removed awaiting change review labels Mar 24, 2025

markshannon approved these changes Mar 26, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

gh-130704: Strength reduce `LOAD_FAST{_LOAD_FAST}` #130708

gh-130704: Strength reduce `LOAD_FAST{_LOAD_FAST}` #130708

mpage commented Feb 28, 2025 •

edited

Loading

mpage commented Mar 18, 2025

markshannon commented Mar 20, 2025

colesbury Mar 20, 2025

mpage Mar 20, 2025

markshannon Mar 20, 2025

markshannon Mar 21, 2025

mpage commented Mar 20, 2025

markshannon left a comment

markshannon Mar 20, 2025

mpage Mar 20, 2025

markshannon Mar 20, 2025

markshannon Mar 20, 2025

mpage Mar 20, 2025

markshannon Mar 20, 2025

markshannon Mar 20, 2025

markshannon Mar 20, 2025

markshannon Mar 20, 2025

markshannon Mar 20, 2025

iritkatriel Mar 20, 2025

mpage Mar 20, 2025

iritkatriel Mar 20, 2025

mpage Mar 20, 2025 •

edited

Loading

markshannon commented Mar 21, 2025

mpage commented Mar 24, 2025

colesbury left a comment

markshannon left a comment

markshannon Mar 26, 2025

markshannon Mar 26, 2025 •

edited

Loading

	# Make sure theres > 1 reference to code_object so that the
	# Make sure there's > 1 reference to code_object so that the

	("STORE_FAST_STORE_FAST", 0 << 4 \| 1, 4),
	("STORE_FAST_STORE_FAST", ((0 << 4) \| 1), 4),

	* Strength reduce LOAD_FAST{_LOAD_FAST} instructions into weaker variants that
	* Strength reduce LOAD_FAST{_LOAD_FAST} instructions into faster variants that

gh-130704: Strength reduce LOAD_FAST{_LOAD_FAST} #130708

Are you sure you want to change the base?

gh-130704: Strength reduce LOAD_FAST{_LOAD_FAST} #130708

Conversation

mpage commented Feb 28, 2025 • edited Loading

mpage commented Mar 18, 2025

markshannon commented Mar 20, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mpage commented Mar 20, 2025

markshannon left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mpage Mar 20, 2025 • edited Loading

Choose a reason for hiding this comment

markshannon commented Mar 21, 2025

mpage commented Mar 24, 2025

colesbury left a comment

Choose a reason for hiding this comment

markshannon left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

markshannon Mar 26, 2025 • edited Loading

Choose a reason for hiding this comment

gh-130704: Strength reduce `LOAD_FAST{_LOAD_FAST}` #130708

gh-130704: Strength reduce `LOAD_FAST{_LOAD_FAST}` #130708

mpage commented Feb 28, 2025 •

edited

Loading

mpage Mar 20, 2025 •

edited

Loading

markshannon Mar 26, 2025 •

edited

Loading