[mypyc] Borrow references during chained attribute access (#12805)

If we have multiple native attribute access operations in succession, we
can borrow the temporaries. This avoids an incref and decref. For
example, when evaluating x.y.z, we don't need to incref the result
of x.y.

We need to make sure that the objects from which we borrow
values are not freed too early by adding keep_alive ops.

This is part of a wider reference counting optimization workstream.
All the improvements together produced around 5% performance
improvement in the richards benchmark.

In carefully constructed microbenchmarks 50+% improvements are
possible.

Author: JukkaL

mypyc/codegen/emitfunc.py

@@ -333,7 +333,7 @@ def visit_get_attr(self, op: GetAttr) -> None:
                         'PyErr_SetString({}, "attribute {} of {} undefined");'.format(
                             exc_class, repr(op.attr), repr(cl.name)))
 
-            if attr_rtype.is_refcounted:
+            if attr_rtype.is_refcounted and not op.is_borrowed:
                 if not merged_branch and not always_defined:
                     self.emitter.emit_line('} else {')
                 self.emitter.emit_inc_ref(dest, attr_rtype)

mypyc/ir/ops.py

@@ -599,13 +599,14 @@ class GetAttr(RegisterOp):
 
     error_kind = ERR_MAGIC
 
-    def __init__(self, obj: Value, attr: str, line: int) -> None:
+    def __init__(self, obj: Value, attr: str, line: int, *, borrow: bool = False) -> None:
         super().__init__(line)
         self.obj = obj
         self.attr = attr
         assert isinstance(obj.type, RInstance), 'Attribute access not supported: %s' % obj.type
         self.class_type = obj.type
         self.type = obj.type.attr_type(attr)
+        self.is_borrowed = borrow
 
     def sources(self) -> List[Value]:
         return [self.obj]

mypyc/ir/pprint.py

@@ -77,7 +77,11 @@ def visit_load_literal(self, op: LoadLiteral) -> str:
         return self.format('%r = %s%s', op, prefix, repr(op.value))
 
     def visit_get_attr(self, op: GetAttr) -> str:
-        return self.format('%r = %r.%s', op, op.obj, op.attr)
+        if op.is_borrowed:
+            borrow = 'borrow '
+        else:
+            borrow = ''
+        return self.format('%r = %s%r.%s', op, borrow, op.obj, op.attr)
 
     def visit_set_attr(self, op: SetAttr) -> str:
         if op.is_init:

mypyc/irbuild/builder.py

