JIT: Move executor to a register

Author: diegorusso

Include/internal/pycore_jit.h

@@ -19,7 +19,7 @@ extern "C" {
 #ifdef _Py_JIT
 
 typedef _Py_CODEUNIT *(*jit_func)(
-    _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate,
+    _PyExecutorObject *executor, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate,
     _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2
 );
 

Python/bytecodes.c

@@ -5309,7 +5309,7 @@ dummy_func(
             assert(current_executor == (_PyExecutorObject*)executor);
 #endif
             assert(tstate->jit_exit == NULL || tstate->jit_exit->executor == current_executor);
-            tstate->current_executor = (PyObject *)executor;
+            tstate->current_executor = (PyObject *)current_executor;
             if (!current_executor->vm_data.valid) {
                 assert(tstate->jit_exit->executor == current_executor);
                 assert(tstate->current_executor == executor);

Python/executor_cases.c.h

@@ -19,8 +19,6 @@ class HoleValue(enum.Enum):
     CODE = enum.auto()
     # The base address of the read-only data for this uop:
     DATA = enum.auto()
-    # The address of the current executor (exposed as _JIT_EXECUTOR):
-    EXECUTOR = enum.auto()
     # The base address of the "global" offset table located in the read-only data.
     # Shouldn't be present in the final stencils, since these are all replaced with
     # equivalent DATA values:
@@ -108,7 +106,6 @@ class HoleValue(enum.Enum):
 _HOLE_EXPRS = {
     HoleValue.CODE: "(uintptr_t)code",
     HoleValue.DATA: "(uintptr_t)data",
-    HoleValue.EXECUTOR: "(uintptr_t)executor",
     HoleValue.GOT: "",
     # These should all have been turned into DATA values by process_relocations:
     HoleValue.OPARG: "instruction->oparg",

Tools/jit/_stencils.py

@@ -9,5 +9,5 @@ typedef jit_func __attribute__((preserve_none)) jit_func_preserve_none;
 
 #define DECLARE_TARGET(NAME)                     \
     _Py_CODEUNIT *__attribute__((preserve_none, visibility("hidden"))) \
-    NAME(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate, \
+    NAME(_PyExecutorObject *executor, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate, \
     _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2);

Tools/jit/jit.h

@@ -12,5 +12,6 @@ _JIT_ENTRY(
 ) {
     // Note that this is *not* a tail call
     jit_func_preserve_none jitted = (jit_func_preserve_none)exec->jit_code;
-    return jitted(frame, stack_pointer, tstate, PyStackRef_ZERO_BITS, PyStackRef_ZERO_BITS, PyStackRef_ZERO_BITS);
+    return jitted(exec, frame, stack_pointer, tstate,
+                  PyStackRef_ZERO_BITS, PyStackRef_ZERO_BITS, PyStackRef_ZERO_BITS);
 }

Tools/jit/shim.c

@@ -76,7 +76,7 @@ do {                                                                       \
     OPT_STAT_INC(traces_executed);                                         \
     _PyExecutorObject *_executor = (EXECUTOR);                             \
     jit_func_preserve_none jitted = _executor->jit_code;                   \
-    __attribute__((musttail)) return jitted(frame, stack_pointer, tstate,  \
+    __attribute__((musttail)) return jitted(_executor, frame, stack_pointer, tstate,  \
     _tos_cache0, _tos_cache1, _tos_cache2); \
 } while (0)
 
@@ -100,7 +100,7 @@ do {                                                                       \
 #define PATCH_JUMP(ALIAS)                                                 \
 do {                                                                      \
     DECLARE_TARGET(ALIAS);                                                \
-    __attribute__((musttail)) return ALIAS(frame, stack_pointer, tstate,  \
+    __attribute__((musttail)) return ALIAS(current_executor, frame, stack_pointer, tstate,  \
     _tos_cache0, _tos_cache1, _tos_cache2); \
 } while (0)
 
@@ -120,11 +120,11 @@ do {                                                                      \
 
 __attribute__((preserve_none)) _Py_CODEUNIT *
 _JIT_ENTRY(
-    _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate,
-     _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2
+    _PyExecutorObject *executor, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate,
+    _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2
 ) {
     // Locals that the instruction implementations expect to exist:
-    PATCH_VALUE(_PyExecutorObject *, current_executor, _JIT_EXECUTOR)
+    _PyExecutorObject *current_executor = executor;
     int oparg;
     int uopcode = _JIT_OPCODE;
     _Py_CODEUNIT *next_instr;