GH-139109: Partial reworking of JIT data structures (GH-144105)

* Halve size of buffers by reusing combined trace + optimizer buffers for TOS caching
* Add simple buffer struct for more maintainable handling of buffers
* Decouple JIT structs from thread state struct
* Ensure terminator is added to trace, when optimizer gives up

Author: markshannon

Include/internal/pycore_optimizer.h

@@ -16,12 +16,102 @@ extern "C" {
 #include <stdbool.h>
 
 
+typedef struct _PyJitUopBuffer {
+    _PyUOpInstruction *start;
+    _PyUOpInstruction *next;
+    _PyUOpInstruction *end;
+} _PyJitUopBuffer;
+
+
+typedef struct _JitOptContext {
+    char done;
+    char out_of_space;
+    bool contradiction;
+    // Has the builtins dict been watched?
+    bool builtins_watched;
+    // The current "executing" frame.
+    _Py_UOpsAbstractFrame *frame;
+    _Py_UOpsAbstractFrame frames[MAX_ABSTRACT_FRAME_DEPTH];
+    int curr_frame_depth;
+
+    // Arena for the symbolic types.
+    ty_arena t_arena;
+
+    JitOptRef *n_consumed;
+    JitOptRef *limit;
+    JitOptRef locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
+    _PyJitUopBuffer out_buffer;
+} JitOptContext;
+
+
+static inline void
+uop_buffer_init(_PyJitUopBuffer *trace, _PyUOpInstruction *start, uint32_t size)
+{
+    trace->next = trace->start = start;
+    trace->end = start + size;
+}
+
+static inline _PyUOpInstruction *
+uop_buffer_last(_PyJitUopBuffer *trace)
+{
+    assert(trace->next > trace->start);
+    return trace->next-1;
+}
+
+static inline int
+uop_buffer_length(_PyJitUopBuffer *trace)
+{
+    return (int)(trace->next - trace->start);
+}
+
+static inline int
+uop_buffer_remaining_space(_PyJitUopBuffer *trace)
+{
+    return (int)(trace->end - trace->next);
+}
+
+typedef struct _PyJitTracerInitialState {
+    int stack_depth;
+    int chain_depth;
+    struct _PyExitData *exit;
+    PyCodeObject *code; // Strong
+    PyFunctionObject *func; // Strong
+    struct _PyExecutorObject *executor; // Strong
+    _Py_CODEUNIT *start_instr;
+    _Py_CODEUNIT *close_loop_instr;
+    _Py_CODEUNIT *jump_backward_instr;
+} _PyJitTracerInitialState;
+
+typedef struct _PyJitTracerPreviousState {
+    bool dependencies_still_valid;
+    int instr_oparg;
+    int instr_stacklevel;
+    _Py_CODEUNIT *instr;
+    PyCodeObject *instr_code; // Strong
+    struct _PyInterpreterFrame *instr_frame;
+    _PyBloomFilter dependencies;
+} _PyJitTracerPreviousState;
+
+typedef struct _PyJitTracerTranslatorState {
+    int jump_backward_seen;
+} _PyJitTracerTranslatorState;
+
+typedef struct _PyJitTracerState {
+    bool is_tracing;
+    _PyJitTracerInitialState initial_state;
+    _PyJitTracerPreviousState prev_state;
+    _PyJitTracerTranslatorState translator_state;
+    JitOptContext opt_context;
+    _PyJitUopBuffer code_buffer;
+    _PyJitUopBuffer out_buffer;
+    _PyUOpInstruction uop_array[2 * UOP_MAX_TRACE_LENGTH];
+} _PyJitTracerState;
+
 typedef struct _PyExecutorLinkListNode {
     struct _PyExecutorObject *next;
     struct _PyExecutorObject *previous;
 } _PyExecutorLinkListNode;
 
-
 typedef struct {
     uint8_t opcode;
     uint8_t oparg;
@@ -86,8 +176,8 @@ PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyInterpreterState *interp);
 
 int _Py_uop_analyze_and_optimize(
     _PyThreadStateImpl *tstate,
-    _PyUOpInstruction *trace, int trace_len, int curr_stackentries,
-    _PyBloomFilter *dependencies);
+    _PyUOpInstruction *input, int trace_len, int curr_stackentries,
+    _PyUOpInstruction *output, _PyBloomFilter *dependencies);
 
 extern PyTypeObject _PyUOpExecutor_Type;
 

Include/internal/pycore_optimizer_types.h

@@ -126,27 +126,6 @@ typedef struct ty_arena {
     JitOptSymbol arena[TY_ARENA_SIZE];
 } ty_arena;
 
