[mypyc] Use METH_FASTCALL with "__call__" (#9946)

Allocate a vectorcall function pointer as a struct field for native
classes that include `__call__`, including nested functions. This
lets us use METH_FASTCALL wrapper functions with `__call__`
methods.

See https://www.python.org/dev/peps/pep-0590/ for details of why
we jump through these hoops.

This makes the `nested_func` microbenchmark about 1.5x faster.

Follow-up to #9894.

Author: JukkaL

mypyc/codegen/emitclass.py

@@ -1,10 +1,10 @@
 """Code generation for native classes and related wrappers."""
 
-
 from typing import Optional, List, Tuple, Dict, Callable, Mapping, Set
+
 from mypy.ordered_dict import OrderedDict
 
-from mypyc.common import PREFIX, NATIVE_PREFIX, REG_PREFIX, USE_FASTCALL
+from mypyc.common import PREFIX, NATIVE_PREFIX, REG_PREFIX, USE_FASTCALL, USE_VECTORCALL
 from mypyc.codegen.emit import Emitter, HeaderDeclaration
 from mypyc.codegen.emitfunc import native_function_header
 from mypyc.codegen.emitwrapper import (
@@ -35,7 +35,7 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 
 SLOT_DEFS = {
     '__init__': ('tp_init', lambda c, t, e: generate_init_for_class(c, t, e)),
-    '__call__': ('tp_call', wrapper_slot),
+    '__call__': ('tp_call', lambda c, t, e: generate_call_wrapper(c, t, e)),
     '__str__': ('tp_str', native_slot),
     '__repr__': ('tp_repr', native_slot),
     '__next__': ('tp_iternext', native_slot),
@@ -71,6 +71,15 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 }
 
 
+def generate_call_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
+    if USE_VECTORCALL:
+        # Use vectorcall wrapper if supported (PEP 590).
+        return 'PyVectorcall_Call'
+    else:
+        # On older Pythons use the legacy wrapper.
+        return wrapper_slot(cl, fn, emitter)
+
+
 def generate_slots(cl: ClassIR, table: SlotTable, emitter: Emitter) -> Dict[str, str]:
     fields = OrderedDict()  # type: Dict[str, str]
     # Sort for determinism on Python 3.5
@@ -241,6 +250,10 @@ def emit_line() -> None:
     flags = ['Py_TPFLAGS_DEFAULT', 'Py_TPFLAGS_HEAPTYPE', 'Py_TPFLAGS_BASETYPE']
     if generate_full:
         flags.append('Py_TPFLAGS_HAVE_GC')
+    if cl.has_method('__call__') and USE_VECTORCALL:
+        fields['tp_vectorcall_offset'] = 'offsetof({}, vectorcall)'.format(
+            cl.struct_name(emitter.names))
+        flags.append('_Py_TPFLAGS_HAVE_VECTORCALL')
     fields['tp_flags'] = ' | '.join(flags)
 
     emitter.emit_line("static PyTypeObject {}_template_ = {{".format(emitter.type_struct_name(cl)))
@@ -277,6 +290,8 @@ def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
     lines += ['typedef struct {',
               'PyObject_HEAD',
               'CPyVTableItem *vtable;']
+    if cl.has_method('__call__') and USE_VECTORCALL:
+        lines.append('vectorcallfunc vectorcall;')
     for base in reversed(cl.base_mro):
         if not base.is_trait:
             for attr, rtype in base.attributes.items():
@@ -451,6 +466,10 @@ def generate_setup_for_class(cl: ClassIR,
     else:
         emitter.emit_line('self->vtable = {};'.format(vtable_name))
 
+    if cl.has_method('__call__') and USE_VECTORCALL:
+        name = cl.method_decl('__call__').cname(emitter.names)
+        emitter.emit_line('self->vectorcall = {}{};'.format(PREFIX, name))
+
     for base in reversed(cl.base_mro):
         for attr, rtype in base.attributes.items():
             emitter.emit_line('self->{} = {};'.format(

mypyc/codegen/emitmodule.py

@@ -24,7 +24,7 @@
 from mypyc.irbuild.mapper import Mapper
 from mypyc.common import (
     PREFIX, TOP_LEVEL_NAME, INT_PREFIX, MODULE_PREFIX, RUNTIME_C_FILES, USE_FASTCALL,
-    shared_lib_name,
+    USE_VECTORCALL, shared_lib_name,
 )
 from mypyc.codegen.cstring import encode_as_c_string, encode_bytes_as_c_string
 from mypyc.codegen.emit import EmitterContext, Emitter, HeaderDeclaration
@@ -1071,5 +1071,10 @@ def visit(item: T) -> None:
 
 
 def is_fastcall_supported(fn: FuncIR) -> bool:
-    # TODO: Support METH_FASTCALL for all methods.
-    return USE_FASTCALL and (fn.class_name is None or fn.name not in ('__init__', '__call__'))
+    if fn.class_name is not None:
+        if fn.name == '__call__':
+            # We can use vectorcalls (PEP 590) when supported
+            return USE_VECTORCALL
+        # TODO: Support fastcall for __init__.
+        return USE_FASTCALL and fn.name != '__init__'
+    return USE_FASTCALL

mypyc/codegen/emitwrapper.py

@@ -14,7 +14,7 @@
 
 from mypy.nodes import ARG_POS, ARG_OPT, ARG_NAMED_OPT, ARG_NAMED, ARG_STAR, ARG_STAR2
 
-from mypyc.common import PREFIX, NATIVE_PREFIX, DUNDER_PREFIX
+from mypyc.common import PREFIX, NATIVE_PREFIX, DUNDER_PREFIX, USE_VECTORCALL
 from mypyc.codegen.emit import Emitter
 from mypyc.ir.rtypes import (
     RType, is_object_rprimitive, is_int_rprimitive, is_bool_rprimitive, object_rprimitive
@@ -157,9 +157,13 @@ def generate_wrapper_function(fn: FuncIR,
         arg_ptrs += ['&obj_{}'.format(groups[ARG_STAR2][0].name) if groups[ARG_STAR2] else 'NULL']
     arg_ptrs += ['&obj_{}'.format(arg.name) for arg in reordered_args]
 
+    if fn.name == '__call__' and USE_VECTORCALL:
+        nargs = 'PyVectorcall_NARGS(nargs)'
+    else:
+        nargs = 'nargs'
     emitter.emit_lines(
-        'if (!CPyArg_ParseStackAndKeywords(args, nargs, kwnames, &parser{})) {{'.format(
-            ''.join(', ' + n for n in arg_ptrs)),
+        'if (!CPyArg_ParseStackAndKeywords(args, {}, kwnames, &parser{})) {{'.format(
+            nargs, ''.join(', ' + n for n in arg_ptrs)),
         'return NULL;',
         '}')
     traceback_code = generate_traceback_code(fn, emitter, source_path, module_name)

mypyc/common.py

@@ -50,8 +50,11 @@
 MAX_LITERAL_SHORT_INT = (sys.maxsize >> 1 if not IS_MIXED_32_64_BIT_BUILD
                          else 2**30 - 1)  # type: Final
 
-# We can use faster wrapper functions on Python 3.7+ (fastcall/vectorcall).
-USE_FASTCALL = sys.version_info >= (3, 7)
+# We can use METH_FASTCALL faster wrapper functions on Python 3.7+.
+USE_FASTCALL = sys.version_info >= (3, 7)  # type: Final
+
+# We can use vectorcalls on Python 3.8+ (PEP 590).
+USE_VECTORCALL = sys.version_info >= (3, 8)  # type: Final
 
 # Runtime C library files
 RUNTIME_C_FILES = [

mypyc/irbuild/callable_class.py

@@ -97,6 +97,7 @@ def add_call_to_callable_class(builder: IRBuilder,
     call_fn_ir = FuncIR(call_fn_decl, args, blocks,
                         fn_info.fitem.line, traceback_name=fn_info.fitem.name)
     fn_info.callable_class.ir.methods['__call__'] = call_fn_ir
+    fn_info.callable_class.ir.method_decls['__call__'] = call_fn_decl
     return call_fn_ir