Specialize for calls to simple Python functions with perfectly matchi…

…ng arguments.
python · markshannon · Oct 20, 2021 · Oct 11, 2021 · Oct 13, 2021 · Oct 13, 2021
commit d5add72681b11594f7a919a4a888139976c33b9e
diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
@@ -17,6 +17,7 @@ typedef struct {
     uint8_t original_oparg;
     uint8_t counter;
     uint16_t index;
+    uint32_t version;
 } _PyAdaptiveEntry;
 
 
@@ -308,7 +309,7 @@ int _Py_Specialize_LoadGlobal(PyObject *globals, PyObject *builtins, _Py_CODEUNI
 int _Py_Specialize_LoadMethod(PyObject *owner, _Py_CODEUNIT *instr, PyObject *name, SpecializedCacheEntry *cache);
 int _Py_Specialize_BinarySubscr(PyObject *sub, PyObject *container, _Py_CODEUNIT *instr);
 int _Py_Specialize_BinaryAdd(PyObject *sub, PyObject *container, _Py_CODEUNIT *instr);
-int _Py_Specialize_CallFunction(PyThreadState *tstate, PyObject *callable, int nargs, _Py_CODEUNIT *instr, SpecializedCacheEntry *cache);
+int _Py_Specialize_CallFunction(PyObject *callable, _Py_CODEUNIT *instr, int nargs, SpecializedCacheEntry *cache);
 
 #define PRINT_SPECIALIZATION_STATS 0
 #define PRINT_SPECIALIZATION_STATS_DETAILED 0

diff --git a/Include/opcode.h b/Include/opcode.h
diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -248,6 +248,7 @@ def jabs_op(name, op):
     "STORE_ATTR_SLOT",
     "STORE_ATTR_WITH_HINT",
     "CALL_FUNCTION_ADAPTIVE",
+    "CALL_FUNCTION_PY_SIMPLE",
     # Super instructions
     "LOAD_FAST__LOAD_FAST",
     "STORE_FAST__LOAD_FAST",

@@ -4663,7 +4663,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
             if (cache->adaptive.counter == 0) {
                 PyObject *callable = PEEK(nargs+1);
                 next_instr--;
-                if (_Py_Specialize_CallFunction(tstate, callable, nargs, next_instr, cache) < 0) {
+                if (_Py_Specialize_CallFunction(callable, next_instr, nargs, cache) < 0) {
                     goto error;
                 }
                 DISPATCH();
@@ -4678,6 +4678,36 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
             }
         }
 
+        TARGET(CALL_FUNCTION_PY_SIMPLE) {
+            SpecializedCacheEntry *caches = GET_CACHE();
+            _PyAdaptiveEntry *cache0 = &caches[0].adaptive;
+            int argcount = cache0->original_oparg;
+            PyObject *callable = PEEK(argcount+1);
+            DEOPT_IF(!PyFunction_Check(callable), CALL_FUNCTION);
+            PyFunctionObject *func = (PyFunctionObject *)callable;
+            DEOPT_IF(func->func_version != cache0->version, CALL_FUNCTION);
+            /* PEP 523 */
+            DEOPT_IF(tstate->interp->eval_frame != NULL, CALL_FUNCTION);
+            STAT_INC(CALL_FUNCTION, hit);
+            record_cache_hit(cache0);
+            InterpreterFrame *new_frame = _PyThreadState_PushFrame(
+                tstate, PyFunction_AS_FRAME_CONSTRUCTOR(func), NULL);
+            if (new_frame == NULL) {
+                goto error;
+            }
+            STACK_SHRINK(argcount);
+            for (int i = 0; i < argcount; i++) {
+                new_frame->localsplus[i] = stack_pointer[i];
+            }
+            STACK_SHRINK(1);
+            Py_DECREF(func);
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            new_frame->previous = tstate->frame;
+            new_frame->depth = frame->depth + 1;
+            tstate->frame = frame = new_frame;
+            goto start_frame;
+        }
+
         TARGET(CALL_FUNCTION_EX) {
             PREDICTED(CALL_FUNCTION_EX);
             PyObject *func, *callargs, *kwargs = NULL, *result;
@@ -4946,6 +4976,7 @@ MISS_WITH_CACHE(LOAD_ATTR)
 MISS_WITH_CACHE(STORE_ATTR)
 MISS_WITH_CACHE(LOAD_GLOBAL)
 MISS_WITH_CACHE(LOAD_METHOD)
+MISS_WITH_CACHE(CALL_FUNCTION)
 MISS_WITH_OPARG_COUNTER(BINARY_SUBSCR)
 MISS_WITH_OPARG_COUNTER(BINARY_ADD)
 

diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -457,7 +457,11 @@ _Py_Quicken(PyCodeObject *code) {
 /* Calls */
 #define SPEC_FAIL_BUILTIN_FUNCTION 7
 #define SPEC_FAIL_CLASS 8
-#define SPEC_FAIL_PYTHON_FUNCTION 9
+#define SPEC_FAIL_GENERATOR 9
+#define SPEC_FAIL_COMPLEX_PARAMETERS 10
+#define SPEC_FAIL_WRONG_NUMBER_ARGUMENTS 10
+#define SPEC_FAIL_CO_NOT_OPTIMIZED 11
+#define SPEC_FAIL_FREE_VARS 12
 
 
 static int
@@ -1196,42 +1200,79 @@ _Py_Specialize_BinaryAdd(PyObject *left, PyObject *right, _Py_CODEUNIT *instr)
     return 0;
 }
 
-static int specialize_c_call(
-    PyObject *callable, int nargs, _Py_CODEUNIT *instr, SpecializedCacheEntry *cache)
+static int
+specialize_c_call(
+    PyObject *callable, _Py_CODEUNIT *instr,
+    int nargs, SpecializedCacheEntry *cache)
 {
     SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_BUILTIN_FUNCTION);
     return -1;
 }
 
-static int specialize_class_call(
-    PyObject *callable, int nargs, _Py_CODEUNIT *instr, SpecializedCacheEntry *cache)
+static int
+specialize_class_call(
+    PyObject *callable, _Py_CODEUNIT *instr,
+    int nargs, SpecializedCacheEntry *cache)
 {
     SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_CLASS);
     return -1;
 }
 
