Repository navigation
Top-of-stack caching in the JIT #135379
Description
Activity
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:
- COPY and SWAP were not handled, so those need custom handling.
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 :).
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.
- added a commit that references this issue
on Jun 19, 2025 - added 4 commits that reference this issue
on Jun 19, 2025 11 remaining items
- added 6 commits that reference this issue
on Aug 19, 2025 Done.
@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.
@wjakob that is correct. TOS caching allocates some of the stack values into registers.
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:
we can tailor each version to the number of registers cached:
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.
BINARY_OP_ADD_INTexpects two inputs, but we only have one cached (from theLOAD_FAST_BORROW) so we need to insert aRELOAD:SPILLandRELOADare 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).
See faster-cpython/ideas#711
Linked PRs