Add Brandt's fixes

python · markshannon · Feb 20, 2024 · Jan 10, 2024 · Jan 11, 2024 · Jan 11, 2024
commit cfd3285fd7278f93f277b2e8363fa396e96d84ae
diff --git a/Include/cpython/optimizer.h b/Include/cpython/optimizer.h
@@ -51,6 +51,7 @@ typedef struct _PyExecutorObject {
     _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
     uint32_t exit_count;
     uint32_t code_size;
+    size_t jit_size;
     void *jit_code;
     _PyExitData exits[1];
 } _PyExecutorObject;

@@ -387,12 +387,26 @@ stack_pointer = _PyFrame_GetStackPointer(frame);
 
 /* Tier-switching macros. */
 
+#ifdef _Py_JIT
+#define GOTO_TIER_TWO(EXECUTOR)                            \
+    do {                                                   \
+        jit_func jitted = (EXECUTOR)->jit_code;            \
+        next_instr = jitted(frame, stack_pointer, tstate); \
+        frame = tstate->current_frame;                     \
+        if (next_instr == NULL) {                          \
+            goto resume_with_error;                        \
+        }                                                  \
+        stack_pointer = _PyFrame_GetStackPointer(frame);   \
+        DISPATCH();                                        \
+    } while (0)
+#else
 #define GOTO_TIER_TWO(EXECUTOR) \
 do { \
     next_uop = (EXECUTOR)->trace; \
     assert(next_uop->opcode == _START_EXECUTOR || next_uop->opcode == _COLD_EXIT); \
     goto enter_tier_two; \
 } while (0)
+#endif
 
 #define GOTO_TIER_ONE(TARGET) \
     next_instr = target; \

@@ -323,6 +323,11 @@ _PyJIT_Compile(_PyExecutorObject *executor, _PyUOpInstruction *trace, size_t len
     // Loop again to emit the code:
     char *code = memory;
     char *data = memory + code_size;
+    char *top = code;
+    if (trace[0].opcode == _START_EXECUTOR) {
+        // Don't want to execute this more than once:
+        top += stencil_groups[_START_EXECUTOR].code.body_size;
+    }
     for (size_t i = 0; i < length; i++) {
         _PyUOpInstruction *instruction = &trace[i];
         const StencilGroup *group = &stencil_groups[instruction->opcode];
@@ -335,7 +340,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, _PyUOpInstruction *trace, size_t len
         patches[HoleValue_OPARG] = instruction->oparg;
         patches[HoleValue_OPERAND] = instruction->operand;
         patches[HoleValue_TARGET] = instruction->target;
-        patches[HoleValue_TOP] = (uint64_t)memory;
+        patches[HoleValue_TOP] = (uint64_t)top;
         patches[HoleValue_ZERO] = 0;
         emit(group, patches);
         code += group->code.body_size;
@@ -348,7 +353,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, _PyUOpInstruction *trace, size_t len
         return -1;
     }
     executor->jit_code = memory;
-    executor->code_size = code_size + data_size;
+    executor->jit_size = code_size + data_size;
     return 0;
 }
 
@@ -359,7 +364,7 @@ _PyJIT_Free(_PyExecutorObject *executor)
     size_t size = executor->code_size;
     if (memory) {
         executor->jit_code = NULL;
-        executor->code_size = 0;
+        executor->jit_size = 0;
         if (jit_free(memory, size)) {
             PyErr_WriteUnraisable(NULL);
         }

@@ -914,7 +914,7 @@ make_executor_from_uops(_PyUOpInstruction *buffer, const _PyBloomFilter *depende
 #endif
 #ifdef _Py_JIT
     executor->jit_code = NULL;
-    executor->code_size = 0;
+    executor->jit_size = 0;
     if (_PyJIT_Compile(executor, executor->trace, length+1)) {
         Py_DECREF(executor);
         return NULL;