-typedef struct _JitOptContext {
-    char done;
-    char out_of_space;
-    bool contradiction;
-    // Has the builtins dict been watched?
-    bool builtins_watched;
-    // The current "executing" frame.
-    _Py_UOpsAbstractFrame *frame;
-    _Py_UOpsAbstractFrame frames[MAX_ABSTRACT_FRAME_DEPTH];
-    int curr_frame_depth;
-
-    // Arena for the symbolic types.
-    ty_arena t_arena;
-
-    JitOptRef *n_consumed;
-    JitOptRef *limit;
-    JitOptRef locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
-    _PyUOpInstruction *out_buffer;
-    int out_len;
-} JitOptContext;
-
 
 #ifdef __cplusplus
 }

Include/internal/pycore_tstate.h

@@ -12,7 +12,6 @@ extern "C" {
 #include "pycore_freelist_state.h"  // struct _Py_freelists
 #include "pycore_interpframe_structs.h"  // _PyInterpreterFrame
 #include "pycore_mimalloc.h"        // struct _mimalloc_thread_state
-#include "pycore_optimizer_types.h" // JitOptContext
 #include "pycore_qsbr.h"            // struct qsbr
 #include "pycore_uop.h"             // struct _PyUOpInstruction
 #include "pycore_structs.h"
@@ -24,46 +23,6 @@ struct _gc_thread_state {
 };
 #endif
 
-#if _Py_TIER2
-typedef struct _PyJitTracerInitialState {
-    int stack_depth;
-    int chain_depth;
-    struct _PyExitData *exit;
-    PyCodeObject *code; // Strong
-    PyFunctionObject *func; // Strong
-    struct _PyExecutorObject *executor; // Strong
-    _Py_CODEUNIT *start_instr;
-    _Py_CODEUNIT *close_loop_instr;
-    _Py_CODEUNIT *jump_backward_instr;
-} _PyJitTracerInitialState;
-
-typedef struct _PyJitTracerPreviousState {
-    bool dependencies_still_valid;
-    int code_max_size;
-    int code_curr_size;
-    int instr_oparg;
-    int instr_stacklevel;
-    _Py_CODEUNIT *instr;
-    PyCodeObject *instr_code; // Strong
-    struct _PyInterpreterFrame *instr_frame;
-    _PyBloomFilter dependencies;
-} _PyJitTracerPreviousState;
-
-typedef struct _PyJitTracerTranslatorState {
-    int jump_backward_seen;
-} _PyJitTracerTranslatorState;
-
-typedef struct _PyJitTracerState {
-    bool is_tracing;
-    _PyJitTracerInitialState initial_state;
-    _PyJitTracerPreviousState prev_state;
-    _PyJitTracerTranslatorState translator_state;
-    JitOptContext opt_context;
-    _PyUOpInstruction code_buffer[UOP_MAX_TRACE_LENGTH];
-    _PyUOpInstruction out_buffer[UOP_MAX_TRACE_LENGTH];
-} _PyJitTracerState;
-
-#endif
 
 // Every PyThreadState is actually allocated as a _PyThreadStateImpl. The
 // PyThreadState fields are exposed as part of the C API, although most fields
@@ -141,7 +100,7 @@ typedef struct _PyThreadStateImpl {
     Py_ssize_t reftotal;  // this thread's total refcount operations
 #endif
 #if _Py_TIER2
-    _PyJitTracerState *jit_tracer_state;
+    struct _PyJitTracerState *jit_tracer_state;
 #endif
 } _PyThreadStateImpl;
 

Include/internal/pycore_uop.h

@@ -38,11 +38,10 @@ typedef struct _PyUOpInstruction{
 // This is the length of the trace we translate initially.
 #ifdef Py_DEBUG
     // With asserts, the stencils are a lot larger
-#define UOP_MAX_TRACE_LENGTH 2000
+#define UOP_MAX_TRACE_LENGTH 1000
 #else
-#define UOP_MAX_TRACE_LENGTH 5000
+#define UOP_MAX_TRACE_LENGTH 2500
 #endif
-#define UOP_BUFFER_SIZE (UOP_MAX_TRACE_LENGTH * sizeof(_PyUOpInstruction))
 
 /* Bloom filter with m = 256
  * https://en.wikipedia.org/wiki/Bloom_filter */

Python/ceval_macros.h

@@ -433,7 +433,7 @@ do {                                                   \
         JUMP_TO_LABEL(error);                          \
     }                                                  \
     if (keep_tracing_bit) { \
-        assert(((_PyThreadStateImpl *)tstate)->jit_tracer_state->prev_state.code_curr_size == 2); \
+        assert(uop_buffer_length(&((_PyThreadStateImpl *)tstate)->jit_tracer_state->code_buffer)); \
         ENTER_TRACING(); \
         DISPATCH_NON_TRACING(); \
     } \