From cbd96f96a67bc5c2088457fcd30742644614f00b Mon Sep 17 00:00:00 2001 From: jairov4 <1904410+jairov4@users.noreply.github.com> Date: Mon, 14 Oct 2024 03:51:41 -0500 Subject: [PATCH] [mypyc] Optimize calls to final classes (#17886) Fixes #9612 This change allows to gain more efficiency where classes are annotated with `@final` bypassing entirely the vtable for method calls and property accessors. For example: In ```python @final class Vector: __slots__ = ("_x", "_y") def __init__(self, x: i32, y: i32) -> None: self._x = x self._y = y @property def y(self) -> i32: return self._y def test_vector() -> None: v3 = Vector(1, 2) assert v3.y == 2 ``` The call will produce: ```c ... cpy_r_r6 = CPyDef_Vector___y(cpy_r_r0); ... ``` Instead of: ```c ... cpy_r_r1 = CPY_GET_ATTR(cpy_r_r0, CPyType_Vector, 2, farm_rush___engine___vectors2___VectorObject, int32_t); /* y */ ... ``` (which uses vtable) --- mypyc/codegen/emitclass.py | 1 + mypyc/codegen/emitfunc.py | 42 ++++++++++-------- mypyc/ir/class_ir.py | 10 +++-- mypyc/ir/rtypes.py | 2 +- mypyc/irbuild/ll_builder.py | 4 +- mypyc/irbuild/prepare.py | 6 ++- mypyc/irbuild/util.py | 11 ++++- mypyc/test-data/run-classes.test | 75 ++++++++++++++++++++++++++++++++ 8 files changed, 125 insertions(+), 26 deletions(-) diff --git a/mypyc/codegen/emitclass.py b/mypyc/codegen/emitclass.py index ad95a1b0f323..3ab6932546a6 100644 --- a/mypyc/codegen/emitclass.py +++ b/mypyc/codegen/emitclass.py @@ -571,6 +571,7 @@ def generate_setup_for_class( emitter.emit_line("}") else: emitter.emit_line(f"self->vtable = {vtable_name};") + for i in range(0, len(cl.bitmap_attrs), BITMAP_BITS): field = emitter.bitmap_field(i) emitter.emit_line(f"self->{field} = 0;") diff --git a/mypyc/codegen/emitfunc.py b/mypyc/codegen/emitfunc.py index d945a28d8481..6088fb06dd32 100644 --- a/mypyc/codegen/emitfunc.py +++ b/mypyc/codegen/emitfunc.py @@ -72,6 +72,7 @@ from mypyc.ir.pprint import generate_names_for_ir from mypyc.ir.rtypes import ( RArray, + RInstance, RStruct, RTuple, RType, @@ -362,20 +363,23 @@ def visit_get_attr(self, op: GetAttr) -> None: prefer_method = cl.is_trait and attr_rtype.error_overlap if cl.get_method(op.attr, prefer_method=prefer_method): # Properties are essentially methods, so use vtable access for them. - version = "_TRAIT" if cl.is_trait else "" - self.emit_line( - "%s = CPY_GET_ATTR%s(%s, %s, %d, %s, %s); /* %s */" - % ( - dest, - version, - obj, - self.emitter.type_struct_name(rtype.class_ir), - rtype.getter_index(op.attr), - rtype.struct_name(self.names), - self.ctype(rtype.attr_type(op.attr)), - op.attr, + if cl.is_method_final(op.attr): + self.emit_method_call(f"{dest} = ", op.obj, op.attr, []) + else: + version = "_TRAIT" if cl.is_trait else "" + self.emit_line( + "%s = CPY_GET_ATTR%s(%s, %s, %d, %s, %s); /* %s */" + % ( + dest, + version, + obj, + self.emitter.type_struct_name(rtype.class_ir), + rtype.getter_index(op.attr), + rtype.struct_name(self.names), + self.ctype(rtype.attr_type(op.attr)), + op.attr, + ) ) - ) else: # Otherwise, use direct or offset struct access. attr_expr = self.get_attr_expr(obj, op, decl_cl) @@ -529,11 +533,13 @@ def visit_call(self, op: Call) -> None: def visit_method_call(self, op: MethodCall) -> None: """Call native method.""" dest = self.get_dest_assign(op) - obj = self.reg(op.obj) + self.emit_method_call(dest, op.obj, op.method, op.args) - rtype = op.receiver_type + def emit_method_call(self, dest: str, op_obj: Value, name: str, op_args: list[Value]) -> None: + obj = self.reg(op_obj) + rtype = op_obj.type + assert isinstance(rtype, RInstance) class_ir = rtype.class_ir - name = op.method method = rtype.class_ir.get_method(name) assert method is not None @@ -547,7 +553,7 @@ def visit_method_call(self, op: MethodCall) -> None: if method.decl.kind == FUNC_STATICMETHOD else [f"(PyObject *)Py_TYPE({obj})"] if method.decl.kind == FUNC_CLASSMETHOD else [obj] ) - args = ", ".join(obj_args + [self.reg(arg) for arg in op.args]) + args = ", ".join(obj_args + [self.reg(arg) for arg in op_args]) mtype = native_function_type(method, self.emitter) version = "_TRAIT" if rtype.class_ir.is_trait else "" if is_direct: @@ -567,7 +573,7 @@ def visit_method_call(self, op: MethodCall) -> None: rtype.struct_name(self.names), mtype, args, - op.method, + name, ) ) diff --git a/mypyc/ir/class_ir.py b/mypyc/ir/class_ir.py index 18f3cbcff987..94bf714b28d4 100644 --- a/mypyc/ir/class_ir.py +++ b/mypyc/ir/class_ir.py @@ -93,6 +93,7 @@ def __init__( is_generated: bool = False, is_abstract: bool = False, is_ext_class: bool = True, + is_final_class: bool = False, ) -> None: self.name = name self.module_name = module_name @@ -100,6 +101,7 @@ def __init__( self.is_generated = is_generated self.is_abstract = is_abstract self.is_ext_class = is_ext_class + self.is_final_class = is_final_class # An augmented class has additional methods separate from what mypyc generates. # Right now the only one is dataclasses. self.is_augmented = False @@ -199,7 +201,8 @@ def __repr__(self) -> str: "ClassIR(" "name={self.name}, module_name={self.module_name}, " "is_trait={self.is_trait}, is_generated={self.is_generated}, " - "is_abstract={self.is_abstract}, is_ext_class={self.is_ext_class}" + "is_abstract={self.is_abstract}, is_ext_class={self.is_ext_class}, " + "is_final_class={self.is_final_class}" ")".format(self=self) ) @@ -248,8 +251,7 @@ def has_method(self, name: str) -> bool: def is_method_final(self, name: str) -> bool: subs = self.subclasses() if subs is None: - # TODO: Look at the final attribute! - return False + return self.is_final_class if self.has_method(name): method_decl = self.method_decl(name) @@ -349,6 +351,7 @@ def serialize(self) -> JsonDict: "is_abstract": self.is_abstract, "is_generated": self.is_generated, "is_augmented": self.is_augmented, + "is_final_class": self.is_final_class, "inherits_python": self.inherits_python, "has_dict": self.has_dict, "allow_interpreted_subclasses": self.allow_interpreted_subclasses, @@ -404,6 +407,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR: ir.is_abstract = data["is_abstract"] ir.is_ext_class = data["is_ext_class"] ir.is_augmented = data["is_augmented"] + ir.is_final_class = data["is_final_class"] ir.inherits_python = data["inherits_python"] ir.has_dict = data["has_dict"] ir.allow_interpreted_subclasses = data["allow_interpreted_subclasses"] diff --git a/mypyc/ir/rtypes.py b/mypyc/ir/rtypes.py index fecfaee5ef77..53e3cee74e56 100644 --- a/mypyc/ir/rtypes.py +++ b/mypyc/ir/rtypes.py @@ -64,7 +64,7 @@ class RType: @abstractmethod def accept(self, visitor: RTypeVisitor[T]) -> T: - raise NotImplementedError + raise NotImplementedError() def short_name(self) -> str: return short_name(self.name) diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py index 0c9310e6a5ca..c98136ce06d2 100644 --- a/mypyc/irbuild/ll_builder.py +++ b/mypyc/irbuild/ll_builder.py @@ -1889,7 +1889,7 @@ def primitive_op( # Does this primitive map into calling a Python C API # or an internal mypyc C API function? if desc.c_function_name: - # TODO: Generate PrimitiOps here and transform them into CallC + # TODO: Generate PrimitiveOps here and transform them into CallC # ops only later in the lowering pass c_desc = CFunctionDescription( desc.name, @@ -1908,7 +1908,7 @@ def primitive_op( ) return self.call_c(c_desc, args, line, result_type) - # This primitve gets transformed in a lowering pass to + # This primitive gets transformed in a lowering pass to # lower-level IR ops using a custom transform function. coerced = [] diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 29e06439abdd..40a40b79df49 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -81,7 +81,11 @@ def build_type_map( # references even if there are import cycles. for module, cdef in classes: class_ir = ClassIR( - cdef.name, module.fullname, is_trait(cdef), is_abstract=cdef.info.is_abstract + cdef.name, + module.fullname, + is_trait(cdef), + is_abstract=cdef.info.is_abstract, + is_final_class=cdef.info.is_final, ) class_ir.is_ext_class = is_extension_class(cdef) if class_ir.is_ext_class: diff --git a/mypyc/irbuild/util.py b/mypyc/irbuild/util.py index ed01a59d1214..e27e509ad7fa 100644 --- a/mypyc/irbuild/util.py +++ b/mypyc/irbuild/util.py @@ -27,10 +27,16 @@ UnaryExpr, Var, ) +from mypy.semanal import refers_to_fullname +from mypy.types import FINAL_DECORATOR_NAMES DATACLASS_DECORATORS = {"dataclasses.dataclass", "attr.s", "attr.attrs"} +def is_final_decorator(d: Expression) -> bool: + return refers_to_fullname(d, FINAL_DECORATOR_NAMES) + + def is_trait_decorator(d: Expression) -> bool: return isinstance(d, RefExpr) and d.fullname == "mypy_extensions.trait" @@ -119,7 +125,10 @@ def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]: def is_extension_class(cdef: ClassDef) -> bool: if any( - not is_trait_decorator(d) and not is_dataclass_decorator(d) and not get_mypyc_attr_call(d) + not is_trait_decorator(d) + and not is_dataclass_decorator(d) + and not get_mypyc_attr_call(d) + and not is_final_decorator(d) for d in cdef.decorators ): return False diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index 7c2998874f78..49eb3028c9ee 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -2519,3 +2519,78 @@ class C: def test_final_attribute() -> None: assert C.A == -1 assert C.a == [-1] + +[case testClassWithFinalDecorator] +from typing import final + +@final +class C: + def a(self) -> int: + return 1 + +def test_class_final_attribute() -> None: + assert C().a() == 1 + + +[case testClassWithFinalDecoratorCtor] +from typing import final + +@final +class C: + def __init__(self) -> None: + self.a = 1 + + def b(self) -> int: + return 2 + + @property + def c(self) -> int: + return 3 + +def test_class_final_attribute() -> None: + assert C().a == 1 + assert C().b() == 2 + assert C().c == 3 + +[case testClassWithFinalDecoratorInheritedWithProperties] +from typing import final + +class B: + def a(self) -> int: + return 2 + + @property + def b(self) -> int: + return self.a() + 2 + + @property + def c(self) -> int: + return 3 + +def test_class_final_attribute_basic() -> None: + assert B().a() == 2 + assert B().b == 4 + assert B().c == 3 + +@final +class C(B): + def a(self) -> int: + return 1 + + @property + def b(self) -> int: + return self.a() + 1 + +def fn(cl: B) -> int: + return cl.a() + +def test_class_final_attribute_inherited() -> None: + assert C().a() == 1 + assert fn(C()) == 1 + assert B().a() == 2 + assert fn(B()) == 2 + + assert B().b == 4 + assert C().b == 2 + assert B().c == 3 + assert C().c == 3