gh-143421: Move JitOptContext from stack allocation to per-thread heap allocation to avoid stack overflow.

[cocolato]

• Issue: JIT optimizer cleanups #143421
• JIT: Implement unique reference tracking in Tier 2 for reference count optimizations #143414 (comment)

[Fidget-Spinner]

Pretty close!

[cocolato]

Updated!

[Fidget-Spinner]

Great, thanks!

[markshannon]

Can we avoid unnecessary pointer chasing by embedding the array in the jit state struct, and allocate the entire struct together instead of having part of it embedded in the thread state, and part in a new allocation?

ie. instead of

typedef struct _PyJitOptState {
    struct _JitOptContext *opt_context;
} _PyJitOptState;

typedef struct _PyJitTracerState {
    _PyUOpInstruction *code_buffer;
    _PyJitOptState opt_state;
    _PyJitTracerInitialState initial_state;
    _PyJitTracerPreviousState prev_state;
    _PyJitTracerTranslatorState translator_state;
} _PyJitTracerState;

we would have

typedef struct _PyJitTracerState {
    _PyJitTracerInitialState initial_state;
    _PyJitTracerPreviousState prev_state;
    _PyJitTracerTranslatorState translator_state;
     _JitOptContext opt_context;
    _PyUOpInstruction code_buffer[UOP_MAX_TRACE_LENGTH];
} _PyJitTracerState;

Doing so would also decouple the thread state header from the JIT headers, as only a pointer to an opaque struct would need to be declared in the thread state header.

I'm also concerned about the amount of allocation and freeing required to compile a small trace. Maybe allocate once and only free when the thread is destroyed?

[cocolato]

we would have

typedef struct _PyJitTracerState {
    _PyJitTracerInitialState initial_state;
    _PyJitTracerPreviousState prev_state;
    _PyJitTracerTranslatorState translator_state;
     _JitOptContext opt_context;
    _PyUOpInstruction code_buffer[UOP_MAX_TRACE_LENGTH];
} _PyJitTracerSt

I can try implementing it this way, but I'd like to know if this should be done in a new PR/issue. Perhaps we could use benchmarks to determine potential performance issues from allocating _PyJitTracerState once.

[Fidget-Spinner]

I'm also concerned about the amount of allocation and freeing required to compile a small trace. Maybe allocate once and only free when the thread is destroyed?

Isnt that already being done?

[Fidget-Spinner]

@cocolato I think we can apply Mark's changes, it means we would need no new allocations anymore as its all heap allocated with the _PyThreadStateImpl struct.

[cocolato]

I'm not sure if the current implementation of the new pycore_optimizer_types.h is correct. I'd appreciate any suggestions for improvements.

[Fidget-Spinner]

Very cool, thank you!

[Fidget-Spinner]

@markshannon

This caused a 100% slowdown in bench_thread_pool benchmark https://doesjitgobrrr.com/run/2026-01-08

We need to lazily allocate it and pointer chase, otherwise we're slowing down significantly the spawning of threads.

[cocolato]

@markshannon

This caused a 100% slowdown in bench_thread_pool benchmark https://doesjitgobrrr.com/run/2026-01-08

We need to lazily allocate it and pointer chase, otherwise we're slowing down significantly the spawning of threads.

So should we revert to the first version?4e7a918

[Fidget-Spinner]

So should we revert to the first version?4e7a918

I will have a PR up.

[cocolato]

Sorry, I will run a local benchmark first before making important changes in the next time.

[Fidget-Spinner]

Sorry, I will run a local benchmark first before making important changes in the next time.

No it's fine. That's the point of having a benchmark runner, to catch things like this! It's hard to predict what a change will affect, so that's the point of having benchark runners to save human time.

[Fidget-Spinner]

PR up at #143597