Reserve a register for localsplus

Fidget-Spinner · Dec 9, 2024 · 3a2f8b2 · 3a2f8b2
1 parent 8fa5ece
commit 3a2f8b2
Show file tree

Hide file tree

Showing 8 changed files with 84 additions and 39 deletions.
diff --git a/Include/internal/pycore_jit.h b/Include/internal/pycore_jit.h
@@ -11,7 +11,7 @@ extern "C" {
 
 #ifdef _Py_JIT
 
-typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate);
+typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate, _PyStackRef *frame_localsplus);
 
 int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -250,7 +250,7 @@ dummy_func(
         };
 
         inst(LOAD_FAST_CHECK, (-- value)) {
-            _PyStackRef value_s = GETLOCAL(oparg);
+            _PyStackRef value_s = PEEKLOCAL(oparg);
             if (PyStackRef_IsNull(value_s)) {
                 _PyEval_FormatExcCheckArg(tstate, PyExc_UnboundLocalError,
                     UNBOUNDLOCAL_ERROR_MSG,
@@ -262,21 +262,21 @@ dummy_func(
         }
 
         replicate(8) pure inst(LOAD_FAST, (-- value)) {
-            assert(!PyStackRef_IsNull(GETLOCAL(oparg)));
-            value = PyStackRef_DUP(GETLOCAL(oparg));
+            assert(!PyStackRef_IsNull(PEEKLOCAL(oparg)));
+            value = PyStackRef_DUP(PEEKLOCAL(oparg));
         }
 
         inst(LOAD_FAST_AND_CLEAR, (-- value)) {
-            value = GETLOCAL(oparg);
+            value = PEEKLOCAL(oparg);
             // do not use SETLOCAL here, it decrefs the old value
             GETLOCAL(oparg) = PyStackRef_NULL;
         }
 
         inst(LOAD_FAST_LOAD_FAST, ( -- value1, value2)) {
             uint32_t oparg1 = oparg >> 4;
             uint32_t oparg2 = oparg & 15;
-            value1 = PyStackRef_DUP(GETLOCAL(oparg1));
-            value2 = PyStackRef_DUP(GETLOCAL(oparg2));
+            value1 = PyStackRef_DUP(PEEKLOCAL(oparg1));
+            value2 = PyStackRef_DUP(PEEKLOCAL(oparg2));
         }
 
         family(LOAD_CONST, 0) = {
@@ -1047,6 +1047,7 @@ dummy_func(
             RELOAD_STACK();
             LOAD_IP(frame->return_offset);
             res = temp;
+            SET_LOCALSPLUS(frame);
             LLTRACE_RESUME_FRAME();
         }
 
@@ -1148,6 +1149,7 @@ dummy_func(
                 frame->return_offset = (uint16_t)(INSTRUCTION_SIZE + oparg);
                 assert(gen_frame->previous == NULL);
                 gen_frame->previous = frame;
+                SET_LOCALSPLUS(gen_frame);
                 DISPATCH_INLINED(gen_frame);
             }
             if (PyStackRef_IsNone(v) && PyIter_Check(receiver_o)) {
@@ -1235,6 +1237,7 @@ dummy_func(
             RELOAD_STACK();
             LOAD_IP(1 + INLINE_CACHE_ENTRIES_SEND);
             value = temp;
+            SET_LOCALSPLUS(frame);
             LLTRACE_RESUME_FRAME();
         }
 
@@ -2317,6 +2320,7 @@ dummy_func(
             DEAD(owner);
             new_frame->localsplus[1] = PyStackRef_FromPyObjectNew(name);
             frame->return_offset = INSTRUCTION_SIZE;
+            SET_LOCALSPLUS(new_frame);
             DISPATCH_INLINED(new_frame);
         }
 
@@ -3381,6 +3385,7 @@ dummy_func(
                     ERROR_NO_POP();
                 }
                 frame->return_offset = INSTRUCTION_SIZE;
+                SET_LOCALSPLUS(new_frame);
                 DISPATCH_INLINED(new_frame);
             }
             /* Callable is not a normal Python function */
@@ -3633,6 +3638,7 @@ dummy_func(
             tstate->py_recursion_remaining--;
             LOAD_SP();
             LOAD_IP(0);
+            SET_LOCALSPLUS(frame);
             LLTRACE_RESUME_FRAME();
         }
 
@@ -4256,6 +4262,7 @@ dummy_func(
                 }
                 assert(INSTRUCTION_SIZE == 1 + INLINE_CACHE_ENTRIES_CALL_KW);
                 frame->return_offset = INSTRUCTION_SIZE;
+                SET_LOCALSPLUS(new_frame);
                 DISPATCH_INLINED(new_frame);
             }
             /* Callable is not a normal Python function */
@@ -4523,6 +4530,7 @@ dummy_func(
                     }
                     assert(INSTRUCTION_SIZE == 1);
                     frame->return_offset = 1;
+                    SET_LOCALSPLUS(new_frame);
                     DISPATCH_INLINED(new_frame);
                 }
                 result_o = PyObject_Call(func, callargs, kwargs);
@@ -4588,6 +4596,7 @@ dummy_func(
             LOAD_IP(frame->return_offset);
             RELOAD_STACK();
             res = PyStackRef_FromPyObjectSteal((PyObject *)gen);
+            SET_LOCALSPLUS(frame);
             LLTRACE_RESUME_FRAME();
         }
 

diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
@@ -248,6 +248,7 @@ GETITEM(PyObject *v, Py_ssize_t i) {
 
 #define LOCALS_ARRAY    (frame->localsplus)
 #define GETLOCAL(i)     (frame->localsplus[i])
+#define PEEKLOCAL(i)    (frame->localsplus[i])
 
 /* The SETLOCAL() macro must not DECREF the local variable in-place and
    then store the new value; it must copy the old value to a temporary
@@ -406,7 +407,7 @@ _PyFrame_SetStackPointer(frame, stack_pointer)
 do {                                                   \
     OPT_STAT_INC(traces_executed);                     \
     jit_func jitted = (EXECUTOR)->jit_code;            \
-    next_instr = jitted(frame, stack_pointer, tstate); \
+    next_instr = jitted(frame, stack_pointer, tstate, frame->localsplus); \
     Py_DECREF(tstate->previous_executor);              \
     tstate->previous_executor = NULL;                  \
     frame = tstate->current_frame;                     \
@@ -475,3 +476,5 @@ do { \
 #else
 #define CONVERSION_FAILED(NAME) (0)
 #endif
+
+#define SET_LOCALSPLUS(f) ((void)(f))
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h