Description

markshannon

opened on Jun 11, 2025

Issue body actions

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

See faster-cpython/ideas#711

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

1. Fidget-Spinner commented on Jun 11, 2025

   

   Fidget-Spinner

   on Jun 11, 2025

   Member

   More actions

   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:

   1. COPY and SWAP were not handled, so those need custom handling.
   2. PyStackRef_CLOSE is 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 :).
2. Fidget-Spinner commented on Jun 11, 2025

   

   Fidget-Spinner

   on Jun 11, 2025

   Member

   More actions

   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.
3. markshannon

   added a commit that references this issue on Jun 11, 2025

   GH-135379: Remove types from stack items in code generator. (GH-135384)

   c87b5b2
4. picnixz

   added

   performancePerformance or resource usagePerformance or resource usage

   interpreter-core(Objects, Python, Grammar, and Parser dirs)(Objects, Python, Grammar, and Parser dirs)

   topic-JIT

   type-featureA feature request or enhancementA feature request or enhancement

   on Jun 13, 2025
5. markshannon

   added 2 commits that reference this issue on Jun 17, 2025

   GH-135379: Support limited scalar replacement for replicated uops in …

   8dd8b5c

   GH-135379: Specialize int operations for compact ints only (GH-135668)

   9731dd2
6. encukou

   added a commit that references this issue on Jun 19, 2025

   gh-135379: Add back const cast to _PyLong_IsCompact (GH-135706)

   1b969f6
7. Fidget-Spinner

   added a commit that references this issue on Jun 19, 2025

   gh-135379: Move PyLong_CheckCompact to private header and rename it (G…

   0243260
8. lkollar

   added 4 commits that reference this issue on Jun 19, 2025

   pythonGH-135379: Remove types from stack items in code generator. (py…

   76eb7da

   pythonGH-135379: Support limited scalar replacement for replicated uo…

   a7fe5be

   pythonGH-135379: Specialize int operations for compact ints only (pyt…

   b499320

   pythongh-135379: Add back const cast to _PyLong_IsCompact (pythonGH-1…

   5ee00d3
9. 11 remaining items
10. taegyunkim

    added 3 commits that reference this issue on Aug 4, 2025

    pythongh-135379: Add back const cast to _PyLong_IsCompact (pythonGH-1…

    977dd64

    pythongh-135379: Move PyLong_CheckCompact to private header and renam…

    724adea

    pythongh-135379: fix MSVC warning: conversion from 'stwodigits' to 'd…

    a465a38
11. Agent-Hellboy

    added 6 commits that reference this issue on Aug 19, 2025

    pythonGH-135379: Remove types from stack items in code generator. (py…

    c6e800f

    pythonGH-135379: Support limited scalar replacement for replicated uo…

    b00b926

    pythonGH-135379: Specialize int operations for compact ints only (pyt…

    4a97ea1

    pythongh-135379: Add back const cast to _PyLong_IsCompact (pythonGH-1…

    04f0385

    pythongh-135379: Move PyLong_CheckCompact to private header and renam…

    68ad2d9

    pythongh-135379: fix MSVC warning: conversion from 'stwodigits' to 'd…

    ecad01a
12. markshannon

    added a commit that references this issue on Dec 11, 2025

    GH-135379: Top of stack caching for the JIT. (GH-135465)

    469f191
13. fatelei

    added a commit that references this issue on Dec 12, 2025

    pythonGH-135379: Top of stack caching for the JIT. (pythonGH-135465)

    74f2dd5
14. markshannon commented on Dec 12, 2025

    

    markshannon

    on Dec 12, 2025

    MemberAuthor

    More actions

    Done.
15. markshannon

    closed this as completedon Dec 12, 2025
16. wjakob commented on Jan 6, 2026

    

    wjakob

    on Jan 6, 2026

    Contributor

    More actions

    @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.
17. Fidget-Spinner commented on Jan 6, 2026

    

    Fidget-Spinner

    on Jan 6, 2026

    Member

    More actions

    @wjakob that is correct. TOS caching allocates some of the stack values into registers.

Issue actions