@@ -141,6 +141,8 @@ def __init__(self,
         # can also do quick lookups.
         self.imports: OrderedDict[str, None] = OrderedDict()
 
+        self.can_borrow = False
+
     # High-level control
 
     def set_module(self, module_name: str, module_path: str) -> None:
@@ -152,15 +154,23 @@ def set_module(self, module_name: str, module_path: str) -> None:
         self.module_path = module_path
 
     @overload
-    def accept(self, node: Expression) -> Value: ...
+    def accept(self, node: Expression, *, can_borrow: bool = False) -> Value: ...
 
     @overload
     def accept(self, node: Statement) -> None: ...
 
-    def accept(self, node: Union[Statement, Expression]) -> Optional[Value]:
-        """Transform an expression or a statement."""
+    def accept(self, node: Union[Statement, Expression], *,
+               can_borrow: bool = False) -> Optional[Value]:
+        """Transform an expression or a statement.
+
+        If can_borrow is true, prefer to generate a borrowed reference.
+        Borrowed references are faster since they don't require reference count
+        manipulation, but they are only safe to use in specific contexts.
+        """
         with self.catch_errors(node.line):
             if isinstance(node, Expression):
+                old_can_borrow = self.can_borrow
+                self.can_borrow = can_borrow
                 try:
                     res = node.accept(self.visitor)
                     res = self.coerce(res, self.node_type(node), node.line)
@@ -170,6 +180,9 @@ def accept(self, node: Union[Statement, Expression]) -> Optional[Value]:
                 # from causing more downstream trouble.
                 except UnsupportedException:
                     res = Register(self.node_type(node))
+                self.can_borrow = old_can_borrow
+                if not can_borrow:
+                    self.builder.flush_keep_alives()
                 return res
             else:
                 try:

mypyc/irbuild/expression.py

@@ -21,7 +21,7 @@
     Value, Register, TupleGet, TupleSet, BasicBlock, Assign, LoadAddress, RaiseStandardError
 )
 from mypyc.ir.rtypes import (
-    RTuple, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive
+    RTuple, RInstance, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive
 )
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD
 from mypyc.irbuild.format_str_tokenizer import (
@@ -130,8 +130,19 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
     if isinstance(expr.node, MypyFile) and expr.node.fullname in builder.imports:
         return builder.load_module(expr.node.fullname)
 
-    obj = builder.accept(expr.expr)
+    obj_rtype = builder.node_type(expr.expr)
+    if (isinstance(obj_rtype, RInstance)
+            and obj_rtype.class_ir.is_ext_class
+            and obj_rtype.class_ir.has_attr(expr.name)
+            and not obj_rtype.class_ir.get_method(expr.name)):
+        # Direct attribute access -> can borrow object
+        can_borrow = True
+    else:
+        can_borrow = False
+    obj = builder.accept(expr.expr, can_borrow=can_borrow)
+
     rtype = builder.node_type(expr)
+
     # Special case: for named tuples transform attribute access to faster index access.
     typ = get_proper_type(builder.types.get(expr.expr))
     if isinstance(typ, TupleType) and typ.partial_fallback.type.is_named_tuple:
@@ -142,7 +153,8 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
 
     check_instance_attribute_access_through_class(builder, expr, typ)
 
-    return builder.builder.get_attr(obj, expr.name, rtype, expr.line)
+    borrow = can_borrow and builder.can_borrow
+    return builder.builder.get_attr(obj, expr.name, rtype, expr.line, borrow=borrow)
 
 
 def check_instance_attribute_access_through_class(builder: IRBuilder,

mypyc/irbuild/ll_builder.py

@@ -105,6 +105,9 @@ def __init__(
         self.blocks: List[BasicBlock] = []
         # Stack of except handler entry blocks
         self.error_handlers: List[Optional[BasicBlock]] = [None]
+        # Values that we need to keep alive as long as we have borrowed
+        # temporaries. Use flush_keep_alives() to mark the end of the live range.
+        self.keep_alives: List[Value] = []
 
     # Basic operations
 
@@ -145,6 +148,11 @@ def self(self) -> Register:
         """
         return self.args[0]
 
+    def flush_keep_alives(self) -> None:
+        if self.keep_alives:
+            self.add(KeepAlive(self.keep_alives[:]))
+            self.keep_alives = []
+
     # Type conversions
 
     def box(self, src: Value) -> Value:
@@ -219,11 +227,14 @@ def coerce_nullable(self, src: Value, target_type: RType, line: int) -> Value:
 
     # Attribute access
 
-    def get_attr(self, obj: Value, attr: str, result_type: RType, line: int) -> Value:
+    def get_attr(self, obj: Value, attr: str, result_type: RType, line: int, *,
+                 borrow: bool = False) -> Value:
         """Get a native or Python attribute of an object."""
         if (isinstance(obj.type, RInstance) and obj.type.class_ir.is_ext_class
                 and obj.type.class_ir.has_attr(attr)):
-            return self.add(GetAttr(obj, attr, line))
+            if borrow:
+                self.keep_alives.append(obj)
+            return self.add(GetAttr(obj, attr, line, borrow=borrow))
         elif isinstance(obj.type, RUnion):
             return self.union_get_attr(obj, obj.type, attr, result_type, line)
         else:

mypyc/test-data/irbuild-classes.test

@@ -1240,3 +1240,51 @@ L0:
     r0 = ''
     __mypyc_self__.s = r0
     return 1
+
+[case testBorrowAttribute]
+def f(d: D) -> int:
+    return d.c.x
+
+class C:
+    x: int
+class D:
+    c: C
+[out]
+def f(d):
+    d :: __main__.D
+    r0 :: __main__.C
+    r1 :: int
+L0:
+    r0 = borrow d.c
+    r1 = r0.x
+    keep_alive d
+    return r1
+
+[case testNoBorrowOverPropertyAccess]
+class C:
+    d: D
+class D:
+    @property
+    def e(self) -> E:
+        return E()
+class E:
+    x: int
+def f(c: C) -> int:
+    return c.d.e.x
+[out]
+def D.e(self):
+    self :: __main__.D
+    r0 :: __main__.E
+L0:
+    r0 = E()
+    return r0
+def f(c):
+    c :: __main__.C
+    r0 :: __main__.D
+    r1 :: __main__.E
+    r2 :: int
+L0:
+    r0 = c.d
+    r1 = r0.e
+    r2 = r1.x
+    return r2

mypyc/test-data/refcount.test

@@ -917,3 +917,36 @@ L0:
     r5 = unbox(int, r4)
     dec_ref r4
     return r5
+
+[case testBorrowAttribute]
+def g() -> int:
+    d = D()
+    return d.c.x
+
+def f(d: D) -> int:
+    return d.c.x
+
+class C:
+    x: int
+class D:
+    c: C
+[out]
+def g():
+    r0, d :: __main__.D
+    r1 :: __main__.C
+    r2 :: int
+L0:
+    r0 = D()
+    d = r0
+    r1 = borrow d.c
+    r2 = r1.x
+    dec_ref d
+    return r2
+def f(d):
+    d :: __main__.D
+    r0 :: __main__.C
+    r1 :: int
+L0:
+    r0 = borrow d.c
+    r1 = r0.x
+    return r1