Repository navigation
Top-of-stack caching in the JIT #135379
Description
Unlike #131498 which was a wash for performance, TOS caching in the JIT promises substantial performance improvements. This is because we can create several stencils for each uop, tailored for the number of registers and dynamically vary the number of values cached.
For example, in this code:
LOAD_FAST_BORROW
LOAD_FAST_BORROW
BINARY_OP_ADD_INT
STORE_FAST
we can tailor each version to the number of registers cached:
LOAD_FAST_BORROW_0_1 ( 0 -> 1 registers)
LOAD_FAST_BORROW_1_2 ( 1 -> 2 registers)
BINARY_OP_ADD_INT_2_1 ( 2 -> 1 registers)
STORE_FAST_1_0 ( 1 -> 0 registers)
thus avoiding any memory traffic to and from the stack at all.
The exact number of variants per uop will need to be determined empirically.
Having more stencils allows more freedom when generating code, but excessive numbers of stencils would cause bloat both at runtime, and in any repository containing the stencils.
Spilling and reloading
There will be an upper bound to the number of values cached and some uops may need a minimum number of values in the cache.
To handle those we will need to insert spill and reload uops. Spills will reduce the number of cached values, saving them to the in-memory stack and reloads will do the opposite moving values from the in-memory stack to the cache.
E.g.
LOAD_FAST_BORROW
BINARY_OP_ADD_INT
BINARY_OP_ADD_INT expects two inputs, but we only have one cached (from the LOAD_FAST_BORROW) so we need to insert a RELOAD:
LOAD_FAST_BORROW ( 0 -> 1 registers)
RELOAD_1_2 ( 1 -> 2 registers)
BINARY_OP_ADD_INT ( 2 -> 1 registers)
SPILL and RELOAD are semantic no-ops, and will be generated automatically.
Deferred references
For this to work the code generator must spill any cached values to the in-memory stack when GC could occur. Fortunately, the code generator already does this (as part of the work for #131498).
Linked PRs
- GH-135379: Remove types from stack items in code generator. #135384
- GH-135379: Top of stack caching for the JIT. #135465
- GH-135379: Support limited scalar replacement for replicated uops in the JIT code generator. #135563
- GH-135379: Specialize int operations for compact ints only #135668
- gh-135379: Add back const cast to _PyLong_IsCompact #135706
- gh-135379: Move PyLong_CheckCompact to private header #135707
- gh-135379: fix MSVC warning: conversion from 'stwodigits' to 'digit' #135714
Activity
Fidget-Spinner commentedon Jun 11, 2025
Not to rain on your parade, but I tried this before at https://github.com/python/cpython/compare/main...Fidget-Spinner:cpython:tos-caching-jit?expand=1 . It has almost no speedup in nbody due to two reasons:
PyStackRef_CLOSEis a spill.The first point is easily addresed. The second point isn't. We need to properly eliminate/specialize for refcounts for this to pay off. I'm thinking of doing what you and Brandt suggested of making an op leave its operands on the stack, and specializing a POP_TOP.
That said, I think it's still worth doing, because once we start removing refcounts, this will automatically become a huge win. It will show up as a huge speedup in the refcount removal work, but we'll know it also partially came from here :).
Fidget-Spinner commentedon Jun 11, 2025
The other thing to beware of is that this will significantly blow up stencil count if we're not careful. Which will significantly increase our JIT build times.
GH-135379: Remove types from stack items in code generator. (GH-135384)
GH-135379: Support limited scalar replacement for replicated uops in …
GH-135379: Specialize int operations for compact ints only (GH-135668)
gh-135379: Add back const cast to _PyLong_IsCompact (GH-135706)
gh-135379: Move PyLong_CheckCompact to private header and rename it (G…
pythonGH-135379: Remove types from stack items in code generator. (py…
pythonGH-135379: Support limited scalar replacement for replicated uo…
pythonGH-135379: Specialize int operations for compact ints only (pyt…
pythongh-135379: Add back const cast to _PyLong_IsCompact (pythonGH-1…
11 remaining items
pythongh-135379: Add back const cast to _PyLong_IsCompact (pythonGH-1…
pythongh-135379: Move PyLong_CheckCompact to private header and renam…
pythongh-135379: fix MSVC warning: conversion from 'stwodigits' to 'd…
pythonGH-135379: Remove types from stack items in code generator. (py…
pythonGH-135379: Support limited scalar replacement for replicated uo…
pythonGH-135379: Specialize int operations for compact ints only (pyt…
pythongh-135379: Add back const cast to _PyLong_IsCompact (pythonGH-1…
pythongh-135379: Move PyLong_CheckCompact to private header and renam…
pythongh-135379: fix MSVC warning: conversion from 'stwodigits' to 'd…
GH-135379: Top of stack caching for the JIT. (GH-135465)
pythonGH-135379: Top of stack caching for the JIT. (pythonGH-135465)
markshannon commentedon Dec 12, 2025
Done.
wjakob commentedon Jan 6, 2026
@markshannon: the Python 3.15 documentation page (https://docs.python.org/3.15/whatsnew/3.15.html) links the "Basic register allocation in the JIT" section to this PR, which looks incorrect to me.
Fidget-Spinner commentedon Jan 6, 2026
@wjakob that is correct. TOS caching allocates some of the stack values into registers.