Add JIT support and tests

Author: Fidget-Spinner

Lib/test/test_capi/test_opt.py

@@ -400,6 +400,23 @@ def dummy(x):
         self.assertIn("_PUSH_FRAME", uops)
         self.assertIn("_BINARY_OP_ADD_INT", uops)
 
+    def test_call_py_ex(self):
+        def testfunc(n):
+            def ex_py(*args, **kwargs):
+                return 1
+
+            for _ in range(n):
+                args = (1, 2, 3)
+                kwargs = {}
+                ex_py(*args, **kwargs)
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_PUSH_FRAME", uops)
+
     def test_branch_taken(self):
         def testfunc(n):
             for i in range(n):

Lib/test/test_opcache.py

@@ -591,6 +591,23 @@ def instantiate():
             instantiate()
         self.assert_specialized(instantiate, "CALL_FUNCTION_EX_PY")
 
+    @requires_jit_disabled
+    @requires_specialization_ft
+    def test_specialize_call_function_ex_py_fail(self):
+        def ex_py(*args, **kwargs):
+            return 1
+
+        def instantiate():
+            args = (1, 2, 3)
+            kwargs = {}
+            return ex_py(*args, **kwargs)
+
+        _testinternalcapi.set_vectorcall_nop(ex_py)
+        # Trigger specialization
+        for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+            instantiate()
+        self.assert_no_opcode(instantiate, "CALL_FUNCTION_EX_PY")
+
 def make_deferred_ref_count_obj():
     """Create an object that uses deferred reference counting.
 

Python/optimizer_bytecodes.c

@@ -831,6 +831,17 @@ dummy_func(void) {
         new_frame = PyJitRef_Wrap((JitOptSymbol *)frame_new(ctx, co, 0, NULL, 0));
     }
 
+    op(_PY_FRAME_EX, (func_st, null, callargs_st, kwargs_st -- ex_frame)) {
+        assert((this_instr + 2)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging((this_instr + 2));
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+
+        ex_frame = PyJitRef_Wrap((JitOptSymbol *)frame_new(ctx, co, 0, NULL, 0));
+    }
+
     op(_CHECK_AND_ALLOCATE_OBJECT, (type_version/2, callable, self_or_null, args[oparg] -- callable, self_or_null, args[oparg])) {
         (void)type_version;
         (void)args;