First steps towards POP_JUMP_IF_XXX uops

This adds: - Hand-written uops JUMP_IF_{TRUE,FALSE,NONE,NOT_NONE}. These peek at the top of the stack. The jump target (in superblock space) is absolute. - Hand-written translation for POP_JUMP_IF_TRUE, assuming the jump is unlikely. This can serve as an example for the rest. (Probably adding the others will cause a little refactoring.) Once we implement jump-likelihood profiling, we can implement the jump-unlikely case. Note that the stubs are placed at the end of the trace array. There is a gap between the main superblock and the stubs. We could close the gap, we'd just have to patch up the JUMP uops.
python · gvanrossum · Jul 10, 2023 · Jul 8, 2023 · Jul 9, 2023 · Jul 9, 2023
commit c707cffbbcb77c412438aebca051a83de7e5089c
@@ -2755,7 +2755,8 @@ _PyUopExecute(_PyExecutorObject *executor, _PyInterpreterFrame *frame, PyObject
         operand = self->trace[pc].operand;
         oparg = (int)operand;
         DPRINTF(3,
-                "  uop %s, operand %" PRIu64 ", stack_level %d\n",
+                "%4d: uop %s, operand %" PRIu64 ", stack_level %d\n",
+                pc,
                 opcode < 256 ? _PyOpcode_OpName[opcode] : _PyOpcode_uop_name[opcode],
                 operand,
                 (int)(stack_pointer - _PyFrame_Stackbase(frame)));
@@ -2767,6 +2768,38 @@ _PyUopExecute(_PyExecutorObject *executor, _PyInterpreterFrame *frame, PyObject
 #define ENABLE_SPECIALIZATION 0
 #include "executor_cases.c.h"
 
+            case JUMP_IF_FALSE:
+            {
+                if (Py_IsFalse(stack_pointer[-1])) {
+                    pc = oparg;
+                }
+                break;
+            }
+
+            case JUMP_IF_TRUE:
+            {
+                if (Py_IsTrue(stack_pointer[-1])) {
+                    pc = oparg;
+                }
+                break;
+            }
+
+            case JUMP_IF_NONE:
+            {
+                if (Py_IsNone(stack_pointer[-1])) {
+                    pc = oparg;
+                }
+                break;
+            }
+
+            case JUMP_IF_NOT_NONE:
+            {
+                if (!Py_IsNone(stack_pointer[-1])) {
+                    pc = oparg;
+                }
+                break;
+            }
+
             case SAVE_IP:
             {
                 frame->prev_instr = ip_offset + oparg;

diff --git a/Python/opcode_metadata.h b/Python/opcode_metadata.h
@@ -401,6 +401,14 @@ translate_bytecode_to_trace(
     trace[trace_length].operand = (OPERAND); \
     trace_length++;
 
+#define ADD_TO_STUB(INDEX, OPCODE, OPERAND) \
+    DPRINTF(2, "    ADD_TO_STUB(%d, %s, %" PRIu64 ")\n", \
+            (INDEX), \
+            (OPCODE) < 256 ? _PyOpcode_OpName[(OPCODE)] : _PyOpcode_uop_name[(OPCODE)], \
+            (uint64_t)(OPERAND)); \
+    trace[(INDEX)].opcode = (OPCODE); \
+    trace[(INDEX)].operand = (OPERAND);
+
     DPRINTF(4,
             "Optimizing %s (%s:%d) at byte offset %ld\n",
             PyUnicode_AsUTF8(code->co_qualname),
@@ -409,7 +417,7 @@ translate_bytecode_to_trace(
             2 * (long)(initial_instr - (_Py_CODEUNIT *)code->co_code_adaptive));
 
     for (;;) {
-        ADD_TO_TRACE(SAVE_IP, (int)(instr - (_Py_CODEUNIT *)code->co_code_adaptive));
+        ADD_TO_TRACE(SAVE_IP, instr - (_Py_CODEUNIT *)code->co_code_adaptive);
         int opcode = instr->op.code;
         int oparg = instr->op.arg;
         int extras = 0;
@@ -426,6 +434,25 @@ translate_bytecode_to_trace(
             oparg = (oparg & 0xffffff00) | executor->vm_data.oparg;
         }
         switch (opcode) {
+
+            case POP_JUMP_IF_FALSE:
+            {
+                // Assume jump unlikely (TODO: handle jump likely case)
+                // Reserve 7 entries (2 here, 3 stub, plus SAVE_IP + EXIT_TRACE)
+                if (trace_length + 7 > max_length) {
+                    DPRINTF(1, "Ran out of space for POP_JUMP_IF_FALSE\n");
+                    goto done;
+                }
+                _Py_CODEUNIT *target_instr = instr + 1 + _PyOpcode_Caches[_PyOpcode_Deopt[opcode]] + oparg;
+                max_length -= 3;  // Really the start of the stubs
+                ADD_TO_TRACE(JUMP_IF_FALSE, max_length);
+                ADD_TO_TRACE(POP_TOP, 0);
+                ADD_TO_STUB(max_length, POP_TOP, 0);
+                ADD_TO_STUB(max_length + 1, SAVE_IP, target_instr - (_Py_CODEUNIT *)code->co_code_adaptive);
+                ADD_TO_STUB(max_length + 2, EXIT_TRACE, 0);
+                break;
+            }
+
             default:
             {
                 const struct opcode_macro_expansion *expansion = &_PyOpcode_macro_expansion[opcode];
@@ -538,7 +565,7 @@ uop_optimize(
         return -1;
     }
     executor->base.execute = _PyUopExecute;
-    memcpy(executor->trace, trace, trace_length * sizeof(_PyUOpInstruction));
+    memcpy(executor->trace, trace, _Py_UOP_MAX_TRACE_LENGTH * sizeof(_PyUOpInstruction));
     *exec_ptr = (_PyExecutorObject *)executor;
     return 1;
 }

@@ -1338,12 +1338,19 @@ def write_pseudo_instrs(self) -> None:
     def write_uop_items(self, make_text: typing.Callable[[str, int], str]) -> None:
         """Write '#define XXX NNN' for each uop"""
         counter = 300  # TODO: Avoid collision with pseudo instructions
+
         def add(name: str) -> None:
             nonlocal counter
             self.out.emit(make_text(name, counter))
             counter += 1
+
         add("EXIT_TRACE")
         add("SAVE_IP")
+        add("JUMP_IF_FALSE")
+        add("JUMP_IF_TRUE")
+        add("JUMP_IF_NONE")
+        add("JUMP_IF_NOT_NONE")
+
         for instr in self.instrs.values():
             if instr.kind == "op" and instr.is_viable_uop():
                 add(instr.name)