pythongh-116168 WIP: Remove extra _CHECK_STACK_SPACE uops

Lib/test/test_capi/test_opt.py

@@ -955,6 +955,87 @@ def testfunc(n):
         _, ex = self._run_with_optimizer(testfunc, 16)
         self.assertIsNone(ex)
 
+    def test_remove_extra_stack_space_checks_sequential(self):
+        def dummy12(x):  # co_framesize: 12 (TODO: assert this?)
+            return x - 1
+        def dummy13(y):  # co_framesize: 13 (TODO: assert this?)
+            z = y + 2
+            return y, z
+        def testfunc(n):
+            a = 0
+            for _ in range(n):
+                b = dummy12(7)
+                c, d = dummy13(9)
+                a += b + c + d
+            return a
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertEqual(res, 832)
+        self.assertIsNotNone(ex)
+
+        uops_and_operands = [(opcode, operand) for opcode, _, _, operand in list(ex)]
+        uop_names = [uop[0] for uop in uops_and_operands]
+        self.assertEqual(uop_names.count("_PUSH_FRAME"), 2)
+        self.assertEqual(uop_names.count("_POP_FRAME"), 2)
+        self.assertEqual(uop_names.count("_CHECK_STACK_SPACE"), 1)
+        # sequential calls: max(12, 13) == 13
+        self.assertIn(("_CHECK_STACK_SPACE", 13), uops_and_operands)
+
+    def test_remove_extra_stack_space_checks_nested(self):
+        def dummy12(x):  # co_framesize: 12 (TODO: assert this?)
+            return x + 3
+        def dummy15(y):  # co_framesize: 15 (TODO: assert this?)
+            z = dummy12(y)
+            return y, z
+        def testfunc(n):
+            a = 0
+            for _ in range(n):
+                b, c = dummy15(2)
+                a += b + c
+            return a
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertEqual(res, 224)
+        self.assertIsNotNone(ex)
+
+        uops_and_operands = [(opcode, operand) for opcode, _, _, operand in list(ex)]
+        uop_names = [uop[0] for uop in uops_and_operands]
+        self.assertEqual(uop_names.count("_PUSH_FRAME"), 2)
+        self.assertEqual(uop_names.count("_POP_FRAME"), 2)
+        self.assertEqual(uop_names.count("_CHECK_STACK_SPACE"), 1)
+        # nested calls: 15 + 12 == 27
+        self.assertIn(("_CHECK_STACK_SPACE", 27), uops_and_operands)
+
+
+    def test_remove_extra_stack_space_checks_many_calls(self):
+        def dummy12(x):  # co_framesize: 12 (TODO: assert this?)
+            return x + 3
+        def dummy13(y):  # co_framesize: 13 (TODO: assert this?)
+            z = y + 2
+            return y, z
+        def dummy18(y):  # co_framesize: 18 (TODO: assert this?)
+            z = dummy12(y)
+            x, w = dummy13(z)
+            return z, x, w
+        def testfunc(n):
+            a = 0
+            for _ in range(n):
+                b = dummy12(5)
+                c, d, e = dummy18(2)
+                a += b + c + d + e
+            return a
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertEqual(res, 800)
+        self.assertIsNotNone(ex)
+
+        uops_and_operands = [(opcode, operand) for opcode, _, _, operand in list(ex)]
+        uop_names = [uop[0] for uop in uops_and_operands]
+        self.assertEqual(uop_names.count("_PUSH_FRAME"), 4)
+        self.assertEqual(uop_names.count("_POP_FRAME"), 4)
+        self.assertEqual(uop_names.count("_CHECK_STACK_SPACE"), 1)
+        # max(12, 18 + max(12, 13)) == 31
+        self.assertIn(("_CHECK_STACK_SPACE", 31), uops_and_operands)
 
 if __name__ == "__main__":
     unittest.main()

Python/optimizer.c

@@ -769,6 +769,12 @@ translate_bytecode_to_trace(
                                 /* Set the operand to the callee's function object,
                                  * to assist optimization passes */
                                 ADD_TO_TRACE(uop, oparg, (uintptr_t)new_func, target);
+                                // tell the previous _CHECK_STACK_SPACE how much space we need
+                                PyCodeObject *new_func_co = (PyCodeObject *)new_func->func_code;
+                                int last_check_stack_idx = trace_length - 4;
+                                assert(last_check_stack_idx >= 0);
+                                assert(trace[last_check_stack_idx].opcode == _CHECK_STACK_SPACE);
+                                trace[last_check_stack_idx].operand = new_func_co->co_framesize;
                                 code = new_code;
                                 func = new_func;
                                 instr = _PyCode_CODE(code);

Python/optimizer_analysis.c

@@ -555,6 +555,17 @@ peephole_opt(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer, int buffer_s
     }
 }
 
+static void
+remove_extra_stack_space_checks(
+    _PyInterpreterFrame *frame,
+    _PyUOpInstruction *buffer,
+    int buffer_size
+)
+{
+    // TODO: implement
+    return;
+}
+
 //  0 - failure, no error raised, just fall back to Tier 1
 // -1 - failure, and raise error
 //  1 - optimizer success
@@ -590,6 +601,9 @@ _Py_uop_analyze_and_optimize(
 
     remove_unneeded_uops(buffer, buffer_size);
 
+    // TODO: is this the best place for this call?
+    remove_extra_stack_space_checks(frame, buffer, buffer_size);
+
     OPT_STAT_INC(optimizer_successes);
     return 1;
 not_ready:

Python/optimizer_bytecodes.c

@@ -538,6 +538,10 @@ dummy_func(void) {
         }
     }
 
+    // op(_CHCEK_STACK_SPACE, (unused/2, unused[oparg] -- unused/2, unused[oparg]) {
+    //     // TODO: implement tier2 behavior. The operand should contain the necessary co_framesize.
+    // }
+
     op(_INIT_CALL_PY_EXACT_ARGS, (callable, self_or_null, args[oparg] -- new_frame: _Py_UOpsAbstractFrame *)) {
         int argcount = oparg;