-static int specialize_py_call(PyThreadState *tstate,
-    PyObject *callable, int nargs, _Py_CODEUNIT *instr, SpecializedCacheEntry *cache)
+static int
+specialize_py_call(
+    PyFunctionObject *func, _Py_CODEUNIT *instr,
+    int nargs, SpecializedCacheEntry *cache)
 {
-    SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_PYTHON_FUNCTION);
-    return -1;
+    /* Exclude generator or coroutines for now */
+    PyCodeObject *code = (PyCodeObject *)func->func_code;
+    int flags = code->co_flags;
+    if (flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_GENERATOR);
+        return -1;
+    }
+    if ((flags & (CO_VARKEYWORDS | CO_VARARGS)) || code->co_kwonlyargcount) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_COMPLEX_PARAMETERS);
+        return -1;
+    }
+    if (code->co_argcount != nargs) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_WRONG_NUMBER_ARGUMENTS);
+        return -1;
+    }
+    if ((flags & CO_OPTIMIZED) == 0) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_CO_NOT_OPTIMIZED);
+        return -1;
+    }
+    if (code->co_nfreevars) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_FREE_VARS);
+        return -1;
+    }
+    _PyAdaptiveEntry *cache0 = &cache->adaptive;
+    int version = _PyFunction_GetVersionForCurrentState(func);
+    if (version == 0) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_OUT_OF_VERSIONS);
+        return -1;
+    }
+    cache0->version = version;
+    *instr = _Py_MAKECODEUNIT(CALL_FUNCTION_PY_SIMPLE, _Py_OPARG(*instr));
+    return 0;
 }
 
+
 int
 _Py_Specialize_CallFunction(
-    PyThreadState *tstate, PyObject *callable, int nargs,
-    _Py_CODEUNIT *instr, SpecializedCacheEntry *cache)
+    PyObject *callable, _Py_CODEUNIT *instr,
+    int nargs, SpecializedCacheEntry *cache)
 {
     _PyAdaptiveEntry *cache0 = &cache->adaptive;
     int fail;
     if (PyCFunction_CheckExact(callable)) {
-        fail = specialize_c_call(callable, nargs, instr, cache);
+        fail = specialize_c_call(callable, instr, nargs, cache);
     }
     else if (PyFunction_Check(callable)) {
-        fail = specialize_py_call(tstate, callable, nargs, instr, cache);
+        fail = specialize_py_call((PyFunctionObject *)callable, instr, nargs, cache);
     }
     else if (PyType_Check(callable)) {
-        fail = specialize_class_call(callable, nargs, instr, cache);
+        fail = specialize_class_call(callable, instr, nargs, cache);
     }
     else {
         SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_OTHER);