WIP cache freed generator object instance to speed up allocation

Author: JukkaL

mypyc/codegen/emit.py

@@ -1115,6 +1115,31 @@ def emit_gc_clear(self, target: str, rtype: RType) -> None:
         else:
             assert False, "emit_gc_clear() not implemented for %s" % repr(rtype)
 
+    def emit_reuse_clear(self, target: str, rtype: RType) -> None:
+        """Emit attribute clear before object is added into freelist.
+
+        Assume that 'target' represents a C expression that refers to a
+        struct member, such as 'self->x'.
+
+        Unlike emit_gc_clear(), initialize attribute value to match a freshly
+        allocated object.
+        """
+        if isinstance(rtype, RTuple):
+            for i, item_type in enumerate(rtype.types):
+                self.emit_reuse_clear(f"{target}.f{i}", item_type)
+        elif not rtype.is_refcounted:
+            self.emit_line(f"{target} = {rtype.c_undefined};")
+        elif isinstance(rtype, RPrimitive) and rtype.name == "builtins.int":
+            self.emit_line(f"if (CPyTagged_CheckLong({target})) {{")
+            self.emit_line(f"CPyTagged __tmp = {target};")
+            self.emit_line(f"{target} = {self.c_undefined_value(rtype)};")
+            self.emit_line("Py_XDECREF(CPyTagged_LongAsObject(__tmp));")
+            self.emit_line("} else {")
+            self.emit_line(f"{target} = {self.c_undefined_value(rtype)};")
+            self.emit_line("}")
+        else:
+            self.emit_gc_clear(target, rtype)
+
     def emit_traceback(
         self, source_path: str, module_name: str, traceback_entry: tuple[str, int]
     ) -> None:

mypyc/codegen/emitclass.py

@@ -186,6 +186,16 @@ def generate_class_type_decl(
         )
 
 
+def generate_class_reuse(
+    cl: ClassIR, c_emitter: Emitter, external_emitter: Emitter, emitter: Emitter
+) -> None:
+    assert cl.reuse_freed_instance
+    context = c_emitter.context
+    name = cl.name_prefix(c_emitter.names) + "_free_instance"
+    struct_name = cl.struct_name(c_emitter.names)
+    context.declarations[name] = HeaderDeclaration(f"{struct_name} *{name};", needs_export=True)
+
+
 def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     """Generate C code for a class.
 
@@ -557,7 +567,17 @@ def generate_setup_for_class(
     emitter.emit_line("static PyObject *")
     emitter.emit_line(f"{func_name}(PyTypeObject *type)")
     emitter.emit_line("{")
-    emitter.emit_line(f"{cl.struct_name(emitter.names)} *self;")
+    struct_name = cl.struct_name(emitter.names)
+    emitter.emit_line(f"{struct_name} *self;")
+
+    prefix = cl.name_prefix(emitter.names)
+    if cl.reuse_freed_instance:
+        emitter.emit_line(f"if ({prefix}_free_instance != NULL) {{")
+        emitter.emit_line(f"self = {prefix}_free_instance;")
+        emitter.emit_line(f"{prefix}_free_instance = NULL;")
+        emitter.emit_line("return (PyObject *)self;")
+        emitter.emit_line("}")
+
     emitter.emit_line(f"self = ({cl.struct_name(emitter.names)} *)type->tp_alloc(type, 0);")
     emitter.emit_line("if (self == NULL)")
     emitter.emit_line("    return NULL;")
@@ -786,6 +806,8 @@ def generate_dealloc_for_class(
         emitter.emit_line("if (!PyObject_GC_IsFinalized((PyObject *)self)) {")
         emitter.emit_line("Py_TYPE(self)->tp_finalize((PyObject *)self);")
         emitter.emit_line("}")
+    if cl.reuse_freed_instance:
+        emit_reuse_dealloc(cl, emitter)
     emitter.emit_line("PyObject_GC_UnTrack(self);")
     # The trashcan is needed to handle deep recursive deallocations
     emitter.emit_line(f"CPy_TRASHCAN_BEGIN(self, {dealloc_func_name})")
@@ -795,6 +817,22 @@ def generate_dealloc_for_class(
     emitter.emit_line("}")
 
 
+def emit_reuse_dealloc(cl: ClassIR, emitter: Emitter) -> None:
+    prefix = cl.name_prefix(emitter.names)
+    emitter.emit_line(f"if ({prefix}_free_instance == NULL) {{")
+    emitter.emit_line(f"{prefix}_free_instance = self;")
+    emitter.emit_line("Py_INCREF(self);")
+
+    # TODO: emit_clear_bitmaps(cl, emitter)
+
+    for base in reversed(cl.base_mro):
+        for attr, rtype in base.attributes.items():
+            emitter.emit_reuse_clear(f"self->{emitter.attr(attr)}", rtype)
+
+    emitter.emit_line("return;")
+    emitter.emit_line("}")
+
+
 def generate_finalize_for_class(
     del_method: FuncIR, finalize_func_name: str, emitter: Emitter
 ) -> None:

mypyc/codegen/emitmodule.py

@@ -29,7 +29,7 @@
 from mypy.util import hash_digest, json_dumps
 from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import Emitter, EmitterContext, HeaderDeclaration, c_array_initializer
-from mypyc.codegen.emitclass import generate_class, generate_class_type_decl
+from mypyc.codegen.emitclass import generate_class, generate_class_reuse, generate_class_type_decl
 from mypyc.codegen.emitfunc import generate_native_function, native_function_header
 from mypyc.codegen.emitwrapper import (
     generate_legacy_wrapper_function,
@@ -609,6 +609,8 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
             self.declare_finals(module_name, module.final_names, declarations)
             for cl in module.classes:
                 generate_class_type_decl(cl, emitter, ext_declarations, declarations)
+                if cl.reuse_freed_instance:
+                    generate_class_reuse(cl, emitter, ext_declarations, declarations)
             self.declare_type_vars(module_name, module.type_var_names, declarations)
             for fn in module.functions:
                 generate_function_declaration(fn, declarations)

mypyc/ir/class_ir.py

@@ -204,6 +204,11 @@ def __init__(
         # If this is a generator environment class, what is the actual method for it
         self.env_user_function: FuncIR | None = None
 
+        # If True, keep one freed, cleared instance available for immediate reuse to
+        # speed up allocations. This helps if many objects are freed quickly, before
+        # other instances of the same class are allocated.
+        self.reuse_freed_instance = False
+
     def __repr__(self) -> str:
         return (
             "ClassIR("

mypyc/irbuild/generator.py

@@ -156,6 +156,7 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
     name = f"{builder.fn_info.namespaced_name()}_gen"
 
     generator_class_ir = ClassIR(name, builder.module_name, is_generated=True, is_final_class=True)
+    generator_class_ir.reuse_freed_instance = True
     if builder.fn_info.can_merge_generator_and_env_classes():
         builder.fn_info.env_class = generator_class_ir
     else: